Add gym-aloha, rename simxarm -> xarm, refactor

lerobot/common/datasets/factory.py

@@ -19,10 +19,10 @@ def make_dataset(
     normalize=True,
     stats_path=None,
 ):
-    if cfg.env.name == "simxarm":
-        from lerobot.common.datasets.simxarm import SimxarmDataset
+    if cfg.env.name == "xarm":
+        from lerobot.common.datasets.xarm import XarmDataset
 
-        clsfunc = SimxarmDataset
+        clsfunc = XarmDataset
 
     elif cfg.env.name == "pusht":
         from lerobot.common.datasets.pusht import PushtDataset

lerobot/common/datasets/xarm.py

@@ -24,7 +24,7 @@ def download(raw_dir):
     zip_path.unlink()
 
 
-class SimxarmDataset(torch.utils.data.Dataset):
+class XarmDataset(torch.utils.data.Dataset):
     available_datasets = [
         "xarm_lift_medium",
     ]

lerobot/common/envs/factory.py

@@ -1,3 +1,5 @@
+import importlib
+
 import gymnasium as gym
 
 
@@ -8,43 +10,28 @@ def make_env(cfg, num_parallel_envs=0) -> gym.Env | gym.vector.SyncVectorEnv:
     """
     kwargs = {
         "obs_type": "pixels_agent_pos",
+        "render_mode": "rgb_array",
         "max_episode_steps": cfg.env.episode_length,
         "visualization_width": 384,
         "visualization_height": 384,
     }
 
-    if cfg.env.name == "simxarm":
-        import gym_xarm  # noqa: F401
+    package_name = f"gym_{cfg.env.name}"
 
-        assert cfg.env.task == "lift"
-        env_fn = lambda: gym.make(  # noqa: E731
-            "gym_xarm/XarmLift-v0",
-            **kwargs,
+    try:
+        importlib.import_module(package_name)
+    except ModuleNotFoundError as e:
+        print(
+            f"{package_name} is not installed. Please install it with `pip install 'lerobot[{cfg.env.name}]'`"
         )
-    elif cfg.env.name == "pusht":
-        import gym_pusht  # noqa: F401
+        raise e
 
-        # assert kwargs["seed"] > 200, "Seed 0-200 are used for the demonstration dataset, so we don't want to seed the eval env with this range."
-        env_fn = lambda: gym.make(  # noqa: E731
-            "gym_pusht/PushTPixels-v0",
-            **kwargs,
-        )
-    elif cfg.env.name == "aloha":
-        from lerobot.common.envs import aloha as gym_aloha  # noqa: F401
-
-        kwargs["task"] = cfg.env.task
-
-        env_fn = lambda: gym.make(  # noqa: E731
-            "gym_aloha/AlohaInsertion-v0",
-            **kwargs,
-        )
-    else:
-        raise ValueError(cfg.env.name)
+    handle = f"{package_name}/{cfg.env.handle}"
 
     if num_parallel_envs == 0:
         # non-batched version of the env that returns an observation of shape (c)
-        env = env_fn()
+        env = gym.make(handle, **kwargs)
     else:
         # batched version of the env that returns an observation of shape (b, c)
-        env = gym.vector.SyncVectorEnv([env_fn for _ in range(num_parallel_envs)])
+        env = gym.vector.SyncVectorEnv([lambda: gym.make(handle, **kwargs) for _ in range(num_parallel_envs)])
     return env