Simplify configs (#550)

Co-authored-by: Remi <remi.cadene@huggingface.co>
Co-authored-by: HUANG TZU-CHUN <137322177+tc-huang@users.noreply.github.com>

tests/test_datasets.py

@@ -43,9 +43,14 @@ from lerobot.common.datasets.utils import (
     hf_transform_to_torch,
     unflatten_dict,
 )
-from lerobot.common.utils.utils import init_hydra_config, seeded_context
+from lerobot.common.envs.factory import make_env_config
+from lerobot.common.policies.factory import make_policy_config
+from lerobot.common.robot_devices.robots.utils import make_robot
+from lerobot.common.utils.utils import seeded_context
+from lerobot.configs.default import DatasetConfig
+from lerobot.configs.train import TrainPipelineConfig
 from tests.fixtures.constants import DUMMY_REPO_ID
-from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, make_robot
+from tests.utils import DEVICE, require_x86_64_kernel
 
 
 def test_same_attributes_defined(lerobot_dataset_factory, tmp_path):
@@ -98,11 +103,13 @@ def test_dataset_initialization(lerobot_dataset_factory, tmp_path):
 # - [ ] test smaller methods
 
 
-@pytest.mark.skip("TODO after v2 migration / removing hydra")
 @pytest.mark.parametrize(
     "env_name, repo_id, policy_name",
-    lerobot.env_dataset_policy_triplets
-    + [("aloha", ["lerobot/aloha_sim_insertion_human", "lerobot/aloha_sim_transfer_cube_human"], "act")],
+    # Single dataset
+    lerobot.env_dataset_policy_triplets,
+    # Multi-dataset
+    # TODO after fix multidataset
+    # + [("aloha", ["lerobot/aloha_sim_insertion_human", "lerobot/aloha_sim_transfer_cube_human"], "act")],
 )
 def test_factory(env_name, repo_id, policy_name):
     """
@@ -110,15 +117,14 @@ def test_factory(env_name, repo_id, policy_name):
         - we can create a dataset with the factory.
         - for a commonly used set of data keys, the data dimensions are correct.
     """
-    cfg = init_hydra_config(
-        DEFAULT_CONFIG_PATH,
-        overrides=[
-            f"env={env_name}",
-            f"dataset_repo_id={repo_id}",
-            f"policy={policy_name}",
-            f"device={DEVICE}",
-        ],
+    cfg = TrainPipelineConfig(
+        # TODO(rcadene, aliberts): remove dataset download
+        dataset=DatasetConfig(repo_id=repo_id, episodes=[0]),
+        env=make_env_config(env_name),
+        policy=make_policy_config(policy_name),
+        device=DEVICE,
     )
+
     dataset = make_dataset(cfg)
     delta_timestamps = dataset.delta_timestamps
     camera_keys = dataset.meta.camera_keys
@@ -171,8 +177,8 @@ def test_factory(env_name, repo_id, policy_name):
 
 
 # TODO(alexander-soare): If you're hunting for savings on testing time, this takes about 5 seconds.
-@pytest.mark.skip("TODO after v2 migration / removing hydra")
-def test_multilerobotdataset_frames():
+@pytest.mark.skip("TODO after fix multidataset")
+def test_multidataset_frames():
     """Check that all dataset frames are incorporated."""
     # Note: use the image variants of the dataset to make the test approx 3x faster.
     # Note: We really do need three repo_ids here as at some point this caught an issue with the chaining
@@ -205,14 +211,14 @@ def test_multilerobotdataset_frames():
 
 
 # TODO(aliberts, rcadene): Refactor and move this to a tests/test_compute_stats.py
-@pytest.mark.skip("TODO after v2 migration / removing hydra")
 def test_compute_stats_on_xarm():
     """Check that the statistics are computed correctly according to the stats_patterns property.
 
     We compare with taking a straight min, mean, max, std of all the data in one pass (which we can do
     because we are working with a small dataset).
     """
-    dataset = LeRobotDataset("lerobot/xarm_lift_medium")
+    # TODO(rcadene, aliberts): remove dataset download
+    dataset = LeRobotDataset("lerobot/xarm_lift_medium", episodes=[0])
 
     # reduce size of dataset sample on which stats compute is tested to 10 frames
     dataset.hf_dataset = dataset.hf_dataset.select(range(10))
@@ -289,7 +295,6 @@ def test_flatten_unflatten_dict():
     assert json.dumps(original_d, sort_keys=True) == json.dumps(d, sort_keys=True), f"{original_d} != {d}"
 
 
-@pytest.mark.skip("TODO after v2 migration / removing hydra")
 @pytest.mark.parametrize(
     "repo_id",
     [
@@ -301,11 +306,12 @@ def test_flatten_unflatten_dict():
         # "lerobot/cmu_stretch",
     ],
 )
-# TODO(rcadene, aliberts): all these tests fail locally on Mac M1, but not on Linux
+@require_x86_64_kernel
 def test_backward_compatibility(repo_id):
     """The artifacts for this test have been generated by `tests/scripts/save_dataset_to_safetensors.py`."""
 
-    dataset = LeRobotDataset(repo_id)
+    # TODO(rcadene, aliberts): remove dataset download
+    dataset = LeRobotDataset(repo_id, episodes=[0])
 
     test_dir = Path("tests/data/save_dataset_to_safetensors") / repo_id
 
@@ -318,6 +324,11 @@ def test_backward_compatibility(repo_id):
         new_frame.pop("language_instruction", None)
         old_frame.pop("language_instruction", None)
 
+        # Remove task_index to allow for backward compatibility
+        # TODO(rcadene): remove when new features have been generated
+        if "task_index" not in old_frame:
+            del new_frame["task_index"]
+
         new_keys = set(new_frame.keys())
         old_keys = set(old_frame.keys())
         assert new_keys == old_keys, f"{new_keys=} and {old_keys=} are not the same"
@@ -361,8 +372,8 @@ def test_backward_compatibility(repo_id):
     # load_and_compare(i - 1)
 
 
-@pytest.mark.skip("TODO after v2 migration / removing hydra")
-def test_aggregate_stats():
+@pytest.mark.skip("TODO after fix multidataset")
+def test_multidataset_aggregate_stats():
     """Makes 3 basic datasets and checks that aggregate stats are computed correctly."""
     with seeded_context(0):
         data_a = torch.rand(30, dtype=torch.float32)