feat(dataset-viz): render generic custom scalar columns

Salvage of huggingface/lerobot#3918 by @nathon-lee — rebased to main.

lerobot-dataset-viz only logged known scalars; custom float/bool fields
now appear in the Rerun blueprint and frame stream when shape is scalar.

refactor(viz): add new non-default colums viz

Co-authored-by: nathon-lee <leejianwoo@gmail.com>

docs/source/hil_data_collection.mdx

@@ -59,7 +59,6 @@ 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 the leader holder using 1 M2x6mm screw.
+- Attach the handle to motor 5 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.
 

docs/source/using_dataset_tools.mdx

@@ -11,10 +11,9 @@ LeRobot provides several utilities for manipulating datasets:
 3. **Merge Datasets** - Combine multiple datasets into one. The datasets must have identical features, and episodes are concatenated in the order specified in `repo_ids`
 4. **Add Features** - Add new features to a dataset
 5. **Remove Features** - Remove features from a dataset
-6. **Modify Tasks** - Change the natural-language task descriptions associated with episodes
-7. **Convert to Video** - Convert image-based datasets to video format for efficient storage (RGB and depth cameras are encoded with separate encoders)
-8. **Re-encode Videos** - Re-encode an existing video dataset's RGB and/or depth streams with new encoder settings
-9. **Show the Info of Datasets** - Show the summary of datasets information such as number of episode etc.
+6. **Convert to Video** - Convert image-based datasets to video format for efficient storage (RGB and depth cameras are encoded with separate encoders)
+7. **Re-encode Videos** - Re-encode an existing video dataset's RGB and/or depth streams with new encoder settings
+8. **Show the Info of Datasets** - Show the summary of datasets information such as number of episode etc.
 
 The core implementation is in `lerobot.datasets.dataset_tools`.
 An example script detailing how to use the tools API is available in `examples/dataset/use_dataset_tools.py`.
@@ -90,53 +89,6 @@ lerobot-edit-dataset \
     --operation.feature_names "['observation.images.top']"
 ```
 
-#### Modify Tasks
-
-Change the natural-language task descriptions attached to episodes. This is useful for fixing typos, standardizing wording, or re-labeling episodes.
-
-> [!WARNING]
-> `modify_tasks` modifies the dataset **in-place** (updating `meta/tasks.parquet`, the `task_index` column in the data files, the `tasks` column in the episode metadata, and `total_tasks` in `meta/info.json`). The `--new_repo_id` and `--new_root` parameters are ignored for this operation.
-
-```bash
-# Set a single task for all episodes
-lerobot-edit-dataset \
-    --repo_id lerobot/pusht \
-    --operation.type modify_tasks \
-    --operation.new_task "Pick up the cube and place it"
-
-# Set different tasks for specific episodes
-lerobot-edit-dataset \
-    --repo_id lerobot/pusht \
-    --operation.type modify_tasks \
-    --operation.episode_tasks '{"0": "Task A", "1": "Task B", "2": "Task A"}'
-
-# Replace existing task strings wherever they appear
-lerobot-edit-dataset \
-    --repo_id lerobot/pusht \
-    --operation.type modify_tasks \
-    --operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}'
-
-# Combine modes in a single run
-lerobot-edit-dataset \
-    --repo_id lerobot/pusht \
-    --operation.type modify_tasks \
-    --operation.new_task "Default task" \
-    --operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}' \
-    --operation.episode_tasks '{"5": "Special task for episode 5"}'
-```
-
-**Parameters:**
-
-- `new_task`: A single task string used as the default for episodes not otherwise covered.
-- `episode_tasks`: Mapping from episode index to task string.
-- `task_replacements`: Mapping from existing task strings to their replacements, applied to episodes whose current task matches a key. Every key must be an existing task in the dataset.
-
-The modes can be combined in a single run. Per episode, the task is resolved with the following precedence:
-
-`episode_tasks` > `task_replacements` > `new_task` > original task
-
-At least one of `new_task`, `episode_tasks`, or `task_replacements` must be specified. An episode that ends up with no task raises an error.
-
 #### Convert to Video
 
 Convert an image-based dataset to video format, creating a new LeRobotDataset where images are stored as videos. This is useful for reducing storage requirements and improving data loading performance. The new dataset will have the exact same structure as the original, but with images encoded as MP4 videos in the proper LeRobot format.

pyproject.toml

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

src/lerobot/configs/policies.py

@@ -163,10 +163,8 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):  # type: igno
         return None
 
     def _save_pretrained(self, save_directory: Path) -> None:
-        # 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)
+        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
+            draccus.dump(self, f, indent=4)
 
     @classmethod
     def from_pretrained(

src/lerobot/configs/rewards.py

@@ -103,10 +103,8 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
         pass
 
     def _save_pretrained(self, save_directory: Path) -> None:
-        # 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)
+        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
+            draccus.dump(self, f, indent=4)
 
     @classmethod
     def from_pretrained(

src/lerobot/configs/train.py

@@ -194,11 +194,7 @@ class TrainPipelineConfig(HubMixin):
             )
 
         if Path(config_path).resolve().exists():
-            # `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
+            policy_dir = Path(config_path).parent
             self.checkpoint_path = policy_dir.parent
         elif self.job.is_remote:
             return

src/lerobot/datasets/dataset_tools.py

@@ -1435,18 +1435,15 @@ def modify_tasks(
     dataset: LeRobotDataset,
     new_task: str | None = None,
     episode_tasks: dict[int, str] | None = None,
-    task_replacements: dict[str, str] | None = None,
 ) -> LeRobotDataset:
     """Modify tasks in a LeRobotDataset.
 
     This function allows you to either:
     1. Set a single task for the entire dataset (using `new_task`)
     2. Set specific tasks for specific episodes (using `episode_tasks`)
-    3. Replace existing task strings wherever they appear (using `task_replacements`)
 
-    Per episode, the task is resolved with precedence:
-    `episode_tasks` > `task_replacements` > `new_task` > original task. An episode that ends
-    up with no task (none of the above apply and it had no original task) raises an error.
+    You can combine both: `new_task` sets the default, and `episode_tasks` overrides
+    specific episodes.
 
     The dataset is modified in-place, updating only the task-related files:
     - meta/tasks.parquet
@@ -1456,14 +1453,11 @@ def modify_tasks(
 
     Args:
         dataset: The source LeRobotDataset to modify.
-        new_task: Default task applied to any episode not covered by `episode_tasks` or a
-            matching `task_replacements` entry.
-        episode_tasks: Optional dict mapping episode indices to task strings. Takes precedence
-            over both `task_replacements` and `new_task`.
-        task_replacements: Optional dict mapping existing task strings to new ones. Applied to
-            episodes whose current task matches a key. Every key must be an existing task.
+        new_task: A single task string to apply to all episodes. If None and episode_tasks
+            is also None, raises an error.
+        episode_tasks: Optional dict mapping episode indices to their task strings.
+            Overrides `new_task` for specific episodes.
 
-    At least one of `new_task`, `episode_tasks`, or `task_replacements` must be provided.
 
     Examples:
         Set a single task for all episodes:
@@ -1481,17 +1475,11 @@ def modify_tasks(
                 new_task="Default task",
                 episode_tasks={5: "Special task for episode 5"}
             )
-
-        Replace existing task strings in-place:
-            dataset = modify_tasks(
-                dataset,
-                task_replacements={"Pick up the cube": "Lift the cube"}
-            )
     """
-    if not new_task and not episode_tasks and not task_replacements:
-        raise ValueError("Must specify at least one of new_task, episode_tasks, or task_replacements")
+    if new_task is None and episode_tasks is None:
+        raise ValueError("Must specify at least one of new_task or episode_tasks")
 
-    if episode_tasks:
+    if episode_tasks is not None:
         valid_indices = set(range(dataset.meta.total_episodes))
         invalid = set(episode_tasks.keys()) - valid_indices
         if invalid:
@@ -1501,29 +1489,19 @@ def modify_tasks(
     if dataset.meta.episodes is None:
         dataset.meta.episodes = load_episodes(dataset.root)
 
-    if task_replacements:
-        current_tasks = set(dataset.meta.tasks.index)
-        invalid_tasks = set(task_replacements) - current_tasks
-        if invalid_tasks:
-            raise ValueError(f"Task replacements reference unknown tasks: {sorted(invalid_tasks)}")
-
     # Build the mapping from episode index to task string
     episode_to_task: dict[int, str] = {}
     for ep_idx in range(dataset.meta.total_episodes):
-        original_tasks = dataset.meta.episodes[ep_idx]["tasks"]
-        original_task = original_tasks[0] if original_tasks else None
-
         if episode_tasks and ep_idx in episode_tasks:
             episode_to_task[ep_idx] = episode_tasks[ep_idx]
-        elif task_replacements and original_task in task_replacements:
-            episode_to_task[ep_idx] = task_replacements[original_task]
-        elif new_task:
+        elif new_task is not None:
             episode_to_task[ep_idx] = new_task
-        elif original_task:
-            # Keep original task if not overridden and no default provided
-            episode_to_task[ep_idx] = original_task
         else:
-            raise ValueError(f"Episode {ep_idx} has no task; provide new_task or episode_tasks")
+            # Keep original task if not overridden and no default provided
+            original_tasks = dataset.meta.episodes[ep_idx]["tasks"]
+            if not original_tasks:
+                raise ValueError(f"Episode {ep_idx} has no tasks and no default task was provided")
+            episode_to_task[ep_idx] = original_tasks[0]
 
     # Collect all unique tasks and create new task mapping
     unique_tasks = sorted(set(episode_to_task.values()))

src/lerobot/policies/evo1/modeling_evo1.py

@@ -42,9 +42,6 @@ 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,9 +68,6 @@ 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,9 +520,6 @@ 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,9 +749,6 @@ 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,9 +714,6 @@ 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,10 +249,6 @@ 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,9 +145,6 @@ 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,23 +168,14 @@ class SmolVLMWithExpertModel(nn.Module):
                 last_layers.append(self.num_vlm_layers - 2)
             frozen_layers = [
                 "lm_head",
-                "text_model.norm.weight",
+                "text_model.model.norm.weight",
             ]
             for layer in last_layers:
-                frozen_layers.append(f"text_model.layers.{layer}.")
+                frozen_layers.append(f"text_model.model.layers.{layer}.")
 
-            unmatched_patterns = set(frozen_layers)
             for name, params in self.vlm.named_parameters():
-                matched_patterns = [k for k in frozen_layers if k in name]
-                if matched_patterns:
+                if any(k in name for k in frozen_layers):
                     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/rl/algorithms/configs.py

@@ -16,7 +16,6 @@ from __future__ import annotations
 
 import abc
 import builtins
-import json
 import logging
 import os
 from dataclasses import dataclass, field
@@ -79,10 +78,8 @@ class RLAlgorithmConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
 
     def _save_pretrained(self, save_directory: Path) -> None:
         """Serialize this config as ``config.json`` inside ``save_directory``."""
-        # 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)
+        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
+            draccus.dump(self, f, indent=4)
 
     @classmethod
     def from_pretrained(

src/lerobot/rollout/context.py

@@ -57,7 +57,6 @@ 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:
@@ -227,12 +226,6 @@ 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,7 +22,6 @@ way via ``notify_observation``.
 
 from __future__ import annotations
 
-import inspect
 import logging
 import math
 import time
@@ -63,23 +62,6 @@ _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,6 +87,11 @@ 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__)
@@ -153,6 +158,8 @@ 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))
 
@@ -244,6 +251,8 @@ 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:
@@ -287,6 +296,10 @@ 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_edit_dataset.py

@@ -127,12 +127,6 @@ Modify tasks - set default task with overrides for specific episodes (WARNING: m
         --operation.new_task "Default task" \
         --operation.episode_tasks '{"5": "Special task for episode 5"}'
 
-Modify tasks - replace existing task strings in-place (WARNING: modifies in-place):
-    lerobot-edit-dataset \
-        --repo_id lerobot/pusht \
-        --operation.type modify_tasks \
-        --operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}'
-
 Convert image dataset to video format and save locally:
     lerobot-edit-dataset \
         --repo_id lerobot/pusht_image \
@@ -309,7 +303,6 @@ class RemoveFeatureConfig(OperationConfig):
 class ModifyTasksConfig(OperationConfig):
     new_task: str | None = None
     episode_tasks: dict[str, str] | None = None
-    task_replacements: dict[str, str] | None = None
 
 
 @OperationConfig.register_subclass("convert_image_to_video")
@@ -558,12 +551,9 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
 
     new_task = cfg.operation.new_task
     episode_tasks_raw = cfg.operation.episode_tasks
-    task_replacements = cfg.operation.task_replacements
 
-    if new_task is None and episode_tasks_raw is None and task_replacements is None:
-        raise ValueError(
-            "Must specify at least one of new_task, episode_tasks, or task_replacements for modify_tasks operation"
-        )
+    if new_task is None and episode_tasks_raw is None:
+        raise ValueError("Must specify at least one of new_task or episode_tasks for modify_tasks operation")
 
     if cfg.new_repo_id is not None or cfg.new_root is not None:
         logging.warning(
@@ -583,14 +573,11 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
         logging.info(f"  Default task: '{new_task}'")
     if episode_tasks:
         logging.info(f"  Episode-specific tasks: {episode_tasks}")
-    if task_replacements:
-        logging.info(f"  Task replacements: {task_replacements}")
 
     modified_dataset = modify_tasks(
         dataset,
         new_task=new_task,
         episode_tasks=episode_tasks,
-        task_replacements=task_replacements,
     )
 
     logging.info(f"Dataset modified at {dataset.root}")

src/lerobot/scripts/lerobot_train.py

@@ -348,6 +348,7 @@ 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,
             },
@@ -355,17 +356,11 @@ 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,15 +67,7 @@ class BiSOLeader(BimanualMixin, Teleoperator):
 
     @cached_property
     def feedback_features(self) -> dict[str, type]:
-        # 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()},
-        }
+        return {}
 
     def setup_motors(self) -> None:
         self.left_arm.setup_motors()
@@ -95,43 +87,6 @@ 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:
-        """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)
+        # TODO: Implement force feedback
+        raise NotImplementedError

src/lerobot/utils/dataset_visualization_utils.py

@@ -0,0 +1,99 @@
+# 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,6 +23,7 @@ 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
@@ -41,6 +42,7 @@ 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}``
@@ -405,6 +407,24 @@ 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,
@@ -452,6 +472,7 @@ 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:
@@ -500,17 +521,14 @@ def serve_foxglove_dataset_playback(
             channels=channels,
             log_time=log_time,
         )
-        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)
+        episode_scalars, extra_scalars = _dataset_frame_scalar_groups(sample, extra_scalar_keys)
         _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,8 +85,6 @@ 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
 
 
@@ -97,8 +95,6 @@ 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,27 +66,3 @@ 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_dataset_tools.py

@@ -1125,61 +1125,9 @@ def test_modify_tasks_default_with_overrides(sample_dataset):
             assert ep_data["tasks"][0] == default_task
 
 
-def test_modify_tasks_replacements(sample_dataset):
-    """Test replacing task strings based on their current values."""
-    modified_dataset = modify_tasks(
-        sample_dataset,
-        task_replacements={
-            "task_0": "Pick the cube",
-            "task_1": "Place the cube",
-        },
-    )
-
-    assert len(modified_dataset.meta.tasks) == 2
-    assert "Pick the cube" in modified_dataset.meta.tasks.index
-    assert "Place the cube" in modified_dataset.meta.tasks.index
-
-    for ep_idx in range(5):
-        expected_task = "Pick the cube" if ep_idx % 2 == 0 else "Place the cube"
-        assert modified_dataset.meta.episodes[ep_idx]["tasks"][0] == expected_task
-
-
-def test_modify_tasks_replacements_with_episode_overrides(sample_dataset):
-    """Test that explicit episode overrides take precedence over replacements."""
-    modified_dataset = modify_tasks(
-        sample_dataset,
-        task_replacements={
-            "task_0": "Pick the cube",
-            "task_1": "Place the cube",
-        },
-        episode_tasks={1: "Inspect the cube"},
-    )
-
-    assert modified_dataset.meta.episodes[0]["tasks"][0] == "Pick the cube"
-    assert modified_dataset.meta.episodes[1]["tasks"][0] == "Inspect the cube"
-    assert modified_dataset.meta.episodes[3]["tasks"][0] == "Place the cube"
-    assert len(modified_dataset.meta.tasks) == 3
-
-
-def test_modify_tasks_default_task_and_replacements(sample_dataset):
-    """Test that new_task acts as the default for episodes not matched by task_replacements."""
-    modified_dataset = modify_tasks(
-        sample_dataset,
-        new_task="Default task",
-        task_replacements={"task_0": "Pick the cube"},
-    )
-
-    for ep_idx in range(5):
-        expected_task = "Pick the cube" if ep_idx % 2 == 0 else "Default task"
-        assert modified_dataset.meta.episodes[ep_idx]["tasks"][0] == expected_task
-    assert len(modified_dataset.meta.tasks) == 2
-
-
 def test_modify_tasks_no_task_specified(sample_dataset):
     """Test error when no task is specified."""
-    with pytest.raises(
-        ValueError, match="Must specify at least one of new_task, episode_tasks, or task_replacements"
-    ):
+    with pytest.raises(ValueError, match="Must specify at least one of new_task or episode_tasks"):
         modify_tasks(sample_dataset)
 
 
@@ -1189,12 +1137,6 @@ def test_modify_tasks_invalid_episode_indices(sample_dataset):
         modify_tasks(sample_dataset, episode_tasks={10: "Task", 20: "Task"})
 
 
-def test_modify_tasks_invalid_task_replacements(sample_dataset):
-    """Test error when task replacements refer to unknown task strings."""
-    with pytest.raises(ValueError, match="Task replacements reference unknown tasks"):
-        modify_tasks(sample_dataset, task_replacements={"missing_task": "New task"})
-
-
 def test_modify_tasks_updates_info_json(sample_dataset):
     """Test that total_tasks is updated in info.json."""
     episode_tasks = {0: "Task A", 1: "Task B", 2: "Task C", 3: "Task A", 4: "Task B"}

tests/datasets/test_visualize_dataset.py

@@ -13,11 +13,78 @@
 # 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")
@@ -33,3 +100,93 @@ 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/scripts/test_edit_dataset_parsing.py

@@ -105,23 +105,6 @@ class TestOperationTypeParsing:
         resolved_name = OperationConfig.get_choice_name(type(cfg.operation))
         assert resolved_name == type_name
 
-    def test_modify_tasks_replacements_args_parse(self):
-        cfg = parse_cfg(
-            [
-                "--repo_id",
-                "test/repo",
-                "--operation.type",
-                "modify_tasks",
-                "--operation.task_replacements",
-                '{"task_0": "pick cube", "task_1": "place cube"}',
-            ]
-        )
-        assert isinstance(cfg.operation, ModifyTasksConfig)
-        assert cfg.operation.task_replacements == {
-            "task_0": "pick cube",
-            "task_1": "place cube",
-        }
-
 
 class TestDepthEncoderParsing:
     """Test that the depth encoder is exposed and parsed for video operations."""

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, OBS_STATE
+from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
 
 
 def test_foxglove_safe_name_collapses_dots():
@@ -93,6 +93,24 @@ 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,17 +73,6 @@ 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 = 2
+revision = 3
 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,17 +1359,18 @@ sdist = { url = "https://files.pythonhosted.org/packages/a2/55/8f8cab2afd404cf57
 
 [[package]]
 name = "draccus"
-version = "0.11.6"
+version = "0.10.0"
 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/df/b0/cc399719e0cf6fea451f155798e19ea8c5e9c838b701d1565f29ea626c18/draccus-0.11.6.tar.gz", hash = "sha256:d134f576a1f4febd93c6b200df7f92e5febffe9efdc0a5f381bf50c2e0568a39", size = 67940, upload-time = "2026-06-12T13:21:19.999Z" }
+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" }
 wheels = [
-    { 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" },
+    { 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" },
 ]
 
 [[package]]
@@ -3288,7 +3289,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.11.6,<0.12.0" },
+    { name = "draccus", specifier = "==0.10.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" },
@@ -5635,6 +5636,18 @@ 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"