scripts/train_push_unit2.py

from __future__ import annotations

import datetime as dt
import json
import pickle
import random
import shutil
from pathlib import Path

import gymnasium as gym
import imageio
import numpy as np
from huggingface_hub import HfApi
from huggingface_hub.errors import HfHubHTTPError
from huggingface_hub.repocard import metadata_eval_result, metadata_save
from tqdm import tqdm


USERNAME = "Sami94"
STUDENT_NAME = "Sami Chellia"
OUTPUT_DIR = Path("artifacts/unit2")


def initialize_q_table(state_space: int, action_space: int) -> np.ndarray:
    return np.zeros((state_space, action_space))


def greedy_policy(qtable: np.ndarray, state: int) -> int:
    return int(np.argmax(qtable[state][:]))


def epsilon_greedy_policy(qtable: np.ndarray, state: int, epsilon: float, env: gym.Env) -> int:
    if random.uniform(0, 1) > epsilon:
        return greedy_policy(qtable, state)
    return int(env.action_space.sample())


def train(
    *,
    n_training_episodes: int,
    learning_rate: float,
    min_epsilon: float,
    max_epsilon: float,
    decay_rate: float,
    gamma: float,
    env: gym.Env,
    max_steps: int,
    qtable: np.ndarray,
) -> np.ndarray:
    for episode in tqdm(range(n_training_episodes), desc=f"training {env.spec.id}"):
        epsilon = min_epsilon + (max_epsilon - min_epsilon) * np.exp(-decay_rate * episode)
        state, _ = env.reset()

        for _ in range(max_steps):
            action = epsilon_greedy_policy(qtable, int(state), epsilon, env)
            new_state, reward, terminated, truncated, _ = env.step(action)
            new_state = int(new_state)

            qtable[int(state)][action] += learning_rate * (
                reward + gamma * np.max(qtable[new_state]) - qtable[int(state)][action]
            )

            if terminated or truncated:
                break
            state = new_state

    return qtable


def evaluate_agent(
    env: gym.Env,
    max_steps: int,
    n_eval_episodes: int,
    qtable: np.ndarray,
    seed: list[int],
) -> tuple[float, float]:
    episode_rewards: list[float] = []
    for episode in tqdm(range(n_eval_episodes), desc=f"evaluating {env.spec.id}"):
        if seed:
            state, _ = env.reset(seed=seed[episode])
        else:
            state, _ = env.reset()
        total_rewards_ep = 0.0

        for _ in range(max_steps):
            action = greedy_policy(qtable, int(state))
            new_state, reward, terminated, truncated, _ = env.step(action)
            total_rewards_ep += float(reward)
            if terminated or truncated:
                break
            state = new_state

        episode_rewards.append(total_rewards_ep)

    return float(np.mean(episode_rewards)), float(np.std(episode_rewards))


def record_video(env: gym.Env, qtable: np.ndarray, out_path: Path, fps: int = 1) -> None:
    images = []
    terminated = False
    truncated = False
    state, _ = env.reset(seed=random.randint(0, 500))
    images.append(env.render())

    while not (terminated or truncated):
        action = greedy_policy(qtable, int(state))
        state, _, terminated, truncated, _ = env.step(action)
        images.append(env.render())

    imageio.mimsave(out_path, [np.array(img) for img in images], fps=fps)


def push_to_hub(repo_id: str, model: dict, env: gym.Env, local_repo_path: Path) -> str:
    _, repo_name = repo_id.split("/")
    api = HfApi()

    if local_repo_path.exists():
        shutil.rmtree(local_repo_path)
    local_repo_path.mkdir(parents=True, exist_ok=True)

    if env.spec.kwargs.get("map_name"):
        model["map_name"] = env.spec.kwargs.get("map_name")
        if env.spec.kwargs.get("is_slippery", "") is False:
            model["slippery"] = False

    with (local_repo_path / "q-learning.pkl").open("wb") as f:
        pickle.dump(model, f)

    mean_reward, std_reward = evaluate_agent(
        env,
        model["max_steps"],
        model["n_eval_episodes"],
        model["qtable"],
        model["eval_seed"],
    )

    evaluate_data = {
        "env_id": model["env_id"],
        "mean_reward": mean_reward,
        "std_reward": std_reward,
        "n_eval_episodes": model["n_eval_episodes"],
        "eval_datetime": dt.datetime.now().isoformat(),
        "student": STUDENT_NAME,
        "hf_username": USERNAME,
    }
    (local_repo_path / "results.json").write_text(json.dumps(evaluate_data, indent=2), encoding="utf-8")

    env_name = model["env_id"]
    if env.spec.kwargs.get("map_name"):
        env_name += "-" + env.spec.kwargs.get("map_name")
    if env.spec.kwargs.get("is_slippery", "") is False:
        env_name += "-no_slippery"

    metadata = {
        "tags": [
            env_name,
            "q-learning",
            "reinforcement-learning",
            "custom-implementation",
            "huggingface-deep-rl-course",
        ]
    }
    eval_metadata = metadata_eval_result(
        model_pretty_name=repo_name,
        task_pretty_name="reinforcement-learning",
        task_id="reinforcement-learning",
        metrics_pretty_name="mean_reward",
        metrics_id="mean_reward",
        metrics_value=f"{mean_reward:.2f} +/- {std_reward:.2f}",
        dataset_pretty_name=env_name,
        dataset_id=env_name,
    )
    metadata = {**metadata, **eval_metadata}

    readme_path = local_repo_path / "README.md"
    readme_path.write_text(
        f"""# Q-Learning Agent for {env_name}

Student: {STUDENT_NAME}
Hugging Face username: {USERNAME}

This repository contains a Q-Learning agent trained for the Hugging Face Deep RL course.

Mean reward: {mean_reward:.2f} +/- {std_reward:.2f}

```python
from huggingface_hub import hf_hub_download
import pickle

model_path = hf_hub_download(repo_id="{repo_id}", filename="q-learning.pkl")
with open(model_path, "rb") as f:
    model = pickle.load(f)
```
""",
        encoding="utf-8",
    )
    metadata_save(readme_path, metadata)

    record_video(env, model["qtable"], local_repo_path / "replay.mp4")

    try:
        repo_url = api.create_repo(repo_id=repo_id, exist_ok=True)
        api.upload_folder(repo_id=repo_id, folder_path=local_repo_path, path_in_repo=".")
        return str(repo_url)
    except HfHubHTTPError as exc:
        print(f"Hub push failed for {repo_id}: {exc}")
        print(f"Local artifacts saved in {local_repo_path.resolve()}")
        return f"LOCAL_ONLY:{local_repo_path.resolve()}"


def run_frozenlake() -> dict:
    env = gym.make("FrozenLake-v1", map_name="4x4", is_slippery=False, render_mode="rgb_array")
    qtable = initialize_q_table(env.observation_space.n, env.action_space.n)
    params = {
        "n_training_episodes": 10000,
        "learning_rate": 0.7,
        "n_eval_episodes": 100,
        "env_id": "FrozenLake-v1",
        "max_steps": 99,
        "gamma": 0.95,
        "eval_seed": [],
        "max_epsilon": 1.0,
        "min_epsilon": 0.05,
        "decay_rate": 0.0005,
    }
    qtable = train(env=env, qtable=qtable, **{k: params[k] for k in [
        "n_training_episodes",
        "learning_rate",
        "min_epsilon",
        "max_epsilon",
        "decay_rate",
        "gamma",
        "max_steps",
    ]})
    model = {**params, "qtable": qtable}
    repo_id = f"{USERNAME}/q-FrozenLake-v1-4x4-noSlippery"
    url = push_to_hub(repo_id, model, env, OUTPUT_DIR / "q-FrozenLake-v1-4x4-noSlippery")
    env.close()
    return {"repo_id": repo_id, "url": url}


def run_taxi() -> dict:
    env = gym.make("Taxi-v3", render_mode="rgb_array")
    qtable = initialize_q_table(env.observation_space.n, env.action_space.n)
    params = {
        "n_training_episodes": 25000,
        "learning_rate": 0.7,
        "n_eval_episodes": 100,
        "eval_seed": [
            16, 54, 165, 177, 191, 191, 120, 80, 149, 178, 48, 38, 6, 125, 174,
            73, 50, 172, 100, 148, 146, 6, 25, 40, 68, 148, 49, 167, 9, 97,
            164, 176, 61, 7, 54, 55, 161, 131, 184, 51, 170, 12, 120, 113,
            95, 126, 51, 98, 36, 135, 54, 82, 45, 95, 89, 59, 95, 124, 9,
            113, 58, 85, 51, 134, 121, 169, 105, 21, 30, 11, 50, 65, 12, 43,
            82, 145, 152, 97, 106, 55, 31, 85, 38, 112, 102, 168, 123, 97,
            21, 83, 158, 26, 80, 63, 5, 81, 32, 11, 28, 148,
        ],
        "env_id": "Taxi-v3",
        "max_steps": 99,
        "gamma": 0.95,
        "max_epsilon": 1.0,
        "min_epsilon": 0.05,
        "decay_rate": 0.005,
    }
    qtable = train(env=env, qtable=qtable, **{k: params[k] for k in [
        "n_training_episodes",
        "learning_rate",
        "min_epsilon",
        "max_epsilon",
        "decay_rate",
        "gamma",
        "max_steps",
    ]})
    model = {**params, "qtable": qtable}
    repo_id = f"{USERNAME}/q-Taxi-v3"
    url = push_to_hub(repo_id, model, env, OUTPUT_DIR / "q-Taxi-v3")
    env.close()
    return {"repo_id": repo_id, "url": url}


def main() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    results = [run_frozenlake(), run_taxi()]
    (OUTPUT_DIR / "pushed_repos.json").write_text(json.dumps(results, indent=2), encoding="utf-8")
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()