Merge branch 'main' into feat/fix_soft_prompt_typo

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>

docs/source/cameras.mdx

@@ -164,8 +164,8 @@ includes the range reported by the sensor. Requesting an unsupported control als
 Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
 require `use_rgb=True`.
 
-Manual color controls require a dedicated RGB module. Cameras without one, such as the RealSense
-D405, do not support them and raise an error at connection time.
+On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual
+exposure or gain also affects the depth stream.
 
 </hfoption>
 </hfoptions>

docs/source/env_processor.mdx

@@ -88,20 +88,6 @@ policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)
 
 The same policy can work with different environment processors, and the same environment processor can work with different policies:
 
-````python
-# Use SmolVLA policy with LIBERO environment
-# Use SmolVLA policy with LIBERO environment
-libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
-    env_cfg=libero_cfg,
-    policy_cfg=smolvla_cfg,
-)
-smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
-# Or use ACT policy with the same LIBERO environment
-libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
-    env_cfg=libero_cfg,
-    policy_cfg=act_cfg,
-)
-act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
 ```python
 # Use SmolVLA policy with LIBERO environment
 libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
@@ -116,6 +102,7 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
     policy_cfg=act_cfg,
 )
 act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
+```
 
 ### 3. **Easier Experimentation**
 
@@ -145,7 +132,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
         state = torch.cat([eef_pos, eef_axisangle, eef_vel,
                           gripper_pos, gripper_vel], dim=-1)  # 14D
         return state
-````
+```
 
 ### 4. **Cleaner Environment Code**
 

docs/source/hil_data_collection.mdx

@@ -59,6 +59,7 @@ The `lerobot-rollout --strategy.type=dagger` mode requires **teleoperators with
 
 - `bi_openarm_mini` - Bimanual OpenArm Mini
 - `so_leader` - SO100 / SO101 leader arm
+- `bi_so_leader` - Bimanual SO100 / SO101 leader arms
 
 > [!IMPORTANT]
 > The provided commands default to `bi_openarm_follower` + `bi_openarm_mini`.

docs/source/hope_jr.mdx

@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.
 
 ### Record
 
-This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
+This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data).
 
 ```bash
 lerobot-record \

docs/source/integrate_hardware.mdx

@@ -18,7 +18,7 @@ If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus inter
 - [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) – for controlling Dynamixel servos
 
 Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
-For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so101_follower/so101_follower.py)
+For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so_follower.py)
 
 Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):
 

docs/source/lekiwi.mdx

@@ -51,7 +51,7 @@ In addition to these instructions, you need to install the Feetech SDK & ZeroMQ
 pip install -e ".[lekiwi]"
 ```
 
-Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:.
+Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖.
 Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
 
 # Step-by-Step Assembly Instructions

docs/source/pi0fast.mdx

@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs
 
 ## Reproducing π₀Fast results
 
-We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40kk steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
+We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
 
 The finetuned model can be found here:
 

docs/source/smolvla.mdx

@@ -93,7 +93,7 @@ lerobot-train --help
 
 ## Evaluate the finetuned model and run it in real-time
 
-Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
+Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots#record-a-dataset).
 Once you are logged in, you can run inference in your setup by doing:
 
 ```bash

docs/source/so101.mdx

@@ -338,7 +338,7 @@ It is advisable to install one 3-pin cable in the motor after placing them befor
 <hfoption id="Leader">
 
 - Mount the leader holder onto the wrist and secure it with 4 M3x6mm screws.
-- Attach the handle to motor 5 using 1 M2x6mm screw.
+- Attach the handle to the leader holder using 1 M2x6mm screw.
 - Insert the gripper motor, secure it with 2 M2x6mm screws on each side, attach a motor horn using a M3x6mm horn screw.
 - Attach the follower trigger with 4 M3x6mm screws.
 

pyproject.toml

@@ -67,7 +67,7 @@ dependencies = [
     "einops>=0.8.0,<0.9.0",
 
     # Config & Hub
-    "draccus==0.10.0", # TODO: Relax version constraint
+    "draccus>=0.11.6,<0.12.0",
     "huggingface-hub>=1.0.0,<2.0.0",
     "requests>=2.32.0,<3.0.0",
 

src/lerobot/cameras/realsense/camera_realsense.py

@@ -365,12 +365,11 @@ class RealSenseCamera(Camera):
         return self._async_read(timeout_ms=10000, read_depth=read_depth)
 
     def _get_color_sensor(self) -> "rs.sensor":
-        """Returns the dedicated "RGB Camera" sensor that controls the color stream.
+        """Returns the sensor that controls the color stream.
 
-        Manual color controls are only applied to a dedicated RGB module. Cameras
-        without one (e.g. the D405, whose color stream comes from the shared
-        "Stereo Module") are unsupported, so we never fall back to another sensor
-        to avoid altering the depth stream.
+        Most RealSense cameras expose "RGB Camera" for color. The D405 has no
+        separate RGB module — its color stream comes from "Stereo Module".
+        We try RGB Camera first, then fall back to Stereo Module.
         """
         if self.rs_profile is None:
             raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
@@ -378,14 +377,12 @@ class RealSenseCamera(Camera):
         device = self.rs_profile.get_device()
         sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
 
-        if "RGB Camera" in sensors:
-            return sensors["RGB Camera"]
+        for name in ("RGB Camera", "Stereo Module"):
+            if name in sensors:
+                return sensors[name]
 
         available = list(sensors.keys())
-        raise RuntimeError(
-            f"{self}: manual color controls require a dedicated 'RGB Camera' module, which this camera does not have. ",
-            f"Available sensors: {available}.",
-        )
+        raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}")
 
     def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
         """Sets a sensor option, re-raising range errors with actionable diagnostics."""

src/lerobot/configs/policies.py

@@ -163,8 +163,10 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):  # type: igno
         return None
 
     def _save_pretrained(self, save_directory: Path) -> None:
-        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
-            draccus.dump(self, f, indent=4)
+        # Encode against the base class so draccus includes the choice "type" key,
+        # which `from_pretrained` needs to resolve the concrete subclass.
+        with open(save_directory / CONFIG_NAME, "w") as f:
+            json.dump(draccus.encode(self, PreTrainedConfig), f, indent=4)
 
     @classmethod
     def from_pretrained(

src/lerobot/configs/rewards.py

@@ -103,8 +103,10 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
         pass
 
     def _save_pretrained(self, save_directory: Path) -> None:
-        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
-            draccus.dump(self, f, indent=4)
+        # Encode against the base class so draccus includes the choice "type" key,
+        # which `from_pretrained` needs to resolve the concrete subclass.
+        with open(save_directory / CONFIG_NAME, "w") as f:
+            json.dump(draccus.encode(self, RewardModelConfig), f, indent=4)
 
     @classmethod
     def from_pretrained(

src/lerobot/configs/train.py

@@ -194,7 +194,11 @@ class TrainPipelineConfig(HubMixin):
             )
 
         if Path(config_path).resolve().exists():
-            policy_dir = Path(config_path).parent
+            # `config_path` may point at the checkpoint's train_config.json or at its
+            # pretrained_model/ directory (both documented above) — resolve either to
+            # the pretrained_model/ directory.
+            config_path_obj = Path(config_path)
+            policy_dir = config_path_obj.parent if config_path_obj.is_file() else config_path_obj
             self.checkpoint_path = policy_dir.parent
         elif self.job.is_remote:
             return

src/lerobot/policies/evo1/modeling_evo1.py

@@ -42,6 +42,9 @@ class Evo1Policy(PreTrainedPolicy):
     config_class = Evo1Config
     name = "evo1"
 
+    def supports_rtc(self) -> bool:
+        return True
+
     def __init__(self, config: Evo1Config, *, vlm_hub_kwargs: dict | None = None, **kwargs):
         super().__init__(config)
         config.validate_features()

src/lerobot/policies/groot/modeling_groot.py

@@ -68,6 +68,9 @@ class GrootPolicy(PreTrainedPolicy):
     name = "groot"
     config_class = GrootConfig
 
+    def supports_rtc(self) -> bool:
+        return True
+
     def __init__(self, config: GrootConfig, **kwargs):
         """Initialize Groot policy wrapper."""
         require_package("transformers", extra="groot")

src/lerobot/policies/molmoact2/modeling_molmoact2.py

@@ -520,6 +520,9 @@ class MolmoAct2Policy(PreTrainedPolicy):
     config_class = MolmoAct2Config
     name = "molmoact2"
 
+    def supports_rtc(self) -> bool:
+        return self.config.inference_action_mode == "continuous"
+
     def __init__(
         self,
         config: MolmoAct2Config,

src/lerobot/policies/pi0/modeling_pi0.py

@@ -749,6 +749,9 @@ class PI0Policy(PreTrainedPolicy):
     config_class = PI0Config
     name = "pi0"
 
+    def supports_rtc(self) -> bool:
+        return True
+
     def __init__(
         self,
         config: PI0Config,

src/lerobot/policies/pi05/modeling_pi05.py

@@ -714,6 +714,9 @@ class PI05Policy(PreTrainedPolicy):
     config_class = PI05Config
     name = "pi05"
 
+    def supports_rtc(self) -> bool:
+        return True
+
     def __init__(
         self,
         config: PI05Config,

src/lerobot/policies/pretrained.py

@@ -249,6 +249,10 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         """
         raise NotImplementedError
 
+    def supports_rtc(self) -> bool:
+        """Whether this policy implements Real-Time Chunking inference semantics."""
+        return False
+
     # TODO(aliberts, rcadene): split into 'forward' and 'compute_loss'?
     @abc.abstractmethod
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -145,6 +145,9 @@ class SmolVLAPolicy(PreTrainedPolicy):
     config_class = SmolVLAConfig
     name = "smolvla"
 
+    def supports_rtc(self) -> bool:
+        return True
+
     def __init__(
         self,
         config: SmolVLAConfig,

src/lerobot/policies/smolvla/smolvlm_with_expert.py

@@ -168,14 +168,23 @@ class SmolVLMWithExpertModel(nn.Module):
                 last_layers.append(self.num_vlm_layers - 2)
             frozen_layers = [
                 "lm_head",
-                "text_model.model.norm.weight",
+                "text_model.norm.weight",
             ]
             for layer in last_layers:
-                frozen_layers.append(f"text_model.model.layers.{layer}.")
+                frozen_layers.append(f"text_model.layers.{layer}.")
 
+            unmatched_patterns = set(frozen_layers)
             for name, params in self.vlm.named_parameters():
-                if any(k in name for k in frozen_layers):
+                matched_patterns = [k for k in frozen_layers if k in name]
+                if matched_patterns:
                     params.requires_grad = False
+                    unmatched_patterns.difference_update(matched_patterns)
+            if unmatched_patterns:
+                raise RuntimeError(
+                    "Some frozen layer patterns matched no VLM parameters, so the corresponding layers "
+                    "would silently remain trainable (parameter naming may have changed in transformers): "
+                    f"{sorted(unmatched_patterns)}"
+                )
         # To avoid unused params issue with distributed training
         for name, params in self.lm_expert.named_parameters():
             if "lm_head" in name:

src/lerobot/policies/xvla/modeling_xvla.py

@@ -150,7 +150,7 @@ class XVLAModel(nn.Module):
         # Freeze or unfreeze policy transformer
         if not self.config.train_policy_transformer:
             for name, param in self.transformer.named_parameters():
-                if "soft_prompts" not in name:
+                if "soft_prompt" not in name:
                     param.requires_grad = False
 
         # Freeze or unfreeze soft prompts

src/lerobot/rl/algorithms/configs.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 
 import abc
 import builtins
+import json
 import logging
 import os
 from dataclasses import dataclass, field
@@ -78,8 +79,10 @@ class RLAlgorithmConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
 
     def _save_pretrained(self, save_directory: Path) -> None:
         """Serialize this config as ``config.json`` inside ``save_directory``."""
-        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
-            draccus.dump(self, f, indent=4)
+        # Encode against the base class so draccus includes the choice "type" key,
+        # which `from_pretrained` needs to resolve the concrete subclass.
+        with open(save_directory / CONFIG_NAME, "w") as f:
+            json.dump(draccus.encode(self, RLAlgorithmConfig), f, indent=4)
 
     @classmethod
     def from_pretrained(

src/lerobot/rollout/context.py

@@ -57,6 +57,7 @@ from .inference import (
     SyncInferenceConfig,
     create_inference_engine,
 )
+from .inference.rtc import supports_rtc_inference
 from .robot_wrapper import ThreadSafeRobot
 
 if TYPE_CHECKING or _peft_available:
@@ -226,6 +227,12 @@ def build_rollout_context(
     policy = _load_pretrained_policy(policy_config)
 
     if is_rtc:
+        if not supports_rtc_inference(policy):
+            raise ValueError(
+                f"RTC inference is not supported by policy type '{policy_config.type}': "
+                "the policy must implement RTC semantics and predict_action_chunk must accept "
+                "inference_delay and prev_chunk_left_over. Use '--inference.type=sync' instead."
+            )
         policy.config.rtc_config = cfg.inference.rtc
         if hasattr(policy, "init_rtc_processor"):
             policy.init_rtc_processor()

src/lerobot/rollout/inference/rtc.py

@@ -22,6 +22,7 @@ way via ``notify_observation``.
 
 from __future__ import annotations
 
+import inspect
 import logging
 import math
 import time
@@ -62,6 +63,23 @@ _RTC_JOIN_TIMEOUT_S: float = 3.0
 # ---------------------------------------------------------------------------
 
 
+def supports_rtc_inference(policy: PreTrainedPolicy) -> bool:
+    """Whether a policy declares RTC support and accepts the RTC call shape."""
+    supports_rtc = getattr(policy, "supports_rtc", None)
+    if not callable(supports_rtc) or not supports_rtc():
+        return False
+
+    try:
+        inspect.signature(policy.predict_action_chunk).bind(
+            object(),
+            inference_delay=0,
+            prev_chunk_left_over=None,
+        )
+    except (TypeError, ValueError):
+        return False
+    return True
+
+
 def _normalize_prev_actions_length(prev_actions: torch.Tensor, target_steps: int) -> torch.Tensor:
     """Pad or truncate RTC prefix actions to a fixed length for stable compiled inference."""
     if prev_actions.ndim != 2:

src/lerobot/scripts/lerobot_train.py

@@ -348,7 +348,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         preprocessor_overrides = {
             "device_processor": {"device": device.type},
             "normalizer_processor": {
-                "stats": dataset.meta.stats,
                 "features": {**policy.config.input_features, **policy.config.output_features},
                 "norm_map": policy.config.normalization_mapping,
             },
@@ -356,11 +355,17 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         }
         postprocessor_overrides = {
             "unnormalizer_processor": {
-                "stats": dataset.meta.stats,
                 "features": policy.config.output_features,
                 "norm_map": policy.config.normalization_mapping,
             },
         }
+        # On resume, the checkpoint's saved processor stats are authoritative: they may have
+        # been adapted by the policy (e.g. EVO1 pads state/action stats to max_state_dim),
+        # and force-feeding raw dataset stats over them crashes normalization (#4006).
+        # This mirrors the `dataset_stats` kwarg above, which is also skipped on resume.
+        if not cfg.resume:
+            preprocessor_overrides["normalizer_processor"]["stats"] = dataset.meta.stats
+            postprocessor_overrides["unnormalizer_processor"]["stats"] = dataset.meta.stats
         if getattr(active_cfg, "use_relative_actions", False):
             preprocessor_overrides["relative_actions_processor"] = {
                 "enabled": True,

src/lerobot/teleoperators/bi_so_leader/bi_so_leader.py

@@ -67,7 +67,15 @@ class BiSOLeader(BimanualMixin, Teleoperator):
 
     @cached_property
     def feedback_features(self) -> dict[str, type]:
-        return {}
+        # Bimanual teleop has feedback (can be actuated for handover).
+        # Return the same structure as action_features for consistency with left/right arms.
+        left_arm_features = self.left_arm.feedback_features
+        right_arm_features = self.right_arm.feedback_features
+
+        return {
+            **{f"left_{k}": v for k, v in left_arm_features.items()},
+            **{f"right_{k}": v for k, v in right_arm_features.items()},
+        }
 
     def setup_motors(self) -> None:
         self.left_arm.setup_motors()
@@ -87,6 +95,43 @@ class BiSOLeader(BimanualMixin, Teleoperator):
 
         return action_dict
 
+    def enable_torque(self) -> None:
+        """Enable torque on both leader arms for smooth handover."""
+        self.left_arm.enable_torque()
+        self.right_arm.enable_torque()
+
+    def disable_torque(self) -> None:
+        """Disable torque on both leader arms to allow human control."""
+        self.left_arm.disable_torque()
+        self.right_arm.disable_torque()
+
+    @check_if_not_connected
     def send_feedback(self, feedback: dict[str, float]) -> None:
-        # TODO: Implement force feedback
-        raise NotImplementedError
+        """Route bimanual feedback to left and right arms with proper prefix stripping.
+
+        Receives feedback dict with keys like: left_shoulder_pan.pos, right_shoulder_pan.pos, ...
+        Splits and routes to each arm by removing the prefix.
+
+        This enables DAgger smooth handover: when transitioning from policy control to human
+        intervention, both leader arms are commanded to the follower's current pose to avoid
+        discontinuities.
+        """
+        # Split feedback by arm prefix
+        left_feedback = {}
+        right_feedback = {}
+
+        for key, value in feedback.items():
+            if key.startswith("left_"):
+                # Strip "left_" prefix and pass to left arm
+                stripped_key = key[5:]  # len("left_") == 5
+                left_feedback[stripped_key] = value
+            elif key.startswith("right_"):
+                # Strip "right_" prefix and pass to right arm
+                stripped_key = key[6:]  # len("right_") == 6
+                right_feedback[stripped_key] = value
+
+        # Send to each arm
+        if left_feedback:
+            self.left_arm.send_feedback(left_feedback)
+        if right_feedback:
+            self.right_arm.send_feedback(right_feedback)

src/lerobot/utils/random_utils.py

@@ -85,6 +85,8 @@ def serialize_torch_rng_state() -> dict[str, torch.Tensor]:
     torch_rng_state_dict = {"torch_rng_state": torch.get_rng_state()}
     if torch.cuda.is_available():
         torch_rng_state_dict["torch_cuda_rng_state"] = torch.cuda.get_rng_state()
+    if torch.backends.mps.is_available():
+        torch_rng_state_dict["torch_mps_rng_state"] = torch.mps.get_rng_state()
     return torch_rng_state_dict
 
 
@@ -95,6 +97,8 @@ def deserialize_torch_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None
     torch.set_rng_state(rng_state_dict["torch_rng_state"])
     if torch.cuda.is_available() and "torch_cuda_rng_state" in rng_state_dict:
         torch.cuda.set_rng_state(rng_state_dict["torch_cuda_rng_state"])
+    if torch.backends.mps.is_available() and "torch_mps_rng_state" in rng_state_dict:
+        torch.mps.set_rng_state(rng_state_dict["torch_mps_rng_state"])
 
 
 def serialize_rng_state() -> dict[str, torch.Tensor]:

tests/cameras/test_realsense.py

@@ -322,16 +322,15 @@ def test_get_color_sensor_prefers_rgb_camera():
     assert camera._get_color_sensor() is rgb
 
 
-def test_get_color_sensor_raises_without_dedicated_rgb_module():
-    """D405 has no separate RGB module; we refuse to touch the shared Stereo Module."""
+def test_get_color_sensor_falls_back_to_stereo_module():
+    """D405 has no separate RGB module; color comes from Stereo Module."""
     config = RealSenseCameraConfig(serial_number_or_name="042")
     camera = RealSenseCamera(config)
 
     stereo = _make_mock_sensor("Stereo Module")
     _attach_mock_color_sensor(camera, stereo)
 
-    with pytest.raises(RuntimeError, match="dedicated 'RGB Camera' module"):
-        camera._get_color_sensor()
+    assert camera._get_color_sensor() is stereo
 
 
 def test_get_color_sensor_raises_with_available_sensors():

tests/configs/test_resume_from_hub.py

@@ -66,3 +66,27 @@ def test_from_pretrained_raises_when_no_root_config_and_no_checkpoints(monkeypat
 
     with pytest.raises(FileNotFoundError, match="train_config.json not found"):
         TrainPipelineConfig.from_pretrained("user/empty-repo")
+
+
+@pytest.mark.parametrize("pass_dir", [False, True])
+def test_resolve_resume_checkpoint_accepts_file_or_pretrained_model_dir(tmp_path, monkeypatch, pass_dir):
+    """`--config_path` may point at the checkpoint's train_config.json or at its
+    pretrained_model/ directory; both must resolve `policy.pretrained_path` to the
+    pretrained_model/ directory (regression test for the directory case, which
+    previously resolved one level too high and failed on model.safetensors)."""
+    pretrained_dir = tmp_path / "checkpoints" / "000002" / "pretrained_model"
+    pretrained_dir.mkdir(parents=True)
+    (pretrained_dir / "train_config.json").touch()
+    target = pretrained_dir if pass_dir else pretrained_dir / "train_config.json"
+
+    from lerobot.policies.act.configuration_act import ACTConfig
+
+    cfg = tc.draccus.parse(TrainPipelineConfig, args=["--dataset.repo_id", "u/d"])
+    cfg.policy = ACTConfig()
+    cfg.resume = True
+    monkeypatch.setattr(tc.parser, "parse_arg", lambda name: str(target) if name == "config_path" else None)
+
+    cfg._resolve_resume_checkpoint()
+
+    assert cfg.policy.pretrained_path == pretrained_dir
+    assert cfg.checkpoint_path == pretrained_dir.parent

tests/utils/test_random_utils.py

@@ -73,6 +73,17 @@ def test_serialize_deserialize_torch_rng(fixed_seed):
     assert val2 == val3
 
 
+@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
+def test_serialize_deserialize_torch_rng_mps(fixed_seed):
+    _ = torch.rand(1, device="mps").item()
+    st = serialize_torch_rng_state()
+    assert "torch_mps_rng_state" in st
+    val2 = torch.rand(1, device="mps").item()
+    deserialize_torch_rng_state(st)
+    val3 = torch.rand(1, device="mps").item()
+    assert val2 == val3
+
+
 def test_serialize_deserialize_rng(fixed_seed):
     # Generate one from each library
     _ = random.random()

uv.lock

@@ -1,5 +1,5 @@
 version = 1
-revision = 3
+revision = 2
 requires-python = ">=3.12"
 resolution-markers = [
     "(python_full_version >= '3.15' and platform_machine == 'AMD64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux')",
@@ -1359,18 +1359,17 @@ sdist = { url = "https://files.pythonhosted.org/packages/a2/55/8f8cab2afd404cf57
 
 [[package]]
 name = "draccus"
-version = "0.10.0"
+version = "0.11.6"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "mergedeep" },
     { name = "pyyaml" },
-    { name = "pyyaml-include" },
     { name = "toml" },
     { name = "typing-inspect" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/4e/e2/f5012fda17ee5d1eaf3481b6ca3e11dffa5348e5e08ab745538fdc8041bb/draccus-0.10.0.tar.gz", hash = "sha256:8dd08304219becdcd66cd16058ba98e9c3e6b7bfe48ccb9579dae39f8d37ae19", size = 62243, upload-time = "2025-02-05T07:27:48.182Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/df/b0/cc399719e0cf6fea451f155798e19ea8c5e9c838b701d1565f29ea626c18/draccus-0.11.6.tar.gz", hash = "sha256:d134f576a1f4febd93c6b200df7f92e5febffe9efdc0a5f381bf50c2e0568a39", size = 67940, upload-time = "2026-06-12T13:21:19.999Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/c4/9a/a83083b230d352ee5d205757b74006dbe084448ca45e3bc5ca99215b1e55/draccus-0.10.0-py3-none-any.whl", hash = "sha256:90243418ae0e9271c390a59cafb6acfd37001193696ed36fcc8525f791a83282", size = 71783, upload-time = "2025-02-05T07:27:46.1Z" },
+    { url = "https://files.pythonhosted.org/packages/18/f1/d56bef4563d1cfaf55dd58de298a456587b9fe0a85dabb55768d44ace8f7/draccus-0.11.6-py3-none-any.whl", hash = "sha256:1cf3f37c64766e0f17d757099b388e762f1d22c772cc2376b4e366b515fb5737", size = 85449, upload-time = "2026-06-12T13:21:18.83Z" },
 ]
 
 [[package]]
@@ -3289,7 +3288,7 @@ requires-dist = [
     { name = "deepdiff", marker = "extra == 'deepdiff-dep'", specifier = ">=7.0.1,<9.0.0" },
     { name = "diffusers", marker = "extra == 'diffusers-dep'", specifier = ">=0.38.0,<0.40.0" },
     { name = "dm-tree", marker = "extra == 'groot'", specifier = ">=0.1.8,<1.0.0" },
-    { name = "draccus", specifier = "==0.10.0" },
+    { name = "draccus", specifier = ">=0.11.6,<0.12.0" },
     { name = "dynamixel-sdk", marker = "extra == 'dynamixel'", specifier = ">=3.7.31,<3.9.0" },
     { name = "einops", specifier = ">=0.8.0,<0.9.0" },
     { name = "faker", marker = "extra == 'sarm'", specifier = ">=33.0.0,<35.0.0" },
@@ -5636,18 +5635,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" },
 ]
 
-[[package]]
-name = "pyyaml-include"
-version = "1.4.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "pyyaml" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/7f/be/2d07ad85e3d593d69640876a8686eae2c533db8cb7bf298d25c421b4d2d5/pyyaml-include-1.4.1.tar.gz", hash = "sha256:1a96e33a99a3e56235f5221273832464025f02ff3d8539309a3bf00dec624471", size = 20592, upload-time = "2024-03-25T14:56:43.748Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/d5/ca/6a2cc3a73170d10b5af1f1613baa2ed1f8f46f62dd0bfab2bffd2c2fe260/pyyaml_include-1.4.1-py3-none-any.whl", hash = "sha256:323c7f3a19c82fbc4d73abbaab7ef4f793e146a13383866831631b26ccc7fb00", size = 19079, upload-time = "2024-03-25T14:56:41.274Z" },
-]
-
 [[package]]
 name = "pyzmq"
 version = "27.1.0"