fix tests

tests/test_datasets.py

@@ -6,12 +6,17 @@ from .utils import init_config
 
 
 @pytest.mark.parametrize(
-    "env_name",
+    "env_name,dataset_id",
     [
-        "simxarm",
-        "pusht",
+        # TODO(rcadene): simxarm is depreciated for now
+        # ("simxarm", "lift"),
+        ("pusht", "pusht"),
+        ("aloha", "sim_insertion_human"),
+        ("aloha", "sim_insertion_scripted"),
+        ("aloha", "sim_transfer_cube_human"),
+        ("aloha", "sim_transfer_cube_scripted"),
     ],
 )
-def test_factory(env_name):
-    cfg = init_config(overrides=[f"env={env_name}"])
+def test_factory(env_name, dataset_id):
+    cfg = init_config(overrides=[f"env={env_name}", f"env.task={dataset_id}"])
     offline_buffer = make_offline_buffer(cfg)

tests/test_envs.py

@@ -36,6 +36,7 @@ def print_spec_rollout(env):
     print("data from rollout:", simple_rollout(100))
 
 
+@pytest.mark.skip(reason="Simxarm is deprecated")
 @pytest.mark.parametrize(
     "task,from_pixels,pixels_only",
     [
@@ -80,7 +81,7 @@ def test_pusht(from_pixels, pixels_only):
 @pytest.mark.parametrize(
     "env_name",
     [
-        "simxarm",
+        # "simxarm",
         "pusht",
     ],
 )