Merge remote-tracking branch 'upstream/main' into feature/add-multitask-dit

# Conflicts:
#	pyproject.toml
#	src/lerobot/policies/pi_gemma.py
#	tests/policies/pi0_fast/test_pi0_fast_original_vs_lerobot.py
#	tests/policies/pi0_pi05/test_pi0.py
#	tests/policies/pi0_pi05/test_pi05.py

tests/artifacts/policies/pusht_diffusion_/actions.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:19eaaa85f66ba4aa6388dbb83819ffad6ea4363247208f871a8dc385689f6fc8
+oid sha256:54aecbc1af72a4cd5e9261492f5e7601890517516257aacdf2a0ffb3ce281f1b
 size 992

tests/artifacts/policies/pusht_diffusion_/grad_stats.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:227296eaeeb54acdc3dae2eb8af3d4d08fb87e245337624447140b1e91cfd002
+oid sha256:88a9c3775a2aa1e90a08850521970070a4fcf0f6b82aab43cd8ccc5cf77e0013
 size 47424

tests/artifacts/policies/pusht_diffusion_/output_dict.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:271b00cb2f0cd5fd26b1d53463638e3d1a6e92692ec625fcffb420ca190869e5
+oid sha256:91a2635e05a75fe187a5081504c5f35ce3417378813fa2deaf9ca4e8200e1819
 size 68

tests/artifacts/policies/pusht_diffusion_/param_stats.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:778fddbbaa64248cee35cb377c02cc2b6076f7ce5855146de677128900617ddf
+oid sha256:645bff922ac7bea63ad018ebf77c303c0e4cd2c1c0dc5ef3192865281bef3dc6
 size 47424

tests/fixtures/dataset_factories.py

@@ -222,7 +222,7 @@ def tasks_factory():
     def _create_tasks(total_tasks: int = 3) -> pd.DataFrame:
         ids = list(range(total_tasks))
         tasks = [f"Perform action {i}." for i in ids]
-        df = pd.DataFrame({"task_index": ids}, index=tasks)
+        df = pd.DataFrame({"task_index": ids}, index=pd.Index(tasks, name="task"))
         return df
 
     return _create_tasks

tests/policies/pi0_fast/test_pi0_fast_original_vs_lerobot.py

@@ -17,7 +17,6 @@
 """Test script to verify PI0Fast policy integration with LeRobot vs the original implementation"""
 # ruff: noqa: E402
 
-import os
 import random
 from copy import deepcopy
 from typing import Any
@@ -28,10 +27,6 @@ import torch
 
 pytest.importorskip("transformers")
 pytest.importorskip("scipy")
-pytestmark = pytest.mark.skipif(
-    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
-    reason="This test requires accepting the model license",
-)
 
 from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig
 from lerobot.policies.pi0_fast.modeling_pi0_fast import PI0FastPolicy
@@ -53,22 +48,23 @@ DUMMY_STATE_DIM = 20
 IMAGE_HEIGHT = 224
 IMAGE_WIDTH = 224
 NUM_VIEWS = 2  # Number of camera views
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-MODEL_PATH_LEROBOT = "jadechoghari/pi0fast-base"
+DEVICE = "cuda"
+MODEL_PATH_LEROBOT = "lerobot/pi0fast-base"
 
 # Expected action token shape: (batch_size, max_decoding_steps)
 EXPECTED_ACTION_TOKENS_SHAPE = (1, 2)
 
 # Expected first 5 action tokens (for reproducibility check)
-EXPECTED_ACTION_TOKENS_FIRST_5 = torch.tensor([255657, 255425])
+EXPECTED_ACTION_TOKENS_FIRST_5 = torch.tensor([255020, 255589])
 
 # Expected actions after detokenization
 EXPECTED_ACTIONS_SHAPE = (1, 2, 32)  # (batch_size, n_action_steps, action_dim)
 EXPECTED_ACTIONS_MEAN = 0.046403881162405014
 EXPECTED_ACTIONS_STD = 0.2607129216194153
-EXPECTED_ACTIONS_FIRST_5 = torch.tensor([-0.0707, 1.4849, 0.0000, 0.0000, 0.0000])
+EXPECTED_ACTIONS_FIRST_5 = torch.tensor([0.0000, 0.3536, 0.0707, 0.0000, 0.0000])
 
 
+@require_cuda
 def set_seed_all(seed: int):
     """Set random seed for all RNG sources to ensure reproducibility."""
     random.seed(seed)
@@ -85,6 +81,7 @@ def set_seed_all(seed: int):
     torch.use_deterministic_algorithms(True, warn_only=True)
 
 
+@require_cuda
 def instantiate_lerobot_pi0_fast(
     from_pretrained: bool = False,
     model_path: str = MODEL_PATH_LEROBOT,
@@ -127,6 +124,7 @@ def instantiate_lerobot_pi0_fast(
     return policy, preprocessor, postprocessor
 
 
+@require_cuda
 def create_dummy_data(device=DEVICE):
     """Create dummy data for testing both implementations."""
     batch_size = 1
@@ -158,22 +156,25 @@ def create_dummy_data(device=DEVICE):
 
 # Pytest fixtures
 @pytest.fixture(scope="module")
+@require_cuda
 def pi0_fast_components():
     """Fixture to instantiate and provide all PI0Fast components for tests."""
     print(f"\nTesting with DEVICE='{DEVICE}'")
     print("\n[Setup] Instantiating LeRobot PI0Fast policy...")
     policy_obj, preprocessor_obj, postprocessor_obj = instantiate_lerobot_pi0_fast(from_pretrained=True)
     print("Model loaded successfully")
-    yield policy_obj, preprocessor_obj, postprocessor_obj
+    return policy_obj, preprocessor_obj, postprocessor_obj
 
 
 @pytest.fixture(scope="module")
+@require_cuda
 def policy(pi0_fast_components):
     """Fixture to provide the PI0Fast policy for tests."""
     return pi0_fast_components[0]
 
 
 @pytest.fixture(scope="module")
+@require_cuda
 def preprocessor(pi0_fast_components):
     """Fixture to provide the PI0Fast preprocessor for tests."""
     return pi0_fast_components[1]

tests/policies/pi0_pi05/test_pi0.py

@@ -16,17 +16,8 @@
 
 """Test script to verify PI0 policy integration with LeRobot, only meant to be run locally!"""
 
-import os
-
-import pytest
 import torch
 
-# Skip this entire module in CI
-pytestmark = pytest.mark.skipif(
-    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
-    reason="This test requires accepting the model license",
-)
-
 from lerobot.policies.factory import make_policy_config  # noqa: E402
 from lerobot.policies.pi0 import (  # noqa: E402
     PI0Config,

tests/policies/pi0_pi05/test_pi05.py

@@ -16,25 +16,15 @@
 
 """Test script to verify PI0.5 (pi05) support in PI0 policy, only meant to be run locally!"""
 
-import os
-
-import pytest
 import torch
 
-from lerobot.utils.random_utils import set_seed
-
-# Skip this entire module in CI
-pytestmark = pytest.mark.skipif(
-    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
-    reason="This test requires accepting the model license",
-)
-
 from lerobot.policies.factory import make_policy_config  # noqa: E402
 from lerobot.policies.pi05 import (  # noqa: E402
     PI05Config,
     PI05Policy,
     make_pi05_pre_post_processors,  # noqa: E402
 )
+from lerobot.utils.random_utils import set_seed
 from tests.utils import require_cuda  # noqa: E402
 
 

tests/policies/pi0_pi05/test_pi05_rtc.py

@@ -24,9 +24,10 @@ import torch
 # Skip this entire module in CI
 pytestmark = pytest.mark.skipif(
     os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
-    reason="This test requires local OpenPI installation and is not meant for CI",
+    reason="TODO: This test seems to hang the CI",
 )
 
+
 from lerobot.configs.types import FeatureType, PolicyFeature, RTCAttentionSchedule  # noqa: E402
 from lerobot.policies.pi05 import PI05Config, PI05Policy, make_pi05_pre_post_processors  # noqa: E402
 from lerobot.policies.rtc.configuration_rtc import RTCConfig  # noqa: E402

tests/policies/pi0_pi05/test_pi0_rtc.py

@@ -24,9 +24,10 @@ import torch
 # Skip this entire module in CI
 pytestmark = pytest.mark.skipif(
     os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
-    reason="This test requires local OpenPI installation and is not meant for CI",
+    reason="TODO: This test seems to hang the CI",
 )
 
+
 from lerobot.configs.types import FeatureType, PolicyFeature, RTCAttentionSchedule  # noqa: E402
 from lerobot.policies.pi0 import PI0Config, PI0Policy, make_pi0_pre_post_processors  # noqa: E402
 from lerobot.policies.rtc.configuration_rtc import RTCConfig  # noqa: E402
@@ -88,6 +89,7 @@ def test_pi0_rtc_initialization_without_rtc_config():
     print("✓ PI0 RTC initialization without RTC config: Test passed")
 
 
+@require_cuda
 def test_pi0_rtc_inference_with_prev_chunk():
     """Test PI0 policy inference with RTC and previous chunk."""
     set_seed(42)

tests/policies/wall_x/test_wallx.py

@@ -29,8 +29,10 @@ from lerobot.policies.wall_x import WallXConfig  # noqa: E402
 from lerobot.policies.wall_x.modeling_wall_x import WallXPolicy  # noqa: E402
 from lerobot.policies.wall_x.processor_wall_x import make_wall_x_pre_post_processors  # noqa: E402
 from lerobot.utils.random_utils import set_seed  # noqa: E402
+from tests.utils import require_cuda  # noqa: E402
 
 
+@require_cuda
 def test_policy_instantiation():
     # Create config
     set_seed(42)
@@ -115,6 +117,7 @@ def test_policy_instantiation():
         raise
 
 
+@require_cuda
 def test_config_creation():
     """Test policy config creation through factory."""
     try:
@@ -126,8 +129,3 @@ def test_config_creation():
     except Exception as e:
         print(f"Config creation failed: {e}")
         raise
-
-
-if __name__ == "__main__":
-    test_policy_instantiation()
-    test_config_creation()

tests/scripts/test_edit_dataset_parsing.py

@@ -27,6 +27,7 @@ from lerobot.scripts.lerobot_edit_dataset import (
     OperationConfig,
     RemoveFeatureConfig,
     SplitConfig,
+    _validate_config,
 )
 
 
@@ -51,11 +52,23 @@ class TestOperationTypeParsing:
         ],
     )
     def test_operation_type_resolves_correct_class(self, type_name, expected_cls):
-        cfg = parse_cfg(["--repo_id", "test/repo", "--operation.type", type_name])
+        cfg = parse_cfg(
+            ["--repo_id", "test/repo", "--new_repo_id", "test/merged", "--operation.type", type_name]
+        )
         assert isinstance(cfg.operation, expected_cls), (
             f"Expected {expected_cls.__name__}, got {type(cfg.operation).__name__}"
         )
 
+    def test_merge_requires_new_repo_id(self):
+        cfg = parse_cfg(["--operation.type", "merge"])
+        with pytest.raises(ValueError, match="--new_repo_id is required for merge"):
+            _validate_config(cfg)
+
+    def test_non_merge_requires_repo_id(self):
+        cfg = parse_cfg(["--operation.type", "delete_episodes"])
+        with pytest.raises(ValueError, match="--repo_id is required for delete_episodes"):
+            _validate_config(cfg)
+
     @pytest.mark.parametrize(
         "type_name, expected_cls",
         [
@@ -69,6 +82,8 @@ class TestOperationTypeParsing:
         ],
     )
     def test_get_choice_name_roundtrips(self, type_name, expected_cls):
-        cfg = parse_cfg(["--repo_id", "test/repo", "--operation.type", type_name])
+        cfg = parse_cfg(
+            ["--repo_id", "test/repo", "--new_repo_id", "test/merged", "--operation.type", type_name]
+        )
         resolved_name = OperationConfig.get_choice_name(type(cfg.operation))
         assert resolved_name == type_name

tests/utils/test_logging_utils.py

@@ -24,6 +24,11 @@ def mock_metrics():
     return {"loss": AverageMeter("loss", ":.3f"), "accuracy": AverageMeter("accuracy", ":.2f")}
 
 
+class MockAccelerator:
+    def __init__(self, num_processes: int):
+        self.num_processes = num_processes
+
+
 def test_average_meter_initialization():
     meter = AverageMeter("loss", ":.2f")
     assert meter.name == "loss"
@@ -82,6 +87,37 @@ def test_metrics_tracker_step(mock_metrics):
     assert tracker.epochs == tracker.samples / 1000
 
 
+def test_metrics_tracker_initialization_with_accelerator(mock_metrics):
+    tracker = MetricsTracker(
+        batch_size=32,
+        num_frames=1000,
+        num_episodes=50,
+        metrics=mock_metrics,
+        initial_step=10,
+        accelerator=MockAccelerator(num_processes=2),
+    )
+    assert tracker.steps == 10
+    assert tracker.samples == 10 * 32 * 2
+    assert tracker.episodes == tracker.samples / (1000 / 50)
+    assert tracker.epochs == tracker.samples / 1000
+
+
+def test_metrics_tracker_step_with_accelerator(mock_metrics):
+    tracker = MetricsTracker(
+        batch_size=32,
+        num_frames=1000,
+        num_episodes=50,
+        metrics=mock_metrics,
+        initial_step=5,
+        accelerator=MockAccelerator(num_processes=2),
+    )
+    tracker.step()
+    assert tracker.steps == 6
+    assert tracker.samples == (5 * 32 * 2) + (32 * 2)
+    assert tracker.episodes == tracker.samples / (1000 / 50)
+    assert tracker.epochs == tracker.samples / 1000
+
+
 def test_metrics_tracker_getattr(mock_metrics):
     tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=mock_metrics)
     assert tracker.loss == mock_metrics["loss"]