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

tests/test_available.py

@@ -19,7 +19,7 @@ import lerobot
 # from gym_pusht.envs import PushtEnv
 # from gym_xarm.envs import SimxarmEnv
 
-# from lerobot.common.datasets.simxarm import SimxarmDataset
+# from lerobot.common.datasets.xarm import SimxarmDataset
 # from lerobot.common.datasets.aloha import AlohaDataset
 # from lerobot.common.datasets.pusht import PushtDataset
 

tests/test_datasets.py

@@ -11,7 +11,7 @@ from .utils import DEVICE, DEFAULT_CONFIG_PATH
 @pytest.mark.parametrize(
     "env_name,dataset_id,policy_name",
     [
-        ("simxarm", "xarm_lift_medium", "tdmpc"),
+        ("xarm", "xarm_lift_medium", "tdmpc"),
         ("pusht", "pusht", "diffusion"),
         ("aloha", "aloha_sim_insertion_human", "act"),
         ("aloha", "aloha_sim_insertion_scripted", "act"),

tests/test_envs.py

@@ -1,3 +1,4 @@
+import importlib
 import pytest
 import torch
 from lerobot.common.datasets.factory import make_dataset
@@ -13,49 +14,25 @@ from .utils import DEVICE, DEFAULT_CONFIG_PATH
 
 
 @pytest.mark.parametrize(
-    "env_task, obs_type",
+    "env_name, handle, obs_type",
     [
         # ("AlohaInsertion-v0", "state"),
-        ("AlohaInsertion-v0", "pixels"),
-        ("AlohaInsertion-v0", "pixels_agent_pos"),
-        ("AlohaTransferCube-v0", "pixels"),
-        ("AlohaTransferCube-v0", "pixels_agent_pos"),
+        ("aloha", "AlohaInsertion-v0", "pixels"),
+        ("aloha", "AlohaInsertion-v0", "pixels_agent_pos"),
+        ("aloha", "AlohaTransferCube-v0", "pixels"),
+        ("aloha", "AlohaTransferCube-v0", "pixels_agent_pos"),
+        ("xarm", "XarmLift-v0", "state"),
+        ("xarm", "XarmLift-v0", "pixels"),
+        ("xarm", "XarmLift-v0", "pixels_agent_pos"),
+        ("pusht", "PushT-v0", "state"),
+        ("pusht", "PushT-v0", "pixels"),
+        ("pusht", "PushT-v0", "pixels_agent_pos"),
     ],
 )
-def test_aloha(env_task, obs_type):
-    from lerobot.common.envs import aloha as gym_aloha  # noqa: F401
-    env = gym.make(f"gym_aloha/{env_task}", obs_type=obs_type)
-    check_env(env.unwrapped)
-
-
-
-@pytest.mark.parametrize(
-    "env_task, obs_type",
-    [
-        ("XarmLift-v0", "state"),
-        ("XarmLift-v0", "pixels"),
-        ("XarmLift-v0", "pixels_agent_pos"),
-        # TODO(aliberts): Add gym_xarm other tasks
-    ],
-)
-def test_xarm(env_task, obs_type):
-    import gym_xarm  # noqa: F401
-    env = gym.make(f"gym_xarm/{env_task}", obs_type=obs_type)
-    check_env(env.unwrapped)
-
-
-
-@pytest.mark.parametrize(
-    "env_task, obs_type",
-    [
-        ("PushTPixels-v0", "state"),
-        ("PushTPixels-v0", "pixels"),
-        ("PushTPixels-v0", "pixels_agent_pos"),
-    ],
-)
-def test_pusht(env_task, obs_type):
-    import gym_pusht  # noqa: F401
-    env = gym.make(f"gym_pusht/{env_task}", obs_type=obs_type)
+def test_env(env_name, handle, obs_type):
+    package_name = f"gym_{env_name}"
+    importlib.import_module(package_name)
+    env = gym.make(f"{package_name}/{handle}", obs_type=obs_type)
     check_env(env.unwrapped)
 
 
@@ -63,7 +40,7 @@ def test_pusht(env_task, obs_type):
     "env_name",
     [
         "pusht",
-        "simxarm",
+        "xarm",
         "aloha",
     ],
 )
@@ -76,7 +53,7 @@ def test_factory(env_name):
     dataset = make_dataset(cfg)
 
     env = make_env(cfg, num_parallel_envs=1)
-    obs, info = env.reset()
+    obs, _ = env.reset()
     obs = preprocess_observation(obs, transform=dataset.transform)
     for key in dataset.image_keys:
         img = obs[key]

tests/test_policies.py

@@ -12,15 +12,15 @@ from .utils import DEVICE, DEFAULT_CONFIG_PATH
 @pytest.mark.parametrize(
     "env_name,policy_name,extra_overrides",
     [
-        ("simxarm", "tdmpc", ["policy.mpc=true"]),
+        ("xarm", "tdmpc", ["policy.mpc=true"]),
         ("pusht", "tdmpc", ["policy.mpc=false"]),
         ("pusht", "diffusion", []),
         # ("aloha", "act", ["env.task=sim_insertion", "dataset_id=aloha_sim_insertion_human"]),
         #("aloha", "act", ["env.task=sim_insertion", "dataset_id=aloha_sim_insertion_scripted"]),
         #("aloha", "act", ["env.task=sim_transfer_cube", "dataset_id=aloha_sim_transfer_cube_human"]),
         #("aloha", "act", ["env.task=sim_transfer_cube", "dataset_id=aloha_sim_transfer_cube_scripted"]),
-        # TODO(aliberts): simxarm not working with diffusion
-        # ("simxarm", "diffusion", []),
+        # TODO(aliberts): xarm not working with diffusion
+        # ("xarm", "diffusion", []),
     ],
 )
 def test_policy(env_name, policy_name, extra_overrides):