src/backend.py

import fnmatch
import importlib
import json
import os
import random
import shutil
import sys
import time
import zipfile
from pathlib import Path
from typing import Optional

import numpy as np
import rl_zoo3.import_envs  # noqa: F401 pylint: disable=unused-import
import torch as th
import yaml
from datasets import load_dataset
from huggingface_hub import HfApi
from huggingface_hub.utils import EntryNotFoundError
from huggingface_sb3 import EnvironmentName, ModelName, ModelRepoId, load_from_hub
from requests.exceptions import HTTPError
from rl_zoo3 import ALGOS, create_test_env, get_latest_run_id, get_saved_hyperparams
from rl_zoo3.exp_manager import ExperimentManager
from rl_zoo3.utils import get_model_path
from stable_baselines3.common.utils import set_random_seed

from src.logging import setup_logger

ALL_ENV_IDS = [
    "AdventureNoFrameskip-v4",
    "AirRaidNoFrameskip-v4",
    "AlienNoFrameskip-v4",
    "AmidarNoFrameskip-v4",
    "AssaultNoFrameskip-v4",
    "AsterixNoFrameskip-v4",
    "AsteroidsNoFrameskip-v4",
    "AtlantisNoFrameskip-v4",
    "BankHeistNoFrameskip-v4",
    "BattleZoneNoFrameskip-v4",
    "BeamRiderNoFrameskip-v4",
    "BerzerkNoFrameskip-v4",
    "BowlingNoFrameskip-v4",
    "BoxingNoFrameskip-v4",
    "BreakoutNoFrameskip-v4",
    "CarnivalNoFrameskip-v4",
    "CentipedeNoFrameskip-v4",
    "ChopperCommandNoFrameskip-v4",
    "CrazyClimberNoFrameskip-v4",
    "DefenderNoFrameskip-v4",
    "DemonAttackNoFrameskip-v4",
    "DoubleDunkNoFrameskip-v4",
    "ElevatorActionNoFrameskip-v4",
    "EnduroNoFrameskip-v4",
    "FishingDerbyNoFrameskip-v4",
    "FreewayNoFrameskip-v4",
    "FrostbiteNoFrameskip-v4",
    "GopherNoFrameskip-v4",
    "GravitarNoFrameskip-v4",
    "HeroNoFrameskip-v4",
    "IceHockeyNoFrameskip-v4",
    "JamesbondNoFrameskip-v4",
    "JourneyEscapeNoFrameskip-v4",
    "KangarooNoFrameskip-v4",
    "KrullNoFrameskip-v4",
    "KungFuMasterNoFrameskip-v4",
    "MontezumaRevengeNoFrameskip-v4",
    "MsPacmanNoFrameskip-v4",
    "NameThisGameNoFrameskip-v4",
    "PhoenixNoFrameskip-v4",
    "PitfallNoFrameskip-v4",
    "PongNoFrameskip-v4",
    "PooyanNoFrameskip-v4",
    "PrivateEyeNoFrameskip-v4",
    "QbertNoFrameskip-v4",
    "RiverraidNoFrameskip-v4",
    "RoadRunnerNoFrameskip-v4",
    "RobotankNoFrameskip-v4",
    "SeaquestNoFrameskip-v4",
    "SkiingNoFrameskip-v4",
    "SolarisNoFrameskip-v4",
    "SpaceInvadersNoFrameskip-v4",
    "StarGunnerNoFrameskip-v4",
    "TennisNoFrameskip-v4",
    "TimePilotNoFrameskip-v4",
    "TutankhamNoFrameskip-v4",
    "UpNDownNoFrameskip-v4",
    "VentureNoFrameskip-v4",
    "VideoPinballNoFrameskip-v4",
    "WizardOfWorNoFrameskip-v4",
    "YarsRevengeNoFrameskip-v4",
    "ZaxxonNoFrameskip-v4",
    # Box2D
    "BipedalWalker-v3",
    "BipedalWalkerHardcore-v3",
    "CarRacing-v2",
    "LunarLander-v2",
    "LunarLanderContinuous-v2",
    # Toy text
    "Blackjack-v1",
    "CliffWalking-v0",
    "FrozenLake-v1",
    "FrozenLake8x8-v1",
    # Classic control
    "Acrobot-v1",
    "CartPole-v1",
    "MountainCar-v0",
    "MountainCarContinuous-v0",
    "Pendulum-v1",
    # MuJoCo
    "Ant-v4",
    "HalfCheetah-v4",
    "Hopper-v4",
    "Humanoid-v4",
    "HumanoidStandup-v4",
    "InvertedDoublePendulum-v4",
    "InvertedPendulum-v4",
    "Pusher-v4",
    "Reacher-v4",
    "Swimmer-v4",
    "Walker2d-v4",
    # PyBullet
    "AntBulletEnv-v0",
    "HalfCheetahBulletEnv-v0",
    "HopperBulletEnv-v0",
    "HumanoidBulletEnv-v0",
    "InvertedDoublePendulumBulletEnv-v0",
    "InvertedPendulumSwingupBulletEnv-v0",
    "MinitaurBulletEnv-v0",
    "ReacherBulletEnv-v0",
    "Walker2DBulletEnv-v0",
]


def download_from_hub(
    algo: str,
    env_name: EnvironmentName,
    exp_id: int,
    folder: str,
    organization: str,
    repo_name: Optional[str] = None,
    force: bool = False,
) -> None:
    """
    Try to load a model from the Huggingface hub
    and save it following the RL Zoo structure.
    Default repo name is {organization}/{algo}-{env_id}
    where repo_name = {algo}-{env_id}

    :param algo: Algorithm
    :param env_name: Environment name
    :param exp_id: Experiment id
    :param folder: Log folder
    :param organization: Huggingface organization
    :param repo_name: Overwrite default repository name
    :param force: Allow overwritting the folder
        if it already exists.
    """

    model_name = ModelName(algo, env_name)

    if repo_name is None:
        repo_name = model_name  # Note: model name is {algo}-{env_name}

    # Note: repo id is {organization}/{repo_name}
    repo_id = ModelRepoId(organization, repo_name)
    logger.info(f"Downloading from https://huggingface.co/{repo_id}")

    checkpoint = load_from_hub(repo_id, model_name.filename)
    try:
        config_path = load_from_hub(repo_id, "config.yml")
    except EntryNotFoundError:  # hotfix for old models
        config_path = load_from_hub(repo_id, "config.json")
        with open(config_path, "r") as f:
            config = json.load(f)
        config_path = config_path.replace(".json", ".yml")
        with open(config_path, "w") as f:
            yaml.dump(config, f)

    # If VecNormalize, download
    try:
        vec_normalize_stats = load_from_hub(repo_id, "vec_normalize.pkl")
    except HTTPError:
        logger.info("No normalization file")
        vec_normalize_stats = None

    try:
        saved_args = load_from_hub(repo_id, "args.yml")
    except EntryNotFoundError:
        logger.info("No args file")
        saved_args = None

    try:
        env_kwargs = load_from_hub(repo_id, "env_kwargs.yml")
    except EntryNotFoundError:
        logger.info("No env_kwargs file")
        env_kwargs = None

    try:
        train_eval_metrics = load_from_hub(repo_id, "train_eval_metrics.zip")
    except EntryNotFoundError:
        logger.info("No train_eval_metrics file")
        train_eval_metrics = None

    if exp_id == 0:
        exp_id = get_latest_run_id(os.path.join(folder, algo), env_name) + 1
    # Sanity checks
    if exp_id > 0:
        log_path = os.path.join(folder, algo, f"{env_name}_{exp_id}")
    else:
        log_path = os.path.join(folder, algo)

    # Check that the folder does not exist
    log_folder = Path(log_path)
    if log_folder.is_dir():
        if force:
            logger.info(f"The folder {log_path} already exists, overwritting")
            # Delete the current one to avoid errors
            shutil.rmtree(log_path)
        else:
            raise ValueError(
                f"The folder {log_path} already exists, use --force to overwrite it, " "or choose '--exp-id 0' to create a new folder"
            )

    logger.info(f"Saving to {log_path}")
    # Create folder structure
    os.makedirs(log_path, exist_ok=True)
    config_folder = os.path.join(log_path, env_name)
    os.makedirs(config_folder, exist_ok=True)

    # Copy config files and saved stats
    shutil.copy(checkpoint, os.path.join(log_path, f"{env_name}.zip"))
    if saved_args is not None:
        shutil.copy(saved_args, os.path.join(config_folder, "args.yml"))
    shutil.copy(config_path, os.path.join(config_folder, "config.yml"))
    if env_kwargs is not None:
        shutil.copy(env_kwargs, os.path.join(config_folder, "env_kwargs.yml"))
    if vec_normalize_stats is not None:
        shutil.copy(vec_normalize_stats, os.path.join(config_folder, "vecnormalize.pkl"))

    # Extract monitor file and evaluation file
    if train_eval_metrics is not None:
        with zipfile.ZipFile(train_eval_metrics, "r") as zip_ref:
            zip_ref.extractall(log_path)


def pattern_match(patterns, source_list):
    if isinstance(patterns, str):
        patterns = [patterns]

    env_ids = set()
    for pattern in patterns:
        for matching in fnmatch.filter(source_list, pattern):
            env_ids.add(matching)
    return sorted(list(env_ids))


def evaluate(
    user_id,
    repo_name,
    env="CartPole-v1",
    folder="rl-trained-agents",
    algo="ppo",
    # n_timesteps=1000,
    n_episodes=50,
    num_threads=-1,
    n_envs=1,
    exp_id=0,
    verbose=1,
    no_render=False,
    deterministic=False,
    device="auto",
    load_best=False,
    load_checkpoint=None,
    load_last_checkpoint=False,
    stochastic=False,
    norm_reward=False,
    seed=0,
    reward_log="",
    gym_packages=[],
    env_kwargs=None,
    custom_objects=False,
    progress=False,
):
    """
    Enjoy trained agent

    :param env: (str) Environment ID
    :param folder: (str) Log folder
    :param algo: (str) RL Algorithm
    :param n_timesteps: (int) Number of timesteps
    :param num_threads: (int) Number of threads for PyTorch
    :param n_envs: (int) Number of environments
    :param exp_id: (int) Experiment ID (default: 0: latest, -1: no exp folder)
    :param verbose: (int) Verbose mode (0: no output, 1: INFO)
    :param no_render: (bool) Do not render the environment (useful for tests)
    :param deterministic: (bool) Use deterministic actions
    :param device: (str) PyTorch device to be use (ex: cpu, cuda...)
    :param load_best: (bool) Load best model instead of last model if available
    :param load_checkpoint: (int) Load checkpoint instead of last model if available
    :param load_last_checkpoint: (bool) Load last checkpoint instead of last model if available
    :param stochastic: (bool) Use stochastic actions
    :param norm_reward: (bool) Normalize reward if applicable (trained with VecNormalize)
    :param seed: (int) Random generator seed
    :param reward_log: (str) Where to log reward
    :param gym_packages: (List[str]) Additional external Gym environment package modules to import
    :param env_kwargs: (Dict[str, Any]) Optional keyword argument to pass to the env constructor
    :param custom_objects: (bool) Use custom objects to solve loading issues
    :param progress: (bool) if toggled, display a progress bar using tqdm and rich
    """

    # Going through custom gym packages to let them register in the global registory
    for env_module in gym_packages:
        importlib.import_module(env_module)

    env_name = EnvironmentName(env)

    # try:
    #     _, model_path, log_path = get_model_path(
    #         exp_id,
    #         folder,
    #         algo,
    #         env_name,
    #         load_best,
    #         load_checkpoint,
    #         load_last_checkpoint,
    #     )
    # except (AssertionError, ValueError) as e:
    #     # Special case for rl-trained agents
    #     # auto-download from the hub
    #     if "rl-trained-agents" not in folder:
    #         raise e
    #     else:
    #         logger.info("Pretrained model not found, trying to download it from sb3 Huggingface hub: https://huggingface.co/sb3")
    # Auto-download
    download_from_hub(
        algo=algo,
        env_name=env_name,
        exp_id=exp_id,
        folder=folder,
        # organization="sb3",
        organization=user_id,
        # repo_name=None,
        repo_name=repo_name,
        force=False,
    )
    # Try again
    _, model_path, log_path = get_model_path(
        exp_id,
        folder,
        algo,
        env_name,
        load_best,
        load_checkpoint,
        load_last_checkpoint,
    )

    logger.info(f"Loading {model_path}")

    # Off-policy algorithm only support one env for now
    off_policy_algos = ["qrdqn", "dqn", "ddpg", "sac", "her", "td3", "tqc"]

    set_random_seed(seed)

    if num_threads > 0:
        if verbose > 1:
            logger.info(f"Setting torch.num_threads to {num_threads}")
        th.set_num_threads(num_threads)

    is_atari = ExperimentManager.is_atari(env_name.gym_id)
    is_minigrid = ExperimentManager.is_minigrid(env_name.gym_id)

    stats_path = os.path.join(log_path, env_name)
    hyperparams, maybe_stats_path = get_saved_hyperparams(stats_path, norm_reward=norm_reward, test_mode=True)

    # load env_kwargs if existing
    env_kwargs = {}
    args_path = os.path.join(log_path, env_name, "args.yml")
    if os.path.isfile(args_path):
        with open(args_path) as f:
            loaded_args = yaml.load(f, Loader=yaml.UnsafeLoader)
            if loaded_args["env_kwargs"] is not None:
                env_kwargs = loaded_args["env_kwargs"]
    # overwrite with command line arguments
    if env_kwargs is not None:
        env_kwargs.update(env_kwargs)

    log_dir = reward_log if reward_log != "" else None

    env = create_test_env(
        env_name.gym_id,
        n_envs=n_envs,
        stats_path=maybe_stats_path,
        seed=seed,
        log_dir=log_dir,
        should_render=not no_render,
        hyperparams=hyperparams,
        env_kwargs=env_kwargs,
    )

    kwargs = dict(seed=seed)
    if algo in off_policy_algos:
        # Dummy buffer size as we don't need memory to enjoy the trained agent
        kwargs.update(dict(buffer_size=1))
        # Hack due to breaking change in v1.6
        # handle_timeout_termination cannot be at the same time
        # with optimize_memory_usage
        if "optimize_memory_usage" in hyperparams:
            kwargs.update(optimize_memory_usage=False)

    # Check if we are running python 3.8+
    # we need to patch saved model under python 3.6/3.7 to load them
    newer_python_version = sys.version_info.major == 3 and sys.version_info.minor >= 8

    custom_objects = {}
    if newer_python_version or custom_objects:
        custom_objects = {
            "learning_rate": 0.0,
            "lr_schedule": lambda _: 0.0,
            "clip_range": lambda _: 0.0,
        }

    if "HerReplayBuffer" in hyperparams.get("replay_buffer_class", ""):
        kwargs["env"] = env

    model = ALGOS[algo].load(model_path, custom_objects=custom_objects, device=device, **kwargs)
    obs = env.reset()

    # Deterministic by default except for atari games
    stochastic = stochastic or (is_atari or is_minigrid) and not deterministic
    deterministic = not stochastic

    episode_reward = 0.0
    episode_rewards, episode_lengths = [], []
    ep_len = 0
    # For HER, monitor success rate
    successes = []
    lstm_states = None
    episode_start = np.ones((env.num_envs,), dtype=bool)

    # generator = range(n_timesteps)
    # if progress:
    #     if tqdm is None:
    #         raise ImportError("Please install tqdm and rich to use the progress bar")
    #     generator = tqdm(generator)

    try:
        # for _ in generator:
        while len(episode_rewards) < n_episodes:
            action, lstm_states = model.predict(
                obs,  # type: ignore[arg-type]
                state=lstm_states,
                episode_start=episode_start,
                deterministic=deterministic,
            )
            obs, reward, done, infos = env.step(action)

            episode_start = done

            if not no_render:
                env.render("human")

            episode_reward += reward[0]
            ep_len += 1

            if n_envs == 1:
                # For atari the return reward is not the atari score
                # so we have to get it from the infos dict
                if is_atari and infos is not None and verbose >= 1:
                    episode_infos = infos[0].get("episode")
                    if episode_infos is not None:
                        logger.info(f"Atari Episode Score: {episode_infos['r']:.2f}")
                        logger.info(f"Atari Episode Length {episode_infos['l']}")
                        episode_rewards.append(episode_infos["r"])
                        episode_lengths.append(episode_infos["l"])

                if done and not is_atari and verbose > 0:
                    # NOTE: for env using VecNormalize, the mean reward
                    # is a normalized reward when `--norm_reward` flag is passed
                    logger.info(f"Episode Reward: {episode_reward:.2f}")
                    logger.info(f"Episode Length {ep_len}")
                    episode_rewards.append(episode_reward)
                    episode_lengths.append(ep_len)
                    episode_reward = 0.0
                    ep_len = 0

                # Reset also when the goal is achieved when using HER
                if done and infos[0].get("is_success") is not None:
                    if verbose > 1:
                        logger.info(f"Success? {infos[0].get('is_success', False)}")

                    if infos[0].get("is_success") is not None:
                        successes.append(infos[0].get("is_success", False))
                        episode_reward, ep_len = 0.0, 0

    except KeyboardInterrupt:
        pass

    if verbose > 0 and len(successes) > 0:
        logger.info(f"Success rate: {100 * np.mean(successes):.2f}%")

    if verbose > 0 and len(episode_rewards) > 0:
        logger.info(f"{len(episode_rewards)} Episodes")
        logger.info(f"Mean reward: {np.mean(episode_rewards):.2f} +/- {np.std(episode_rewards):.2f}")

    if verbose > 0 and len(episode_lengths) > 0:
        logger.info(f"Mean episode length: {np.mean(episode_lengths):.2f} +/- {np.std(episode_lengths):.2f}")

    env.close()

    return episode_rewards


logger = setup_logger(__name__)

API = HfApi(token=os.environ.get("TOKEN"))
RESULTS_REPO = "open-rl-leaderboard/results_v2"


def _backend_routine():
    # List only the text classification models
    sb3_models = [(model.modelId, model.sha) for model in API.list_models(filter=["reinforcement-learning", "stable-baselines3"])]
    logger.info(f"Found {len(sb3_models)} SB3 models")
    dataset = load_dataset(RESULTS_REPO, split="train", download_mode="force_redownload", verification_mode="no_checks")
    evaluated_models = [("/".join([x["user_id"], x["model_id"]]), x["sha"]) for x in dataset]
    pending_models = list(set(sb3_models) - set(evaluated_models))
    logger.info(f"Found {len(pending_models)} pending models")

    if len(pending_models) == 0:
        return None

    # Shuffle the dataset
    random.shuffle(pending_models)

    # Select a random model
    repo_id, sha = pending_models.pop()
    user_id, model_id = repo_id.split("/")
    row = {"model_id": model_id, "user_id": user_id, "sha": sha}

    # Run an evaluation on the models
    try:
        model_info = API.model_info(repo_id, revision=sha)
        # Extract the environment IDs from the tags (usually only one)
        env_ids = pattern_match(model_info.tags, ALL_ENV_IDS)
    except Exception as e:
        logger.error(f"Error fetching model info for {repo_id}: {e}")
        logger.exception(e)
        env_ids = []

    if len(env_ids) > 0:
        env = env_ids[0]
        logger.info(f"Running evaluation on {user_id}/{model_id}")
        algo = model_info.model_index[0]["name"].lower()

        try:
            episodic_returns = evaluate(user_id, model_id, env, "rl-trained-agents", algo, no_render=True, verbose=1)
            row["status"] = "DONE"
            row["env_id"] = env
            row["episodic_returns"] = episodic_returns
        except Exception as e:
            logger.error(f"Error evaluating {model_id}: {e}")
            logger.exception(e)
            row["status"] = "FAILED"

    else:
        logger.error(f"No environment found for {model_id}")
        row["status"] = "FAILED"

    # load the last version of the dataset
    dataset = load_dataset(RESULTS_REPO, split="train", download_mode="force_redownload", verification_mode="no_checks")
    dataset = dataset.add_item(row)
    dataset.push_to_hub(RESULTS_REPO, split="train", token=API.token)
    time.sleep(60)  # Sleep for 1 minute to avoid rate limiting


def backend_routine():
    try:
        _backend_routine()
    except Exception as e:
        logger.error(f"{e.__class__.__name__}: {str(e)}")
        logger.exception(e)


if __name__ == "__main__":
    while True:
        backend_routine()