update to results_v2

requirements.txt

@@ -11,6 +11,7 @@ free-mujoco-py
 mujoco<=2.3.7
 numpy==1.24.2
 pandas==2.0.0
+pybullet_envs_gymnasium==0.4.0
 python-dateutil==2.8.2
 requests==2.28.2
 rliable==1.0.8

src/backend.py

@@ -1,10 +1,11 @@
-import json
+import fnmatch
 import os
 import random
-import re
-import tempfile
+import time
 
-from huggingface_hub import CommitOperationAdd, HfApi
+import pybullet_envs_gymnasium  # noqa: F401 pylint: disable=unused-import
+from datasets import load_dataset
+from huggingface_hub import HfApi
 
 from src.evaluation import evaluate
 from src.logging import setup_logger
@@ -12,71 +13,182 @@ from src.logging import setup_logger
 logger = setup_logger(__name__)
 
 API = HfApi(token=os.environ.get("TOKEN"))
-RESULTS_REPO = "open-rl-leaderboard/results"
+RESULTS_REPO = "open-rl-leaderboard/results_v2"
+
+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 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 _backend_routine():
     # List only the text classification models
-    rl_models = list(API.list_models(filter="reinforcement-learning"))
+    rl_models = [(model.modelId, model.sha) for model in API.list_models(filter=["reinforcement-learning"])]
     logger.info(f"Found {len(rl_models)} RL models")
-    compatible_models = []
-    for model in rl_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(rl_models) - set(evaluated_models))
+    pending_and_compatible_models = []
+    for model in pending_models:
         filenames = [sib.rfilename for sib in model.siblings]
         if "agent.pt" in filenames:
-            compatible_models.append((model.modelId, model.sha))
+            pending_and_compatible_models.append((model.modelId, model.sha))
 
-    logger.info(f"Found {len(compatible_models)} compatible models")
+    logger.info(f"Found {len(pending_and_compatible_models)} compatible pending models")
 
-    # Get the results
-    pattern = re.compile(r"^[^/]*/[^/]*/[^/]*results_[a-f0-9]+\.json$")
-    filenames = API.list_repo_files(RESULTS_REPO, repo_type="dataset")
-    filenames = [filename for filename in filenames if pattern.match(filename)]
-
-    evaluated_models = set()
-    for filename in filenames:
-        path = API.hf_hub_download(repo_id=RESULTS_REPO, filename=filename, repo_type="dataset")
-        with open(path) as fp:
-            report = json.load(fp)
-        evaluated_models.add((report["config"]["model_id"], report["config"]["model_sha"]))
+    if len(pending_and_compatible_models) == 0:
+        return None
 
-    # Find the models that are not associated with any results
-    pending_models = list(set(compatible_models) - evaluated_models)
-    logger.info(f"Found {len(pending_models)} pending models")
+    # Shuffle the dataset
+    random.shuffle(pending_and_compatible_models)
 
-    if len(pending_models) == 0:
-        return None
+    # Select a random model
+    repo_id, sha = pending_and_compatible_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
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        commits = []
-        model_id, sha = random.choice(pending_models)
-        logger.info(f"Running evaluation on {model_id}")
-        report = {"config": {"model_id": model_id, "model_sha": sha}}
+    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)
+    if len(env_ids) > 0:
+        env_id = env_ids[0]
+        logger.info(f"Running evaluation on {user_id}/{model_id}")
+
         try:
-            evaluations = evaluate(model_id, revision=sha)
+            episodic_returns = evaluate(repo_id, sha, env_id)
+            row["status"] = "DONE"
+            row["env_id"] = env_id
+            row["episodic_returns"] = episodic_returns
         except Exception as e:
-            logger.error(f"Error evaluating {model_id}: {e}")
-            evaluations = None
-
-        if evaluations is not None:
-            report["results"] = evaluations
-            report["status"] = "DONE"
-        else:
-            report["status"] = "FAILED"
-
-        # Update the results
-        dumped = json.dumps(report, indent=2)
-        path_in_repo = f"{model_id}/results_{sha}.json"
-        local_path = os.path.join(tmp_dir, path_in_repo)
-        os.makedirs(os.path.dirname(local_path), exist_ok=True)
-        with open(local_path, "w") as f:
-            f.write(dumped)
-
-        commits.append(CommitOperationAdd(path_in_repo=path_in_repo, path_or_fileobj=local_path))
-
-        API.create_commit(
-            repo_id=RESULTS_REPO, commit_message="Add evaluation results", operations=commits, repo_type="dataset"
-        )
+            logger.error(f"Error evaluating {repo_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.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():

src/evaluation.py

@@ -1,4 +1,3 @@
-import fnmatch
 import os
 from typing import Dict, SupportsFloat
 
@@ -303,35 +302,18 @@ def make(env_id):
     return thunk
 
 
-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(model_id, revision):
-    tags = API.model_info(model_id, revision=revision).tags
-
-    # Extract the environment IDs from the tags (usually only one)
-    env_ids = pattern_match(tags, ALL_ENV_IDS)
-    logger.info(f"Selected environments: {env_ids}")
-
-    results = {}
+def evaluate(repo_id, revision, env_id):
+    tags = API.model_info(repo_id, revision=revision).tags
 
     # Check if the agent exists
     try:
-        agent_path = API.hf_hub_download(repo_id=model_id, filename="agent.pt")
+        agent_path = API.hf_hub_download(repo_id=repo_id, filename="agent.pt")
     except EntryNotFoundError:
         logger.error("Agent not found")
         return None
 
     # Check safety
-    security = next(iter(API.get_paths_info(model_id, "agent.pt", expand=True))).security
+    security = next(iter(API.get_paths_info(repo_id, "agent.pt", expand=True))).security
     if security is None or "safe" not in security:
         logger.warn("Agent safety not available")
         # return None
@@ -341,25 +323,22 @@ def evaluate(model_id, revision):
 
     # Load the agent
     try:
-        agent = torch.jit.load(agent_path).to("cuda")
+        agent = torch.jit.load(agent_path)
     except Exception as e:
         logger.error(f"Error loading agent: {e}")
         return None
 
     # Evaluate the agent on the environments
-    for env_id in env_ids:
-        envs = gym.vector.SyncVectorEnv([make(env_id) for _ in range(1)])
-        observations, _ = envs.reset()
-        episodic_returns = []
-        while len(episodic_returns) < NUM_EPISODES:
-            actions = agent(torch.tensor(observations)).numpy()
-            observations, _, _, _, infos = envs.step(actions)
-            if "final_info" in infos:
-                for info in infos["final_info"]:
-                    if info is None or "episode" not in info:
-                        continue
-                    episodic_returns.append(float(info["episode"]["r"]))
-
-        results[env_id] = {"episodic_returns": episodic_returns}
-        logger.info(f"Environment {env_id}: {np.mean(episodic_returns)} ± {np.std(episodic_returns)}")
-    return results
+    envs = gym.vector.SyncVectorEnv([make(env_id) for _ in range(1)])
+    observations, _ = envs.reset()
+    episodic_returns = []
+    while len(episodic_returns) < NUM_EPISODES:
+        actions = agent(torch.tensor(observations)).numpy()
+        observations, _, _, _, infos = envs.step(actions)
+        if "final_info" in infos:
+            for info in infos["final_info"]:
+                if info is None or "episode" not in info:
+                    continue
+                episodic_returns.append(float(info["episode"]["r"]))
+
+    return episodic_returns