Merge branch 'main' into feat/fix_soft_prompt_typo

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

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/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/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_dataset_viz.py

@@ -87,11 +87,6 @@ import tqdm
 from lerobot.configs import DEPTH_MILLIMETER_UNIT
 from lerobot.datasets import LeRobotDataset
 from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
-from lerobot.utils.dataset_visualization_utils import (
-    get_extra_scalar_keys,
-    is_scalar_like,
-    scalar_to_float,
-)
 from lerobot.utils.utils import init_logging
 
 logger = logging.getLogger(__name__)
@@ -158,8 +153,6 @@ def build_blueprint_from_dataset(dataset: LeRobotDataset):
     for key in (DONE, REWARD, SUCCESS):
         if key in dataset.features:
             views.append(rrb.TimeSeriesView(origin=key, name=key))
-    for key in get_extra_scalar_keys(dataset):
-        views.append(rrb.TimeSeriesView(origin=key, name=key))
 
     return rrb.Blueprint(rrb.Grid(*views))
 
@@ -251,8 +244,6 @@ def visualize_dataset(
         hi = stats["q99"] if "q99" in stats else stats["max"]
         depth_ranges[key] = (float(np.asarray(lo).item()), float(np.asarray(hi).item()))
 
-    extra_scalar_keys = get_extra_scalar_keys(dataset)
-
     first_index = None
     for batch in tqdm.tqdm(dataloader, total=len(dataloader)):
         if first_index is None:
@@ -296,10 +287,6 @@ def visualize_dataset(
             if SUCCESS in batch:
                 rr.log(SUCCESS, rr.Scalars(batch[SUCCESS][i].item()))
 
-            for key in extra_scalar_keys:
-                if key in batch and is_scalar_like(batch[key][i]):
-                    rr.log(key, rr.Scalars(scalar_to_float(batch[key][i])))
-
     # save .rrd locally
     if mode == "local" and save:
         output_dir.mkdir(parents=True, exist_ok=True)

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/dataset_visualization_utils.py

@@ -1,99 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Shared helpers for visualizing scalar features from a LeRobot dataset."""
-
-from collections.abc import Iterable, Mapping
-from typing import TYPE_CHECKING
-
-import numpy as np
-import torch
-
-from .constants import ACTION, DEFAULT_FEATURES, DONE, OBS_STATE, REWARD, SUCCESS
-
-if TYPE_CHECKING:
-    from lerobot.datasets import LeRobotDataset
-
-
-METADATA_KEYS = {*DEFAULT_FEATURES, "task"}
-KNOWN_SCALAR_KEYS = {DONE, REWARD, SUCCESS}
-SCALAR_DTYPE_KINDS = {"b", "i", "u", "f"}
-
-
-def is_scalar_feature(feature: Mapping) -> bool:
-    """Return whether a feature schema describes a numeric or boolean scalar."""
-
-    dtype = feature.get("dtype")
-    if not isinstance(dtype, str):
-        return False
-    try:
-        dtype_kind = np.dtype(dtype).kind
-    except (TypeError, ValueError):
-        return False
-    if dtype_kind not in SCALAR_DTYPE_KINDS:
-        return False
-
-    shape = feature.get("shape")
-    if shape is None:
-        return True
-    if isinstance(shape, int):
-        return shape == 1
-    if not isinstance(shape, (list, tuple)):
-        return False
-    return len(shape) == 0 or (len(shape) == 1 and shape[0] == 1)
-
-
-def get_extra_scalar_keys(dataset: "LeRobotDataset", additional_known_keys: Iterable[str] = ()) -> list[str]:
-    """Return scalar feature keys not handled by the visualizer's standard paths."""
-
-    known_keys = {
-        ACTION,
-        OBS_STATE,
-        *KNOWN_SCALAR_KEYS,
-        *METADATA_KEYS,
-        *additional_known_keys,
-        *dataset.meta.camera_keys,
-    }
-    return [
-        key
-        for key, feature in dataset.features.items()
-        if key not in known_keys and is_scalar_feature(feature)
-    ]
-
-
-def is_scalar_like(value: object) -> bool:
-    """Return whether a runtime value contains exactly one numeric or boolean scalar."""
-
-    if isinstance(value, torch.Tensor):
-        return value.numel() == 1 and not value.is_complex()
-    if isinstance(value, np.ndarray):
-        return value.size == 1 and value.dtype.kind in SCALAR_DTYPE_KINDS
-    return np.isscalar(value) and np.asarray(value).dtype.kind in SCALAR_DTYPE_KINDS
-
-
-def scalar_to_float(value: object) -> float:
-    """Convert a scalar-like tensor, array, or Python value to ``float``."""
-
-    return float(value.item() if hasattr(value, "item") else value)
-
-
-def get_scalar_values(sample: Mapping, keys: Iterable[str]) -> dict[str, float]:
-    """Select and convert scalar-like values from ``sample`` for the requested keys."""
-
-    values = {}
-    for key in keys:
-        value = sample.get(key)
-        if value is not None and is_scalar_like(value):
-            values[key] = scalar_to_float(value)
-    return values

src/lerobot/utils/foxglove_visualization.py

@@ -23,7 +23,6 @@ importing from here directly. Requires the ``viz`` extra (``pip install 'lerobot
 import logging
 import numbers
 import time
-from collections.abc import Iterable, Mapping
 
 import cv2
 import numpy as np
@@ -42,7 +41,6 @@ from .constants import (
     SUCCESS,
     TRUNCATED,
 )
-from .dataset_visualization_utils import get_extra_scalar_keys, get_scalar_values
 from .import_utils import require_package
 
 # Static schema shared by all scalar topics. Each message carries a flat list of ``{label, value}``
@@ -407,24 +405,6 @@ def _frame_to_scalars(sample: dict, key: str, labels: list[str] | None = None) -
     return _labeled_scalars(name, arr.flatten(), labels)
 
 
-def _dataset_frame_scalar_groups(
-    sample: Mapping, extra_scalar_keys: Iterable[str]
-) -> tuple[dict[str, float], dict[str, float]]:
-    """Return standard episode scalars and custom scalar features for one dataset frame."""
-
-    episode_scalars = {}
-    for feature, label in (
-        (DONE, "done"),
-        (TRUNCATED, "truncated"),
-        (REWARD, "reward"),
-        (SUCCESS, "success"),
-    ):
-        value = sample.get(feature)
-        if value is not None:
-            episode_scalars[label] = float(value)
-    return episode_scalars, get_scalar_values(sample, extra_scalar_keys)
-
-
 def serve_foxglove_dataset_playback(
     dataset,
     episode_index: int,
@@ -472,7 +452,6 @@ def serve_foxglove_dataset_playback(
         raise ValueError("Cannot visualize an empty episode.")
     first_ns, last_ns = times_ns[0], times_ns[-1]
     camera_keys = list(dataset.meta.camera_keys)
-    extra_scalar_keys = get_extra_scalar_keys(dataset, additional_known_keys=(TRUNCATED,))
     # Dataset-wide q01/q99 depth bounds (fallback min/max) used to normalize depth to [0, 1].
     depth_ranges: dict[str, tuple[float, float]] = {}
     for key in dataset.meta.depth_keys:
@@ -521,14 +500,17 @@ def serve_foxglove_dataset_playback(
             channels=channels,
             log_time=log_time,
         )
-        episode_scalars, extra_scalars = _dataset_frame_scalar_groups(sample, extra_scalar_keys)
+        episode_scalars = {}
+        for feat, label in (
+            (DONE, "done"),
+            (TRUNCATED, "truncated"),
+            (REWARD, "reward"),
+            (SUCCESS, "success"),
+        ):
+            v = sample.get(feat)
+            if v is not None:
+                episode_scalars[label] = float(v)
         _log_foxglove_scalars("/episode/state", episode_scalars, channels=channels, log_time=log_time)
-        _log_foxglove_scalars(
-            "/episode/extras",
-            extra_scalars,
-            channels=channels,
-            log_time=log_time,
-        )
 
     lock = threading.Lock()
     stop_event = threading.Event()

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/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/datasets/test_visualize_dataset.py

@@ -13,78 +13,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import sys
-from types import SimpleNamespace
-
-import numpy as np
 import pytest
-import torch
 
 pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
 
 from lerobot.scripts.lerobot_dataset_viz import visualize_dataset
-from lerobot.utils import import_utils
-from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS, TRUNCATED
-from lerobot.utils.dataset_visualization_utils import (
-    get_extra_scalar_keys,
-    is_scalar_feature,
-    is_scalar_like,
-    scalar_to_float,
-)
-
-
-class DummyMeta:
-    camera_keys = ["observation.images.front"]
-
-
-class DummyFeatureDataset:
-    meta = DummyMeta()
-    features = {
-        "index": {"dtype": "int64", "shape": [1]},
-        "timestamp": {"dtype": "float32", "shape": [1]},
-        "episode_index": {"dtype": "int64", "shape": [1]},
-        "frame_index": {"dtype": "int64", "shape": [1]},
-        "task_index": {"dtype": "int64", "shape": [1]},
-        ACTION: {"dtype": "float32", "shape": [6]},
-        OBS_STATE: {"dtype": "float32", "shape": [6]},
-        DONE: {"dtype": "bool", "shape": [1]},
-        REWARD: {"dtype": "float32", "shape": [1]},
-        SUCCESS: {"dtype": "bool", "shape": [1]},
-        TRUNCATED: {"dtype": "bool", "shape": [1]},
-        "observation.images.front": {"dtype": "video", "shape": [3, 480, 640]},
-        "q_target": {"dtype": "float32", "shape": [1]},
-        "quality": {"dtype": "double", "shape": [1]},
-        "intervention": {"dtype": "bool", "shape": []},
-        "embedding": {"dtype": "float32", "shape": [32]},
-        "comment": {"dtype": "string", "shape": [1]},
-    }
-
-
-class DummyVizDataset:
-    repo_id = "dummy/custom-scalars"
-    depth_output_unit = "m"
-    meta = SimpleNamespace(camera_keys=[], depth_keys=[], stats=None)
-    features = {
-        "index": {"dtype": "int64", "shape": [1]},
-        "timestamp": {"dtype": "float32", "shape": [1]},
-        ACTION: {"dtype": "float32", "shape": [2], "names": ["x", "y"]},
-        "q_target": {"dtype": "float32", "shape": [1]},
-        "intervention": {"dtype": "bool", "shape": [1]},
-        "embedding": {"dtype": "float32", "shape": [2]},
-    }
-
-    def __len__(self):
-        return 2
-
-    def __getitem__(self, index):
-        return {
-            "index": torch.tensor(index),
-            "timestamp": torch.tensor(index / 10),
-            ACTION: torch.tensor([index, index + 1], dtype=torch.float32),
-            "q_target": torch.tensor([index + 0.5]),
-            "intervention": torch.tensor([index == 0]),
-            "embedding": torch.tensor([index, index + 1], dtype=torch.float32),
-        }
 
 
 @pytest.mark.skip("TODO: add dummy videos")
@@ -100,93 +33,3 @@ def test_visualize_local_dataset(tmp_path, lerobot_dataset_factory):
         output_dir=output_dir,
     )
     assert rrd_path.exists()
-
-
-def test_get_extra_scalar_keys_skips_known_metadata_and_non_scalars():
-    dataset = DummyFeatureDataset()
-
-    assert get_extra_scalar_keys(dataset) == [TRUNCATED, "q_target", "quality", "intervention"]
-    assert get_extra_scalar_keys(dataset, additional_known_keys=(TRUNCATED,)) == [
-        "q_target",
-        "quality",
-        "intervention",
-    ]
-
-
-@pytest.mark.parametrize(
-    ("feature", "expected"),
-    [
-        ({"dtype": "float32", "shape": (1,)}, True),
-        ({"dtype": "double", "shape": [1]}, True),
-        ({"dtype": "bool", "shape": []}, True),
-        ({"dtype": "int64", "shape": 1}, True),
-        ({"dtype": "float32", "shape": (3,)}, False),
-        ({"dtype": "complex64", "shape": [1]}, False),
-        ({"dtype": "datetime64[ns]", "shape": [1]}, False),
-        ({"dtype": "string", "shape": [1]}, False),
-        ({"shape": [1]}, False),
-    ],
-)
-def test_is_scalar_feature(feature, expected):
-    assert is_scalar_feature(feature) is expected
-
-
-@pytest.mark.parametrize(
-    ("value", "expected"),
-    [
-        (torch.tensor(1.5), True),
-        (torch.tensor([1.5]), True),
-        (torch.tensor([1.5, 2.5]), False),
-        (np.array(2.0), True),
-        (np.array([2.0]), True),
-        (np.array([2.0, 3.0]), False),
-        (True, True),
-        ("not numeric", False),
-    ],
-)
-def test_is_scalar_like(value, expected):
-    assert is_scalar_like(value) is expected
-
-
-def test_scalar_to_float():
-    assert scalar_to_float(torch.tensor(3.0)) == 3.0
-    assert scalar_to_float(np.array([4.0])) == 4.0
-
-
-def test_visualize_dataset_logs_extra_scalars(monkeypatch):
-    logged = []
-    initialized = []
-
-    dummy_rrb = SimpleNamespace(
-        Spatial2DView=lambda origin=None, name=None: SimpleNamespace(
-            kind="spatial", origin=origin, name=name
-        ),
-        TimeSeriesView=lambda origin=None, name=None, overrides=None: SimpleNamespace(
-            kind="time_series", origin=origin, name=name, overrides=overrides
-        ),
-        Grid=lambda *views: SimpleNamespace(views=views),
-        Blueprint=lambda root: SimpleNamespace(root=root),
-    )
-    dummy_rr = SimpleNamespace(
-        SeriesLines=lambda names=None: SimpleNamespace(names=names),
-        Scalars=lambda value: SimpleNamespace(value=value),
-        init=lambda *args, **kwargs: initialized.append((args, kwargs)),
-        log=lambda key, entity: logged.append((key, entity)),
-        set_time=lambda *args, **kwargs: None,
-        blueprint=dummy_rrb,
-    )
-    monkeypatch.setitem(sys.modules, "rerun", dummy_rr)
-    monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb)
-    monkeypatch.setattr(import_utils, "require_package", lambda *args, **kwargs: None)
-
-    visualize_dataset(DummyVizDataset(), episode_index=0, batch_size=2)
-
-    logged_keys = [key for key, _ in logged]
-    assert logged_keys.count("q_target") == 2
-    assert logged_keys.count("intervention") == 2
-    assert "embedding" not in logged_keys
-
-    blueprint = initialized[0][1]["default_blueprint"]
-    time_series_origins = {view.origin for view in blueprint.root.views if view.kind == "time_series"}
-    assert {"q_target", "intervention"} <= time_series_origins
-    assert "embedding" not in time_series_origins

tests/utils/test_foxglove_visualization.py

@@ -24,7 +24,7 @@ the functions that talk to the server, so the helpers below run in the base test
 import numpy as np
 
 from lerobot.utils import foxglove_visualization as fv
-from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
+from lerobot.utils.constants import ACTION, OBS_STATE
 
 
 def test_foxglove_safe_name_collapses_dots():
@@ -93,24 +93,6 @@ def test_feature_dim_names_formats():
     assert fv._feature_dim_names({"shape": [2]}) is None
 
 
-def test_dataset_frame_scalar_groups_include_custom_scalars():
-    sample = {
-        DONE: np.array(True),
-        REWARD: np.array(0.5),
-        SUCCESS: np.array(False),
-        "q_target": np.array([0.75]),
-        "intervention": np.array([True]),
-        "embedding": np.array([1.0, 2.0]),
-    }
-
-    episode_scalars, extra_scalars = fv._dataset_frame_scalar_groups(
-        sample, ("q_target", "intervention", "embedding")
-    )
-
-    assert episode_scalars == {"done": 1.0, "reward": 0.5, "success": 0.0}
-    assert extra_scalars == {"q_target": 0.75, "intervention": 1.0}
-
-
 def test_is_scalar():
     assert fv._is_scalar(1.0)
     assert fv._is_scalar(np.float32(2.0))

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"