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

lerobot/__init__.py

@@ -27,7 +27,7 @@ from lerobot.__version__ import __version__  # noqa: F401
 available_envs = [
     "aloha",
     "pusht",
-    "simxarm",
+    "xarm",
 ]
 
 available_tasks_per_env = {
@@ -36,7 +36,7 @@ available_tasks_per_env = {
         "sim_transfer_cube",
     ],
     "pusht": ["pusht"],
-    "simxarm": ["lift"],
+    "xarm": ["lift"],
 }
 
 available_datasets_per_env = {
@@ -47,7 +47,7 @@ available_datasets_per_env = {
         "aloha_sim_transfer_cube_scripted",
     ],
     "pusht": ["pusht"],
-    "simxarm": ["xarm_lift_medium"],
+    "xarm": ["xarm_lift_medium"],
 }
 
 available_datasets = [dataset for env in available_envs for dataset in available_datasets_per_env[env]]

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

lerobot/configs/env/aloha.yaml

@@ -4,7 +4,7 @@ eval_episodes: 50
 eval_freq: 7500
 save_freq: 75000
 log_freq: 250
-# TODO: same as simxarm, need to adjust
+# TODO: same as xarm, need to adjust
 offline_steps: 25000
 online_steps: 25000
 
@@ -14,6 +14,8 @@ dataset_id: aloha_sim_insertion_human
 
 env:
   name: aloha
+  handle: AlohaInsertion-v0
+  # TODO(aliberts): replace task with handle
   task: insertion
   from_pixels: True
   pixels_only: False

lerobot/configs/env/pusht.yaml

@@ -4,7 +4,7 @@ eval_episodes: 50
 eval_freq: 7500
 save_freq: 75000
 log_freq: 250
-# TODO: same as simxarm, need to adjust
+# TODO: same as xarm, need to adjust
 offline_steps: 25000
 online_steps: 25000
 
@@ -14,6 +14,8 @@ dataset_id: pusht
 
 env:
   name: pusht
+  handle: PushT-v0
+  # TODO(aliberts): replace task with handle
   task: pusht
   from_pixels: True
   pixels_only: False

lerobot/configs/env/xarm.yaml

@@ -12,7 +12,9 @@ fps: 15
 dataset_id: xarm_lift_medium
 
 env:
-  name: simxarm
+  name: xarm
+  handle: XarmLift-v0
+  # TODO(aliberts): replace task with handle
   task: lift
   from_pixels: True
   pixels_only: False

lerobot/scripts/train.py

@@ -162,7 +162,7 @@ def train(cfg: dict, out_dir=None, job_name=None):
     logger = Logger(out_dir, job_name, cfg)
 
     log_output_dir(out_dir)
-    logging.info(f"{cfg.env.task=}")
+    logging.info(f"{cfg.env.handle=}")
     logging.info(f"{cfg.offline_steps=} ({format_big_number(cfg.offline_steps)})")
     logging.info(f"{cfg.online_steps=}")
     logging.info(f"{cfg.env.action_repeat=}")