docs(update): updating docs with the task modification features

Part of #2326.

Signed-off-by: 陈伟 <woshei0a0a0a@qq.com>

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.
 

docs/source/using_dataset_tools.mdx

@@ -11,9 +11,10 @@ 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. **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.
+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.
 
 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`.
@@ -89,6 +90,53 @@ 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.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/datasets/dataset_tools.py

@@ -1435,15 +1435,18 @@ 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`)
 
-    You can combine both: `new_task` sets the default, and `episode_tasks` overrides
-    specific episodes.
+    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.
 
     The dataset is modified in-place, updating only the task-related files:
     - meta/tasks.parquet
@@ -1453,11 +1456,14 @@ def modify_tasks(
 
     Args:
         dataset: The source LeRobotDataset to modify.
-        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.
+        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.
 
+    At least one of `new_task`, `episode_tasks`, or `task_replacements` must be provided.
 
     Examples:
         Set a single task for all episodes:
@@ -1475,11 +1481,17 @@ def modify_tasks(
                 new_task="Default task",
                 episode_tasks={5: "Special task for episode 5"}
             )
-    """
-    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 is not None:
+        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 episode_tasks:
         valid_indices = set(range(dataset.meta.total_episodes))
         invalid = set(episode_tasks.keys()) - valid_indices
         if invalid:
@@ -1489,19 +1501,29 @@ 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 new_task is not None:
+        elif task_replacements and original_task in task_replacements:
+            episode_to_task[ep_idx] = task_replacements[original_task]
+        elif new_task:
             episode_to_task[ep_idx] = new_task
-        else:
+        elif original_task:
             # 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]
+            episode_to_task[ep_idx] = original_task
+        else:
+            raise ValueError(f"Episode {ep_idx} has no task; provide new_task or episode_tasks")
 
     # 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,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/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/robots/unitree_g1/config_unitree_g1.py

@@ -68,10 +68,6 @@ class UnitreeG1Config(RobotConfig):
     # Compensates for gravity on the unitree's arms using the arm ik solver
     gravity_compensation: bool = False
 
-    # Locomotion controller class name, e.g. "GrootLocomotionController",
-    # "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
-    # Selecting "SonicWholeBodyController" implicitly switches the robot to the 64-D
-    # latent-token action/observation interface (``motion_token.{i}.pos`` action and a
-    # ``motion_token_state.{i}.pos`` state echo) so ``lerobot-rollout`` can drive a
-    # policy trained on SONIC motion tokens (e.g. nepyope/sonic_walk).
+    # Lower-body controller class name, e.g. "GrootLocomotionController" or
+    # "HolosomaLocomotionController". None disables it.
     controller: str | None = None

src/lerobot/robots/unitree_g1/controllers/__init__.py

@@ -1,27 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 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.
-
-"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
-
-from .gr00t_locomotion import GrootLocomotionController
-from .holosoma_locomotion import HolosomaLocomotionController
-from .sonic_whole_body import SonicWholeBodyController
-
-__all__ = [
-    "GrootLocomotionController",
-    "HolosomaLocomotionController",
-    "SonicWholeBodyController",
-]

src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py

@@ -1,378 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 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.
-
-"""SONIC decoder whole-body controller for the Unitree G1 (token-only).
-
-Pure-Python/ONNX re-implementation of the *decode* half of NVIDIA's SONIC deploy stack.
-The encoder is intentionally absent: a token-output VLA (e.g. ``nepyope/sonic_walk``)
-supplies the 64-D latent ``motion_token`` directly each tick, and the SONIC **decoder**
-maps ``token + recent proprioception history`` to a residual action that is scaled and
-added onto the standing pose (``default_angles``) to produce 50 Hz joint-position targets
-for the robot's PD controller.
-
-Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order) and
-**MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` (in g1_utils)
-convert between them. Quaternions are scalar-first ``(w, x, y, z)``.
-"""
-
-from __future__ import annotations
-
-import json
-import logging
-
-import numpy as np
-import onnx
-import onnxruntime as ort
-from huggingface_hub import hf_hub_download
-
-from ..g1_utils import (
-    ISAACLAB_TO_MUJOCO,
-    MUJOCO_TO_ISAACLAB,
-    G1_29_JointIndex,
-    get_gravity_orientation,
-)
-from ..unitree_g1 import lowstate_to_obs
-
-logger = logging.getLogger(__name__)
-
-# ── Constants (hardware-validated; see the NVIDIA SONIC deploy reference) ──────
-CONTROL_DT = 0.02  # 50 Hz control period (s)
-TOKEN_DIM = 64  # decoder latent size
-
-# SONIC decoder checkpoint: NVIDIA's decoder ONNX re-packaged with its deploy constants
-# (kp/kd PD gains, the standing pose default_angles, and the residual action_scale) embedded
-# in the ONNX metadata; see upload_sonic_decoder.py for provisioning. The runtime loads the
-# model *and* all of these straight from the checkpoint (the Holosoma convention), so no
-# motor-physics math happens at deploy time.
-DEFAULT_SONIC_REPO_ID = "lerobot/sonic_decoder"
-DECODER_FILENAME = "model_decoder.onnx"
-DECODER_INPUT_DIM = 994  # token(64) + 10-frame proprio history + gravity
-
-
-def load_sonic_decoder(repo_id: str = DEFAULT_SONIC_REPO_ID):
-    """Load the SONIC decoder ONNX and its baked-in deploy constants from the checkpoint.
-
-    Returns ``(decoder_session, kp, kd, default_angles, action_scale, neutral_token)``. The
-    gains/pose/scale are (29,) float32 in IsaacLab joint order and ``neutral_token`` is the
-    (64,) float32 idle latent -- all read from the ONNX ``metadata_props`` rather than
-    recomputed/hardcoded at deploy time (mirrors ``holosoma_locomotion.load_policy``).
-    """
-    decoder_path = hf_hub_download(repo_id=repo_id, filename=DECODER_FILENAME)
-    so = ort.SessionOptions()
-    so.log_severity_level = 3  # quiet ORT logs
-    session = ort.InferenceSession(decoder_path, sess_options=so)
-    dec_dim = int(session.get_inputs()[0].shape[1])
-    if dec_dim != DECODER_INPUT_DIM:
-        raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected {DECODER_INPUT_DIM})")
-
-    meta = {p.key: p.value for p in onnx.load(decoder_path, load_external_data=False).metadata_props}
-    required = ("kp", "kd", "default_angles", "action_scale", "neutral_token")
-    missing = [k for k in required if k not in meta]
-    if missing:
-        raise ValueError(
-            f"SONIC decoder ONNX at {repo_id} is missing metadata {missing}; "
-            "re-run upload_sonic_decoder.py to (re)provision the checkpoint."
-        )
-    arr = {k: np.array(json.loads(meta[k]), dtype=np.float32) for k in required}
-    logger.info("Loaded SONIC deploy constants from %s (%d joints)", repo_id, len(arr["kp"]))
-    return session, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"]
-
-
-def _to_mujoco(a):
-    """Apply the ``MUJOCO_TO_ISAACLAB`` gather to a 29-vector (deploy-order reorder).
-
-    NOTE: this returns ``a[MUJOCO_TO_ISAACLAB]``. The ``_mj`` suffixes and the exact
-    permutation direction are a fixed convention validated against the deployed SONIC ONNX
-    policy (the decoder consumes vectors in this order). Do not "correct" the table or
-    rename toward the opposite direction without re-validating on hardware.
-    """
-    return a[MUJOCO_TO_ISAACLAB]
-
-
-# Action-feature prefix for the latent-token interface (see _extract_token_from_action).
-TOKEN_ACTION_PREFIX = "motion_token"  # nosec B105 - feature-key prefix, not a secret
-# Proprio-state prefix for the token interface: the robot echoes the last commanded token
-# here so ``lerobot-rollout`` aggregates it into a 64-D ``observation.state``.
-TOKEN_STATE_PREFIX = "motion_token_state"  # nosec B105 - feature-key prefix, not a secret
-
-
-def token_action_key(i: int) -> str:
-    """Action-dict key for the i-th component of the 64-D SONIC latent token.
-
-    The ``.pos`` suffix is required so the value flows through ``lerobot-rollout``, which
-    only routes ``.pos`` scalar features onto the policy action vector.
-    """
-    return f"{TOKEN_ACTION_PREFIX}.{i}.pos"
-
-
-def token_state_key(i: int) -> str:
-    """Observation key for the i-th component of the 64-D SONIC latent token state."""
-    return f"{TOKEN_STATE_PREFIX}.{i}.pos"
-
-
-# Startup blend duration: over the first control ticks, linearly interpolate every joint
-# from the robot's initial measured pose into the policy's commanded target, so control
-# eases in without a snap on the first command.
-INIT_RAMP_S = 3.0
-
-
-def _extract_token_from_action(action: dict | None) -> np.ndarray | None:
-    """Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or None.
-
-    The token-only interface: the caller supplies the 64-D encoder latent directly (e.g. a
-    token-output VLA's action), which the decoder consumes with the encoder bypassed.
-    Requires the full dense token; a partial one is ignored (returns None).
-    """
-    if not action:
-        return None
-    keys = [token_action_key(i) for i in range(TOKEN_DIM)]
-    if any(key not in action for key in keys):
-        return None
-    return np.fromiter((float(action[key]) for key in keys), dtype=np.float32, count=TOKEN_DIM)
-
-
-class SonicDecoder:
-    """Runs the SONIC decoder ONNX model and owns the proprioception history.
-
-    Each tick it appends the latest robot state to 10-frame history buffers, then maps the
-    supplied 64-D ``token`` + that history to a residual action added onto ``default_angles``.
-    The encoder is bypassed entirely (token supplied by the policy). ``default_angles`` and
-    ``action_scale`` are (29,) float32 in IsaacLab order, loaded from the checkpoint.
-    """
-
-    def __init__(self, decoder, default_angles, action_scale):
-        self.decoder = decoder
-        self.decoder_input = decoder.get_inputs()[0].name
-        self.default_angles = np.asarray(default_angles, np.float32)
-        self.action_scale = np.asarray(action_scale, np.float32)
-        self.default_angles_mj = _to_mujoco(self.default_angles)
-        self.token = np.zeros(TOKEN_DIM, np.float32)
-        self.last_action_mj = np.zeros(29, np.float32)
-        self.h_q_mj = [np.zeros(29, np.float32)] * 10
-        self.h_dq_mj = [np.zeros(29, np.float32)] * 10
-        self.h_ang = [np.zeros(3, np.float32)] * 10
-        self.h_act_mj = [np.zeros(29, np.float32)] * 10
-        self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
-
-    def reset(self):
-        """Clear the token and 10-frame proprioception history.
-
-        ``UnitreeG1.reset()`` relies on this so the first decoder outputs of a new episode
-        are not contaminated by the previous episode's state.
-        """
-        self.token = np.zeros(TOKEN_DIM, np.float32)
-        self.last_action_mj = np.zeros(29, np.float32)
-        self.h_q_mj = [np.zeros(29, np.float32)] * 10
-        self.h_dq_mj = [np.zeros(29, np.float32)] * 10
-        self.h_ang = [np.zeros(3, np.float32)] * 10
-        self.h_act_mj = [np.zeros(29, np.float32)] * 10
-        self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
-
-    def update_history(self, q, dq, ang, quat):
-        """Push the latest proprioception (pos/vel/gyro/orientation) into the 10-frame buffers."""
-        quat = quat / (np.linalg.norm(quat) + 1e-8)
-        q_mj = _to_mujoco(q)
-        dq_mj = _to_mujoco(dq)
-        self.h_q_mj = [q_mj - self.default_angles_mj] + self.h_q_mj[:-1]
-        self.h_dq_mj = [dq_mj] + self.h_dq_mj[:-1]
-        self.h_ang = [ang.copy()] + self.h_ang[:-1]
-        self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
-        self.h_quat = [quat.copy()] + self.h_quat[:-1]
-
-    def build_decoder_obs(self):
-        """Assemble the 994-D decoder input: token + 10-frame proprioception history + gravity."""
-        obs = np.zeros(994, np.float32)
-        off = 0
-        obs[off : off + 64] = self.token
-        off += 64
-        for h, sz in [
-            (list(reversed(self.h_ang)), 3),
-            (list(reversed(self.h_q_mj)), 29),
-            (list(reversed(self.h_dq_mj)), 29),
-            (list(reversed(self.h_act_mj)), 29),
-        ]:
-            for f in range(10):
-                obs[off : off + sz] = h[f]
-                off += sz
-        for q in reversed(self.h_quat):
-            obs[off : off + 3] = get_gravity_orientation(q)
-            off += 3
-        assert off == 994, f"Decoder obs mismatch: {off}"
-        return obs
-
-    def step(self, robot_obs, token, debug=False):
-        """One control tick: read robot obs, decode the supplied token -> joint targets.
-
-        Args:
-            robot_obs: dict with ``<joint>.q``/``.dq`` and ``imu.*`` fields.
-            token: 64-D latent supplied by the policy (encoder bypassed).
-            debug: log action/delta norms.
-
-        Returns:
-            dict of ``<joint>.q`` target positions (rad) in IsaacLab joint order.
-        """
-        self.token = np.asarray(token, np.float32)
-        jnames = [m.name for m in G1_29_JointIndex]
-        q = np.array(
-            [
-                robot_obs.get(f"{n}.q", self.default_angles[m.value])
-                for m, n in zip(G1_29_JointIndex, jnames, strict=False)
-            ],
-            np.float32,
-        )
-        dq = np.array([robot_obs.get(f"{n}.dq", 0.0) for n in jnames], np.float32)
-        quat = np.array(
-            [
-                robot_obs.get("imu.quat.w", 1),
-                robot_obs.get("imu.quat.x", 0),
-                robot_obs.get("imu.quat.y", 0),
-                robot_obs.get("imu.quat.z", 0),
-            ],
-            np.float32,
-        )
-        ang = np.array([robot_obs.get(f"imu.gyro.{a}", 0) for a in "xyz"], np.float32)
-        self.update_history(q, dq, ang, quat)
-        action_mj = (
-            self.decoder.run(None, {self.decoder_input: self.build_decoder_obs().reshape(1, -1)})[0]
-            .squeeze()
-            .astype(np.float32)
-        )
-        self.last_action_mj = action_mj.copy()
-        target = self.default_angles + action_mj[ISAACLAB_TO_MUJOCO] * self.action_scale
-        if debug:
-            delta = target - q
-            logger.debug(
-                "token_norm=%.4f action_norm=%.4f delta_max=%.4f delta_rms=%.4f",
-                np.linalg.norm(self.token),
-                np.linalg.norm(action_mj),
-                np.max(np.abs(delta)),
-                np.sqrt(np.mean(delta**2)),
-            )
-        return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
-
-
-class SonicRuntime:
-    """Loads the SONIC decoder ONNX model and owns the decode controller.
-
-    Token-only deploy: the encoder is bypassed; each tick the decoder consumes a 64-D
-    latent token supplied directly by the policy.
-    """
-
-    def __init__(self):
-        decoder_sess, self.kp, self.kd, default_angles, action_scale, neutral_token = load_sonic_decoder()
-        self.default_angles = default_angles
-        self.neutral_token = neutral_token
-        self.controller = SonicDecoder(decoder_sess, default_angles, action_scale)
-
-    @property
-    def pipeline(self):
-        return self.controller
-
-    def reset(self):
-        self.controller.reset()
-
-    def shutdown(self):
-        pass
-
-
-class SonicWholeBodyController:
-    """Full-body SONIC controller for UnitreeG1's background controller thread."""
-
-    control_dt = CONTROL_DT
-    full_body = True
-
-    def __init__(self):
-        logger.info("Loading SONIC whole-body controller...")
-        self._runtime = SonicRuntime()
-        self.kp = self._runtime.kp
-        self.kd = self._runtime.kd
-        self.controller = self._runtime.controller
-        self._default_angles = self._runtime.default_angles
-        self._neutral_token = self._runtime.neutral_token
-
-        # Startup blend: ease from the robot's initial pose into the first commanded policy
-        # targets over INIT_RAMP_S (captured on the first control tick).
-        self._init_ramp_steps = max(1, round(INIT_RAMP_S / CONTROL_DT))
-        self._init_step = 0
-        self._start_pose: dict[str, float] = {}
-
-        # Token-interface state. ``token_mode`` is set True by the robot whenever a SONIC
-        # whole-body controller is selected (token-driven deploy): the controller then holds a
-        # stable *neutral* token until the first real token arrives, and afterwards holds the
-        # *last* token received between ticks (the async controller runs ~50 Hz while a token
-        # VLA streams ~30 Hz). This lives here (not in the entry-point script) so it applies
-        # uniformly to run_g1_server, lerobot-rollout and the sim replays.
-        self.token_mode = False
-        self._last_token: np.ndarray | None = None
-
-        logger.info("SONIC ready (decoder, 64-D token command path)")
-
-    def _startup_blend(self, obs: dict, out: dict) -> dict:
-        """Ease into policy control at startup: for the first ``INIT_RAMP_S`` seconds,
-        interpolate between the robot's pose captured on the first tick and the policy's
-        live commanded target, so the handoff has no snap.
-
-        ``out`` is the policy's ``<joint>.q`` target dict for this tick; the blend ratio
-        climbs 0->1 over the ramp, after which the raw policy target passes through.
-        """
-        if self._init_step >= self._init_ramp_steps or not out:
-            return out
-        if self._init_step == 0:
-            # Capture the robot's actual pose as the interpolation start point.
-            self._start_pose = {
-                f"{m.name}.q": float(obs.get(f"{m.name}.q", self._default_angles[m.value]))
-                for m in G1_29_JointIndex
-            }
-        self._init_step += 1
-        ratio = min(1.0, self._init_step / self._init_ramp_steps)
-        blended = {
-            k: self._start_pose.get(k, float(tgt)) * (1.0 - ratio) + float(tgt) * ratio
-            for k, tgt in out.items()
-        }
-        if self._init_step >= self._init_ramp_steps:
-            logger.info("SONIC startup blend complete -> full policy control")
-        return blended
-
-    def run_step(self, action: dict, lowstate) -> dict:
-        if lowstate is None:
-            return {}
-        obs = lowstate_to_obs(lowstate)
-
-        # Token-only interface (token-output VLA): a dense 64-D ``motion_token.{i}`` command
-        # is decoded directly, encoder bypassed.
-        token = _extract_token_from_action(action)
-        if token is not None:
-            self._last_token = token
-        elif self._last_token is None and self.token_mode:
-            # Token-driven deploy, but no token has arrived yet: hold the checkpoint's neutral
-            # token, which the decoder maps to a stable, natural standing pose.
-            self._last_token = self._neutral_token.copy()
-        if self._last_token is None:
-            # No token yet and not in token_mode: hold (keep last target).
-            return {}
-        # Either a fresh token this tick or the last one received (held between the ~30 Hz
-        # token stream and the ~50 Hz control loop).
-        return self._startup_blend(obs, self.controller.step(obs, self._last_token))
-
-    def reset(self):
-        self._runtime.reset()
-        self._init_step = 0  # re-run the startup blend after a reset
-        self._start_pose = {}
-        # Drop the held token so token_mode re-seeds the neutral token after a reset.
-        self._last_token = None
-
-    def shutdown(self):
-        self._runtime.shutdown()

src/lerobot/robots/unitree_g1/g1_utils.py

@@ -23,47 +23,6 @@ import numpy as np
 
 NUM_MOTORS = 29
 
-# Joint-order permutations between the two 29-DoF layouts used across the G1 stack:
-# IsaacLab (policy/training order) and MuJoCo (deploy order). ``a[ISAACLAB_TO_MUJOCO]``
-# reorders an IsaacLab-ordered vector into MuJoCo order, and vice-versa.
-ISAACLAB_TO_MUJOCO = np.array(
-    [
-        0,
-        3,
-        6,
-        9,
-        13,
-        17,
-        1,
-        4,
-        7,
-        10,
-        14,
-        18,
-        2,
-        5,
-        8,
-        11,
-        15,
-        19,
-        21,
-        23,
-        25,
-        27,
-        12,
-        16,
-        20,
-        22,
-        24,
-        26,
-        28,
-    ],
-    dtype=np.int32,
-)
-# The two orderings are inverses of each other, so derive one from the other (argsort) to
-# guarantee they can never drift out of sync.
-MUJOCO_TO_ISAACLAB = np.argsort(ISAACLAB_TO_MUJOCO).astype(np.int32)
-
 REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
 REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
 REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
@@ -109,9 +68,8 @@ def make_locomotion_controller(name: str | None):
     if name is None:
         return None
     controllers = {
-        "GrootLocomotionController": "lerobot.robots.unitree_g1.controllers.gr00t_locomotion",
-        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion",
-        "SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
+        "GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
+        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
     }
     module_path = controllers.get(name)
     if module_path is None:

src/lerobot/robots/unitree_g1/gr00t_locomotion.py

@@ -21,7 +21,7 @@ import numpy as np
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from ..g1_utils import (
+from .g1_utils import (
     REMOTE_AXES,
     REMOTE_BUTTONS,
     G1_29_JointIndex,

src/lerobot/robots/unitree_g1/holosoma_locomotion.py

@@ -22,7 +22,7 @@ import onnx
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from ..g1_utils import (
+from .g1_utils import (
     REMOTE_AXES,
     G1_29_JointArmIndex,
     G1_29_JointIndex,

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -34,6 +34,7 @@ from .config_unitree_g1 import UnitreeG1Config
 from .g1_kinematics import G1_29_ArmIK
 from .g1_utils import (
     REMOTE_AXES,
+    REMOTE_KEYS,
     G1_29_JointArmIndex,
     G1_29_JointIndex,
     default_remote_input,
@@ -105,47 +106,6 @@ class G1_29_LowState:  # noqa: N801
     mode_machine: int = 0  # Robot mode
 
 
-def lowstate_to_obs(lowstate) -> dict:
-    """Build a robot observation dict from a Unitree lowstate.
-
-    Shared by ``UnitreeG1.get_observation`` and the SONIC pipeline so the
-    lowstate -> obs mapping lives in exactly one place. Keys match the
-    ``<joint>.q``/``imu.*`` schema consumed across the controllers.
-    """
-    obs: dict = {}
-
-    for motor in G1_29_JointIndex:
-        idx = motor.value
-        obs[f"{motor.name}.q"] = lowstate.motor_state[idx].q
-        obs[f"{motor.name}.dq"] = lowstate.motor_state[idx].dq
-        obs[f"{motor.name}.tau"] = lowstate.motor_state[idx].tau_est
-
-    imu = lowstate.imu_state
-    if imu.gyroscope:
-        obs["imu.gyro.x"] = imu.gyroscope[0]
-        obs["imu.gyro.y"] = imu.gyroscope[1]
-        obs["imu.gyro.z"] = imu.gyroscope[2]
-    if imu.accelerometer:
-        obs["imu.accel.x"] = imu.accelerometer[0]
-        obs["imu.accel.y"] = imu.accelerometer[1]
-        obs["imu.accel.z"] = imu.accelerometer[2]
-    if imu.quaternion:
-        obs["imu.quat.w"] = imu.quaternion[0]
-        obs["imu.quat.x"] = imu.quaternion[1]
-        obs["imu.quat.y"] = imu.quaternion[2]
-        obs["imu.quat.z"] = imu.quaternion[3]
-    if imu.rpy:
-        obs["imu.rpy.roll"] = imu.rpy[0]
-        obs["imu.rpy.pitch"] = imu.rpy[1]
-        obs["imu.rpy.yaw"] = imu.rpy[2]
-
-    wr = getattr(lowstate, "wireless_remote", None)
-    if wr:
-        obs["wireless_remote"] = bytes(wr) if not isinstance(wr, (bytes, bytearray)) else wr
-
-    return obs
-
-
 class UnitreeG1(Robot):
     config_class = UnitreeG1Config
     name = "unitree_g1"
@@ -188,60 +148,22 @@ class UnitreeG1(Robot):
 
         self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
 
-        # Lower-body / whole-body controller loaded dynamically
+        # Lower-body controller loaded dynamically
         self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
 
-        # A SONIC whole-body controller always runs in token mode: it holds a neutral
-        # token until the first real one arrives, then holds the last token between ticks.
-        if self.controller is not None and hasattr(self.controller, "token_mode"):
-            self.controller.token_mode = True
-
         # Controller thread state
         self._controller_thread = None
-        # When set, the controller loop stops publishing low commands so reset() can
-        # drive the joints directly without two publishers fighting (single-publisher).
-        self._controller_paused = threading.Event()
         self._controller_action_lock = threading.Lock()
         self.controller_input = default_remote_input()
         self.controller_output = {}
 
-        # Token-mode state: last 64-D SONIC latent token commanded by the policy,
-        # echoed back as ``observation.state`` so a token-output VLA closes the loop
-        # on its own previous token. Implicit whenever the SONIC whole-body controller
-        # is active. Seeded to zeros; the controller's startup blend eases joints in.
-        self._last_token: np.ndarray | None = None
-        if self._sonic_token:
-            from .controllers.sonic_whole_body import TOKEN_DIM
-
-            self._last_token = np.zeros(TOKEN_DIM, dtype=np.float32)
-
-    @property
-    def _sonic_token(self) -> bool:
-        """Whether the SONIC whole-body decoder is active.
-
-        A SONIC controller consumes a 64-D latent motion token as its action and echoes
-        the last commanded token as ``observation.state``. Keyed purely off the selected
-        controller so the token interface is implicit -- no separate config flag.
-        """
-        return self.config.controller == "SonicWholeBodyController"
-
     def _subscribe_lowstate(self):  # polls robot state @ 250Hz
         while not self._shutdown_event.is_set():
             start_time = time.time()
 
             # Step simulation if in simulation mode
             if self.config.is_simulation and self.sim_env is not None:
-                try:
-                    self.sim_env.step()
-                except ValueError as e:
-                    # Startup race: the sim thread can step once before reset() has
-                    # written a valid base pose, giving a zero-norm pelvis quaternion
-                    # (scipy>=1.11 raises instead of normalizing). Skip and retry so
-                    # the thread survives instead of dying and freezing the sim.
-                    if "zero norm" not in str(e).lower():
-                        raise
-                    time.sleep(self.control_dt)
-                    continue
+                self.sim_env.step()
 
             msg = self.lowstate_subscriber.Read()
             if msg is not None:
@@ -309,38 +231,15 @@ class UnitreeG1(Robot):
                 features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
         return features
 
-    @property
-    def _token_state_ft(self) -> dict[str, type]:
-        """64-D SONIC latent-token proprio state (``motion_token_state.{i}.pos``).
-
-        Exposed only when a SONIC whole-body controller is active; aggregated by the
-        rollout into a 64-D ``observation.state`` (the last token the policy commanded).
-        """
-        if not self._sonic_token:
-            return {}
-        from .controllers.sonic_whole_body import TOKEN_DIM, token_state_key
-
-        return {token_state_key(i): float for i in range(TOKEN_DIM)}
-
     @cached_property
     def observation_features(self) -> dict[str, type | tuple]:
-        return {**self._motors_ft, **self._token_state_ft, **self._cameras_ft}
+        return {**self._motors_ft, **self._cameras_ft}
 
     @cached_property
     def action_features(self) -> dict[str, type]:
-        # No controller configured at all: raw 29-DoF joint teleop.
         if self.controller is None:
             return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
 
-        # Token-output VLA (SONIC decoder): advertise a 64-D latent-token action space
-        # (``motion_token.{i}.pos``) so ``lerobot-rollout`` maps a 64-D policy output
-        # straight onto the decoder, bypassing the encoder.
-        if self._sonic_token:
-            from .controllers.sonic_whole_body import TOKEN_DIM, token_action_key
-
-            return {token_action_key(i): float for i in range(TOKEN_DIM)}
-
-        # Locomotion controllers (GR00T / Holosoma): arm joint targets + joystick axes.
         arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
         remote_features = dict.fromkeys(REMOTE_AXES, float)
         return {**arm_features, **remote_features}
@@ -356,11 +255,6 @@ class UnitreeG1(Robot):
         while not self._shutdown_event.is_set():
             start_time = time.time()
 
-            # Paused during reset() so the reset routine is the sole low-cmd publisher.
-            if self._controller_paused.is_set():
-                time.sleep(control_dt)
-                continue
-
             with self._lowstate_lock:
                 lowstate = self._lowstate
 
@@ -449,9 +343,6 @@ class UnitreeG1(Robot):
 
         self.kp = np.array(self.config.kp, dtype=np.float32)
         self.kd = np.array(self.config.kd, dtype=np.float32)
-        if self.controller is not None and hasattr(self.controller, "kp"):
-            self.kp = np.array(self.controller.kp, dtype=np.float32)
-            self.kd = np.array(self.controller.kd, dtype=np.float32)
 
         for joint in G1_29_JointIndex:
             self.msg.motor_cmd[joint].mode = 1
@@ -500,10 +391,6 @@ class UnitreeG1(Robot):
             if self._controller_thread.is_alive():
                 logger.warning("Controller thread did not stop cleanly")
 
-        # Release controller resources (e.g. SONIC decoder sessions).
-        if self.controller is not None and hasattr(self.controller, "shutdown"):
-            self.controller.shutdown()
-
         # Close simulation environment
         if self.config.is_simulation and self.sim_env is not None:
             try:
@@ -574,15 +461,6 @@ class UnitreeG1(Robot):
         if lowstate.wireless_remote:
             obs["wireless_remote"] = lowstate.wireless_remote
 
-        # Token mode: echo the last commanded latent token as observation.state so a
-        # token-output VLA closes the loop on its own previous token.
-        if self._sonic_token:
-            from .controllers.sonic_whole_body import token_state_key
-
-            token = self._last_token if self._last_token is not None else []
-            for i, v in enumerate(token):
-                obs[token_state_key(i)] = float(v)
-
         # Cameras - read images from ZMQ cameras
         for cam_name, cam in self._cameras.items():
             if getattr(cam, "use_rgb", True):
@@ -595,22 +473,9 @@ class UnitreeG1(Robot):
     def send_action(self, action: RobotAction) -> RobotAction:
         action_to_publish = action
         if self.controller is not None:
-            # SONIC decoder: pull the 64-D latent token out of the action and remember it
-            # for the observation.state echo. The controller thread reads it back from
-            # controller_input (populated below) and decodes it into a 29-DoF command.
-            if self._sonic_token:
-                from .controllers.sonic_whole_body import _extract_token_from_action
-
-                token = _extract_token_from_action(action)
-                if token is not None:
-                    self._last_token = token
-            self._update_controller_action(action)
-            # Full-body controllers (SONIC) own the whole 29-DoF command; nothing to
-            # publish here (the controller thread is the sole publisher).
-            if getattr(self.controller, "full_body", False):
-                return action
             # Controller thread owns legs/waist. Here we only update joystick inputs
             # and publish arm targets from the teleoperator.
+            self._update_controller_action(action)
             arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
             action_to_publish = {
                 key: value
@@ -638,17 +503,11 @@ class UnitreeG1(Robot):
         return action
 
     def _update_controller_action(self, action: RobotAction) -> None:
-        """Update controller input state from an incoming teleop action.
-
-        Controller-agnostic: every value-carrying key (locomotion ``remote.*`` axes or
-        SONIC ``motion_token.*`` values) is forwarded verbatim into ``controller_input``
-        and each controller extracts only the keys it understands. The robot deliberately
-        does not enumerate any controller's key schema here.
-        """
+        """Update controller input state from incoming teleop action."""
         with self._controller_action_lock:
-            for key, value in action.items():
-                if isinstance(key, str) and value is not None:
-                    self.controller_input[key] = value
+            for key in REMOTE_KEYS:
+                if key in action:
+                    self.controller_input[key] = action[key]
 
     @property
     def is_calibrated(self) -> bool:
@@ -678,64 +537,43 @@ class UnitreeG1(Robot):
         if default_positions is None:
             default_positions = np.array(self.config.default_positions, dtype=np.float32)
 
-        # Full-body controllers (SONIC) own the whole 29-DoF command and ignore
-        # ``<joint>.q`` in send_action(), so reset() must publish the default pose
-        # directly. Pause the background controller first so the two aren't both writing
-        # low commands while the robot moves to the default pose.
-        full_body = getattr(self.controller, "full_body", False)
-        paused = False
-        if full_body and self._controller_thread is not None:
-            self._controller_paused.set()
-            paused = True
-            time.sleep(control_dt)  # let any in-flight controller tick settle
+        if self.config.is_simulation and self.sim_env is not None:
+            self.sim_env.reset()
+            self.publish_lowcmd(
+                {f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
+            )
+        else:
+            total_time = 3.0
+            num_steps = int(total_time / control_dt)
 
-        try:
-            if self.config.is_simulation and self.sim_env is not None:
-                self.sim_env.reset()
-                self.publish_lowcmd(
-                    {f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
-                )
-            else:
-                total_time = 3.0
-                num_steps = int(total_time / control_dt)
+            # get current state
+            obs = self.get_observation()
 
-                # get current state
-                obs = self.get_observation()
+            # record current positions
+            init_dof_pos = np.zeros(29, dtype=np.float32)
+            for motor in G1_29_JointIndex:
+                init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
 
-                # record current positions
-                init_dof_pos = np.zeros(29, dtype=np.float32)
+            # Interpolate to default position
+            for step in range(num_steps):
+                start_time = time.time()
+
+                alpha = step / num_steps
+                action_dict = {}
                 for motor in G1_29_JointIndex:
-                    init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
+                    target_pos = default_positions[motor.value]
+                    interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
+                    action_dict[f"{motor.name}.q"] = float(interp_pos)
 
-                # Interpolate to default position
-                for step in range(num_steps):
-                    start_time = time.time()
+                self.send_action(action_dict)
 
-                    alpha = step / num_steps
-                    action_dict = {}
-                    for motor in G1_29_JointIndex:
-                        target_pos = default_positions[motor.value]
-                        interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
-                        action_dict[f"{motor.name}.q"] = float(interp_pos)
+                # Maintain constant control rate
+                elapsed = time.time() - start_time
+                sleep_time = max(0, control_dt - elapsed)
+                time.sleep(sleep_time)
 
-                    # Full-body controllers no-op in send_action(); publish the pose
-                    # directly (arm-only controllers keep the send_action() path).
-                    if full_body:
-                        self.publish_lowcmd(action_dict)
-                    else:
-                        self.send_action(action_dict)
-
-                    # Maintain constant control rate
-                    elapsed = time.time() - start_time
-                    sleep_time = max(0, control_dt - elapsed)
-                    time.sleep(sleep_time)
-
-            # Reset controller internal state (gait phase, obs history, etc.) before
-            # resuming so its buffers reflect the post-reset pose.
-            if self.controller is not None and hasattr(self.controller, "reset"):
-                self.controller.reset()
-        finally:
-            if paused:
-                self._controller_paused.clear()
+        # Reset controller internal state (gait phase, obs history, etc.)
+        if self.controller is not None and hasattr(self.controller, "reset"):
+            self.controller.reset()
 
         logger.info("Reached default position")

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_edit_dataset.py

@@ -127,6 +127,12 @@ 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 \
@@ -303,6 +309,7 @@ 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")
@@ -551,9 +558,12 @@ 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:
-        raise ValueError("Must specify at least one of new_task or episode_tasks for modify_tasks operation")
+    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 cfg.new_repo_id is not None or cfg.new_root is not None:
         logging.warning(
@@ -573,11 +583,14 @@ 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,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/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_dataset_tools.py

@@ -1125,9 +1125,61 @@ 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 or episode_tasks"):
+    with pytest.raises(
+        ValueError, match="Must specify at least one of new_task, episode_tasks, or task_replacements"
+    ):
         modify_tasks(sample_dataset)
 
 
@@ -1137,6 +1189,12 @@ 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/scripts/test_edit_dataset_parsing.py

@@ -105,6 +105,23 @@ 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_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()

upload_sonic_decoder.py

@@ -1,230 +0,0 @@
-#!/usr/bin/env python3
-"""Provision the SONIC decoder checkpoint at ``lerobot/sonic_decoder``.
-
-Takes NVIDIA's ``nvidia/GEAR-SONIC/model_decoder.onnx``, embeds the SONIC deploy constants
-(``kp``/``kd`` PD gains, ``default_angles`` standing pose, the residual ``action_scale``, and
-the ``neutral_token`` idle latent) into the ONNX ``metadata_props`` (the convention Holosoma
-uses for its gains), and pushes the result to ``lerobot/sonic_decoder``. After this runs, the
-runtime loads the decoder *and* every one of these constants straight from the checkpoint --
-no motor-physics math at deploy time, so ``sonic_whole_body.py`` carries none of the
-armature/bandwidth machinery nor any hardcoded deploy constants.
-
-The constants here are derived once from Unitree motor physics (armature + target bandwidth).
-That derivation is intentionally kept in this one-off provisioning script (not the runtime);
-the shared/harmonic helper is a separate PR.
-
-Build only (no network/auth needed if the source ONNX is already cached):
-    python upload_sonic_decoder.py --out ./sonic_decoder
-
-Build + upload:
-    huggingface-cli login          # or export HF_TOKEN=...
-    python upload_sonic_decoder.py --upload
-"""
-
-from __future__ import annotations
-
-import argparse
-import json
-import pathlib
-
-import numpy as np
-import onnx
-from huggingface_hub import hf_hub_download
-
-SRC_REPO_ID = "nvidia/GEAR-SONIC"
-SRC_FILENAME = "model_decoder.onnx"
-DST_REPO_ID = "lerobot/sonic_decoder"
-
-# ── SONIC deploy-constant derivation (provisioning-time only) ─────────────────
-# All constants are (29,) in IsaacLab joint order: legs, waist, arms.
-# kp = armature * w**2, kd = 4 * armature * w, with a x2 factor on the stiff joints
-# (ankles + waist). action_scale = 0.25 * effort / (armature * w**2) is the residual
-# scaling that maps decoder output to a joint-angle delta on top of default_angles.
-NATURAL_FREQ = 10.0 * 2.0 * np.pi
-MOTOR_ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425}
-EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0}
-MOTOR_MODELS = (
-    ["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2
-    + ["7520_14", "5020", "5020"]
-    + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2
-)
-DOUBLE_INDICES = {4, 5, 10, 11, 13, 14}  # ankles + waist
-
-# Nominal standing pose (rad), 29 joints in IsaacLab order. Decoder actions are residuals
-# added on top of this.
-DEFAULT_ANGLES = [
-    -0.312,
-    0.0,
-    0.0,
-    0.669,
-    -0.363,
-    0.0,  # left leg
-    -0.312,
-    0.0,
-    0.0,
-    0.669,
-    -0.363,
-    0.0,  # right leg
-    0.0,
-    0.0,
-    0.0,  # waist
-    0.2,
-    0.2,
-    0.0,
-    0.6,
-    0.0,
-    0.0,
-    0.0,  # left arm
-    0.2,
-    -0.2,
-    0.0,
-    0.6,
-    0.0,
-    0.0,
-    0.0,  # right arm
-]
-
-# Neutral idle token (64-D), held until the first real token arrives. Captured from the
-# encoder while the robot stood idle in sim: the encoder is an FSQ bottleneck (~5 bit/dim,
-# Div(16)), so tokens live on the 1/16 grid. We store the integer FSQ codes and rescale by
-# 1/16 -> an exact on-grid token that decodes to a stable, natural standing pose (unlike the
-# literal all-zero token, which is off-manifold and decodes to a slightly goofy stance).
-NEUTRAL_TOKEN_CODES = [
-    -1,
-    3,
-    1,
-    -1,
-    1,
-    -3,
-    6,
-    1,
-    1,
-    1,
-    -2,
-    -4,
-    -2,
-    0,
-    -3,
-    -1,
-    2,
-    -1,
-    -3,
-    -5,
-    3,
-    1,
-    1,
-    -4,
-    -1,
-    -1,
-    1,
-    -7,
-    0,
-    1,
-    2,
-    -2,
-    5,
-    -2,
-    -2,
-    -4,
-    0,
-    -1,
-    3,
-    -1,
-    0,
-    -5,
-    -1,
-    0,
-    -4,
-    0,
-    0,
-    -1,
-    -1,
-    2,
-    -2,
-    1,
-    3,
-    3,
-    1,
-    0,
-    0,
-    6,
-    0,
-    -7,
-    3,
-    0,
-    2,
-    -2,
-]
-
-
-def compute_kp_kd() -> tuple[list[float], list[float]]:
-    """Return (kp, kd) as plain float lists, (29,) in IsaacLab joint order."""
-
-    def stiffness(k):
-        return MOTOR_ARMATURE[k] * NATURAL_FREQ**2
-
-    def damping(k):
-        return 4.0 * MOTOR_ARMATURE[k] * NATURAL_FREQ
-
-    kp = [(2 if i in DOUBLE_INDICES else 1) * stiffness(k) for i, k in enumerate(MOTOR_MODELS)]
-    kd = [(2 if i in DOUBLE_INDICES else 1) * damping(k) for i, k in enumerate(MOTOR_MODELS)]
-    return kp, kd
-
-
-def compute_action_scale() -> list[float]:
-    """Return the per-joint residual action scale, (29,) in IsaacLab joint order."""
-    return [0.25 * EFFORT[k] / (MOTOR_ARMATURE[k] * NATURAL_FREQ**2) for k in MOTOR_MODELS]
-
-
-def build(out_dir: pathlib.Path) -> pathlib.Path:
-    """Download the source decoder, embed the deploy-constant metadata, save to ``out_dir``."""
-    src = hf_hub_download(repo_id=SRC_REPO_ID, filename=SRC_FILENAME)
-    model = onnx.load(src)
-
-    kp, kd = compute_kp_kd()
-    neutral_token = [c / 16.0 for c in NEUTRAL_TOKEN_CODES]  # FSQ Div(16): codes -> on-grid token
-    meta = {prop.key: prop.value for prop in model.metadata_props}
-    meta["kp"] = json.dumps(kp)
-    meta["kd"] = json.dumps(kd)
-    meta["action_scale"] = json.dumps(compute_action_scale())
-    meta["default_angles"] = json.dumps(DEFAULT_ANGLES)
-    meta["neutral_token"] = json.dumps(neutral_token)
-    # Rewrite metadata_props with the merged dict.
-    del model.metadata_props[:]
-    for key, value in meta.items():
-        model.metadata_props.add(key=key, value=value)
-
-    out_dir.mkdir(parents=True, exist_ok=True)
-    out_path = out_dir / SRC_FILENAME
-    onnx.save(model, out_path)
-    print(f"Wrote {out_path} with kp/kd/action_scale/default_angles/neutral_token metadata.")
-    return out_path
-
-
-def upload(out_path: pathlib.Path) -> None:
-    from huggingface_hub import HfApi
-
-    api = HfApi()
-    api.create_repo(repo_id=DST_REPO_ID, repo_type="model", exist_ok=True)
-    api.upload_file(
-        path_or_fileobj=str(out_path),
-        path_in_repo=SRC_FILENAME,
-        repo_id=DST_REPO_ID,
-        repo_type="model",
-    )
-    print(f"Uploaded {out_path.name} -> {DST_REPO_ID}")
-
-
-def main() -> None:
-    p = argparse.ArgumentParser()
-    p.add_argument("--out", type=pathlib.Path, default=pathlib.Path("./sonic_decoder"))
-    p.add_argument("--upload", action="store_true", help="Push the built ONNX to the hub")
-    args = p.parse_args()
-
-    out_path = build(args.out)
-    if args.upload:
-        upload(out_path)
-
-
-if __name__ == "__main__":
-    main()

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"