Merge branch 'main' into feat/fix_soft_prompt_typo

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

docs/source/hil_data_collection.mdx

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

docs/source/so101.mdx

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

docs/source/unitree_g1.mdx

@@ -8,15 +8,6 @@
 
 The Unitree G1 humanoid is now supported in LeRobot! You can teleoperate, train locomanipulation policies, test in sim, and more. Both 29 and 23 DoF variants are supported.
 
-<Tip>
-  **New: SONIC whole-body control.** The `SonicWholeBodyController` runs NVIDIA's
-  [GEAR-SONIC](https://huggingface.co/nvidia/GEAR-SONIC) decoder on the G1, turning a
-  64-D latent motion token into full-body joint targets at 50 Hz. This lets you drive
-  the robot from a VLA policy trained on SONIC motion tokens (token in → whole-body
-  motion out) with `lerobot-rollout`, in sim or on the physical robot. See
-  [Whole-body control with SONIC](#whole-body-control-with-sonic) below.
-</Tip>
-
 ---
 
 ## Part 1: Getting Started
@@ -68,7 +59,7 @@ lerobot-teleoperate \
   --robot.controller=GrootLocomotionController
 ```
 
-This will launch a [MuJoCo sim instance](https://huggingface.co/lerobot/unitree-g1-mujoco/tree/main) for the G1. You can connect a gamepad to your machine before launching in order to control the robot's locomotion in sim. We support [HolosomaLocomotionController](https://github.com/amazon-far/holosoma), [GrootLocomotionController](https://github.com/NVlabs/GR00T-WholeBodyControl), and [SonicWholeBodyController](https://huggingface.co/nvidia/GEAR-SONIC) via `--robot.controller`.
+This will launch a [MuJoCo sim instance](https://huggingface.co/lerobot/unitree-g1-mujoco/tree/main) for the G1. You can connect a gamepad to your machine before launching in order to control the robot's locomotion in sim. We support both [HolosomaLocomotionController](https://github.com/amazon-far/holosoma) and [GrootLocomotionController](https://github.com/NVlabs/GR00T-WholeBodyControl) via `--robot.controller`.
 
 - Press `9` to release the robot
 - Press `7` / `8` to increase / decrease waist height
@@ -299,52 +290,6 @@ lerobot-rollout \
 
 ---
 
-## Whole-body control with SONIC
-
-The `SonicWholeBodyController` runs NVIDIA's [GEAR-SONIC](https://huggingface.co/nvidia/GEAR-SONIC)
-decoder on the G1. Each 50 Hz tick it consumes a **64-D latent motion token** and emits
-full-body joint targets — the encoder is bypassed, so a policy feeds tokens in and the
-decoder turns them into motion. Before the first token arrives the controller holds a
-neutral (idle) pose.
-
-This makes the G1 drivable by a VLA policy trained to output SONIC motion tokens (token
-as both `observation.state` and `action`, e.g. [`nepyope/sonic_walk`](https://huggingface.co/nepyope/sonic_walk))
-using the standard `lerobot-rollout`. The controller always runs **onboard** the robot;
-the laptop is a thin client that streams tokens and receives camera frames over ZMQ.
-
-**On the robot** — start the server in handshake mode so it instantiates and runs the
-controller onboard against local DDS at full rate:
-
-```bash
-cd ~/lerobot
-python src/lerobot/robots/unitree_g1/run_g1_server.py --handshake --camera
-```
-
-**From your laptop** — run the token policy; selecting `--robot.controller=SonicWholeBodyController`
-implicitly switches the robot to the 64-D latent-token action/observation interface:
-
-```bash
-lerobot-rollout \
-  --policy.path=nepyope/sonic_walk \
-  --policy.device=cuda \
-  --robot.type=unitree_g1 \
-  --robot.is_simulation=false \
-  --robot.robot_ip=<ROBOT_IP> \
-  --robot.controller=SonicWholeBodyController \
-  --robot.cameras='{"ego_view": {"type": "zmq", "server_address": "<ROBOT_IP>", "port": 5555, "camera_name": "head_camera", "width": 640, "height": 480, "fps": 30}}' \
-  --task="walk back and forth" \
-  --duration=1000 \
-  --fps=30
-```
-
-<Tip>
-  SONIC is a token-only decoder in LeRobot: the only input path is the 64-D latent
-  vector. To train your own token policy, expose the 64-D token as the action (a config
-  choice, e.g. `pi05` with a 64-D action dim) — no policy code changes are needed.
-</Tip>
-
----
-
 ## Additional Resources
 
 - [Unitree SDK Documentation](https://github.com/unitreerobotics/unitree_sdk2_python)

pyproject.toml

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

src/lerobot/configs/policies.py

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

src/lerobot/configs/rewards.py

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

src/lerobot/configs/train.py

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

src/lerobot/policies/evo1/modeling_evo1.py

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

src/lerobot/policies/groot/modeling_groot.py

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

src/lerobot/policies/molmoact2/modeling_molmoact2.py

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

src/lerobot/policies/pi0/modeling_pi0.py

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

src/lerobot/policies/pi05/modeling_pi05.py

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

src/lerobot/policies/pretrained.py

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

src/lerobot/policies/smolvla/modeling_smolvla.py

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

src/lerobot/policies/smolvla/smolvlm_with_expert.py

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

src/lerobot/policies/xvla/modeling_xvla.py

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

src/lerobot/rl/algorithms/configs.py

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

src/lerobot/robots/unitree_g1/config_unitree_g1.py

@@ -62,34 +62,12 @@ class UnitreeG1Config(RobotConfig):
     # Socket config for ZMQ bridge
     robot_ip: str = "192.168.123.164"  # default G1 IP
 
-    # Run the locomotion / whole-body controller ONBOARD the robot (policy on the G1
-    # itself, against local DDS at full rate) instead of on the laptop over the ZMQ
-    # socket bridge. In this mode the robot object uses the real Unitree SDK channels
-    # and expects high-level actions (arm targets + joystick axes, or 64-D SONIC
-    # tokens) fed via send_action -- e.g. by run_g1_server's serve_onboard_controller,
-    # which receives them from the laptop over ZMQ. Mutually exclusive with is_simulation.
-    onboard: bool = False
-    # DDS network interface for onboard mode (None = SDK default, matching
-    # run_g1_server.py's ChannelFactoryInitialize(0)).
-    dds_interface: str | None = None
-    # Onboard sub-flags. On a real G1 both are True: the built-in motion services
-    # must be released before we can write lowcmd, and locomotion axes are read from
-    # the physical wireless remote. Against a DDS sim neither applies (no
-    # MotionSwitcher, no physical remote), so set both False so the controller takes
-    # its locomotion axes purely from send_action (ZMQ) input.
-    release_motion_control: bool = True
-    physical_remote: bool = True
-
     # Cameras (ZMQ-based remote cameras)
     cameras: dict[str, CameraConfig] = field(default_factory=dict)
 
     # 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,401 +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 ``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 logging
-
-import numpy as np
-import onnxruntime as ort
-from huggingface_hub import hf_hub_download
-
-from ..g1_utils import (
-    ISAACLAB_TO_MUJOCO,
-    MOTOR_ARMATURE,
-    MUJOCO_TO_ISAACLAB,
-    NATURAL_FREQ,
-    G1_29_JointIndex,
-    compute_pd_gains,
-    get_gravity_orientation,
-    lowstate_to_obs,
-    make_ort_session_options,
-)
-
-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
-
-# Nominal standing pose (rad), 29 joints in IsaacLab order. Decoder actions are residuals
-# added on top of this.
-DEFAULT_ANGLES = np.array(
-    [
-        -0.312, 0.0, 0.0, 0.669, -0.363, 0.0,
-        -0.312, 0.0, 0.0, 0.669, -0.363, 0.0,
-        0.0, 0.0, 0.0,
-        0.2, 0.2, 0.0, 0.6, 0.0, 0.0, 0.0,
-        0.2, -0.2, 0.0, 0.6, 0.0, 0.0, 0.0,
-    ],
-    dtype=np.float32,
-)
-
-# Per-motor torque limits (N·m), used only for SONIC's residual-action scaling. The
-# armature / bandwidth constants and the PD-gain formula are shared (see g1_utils).
-EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0}
-
-
-def _action_scale(k):
-    """Per-motor residual-action scale (maps policy output to joint-angle delta)."""
-    return 0.25 * EFFORT[k] / (MOTOR_ARMATURE[k] * NATURAL_FREQ**2)
-
-
-# Per-joint motor model (IsaacLab order): legs, waist, then arms. Single source of truth
-# for both ACTION_SCALE and compute_kp_kd().
-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
-)
-ACTION_SCALE = np.array([_action_scale(k) for k in MOTOR_MODELS], dtype=np.float32)  # (29,) IsaacLab
-
-
-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]
-
-
-DEFAULT_ANGLES_MUJOCO = _to_mujoco(DEFAULT_ANGLES)
-
-
-# Ankle + waist joint indices (IsaacLab order) that get a x2 stiffness/damping factor.
-_SONIC_DOUBLE = {4, 5, 10, 11, 13, 14}
-
-
-def compute_kp_kd():
-    """SONIC per-joint PD gains (kp, kd), (29,) float32 in IsaacLab joint order."""
-    return compute_pd_gains(MOTOR_MODELS, _SONIC_DOUBLE)
-
-
-# Action-feature prefix for the latent-token interface (see _extract_token_from_action).
-TOKEN_ACTION_PREFIX = "motion_token"
-# 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"
-
-
-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
-
-# Neutral ("zero pose") SONIC token, held by token_mode until the first real token arrives.
-# Captured from the encoder's own output while the robot stood idle in sim: the encoder is
-# an FSQ bottleneck (~5 bit/dim, Div(16)), so its tokens live on the 1/16 grid. We store the
-# integer FSQ codes and rescale by 1/16, giving an exact on-grid token -- unlike the literal
-# all-zero token, which is off the learned manifold and decodes to a slightly goofy stance.
-# This one decodes to a stable, natural standing pose.
-_NEUTRAL_TOKEN_CODES = np.array(
-    [-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],
-    dtype=np.float32,
-)
-NEUTRAL_TOKEN = _NEUTRAL_TOKEN_CODES / 16.0  # FSQ Div(16): integer codes -> on-grid token
-
-
-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).
-    """
-
-    def __init__(self, decoder):
-        self.decoder = decoder
-        self.decoder_input = decoder.get_inputs()[0].name
-        dec_dim = int(decoder.get_inputs()[0].shape[1])
-        if dec_dim != 994:
-            raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected 994)")
-        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 - DEFAULT_ANGLES_MUJOCO] + 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", 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 = DEFAULT_ANGLES + action_mj[ISAACLAB_TO_MUJOCO] * 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_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx")
-
-        so = make_ort_session_options()
-        decoder_sess = ort.InferenceSession(decoder_path, sess_options=so)
-
-        self.kp, self.kd = compute_kp_kd()
-        self.controller = SonicDecoder(decoder_sess)
-
-    @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
-
-        # 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", 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 captured neutral
-            # token (NEUTRAL_TOKEN), which the decoder maps to a stable, natural standing pose.
-            self._last_token = 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,82 +23,11 @@ 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,
-)
-MUJOCO_TO_ISAACLAB = np.array(
-    [
-        0,
-        6,
-        12,
-        1,
-        7,
-        13,
-        2,
-        8,
-        14,
-        3,
-        9,
-        15,
-        22,
-        4,
-        10,
-        16,
-        23,
-        5,
-        11,
-        17,
-        24,
-        18,
-        25,
-        19,
-        26,
-        20,
-        27,
-        21,
-        28,
-    ],
-    dtype=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
 
+
 def default_remote_input() -> dict[str, float]:
     """Return a zeroed-out remote input dict (axes + buttons)."""
     return dict.fromkeys(REMOTE_KEYS, 0.0)
@@ -114,53 +43,6 @@ def get_gravity_orientation(quaternion: list[float] | np.ndarray) -> np.ndarray:
     return gravity_orientation
 
 
-# Unitree motor-model parameters shared by the controllers that derive their PD gains
-# from motor physics rather than hand-tuning (SONIC decoder, Holosoma). NATURAL_FREQ is
-# the target closed-loop stiffness bandwidth (rad/s); MOTOR_ARMATURE is per-model rotor
-# inertia (keys are Unitree motor model names). From these: kp = armature * w**2 and
-# kd = 4 * armature * w, with an optional x2 factor on stiff joints (ankles/waist).
-NATURAL_FREQ = 10.0 * 2.0 * np.pi
-MOTOR_ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425}
-
-
-def compute_pd_gains(motor_models, double_indices=()) -> tuple[np.ndarray, np.ndarray]:
-    """Derive per-joint PD gains (kp, kd) from motor armature and target bandwidth.
-
-    ``motor_models`` is a per-joint sequence of Unitree motor model names (in the
-    controller's own joint order); joints whose index is in ``double_indices`` get a
-    x2 stiffness/damping factor. Returns two (N,) float32 arrays in that same order.
-    """
-    double = set(double_indices)
-
-    def s(k):
-        return MOTOR_ARMATURE[k] * NATURAL_FREQ**2
-
-    def d(k):
-        return 4.0 * MOTOR_ARMATURE[k] * NATURAL_FREQ
-
-    kp = np.array([2 * s(k) if i in double else s(k) for i, k in enumerate(motor_models)], dtype=np.float32)
-    kd = np.array([2 * d(k) if i in double else d(k) for i, k in enumerate(motor_models)], dtype=np.float32)
-    return kp, kd
-
-
-def make_ort_session_options(intra_op_num_threads: int | None = None, inter_op_num_threads: int | None = None):
-    """Build quiet ONNX Runtime SessionOptions, optionally capping the CPU thread pool.
-
-    These tiny MLP policies are latency-bound, not throughput-bound, so letting ORT grab
-    every core starves the real-time control loop / torch policy and causes stutter. Pass
-    1 intra + 1 inter thread for lowest-latency per-step inference.
-    """
-    import onnxruntime as ort
-
-    so = ort.SessionOptions()
-    so.log_severity_level = 3
-    if intra_op_num_threads is not None:
-        so.intra_op_num_threads = intra_op_num_threads
-    if inter_op_num_threads is not None:
-        so.inter_op_num_threads = inter_op_num_threads
-    return so
-
-
 class G1_29_JointArmIndex(IntEnum):
     # Left arm
     kLeftShoulderPitch = 15
@@ -181,55 +63,13 @@ class G1_29_JointArmIndex(IntEnum):
     kRightWristYaw = 28
 
 
-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
-
-
 def make_locomotion_controller(name: str | None):
     """Instantiate a locomotion controller by class name. Returns None if name is 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

@@ -14,8 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import logging
 from collections import deque
 
@@ -23,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,
@@ -70,15 +68,9 @@ def load_groot_policies(
         filename="GR00T-WholeBodyControl-Walk.onnx",
     )
 
-    # Load ONNX policies with a capped thread pool. GR00T runs at 50 Hz in a
-    # background thread alongside the (torch) upper-body policy, IK and sim; letting
-    # ORT grab every core starves those and makes the whole rollout stutter. These
-    # are small MLPs, so 1 thread is both enough and lowest-latency.
-    from ..g1_utils import make_ort_session_options
-
-    so = make_ort_session_options(intra_op_num_threads=1, inter_op_num_threads=1)
-    policy_balance = ort.InferenceSession(balance_path, sess_options=so)
-    policy_walk = ort.InferenceSession(walk_path, sess_options=so)
+    # Load ONNX policies
+    policy_balance = ort.InferenceSession(balance_path)
+    policy_walk = ort.InferenceSession(walk_path)
 
     logger.info("GR00T policies loaded successfully")
 
@@ -204,16 +196,6 @@ class GrootLocomotionController:
         # Transform action back to target joint positions
         target_dof_pos_15 = GROOT_DEFAULT_ANGLES[:15] + self.groot_action * ACTION_SCALE
 
-        # Waist override: an external upper-body IK can command the 3 waist joints
-        # (indices 12/13/14) via ``kWaist{Yaw,Roll,Pitch}.q`` in the action dict. When
-        # present, we substitute the balance policy's waist target so the torso tracks
-        # the IK while the policy keeps only the legs balanced. Single-publisher stays
-        # intact (this thread still owns joints 0-14).
-        for idx in (G1_29_JointIndex.kWaistYaw, G1_29_JointIndex.kWaistRoll, G1_29_JointIndex.kWaistPitch):
-            key = f"{idx.name}.q"
-            if key in action and action[key] is not None:
-                target_dof_pos_15[idx.value] = float(action[key])
-
         # Build action dict
         action_dict = {}
         for i in range(15):

src/lerobot/robots/unitree_g1/holosoma_locomotion.py

@@ -14,19 +14,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-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 (
+from .g1_utils import (
     REMOTE_AXES,
     G1_29_JointArmIndex,
     G1_29_JointIndex,
-    compute_pd_gains,
     get_gravity_orientation,
 )
 
@@ -58,23 +57,12 @@ POLICY_FILES = {
     "ppo": "ppo_g1_29dof.onnx",
 }
 
-# Per-joint motor model in Holosoma's joint order, plus the joints that get a x2
-# stiffness/damping factor. These reproduce the kp/kd that used to be read from the
-# policy's ONNX metadata exactly (both fastsac and ppo), so gains are now derived from
-# the shared motor model (see g1_utils.compute_pd_gains) instead.
-HOLOSOMA_MOTOR_MODELS = (
-    ["7520_14", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2
-    + ["7520_14", "5020", "5020"]
-    + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2
-)
-HOLOSOMA_DOUBLE = {4, 5, 10, 11, 13, 14}
-
 
 def load_policy(
     repo_id: str = DEFAULT_HOLOSOMA_REPO_ID,
     policy_type: str = "fastsac",
 ) -> tuple[ort.InferenceSession, np.ndarray, np.ndarray]:
-    """Load the Holosoma locomotion policy and its motor-model-derived PD gains.
+    """Load Holosoma locomotion policy and extract KP/KD from metadata.
 
     Args:
         repo_id: Hugging Face Hub repo ID
@@ -93,7 +81,16 @@ def load_policy(
     policy = ort.InferenceSession(policy_path)
     logger.info(f"Policy loaded: {policy.get_inputs()[0].shape} → {policy.get_outputs()[0].shape}")
 
-    kp, kd = compute_pd_gains(HOLOSOMA_MOTOR_MODELS, HOLOSOMA_DOUBLE)
+    # Extract KP/KD from ONNX metadata
+    model = onnx.load(policy_path, load_external_data=False)
+    metadata = {prop.key: prop.value for prop in model.metadata_props}
+
+    if "kp" not in metadata or "kd" not in metadata:
+        raise ValueError("ONNX model must contain 'kp' and 'kd' in metadata")
+
+    kp = np.array(json.loads(metadata["kp"]), dtype=np.float32)
+    kd = np.array(json.loads(metadata["kd"]), dtype=np.float32)
+    logger.info(f"Loaded KP/KD from ONNX ({len(kp)} joints)")
 
     return policy, kp, kd
 

src/lerobot/robots/unitree_g1/run_g1_server.py

@@ -22,33 +22,16 @@ This server runs on the robot and forwards:
 - Robot commands (LowCmd) from ZMQ to DDS (from remote clients)
 
 Uses JSON for secure serialization instead of pickle.
-
-Controller-negotiation handshake
---------------------------------
-The first message from a client agrees on which controller the server will run onboard
-(``serve_onboard_controller``); the controller NEVER runs on the laptop client.
-Test the handshake in isolation (no DDS, runs on a laptop) in two terminals::
-
-    # terminal A: handshake-only server
-    python -m lerobot.robots.unitree_g1.run_g1_server --handshake-only
-
-    # terminal B: client proposes a controller
-    python -m lerobot.robots.unitree_g1.run_g1_server \\
-        --handshake-client SonicWholeBodyController --sonic-token-action --server-ip 127.0.0.1
-
-On the real robot, add ``--handshake`` to the normal bridge to require agreement first.
 """
 
 import argparse
 import base64
 import contextlib
 import json
-import signal
 import threading
 import time
 from typing import Any
 
-import numpy as np
 import zmq
 from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
 from unitree_sdk2py.core.channel import ChannelFactoryInitialize, ChannelPublisher, ChannelSubscriber
@@ -67,253 +50,6 @@ LOWCMD_PORT = 6000
 LOWSTATE_PORT = 6001
 NUM_MOTORS = 35
 
-# Onboard high-level channels (serve_onboard_controller): compact actions in, state out.
-ACTION_PORT = 6004
-STATE_PORT = 6005
-
-# Controller-negotiation handshake (REQ/REP). The client's first message agrees on
-# which controller the server will run before any control data flows.
-HANDSHAKE_PORT = 6002
-PROTOCOL_VERSION = 1
-
-# Controllers that can run ONBOARD (must match g1_utils.make_locomotion_controller).
-# ``None`` (a.k.a. "bridge") means no onboard controller: the laptop owns control and
-# streams raw lowcmd over the ZMQ DDS bridge (the legacy run_g1_server behavior).
-VALID_CONTROLLERS = (
-    "GrootLocomotionController",
-    "HolosomaLocomotionController",
-    "SonicWholeBodyController",
-)
-# SONIC latent-token dimensionality (mirrors sonic_whole_body.TOKEN_DIM; kept local so
-# the handshake can run without importing the heavy controller / onnxruntime).
-TOKEN_DIM = 64
-_BRIDGE_ALIASES = {"", "none", "null", "bridge", "raw"}
-
-
-def _normalize_controller(name: str | None) -> str | None:
-    """Map a requested controller name to a canonical value (or None for raw bridge)."""
-    if name is None:
-        return None
-    low = str(name).strip().lower()
-    if low in _BRIDGE_ALIASES:
-        return None
-    for c in VALID_CONTROLLERS:
-        if c.lower() == low:
-            return c
-    raise ValueError(f"Unknown controller {name!r}. Available: {list(VALID_CONTROLLERS)} or 'bridge'")
-
-
-def _capabilities(controller: str | None, sonic_token_action: bool) -> dict[str, Any]:
-    """The interface the server advertises for an agreed controller."""
-    caps: dict[str, Any] = {
-        "controller": controller,
-        "sonic_token_action": bool(sonic_token_action),
-        "protocol": PROTOCOL_VERSION,
-    }
-    if controller is None:
-        # Raw DDS bridge: the laptop runs the controller and streams lowcmd.
-        caps["mode"] = "bridge"
-        caps["lowcmd_port"] = LOWCMD_PORT
-        caps["lowstate_port"] = LOWSTATE_PORT
-    else:
-        # Onboard: the controller runs here; the laptop ships compact high-level actions.
-        caps["mode"] = "onboard"
-        caps["action_port"] = ACTION_PORT
-        caps["state_port"] = STATE_PORT
-        if sonic_token_action:
-            caps["action_space"] = "motion_token"
-            caps["action_dim"] = TOKEN_DIM
-    return caps
-
-
-def negotiate_controller(sock: zmq.Socket, shutdown_event: threading.Event) -> dict[str, Any]:
-    """Server side of the handshake: block on one REP socket until a client sends a
-    valid ``hello``, then reply with the negotiated capabilities and return them.
-
-    Rejects malformed / unknown-controller requests with an error reply and keeps
-    waiting (a rejected client can retry). Honors ``shutdown_event`` so Ctrl-C works.
-    """
-    poller = zmq.Poller()
-    poller.register(sock, zmq.POLLIN)
-    while not shutdown_event.is_set():
-        if not dict(poller.poll(timeout=200)):
-            continue
-        raw = sock.recv()
-        try:
-            hello = json.loads(raw.decode("utf-8"))
-        except (ValueError, UnicodeDecodeError) as e:
-            sock.send_json({"type": "error", "ok": False, "error": f"bad hello: {e}"})
-            continue
-        try:
-            controller = _normalize_controller(hello.get("controller"))
-        except ValueError as e:
-            sock.send_json(
-                {"type": "error", "ok": False, "error": str(e), "available": list(VALID_CONTROLLERS)}
-            )
-            continue
-        reply = {"type": "welcome", "ok": True, **_capabilities(controller, hello.get("sonic_token_action", False))}
-        sock.send_json(reply)
-        return reply
-    raise KeyboardInterrupt
-
-
-def request_controller(
-    server_ip: str,
-    controller: str | None,
-    *,
-    sonic_token_action: bool = False,
-    port: int = HANDSHAKE_PORT,
-    timeout_s: float = 5.0,
-) -> dict[str, Any]:
-    """Client side of the handshake: propose a controller, return the server's agreed
-    capabilities (or raise on rejection / timeout)."""
-    ctx = zmq.Context.instance()
-    sock = ctx.socket(zmq.REQ)
-    sock.setsockopt(zmq.LINGER, 0)
-    sock.setsockopt(zmq.RCVTIMEO, int(timeout_s * 1000))
-    sock.setsockopt(zmq.SNDTIMEO, int(timeout_s * 1000))
-    sock.connect(f"tcp://{server_ip}:{port}")
-    hello = {
-        "type": "hello",
-        "controller": controller,
-        "sonic_token_action": bool(sonic_token_action),
-        "protocol": PROTOCOL_VERSION,
-    }
-    try:
-        sock.send_json(hello)
-        reply = sock.recv_json()
-    except zmq.Again as e:
-        raise TimeoutError(f"no handshake reply from {server_ip}:{port} within {timeout_s}s") from e
-    finally:
-        sock.close(linger=0)
-    if not reply.get("ok"):
-        raise RuntimeError(f"handshake rejected: {reply.get('error')} (available: {reply.get('available')})")
-    return reply
-
-
-def serve_onboard_controller(
-    *,
-    controller: str,
-    sonic_token_action: bool,
-    dds_interface: str | None = None,
-    sim: bool = False,
-    cameras: dict | None = None,
-    camera_fps: int = 30,
-    camera_port: int = 5555,
-    action_port: int = ACTION_PORT,
-    state_port: int = STATE_PORT,
-    state_fps: float = 30.0,
-    stop: threading.Event | None = None,
-) -> None:
-    """Run the negotiated controller ONBOARD -- the single control path on the robot.
-
-    Builds ``UnitreeG1(onboard=True, controller=...)`` so the controller/balance loop runs
-    locally against DDS at full rate (the 50 Hz ``_controller_loop`` thread lives in
-    UnitreeG1), then receives compact high-level actions from the laptop over ZMQ
-    (:action_port), decodes them via the controller, publishes ``observation.state``
-    (:state_port), and optionally serves the ego camera. The controller NEVER runs on the
-    laptop; the laptop (lerobot-rollout thin-client) only ships tokens/axes and reads back
-    state + camera frames.
-    """
-    # Imported lazily: UnitreeG1 imports request_controller from this module, so a
-    # top-level import here would be circular.
-    from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
-    from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
-
-    if stop is None:
-        stop = threading.Event()
-        signal.signal(signal.SIGINT, lambda *_: stop.set())
-        signal.signal(signal.SIGTERM, lambda *_: stop.set())
-
-    cfg = UnitreeG1Config(
-        is_simulation=False,
-        onboard=True,
-        controller=controller,
-        dds_interface=dds_interface,
-        release_motion_control=not sim,
-        physical_remote=not sim,
-        cameras={},
-    )
-
-    # Optional camera server (background daemon thread; independent of DDS).
-    if cameras:
-        camera_server = ImageServer({"fps": camera_fps, "cameras": cameras}, port=camera_port)
-        threading.Thread(target=camera_server.run, daemon=True).start()
-        cam_summary = ", ".join(f"{name}(dev {c['device_id']})" for name, c in cameras.items())
-        print(f"Camera server started on :{camera_port}: {cam_summary}")
-
-    robot = UnitreeG1(cfg)
-    print(f"Connecting onboard robot (controller={controller}, token={sonic_token_action})...")
-    robot.connect()
-
-    ctx = zmq.Context.instance()
-    sock = ctx.socket(zmq.PULL)
-    sock.setsockopt(zmq.CONFLATE, 1)  # only ever act on the freshest command
-    sock.setsockopt(zmq.RCVTIMEO, 200)  # keeps the loop responsive to the stop event
-    sock.bind(f"tcp://0.0.0.0:{action_port}")
-    print(f"Onboard controller live. Waiting for laptop actions on :{action_port} ...")
-    print("Ctrl-C for graceful shutdown.")
-
-    state_sock = None
-    if state_fps > 0:
-        state_sock = ctx.socket(zmq.PUB)
-        state_sock.setsockopt(zmq.SNDHWM, 2)
-        state_sock.setsockopt(zmq.LINGER, 0)
-        state_sock.bind(f"tcp://0.0.0.0:{state_port}")
-        print(f"Publishing observation.state on :{state_port} at {state_fps:.0f} Hz")
-
-        def publish_state() -> None:
-            period = 1.0 / state_fps
-            while not stop.is_set():
-                t0 = time.time()
-                obs = robot.get_observation()
-                if obs:
-                    # Forward every scalar proprio key the robot exposes (29 joint .q, IMU,
-                    # and the SONIC token echo: 64-D motion_token_state.*). Camera arrays are
-                    # streamed separately by the ImageServer, so drop ndarrays here. This
-                    # makes the laptop thin-client a pure relay.
-                    state = {
-                        k: float(v)
-                        for k, v in obs.items()
-                        if isinstance(v, (bool, int, float, np.floating, np.integer))
-                    }
-                    with contextlib.suppress(zmq.Again):
-                        state_sock.send_json(state, zmq.NOBLOCK)
-                time.sleep(max(0.0, period - (time.time() - t0)))
-
-        threading.Thread(target=publish_state, daemon=True).start()
-    else:
-        print("observation.state PUB disabled (state_fps<=0)")
-
-    n = 0
-    try:
-        while not stop.is_set():
-            try:
-                payload = sock.recv()
-            except zmq.Again:
-                continue
-            except zmq.ContextTerminated:
-                break
-
-            try:
-                action = json.loads(payload.decode("utf-8"))
-            except (ValueError, UnicodeDecodeError) as e:
-                print(f"Dropping malformed action: {e}")
-                continue
-
-            robot.send_action(action)
-
-            n += 1
-            if n % 60 == 0:
-                print(f"Applied {n} actions")
-    finally:
-        print("Shutting down onboard controller...")
-        stop.set()
-        if state_sock is not None:
-            with contextlib.suppress(Exception):
-                state_sock.close(linger=0)
-        robot.disconnect()
-
 
 def lowstate_to_dict(msg: hg_LowState) -> dict[str, Any]:
     """Convert LowState SDK message to a JSON-serializable dictionary."""
@@ -424,86 +160,8 @@ def main() -> None:
     parser.add_argument("--camera-width", type=int, default=640, help="Camera width (default: 640)")
     parser.add_argument("--camera-height", type=int, default=480, help="Camera height (default: 480)")
     parser.add_argument("--camera-port", type=int, default=5555, help="Camera ZMQ port (default: 5555)")
-    # Controller-negotiation handshake (first message agrees on the controller).
-    parser.add_argument("--handshake", action="store_true",
-                        help="Wait for a client to negotiate the controller before bridging")
-    parser.add_argument("--handshake-port", type=int, default=HANDSHAKE_PORT,
-                        help=f"Handshake REQ/REP port (default: {HANDSHAKE_PORT})")
-    parser.add_argument("--handshake-only", action="store_true",
-                        help="Run ONLY the handshake server (no DDS/cameras) to test negotiation")
-    parser.add_argument("--handshake-client", default=None, metavar="CONTROLLER",
-                        help="Act as a client: propose CONTROLLER (or 'bridge') to --server-ip and print the reply")
-    parser.add_argument("--server-ip", default="127.0.0.1", help="[--handshake-client] server IP")
-    parser.add_argument("--sonic-token-action", action="store_true",
-                        help="[handshake] negotiate the 64-D SONIC token action interface")
     args = parser.parse_args()
 
-    # --- Isolated handshake test paths (no DDS, safe to run on a laptop) ---
-    if args.handshake_client is not None:
-        controller = None if args.handshake_client.strip().lower() in _BRIDGE_ALIASES else args.handshake_client
-        reply = request_controller(
-            args.server_ip, controller,
-            sonic_token_action=args.sonic_token_action, port=args.handshake_port,
-        )
-        print(json.dumps(reply, indent=2))
-        return
-
-    if args.handshake_only:
-        ctx = zmq.Context.instance()
-        rep = ctx.socket(zmq.REP)
-        rep.bind(f"tcp://0.0.0.0:{args.handshake_port}")
-        print(f"[handshake] server listening on :{args.handshake_port} (no DDS). Ctrl-C to stop.")
-        shutdown = threading.Event()
-        try:
-            while True:
-                reply = negotiate_controller(rep, shutdown)
-                print(f"[handshake] agreed: controller={reply['controller']} mode={reply['mode']} "
-                      f"sonic_token_action={reply['sonic_token_action']}")
-        except KeyboardInterrupt:
-            print("\n[handshake] stopping")
-        finally:
-            rep.close(linger=0)
-            ctx.term()
-        return
-
-    # Controller-negotiation handshake: the client's first message agrees on the
-    # controller, which we then run ONBOARD (the controller NEVER runs on the laptop).
-    # Bridge/None falls through to the legacy raw DDS forward (deprecated laptop control).
-    if args.handshake:
-        ctx = zmq.Context.instance()
-        hs = ctx.socket(zmq.REP)
-        hs.bind(f"tcp://0.0.0.0:{args.handshake_port}")
-        print(f"[handshake] waiting for client controller agreement on :{args.handshake_port} ...")
-        shutdown = threading.Event()
-        try:
-            agreed = negotiate_controller(hs, shutdown)
-        except KeyboardInterrupt:
-            print("[handshake] interrupted before agreement; exiting")
-            hs.close(linger=0)
-            ctx.term()
-            return
-        hs.close(linger=0)
-        if agreed["controller"] is not None:
-            print(f"[handshake] running controller ONBOARD: {agreed['controller']} "
-                  f"(sonic_token_action={agreed['sonic_token_action']})")
-            cameras = None
-            if args.camera:
-                cameras = {
-                    "head_camera": {
-                        "device_id": args.camera_device,
-                        "shape": [args.camera_height, args.camera_width],
-                    }
-                }
-            serve_onboard_controller(
-                controller=agreed["controller"],
-                sonic_token_action=bool(agreed["sonic_token_action"]),
-                cameras=cameras,
-                camera_fps=args.camera_fps,
-                camera_port=args.camera_port,
-            )
-            return
-        print("[handshake] client selected raw DDS bridge (laptop owns control) -> legacy forward.")
-
     # Optionally start camera server in background thread
     camera_thread = None
     if args.camera:
@@ -547,7 +205,6 @@ def main() -> None:
 
     # initialize ZMQ
     ctx = zmq.Context.instance()
-    shutdown_event = threading.Event()
 
     # receive commands from remote client
     lowcmd_sock = ctx.socket(zmq.PULL)
@@ -558,6 +215,7 @@ def main() -> None:
     lowstate_sock.bind(f"tcp://0.0.0.0:{LOWSTATE_PORT}")
 
     state_period = 0.002  # ~500 hz
+    shutdown_event = threading.Event()
 
     # start observation forwarding in background thread
     t_state = threading.Thread(

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -16,8 +16,6 @@
 
 from __future__ import annotations
 
-import contextlib
-import json
 import logging
 import threading
 import time
@@ -28,7 +26,6 @@ from typing import TYPE_CHECKING, Protocol, runtime_checkable
 import numpy as np
 
 from lerobot.cameras import make_cameras_from_configs
-from lerobot.utils.errors import DeviceNotConnectedError
 from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.import_utils import _unitree_sdk_available, require_package
 
@@ -37,10 +34,10 @@ 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,
-    lowstate_to_obs,
     make_locomotion_controller,
 )
 
@@ -50,9 +47,7 @@ if TYPE_CHECKING or _unitree_sdk_available:
         ChannelPublisher as _SDKChannelPublisher,
         ChannelSubscriber as _SDKChannelSubscriber,
     )
-    from unitree_sdk2py.idl.default import (
-        unitree_hg_msg_dds__LowCmd_,
-    )
+    from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
     from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
         LowCmd_ as hg_LowCmd,
         LowState_ as hg_LowState,
@@ -84,14 +79,6 @@ class LocomotionController(Protocol):
 kTopicLowCommand_Debug = "rt/lowcmd"
 kTopicLowState = "rt/lowstate"
 
-# Wireless-remote button byte layout, mapped to the positional button indices the
-# locomotion controllers expect. Used in onboard mode to read the physical Unitree
-# remote from lowstate (mirrors the exo teleoperator's RemoteController).
-_REMOTE_BUTTON_MAP: list[str] = [
-    "RB", "LB", "start", "back", "RT", "LT", "", "",
-    "A", "B", "X", "Y", "up", "right", "down", "left",
-]
-
 
 @dataclass
 class MotorState:
@@ -132,37 +119,24 @@ class UnitreeG1(Robot):
         self.config = config
         self.control_dt = config.control_dt
 
-        # Three mutually-exclusive roles:
-        #   * simulation : local DDS + controller run in-process against a MuJoCo world.
-        #   * onboard    : local DDS + controller run in-process on the robot NX.
-        #   * client     : thin laptop client. No DDS, no controller. It negotiates a
-        #                  controller with ``run_g1_server`` (which runs it onboard),
-        #                  PUSHes high-level actions and reads back state + cameras over
-        #                  ZMQ. The controller *always* runs on the robot, never here.
-        self._client = not config.is_simulation and not config.onboard
-
         # Initialize cameras config (ZMQ-based) - actual connection in connect()
         self._cameras = make_cameras_from_configs(config.cameras)
 
-        # DDS channel classes are only needed by the in-process control roles. The thin
-        # client never touches DDS, so we don't import the socket shim at all.
-        if config.is_simulation or config.onboard:
+        # Import channel classes based on mode
+        if config.is_simulation:
             self._ChannelFactoryInitialize = _SDKChannelFactoryInitialize
             self._ChannelPublisher = _SDKChannelPublisher
             self._ChannelSubscriber = _SDKChannelSubscriber
         else:
-            self._ChannelFactoryInitialize = None
-            self._ChannelPublisher = None
-            self._ChannelSubscriber = None
+            from .unitree_sdk2_socket import (
+                ChannelFactoryInitialize,
+                ChannelPublisher,
+                ChannelSubscriber,
+            )
 
-        # Client-side ZMQ handles / negotiated capabilities (populated in connect()).
-        self._client_action_sock = None
-        self._client_state_sock = None
-        self._client_state_latest: dict[str, float] = {}
-        self._client_caps: dict | None = None
-
-        # Optional arm gravity compensation (feed-forward torque via the arm IK solver).
-        self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
+            self._ChannelFactoryInitialize = ChannelFactoryInitialize
+            self._ChannelPublisher = ChannelPublisher
+            self._ChannelSubscriber = ChannelSubscriber
 
         # Initialize state variables
         self.sim_env = None
@@ -172,69 +146,24 @@ class UnitreeG1(Robot):
         self._shutdown_event = threading.Event()
         self.subscribe_thread = None
 
-        # Lower-body controller loaded dynamically. GUARDRAIL: the controller must never
-        # be built or run on the laptop client -- it always runs onboard (or in sim).
-        if self._client:
-            self.controller: LocomotionController | None = None
-        else:
-            self.controller = make_locomotion_controller(config.controller)
+        self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
 
-            # Token-driven deploy: 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 control ticks.
-            if hasattr(self.controller, "token_mode"):
-                self.controller.token_mode = True
+        # Lower-body controller loaded dynamically
+        self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
 
         # 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 = {}
 
-        # Onboard-only: parser for the physical Unitree wireless remote (read straight
-        # from local lowstate so joystick locomotion works without a laptop round-trip).
-        self._joystick = None
-
-        # 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:
@@ -302,46 +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]:
-        # Role-agnostic: the schema is a pure function of the controller name. The thin
-        # client advertises the same schema as the onboard robot so the exact same
-        # policy output routes straight through.
-
-        # No controller configured at all: raw 29-DoF joint teleop.
-        if self.config.controller is None:
+        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}
@@ -357,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
 
@@ -378,13 +271,6 @@ class UnitreeG1(Robot):
                 with self._controller_action_lock:
                     controller_input = dict(self.controller_input)
 
-                # Onboard: the physical Unitree remote (in local lowstate) takes
-                # priority for locomotion when active; otherwise laptop/ZMQ axes stand.
-                if self.config.onboard:
-                    wl = self._wireless_remote_input(lowstate)
-                    if wl is not None:
-                        controller_input.update(wl)
-
                 # Run controller step
                 controller_action = self.controller.run_step(controller_input, lowstate)
 
@@ -407,163 +293,7 @@ class UnitreeG1(Robot):
     def configure(self) -> None:
         pass
 
-    def _wireless_remote_input(self, lowstate) -> dict | None:
-        """Parse the physical Unitree remote from lowstate into controller inputs.
-
-        Onboard only. Returns None when the remote is idle so the laptop-provided
-        (ZMQ) axes keep control; otherwise the physical remote takes priority.
-        """
-        js = self._joystick
-        if js is None:
-            return None
-        wr = getattr(lowstate, "wireless_remote", None)
-        if not wr or len(wr) < 24:
-            return None
-        try:
-            js.extract(wr)
-        except Exception:  # noqa: BLE001
-            return None
-
-        axes = {
-            "remote.lx": float(js.lx.data),
-            "remote.ly": float(js.ly.data),
-            "remote.rx": float(js.rx.data),
-            "remote.ry": float(js.ry.data),
-        }
-        active = any(abs(v) > 1e-2 for v in axes.values())
-        out = dict(axes)
-        for i, name in enumerate(_REMOTE_BUTTON_MAP):
-            if name:
-                val = float(getattr(js, name).data)
-                out[f"remote.button.{i}"] = val
-                if val:
-                    active = True
-        return out if active else None
-
-    def _release_motion_control(self) -> None:
-        """Release the robot's built-in motion services so we can send raw lowcmd.
-
-        Onboard-only. Mirrors run_g1_server.py: on the real robot the factory
-        locomotion/hand services must relinquish control before our controller can
-        write to ``rt/lowcmd``, otherwise commands are ignored or fought.
-        """
-        from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
-
-        msc = MotionSwitcherClient()
-        msc.SetTimeout(5.0)
-        msc.Init()
-        _, result = msc.CheckMode()
-        while result is not None and "name" in result and result["name"]:
-            logger.info("[UnitreeG1] Releasing built-in mode '%s'...", result["name"])
-            msc.ReleaseMode()
-            _, result = msc.CheckMode()
-            time.sleep(1.0)
-
-    # ------------------------------------------------------------------ #
-    # Thin-client role (laptop): no DDS, no controller. Talks to run_g1_server
-    # over ZMQ. The controller ALWAYS runs onboard; we only relay high-level
-    # actions and read back the state echo + camera frames.
-    # ------------------------------------------------------------------ #
-    def _connect_client(self) -> None:
-        import zmq
-
-        from .run_g1_server import ACTION_PORT, HANDSHAKE_PORT, STATE_PORT, request_controller
-
-        server_ip = self.config.robot_ip
-        if not server_ip:
-            raise ValueError("client mode requires config.robot_ip (the G1 running run_g1_server)")
-
-        # 1) Handshake: agree with the server on which controller it will run onboard.
-        logger.info(
-            "[client] handshaking with %s:%d (controller=%s, token=%s)...",
-            server_ip, HANDSHAKE_PORT, self.config.controller, self._sonic_token,
-        )
-        self._client_caps = request_controller(
-            server_ip,
-            self.config.controller,
-            sonic_token_action=self._sonic_token,
-            port=HANDSHAKE_PORT,
-        )
-        logger.info("[client] server agreed: %s", self._client_caps)
-
-        ctx = zmq.Context.instance()
-
-        # 2) Action PUSH: ship compact high-level actions to the onboard controller.
-        self._client_action_sock = ctx.socket(zmq.PUSH)
-        self._client_action_sock.setsockopt(zmq.SNDHWM, 2)
-        self._client_action_sock.setsockopt(zmq.LINGER, 0)
-        self._client_action_sock.connect(f"tcp://{server_ip}:{ACTION_PORT}")
-
-        # 3) State SUB: read the onboard observation.state echo (last token / joints).
-        self._client_state_sock = ctx.socket(zmq.SUB)
-        self._client_state_sock.setsockopt(zmq.CONFLATE, 1)
-        self._client_state_sock.setsockopt_string(zmq.SUBSCRIBE, "")
-        self._client_state_sock.connect(f"tcp://{server_ip}:{STATE_PORT}")
-
-        # 4) Cameras (ZMQ ImageServer served by run_g1_server) - same as any client.
-        for cam in self._cameras.values():
-            if not cam.is_connected:
-                cam.connect()
-        logger.info("[client] connected: actions ->:%d, state <-:%d, %d camera(s).",
-                    ACTION_PORT, STATE_PORT, len(self._cameras))
-
-    def _recv_client_state(self) -> None:
-        """Drain the state SUB (CONFLATE keeps only the freshest) into the latest cache."""
-        import zmq
-
-        if self._client_state_sock is None:
-            return
-        while True:
-            try:
-                state = self._client_state_sock.recv_json(flags=zmq.NOBLOCK)
-            except zmq.Again:
-                break
-            except (ValueError, zmq.ZMQError):
-                break
-            if isinstance(state, dict):
-                self._client_state_latest = {k: float(v) for k, v in state.items()}
-
-    def _get_observation_client(self) -> RobotObservation:
-        self._recv_client_state()
-        obs: dict = dict(self._client_state_latest)
-        for cam_name, cam in self._cameras.items():
-            if getattr(cam, "use_rgb", True):
-                obs[cam_name] = cam.read_latest()
-            if getattr(cam, "use_depth", False):
-                obs[f"{cam_name}_depth"] = cam.read_latest_depth()
-        return obs
-
-    def _send_action_client(self, action: RobotAction) -> RobotAction:
-        """Relay the raw action straight to the onboard controller. NO processing here:
-        the controller negotiated in the handshake interprets it (token / wb / arm)."""
-        import zmq
-
-        if self._client_action_sock is None:
-            raise DeviceNotConnectedError("UnitreeG1 client is not connected")
-        payload = json.dumps({k: float(v) for k, v in action.items()}).encode("utf-8")
-        with contextlib.suppress(zmq.Again):
-            self._client_action_sock.send(payload, zmq.NOBLOCK)
-        return action
-
-    def _disconnect_client(self) -> None:
-        for sock in (self._client_action_sock, self._client_state_sock):
-            if sock is not None:
-                with contextlib.suppress(Exception):
-                    sock.close(linger=0)
-        self._client_action_sock = None
-        self._client_state_sock = None
-        for cam in self._cameras.values():
-            with contextlib.suppress(Exception):
-                cam.disconnect()
-
     def connect(self, calibrate: bool = True) -> None:  # connect to DDS
-        # Thin-client role: no DDS, no controller. Negotiate the controller with
-        # run_g1_server (which runs it onboard), then open the high-level ZMQ links:
-        # PUSH actions on :ACTION_PORT, SUB state echo on :STATE_PORT, cameras via ZMQ.
-        if self._client:
-            self._connect_client()
-            return
-
         # Initialize DDS channel and simulation environment
         if self.config.is_simulation:
             from lerobot.envs import make_env
@@ -572,28 +302,6 @@ class UnitreeG1(Robot):
             self._env_wrapper = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True)
             # Extract the actual gym env from the dict structure
             self.sim_env = self._env_wrapper["hub_env"][0].envs[0]
-        elif self.config.onboard:
-            # Real robot, controller running onboard against local DDS. Initialize the
-            # real SDK channel factory on the robot's DDS interface and take low-level
-            # control from the built-in services before we start writing lowcmd.
-            if self.config.dds_interface:
-                self._ChannelFactoryInitialize(0, self.config.dds_interface)
-            else:
-                self._ChannelFactoryInitialize(0)
-            # Real robot: hand low-level control over from the built-in services.
-            # A DDS sim has no MotionSwitcher, so this is skipped there.
-            if self.config.release_motion_control:
-                self._release_motion_control()
-            # Real robot: read the physical wireless remote from lowstate for
-            # locomotion. A sim has no physical remote, so leave _joystick=None and
-            # let send_action (ZMQ) drive the locomotion axes instead.
-            if self.config.physical_remote:
-                from unitree_sdk2py.utils.joystick import Joystick
-
-                self._joystick = Joystick()
-                for axis in (self._joystick.lx, self._joystick.ly, self._joystick.rx, self._joystick.ry):
-                    axis.smooth = 1.0
-                    axis.deadzone = 0.0
         else:
             self._ChannelFactoryInitialize(0, config=self.config)
 
@@ -635,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
@@ -645,8 +350,7 @@ class UnitreeG1(Robot):
             self.msg.motor_cmd[joint].kd = self.kd[joint.value]
             self.msg.motor_cmd[joint].q = lowstate.motor_state[joint.value].q
 
-        # Start the 50 Hz controller thread (runs the locomotion/whole-body policy and
-        # publishes low commands to DDS).
+        # Start controller thread if enabled
         if self.controller is not None:
             self._controller_thread = threading.Thread(target=self._controller_loop, daemon=True)
             self._controller_thread.start()
@@ -668,34 +372,12 @@ class UnitreeG1(Robot):
             logger.warning(f"Failed to send zero-torque on disconnect: {e}")
 
     def disconnect(self):
-        if self._client:
-            self._disconnect_client()
-            return
-
-        # Stop the controller loop first so it isn't fighting the shutdown ramp.
-        self._shutdown_event.set()
-        controller_stopped = True
-        if self._controller_thread is not None:
-            # Wait long enough for any in-flight inference tick to finish and the loop
-            # to observe the shutdown flag, so no stray low command is published while
-            # the ramp runs (the shutdown routine must be the single publisher).
-            self._controller_thread.join(timeout=5.0)
-            if self._controller_thread.is_alive():
-                controller_stopped = False
-                logger.error(
-                    "Controller thread did not stop; skipping graceful ramp to avoid "
-                    "concurrent low commands (fail-safe: joints keep last command until exit)"
-                )
-
-        # Put the robot in passive mode (zero-torque) before stopping the rest (real
-        # robot only; the subscribe thread is still alive here to supply the current
-        # pose). Only publish once the controller thread has definitely exited so the
-        # two aren't publishing at once.
-        if not self.config.is_simulation and controller_stopped:
+        # Put robot in passive mode before stopping threads
+        if not self.config.is_simulation:
             self._send_zero_torque()
 
-        if self.controller is not None and hasattr(self.controller, "shutdown"):
-            self.controller.shutdown()
+        # Signal thread to stop and unblock any waits
+        self._shutdown_event.set()
 
         # Wait for subscribe thread to finish
         if self.subscribe_thread is not None:
@@ -703,6 +385,12 @@ class UnitreeG1(Robot):
             if self.subscribe_thread.is_alive():
                 logger.warning("Subscribe thread did not stop cleanly")
 
+        # Wait for controller thread to finish
+        if self._controller_thread is not None:
+            self._controller_thread.join(timeout=2.0)
+            if self._controller_thread.is_alive():
+                logger.warning("Controller thread did not stop cleanly")
+
         # Close simulation environment
         if self.config.is_simulation and self.sim_env is not None:
             try:
@@ -729,25 +417,49 @@ class UnitreeG1(Robot):
             cam.disconnect()
 
     def get_observation(self) -> RobotObservation:
-        if self._client:
-            return self._get_observation_client()
-
         with self._lowstate_lock:
             lowstate = self._lowstate
         if lowstate is None:
             return {}
 
-        # Motors + IMU + wireless remote (shared lowstate -> obs mapping)
-        obs = lowstate_to_obs(lowstate)
+        obs = {}
 
-        # 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
+        # Motors - q, dq, tau for all joints
+        for motor in G1_29_JointIndex:
+            name = motor.name
+            idx = motor.value
+            obs[f"{name}.q"] = lowstate.motor_state[idx].q
+            obs[f"{name}.dq"] = lowstate.motor_state[idx].dq
+            obs[f"{name}.tau"] = lowstate.motor_state[idx].tau_est
 
-            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)
+        # IMU - gyroscope
+        if lowstate.imu_state.gyroscope:
+            obs["imu.gyro.x"] = lowstate.imu_state.gyroscope[0]
+            obs["imu.gyro.y"] = lowstate.imu_state.gyroscope[1]
+            obs["imu.gyro.z"] = lowstate.imu_state.gyroscope[2]
+
+        # IMU - accelerometer
+        if lowstate.imu_state.accelerometer:
+            obs["imu.accel.x"] = lowstate.imu_state.accelerometer[0]
+            obs["imu.accel.y"] = lowstate.imu_state.accelerometer[1]
+            obs["imu.accel.z"] = lowstate.imu_state.accelerometer[2]
+
+        # IMU - quaternion
+        if lowstate.imu_state.quaternion:
+            obs["imu.quat.w"] = lowstate.imu_state.quaternion[0]
+            obs["imu.quat.x"] = lowstate.imu_state.quaternion[1]
+            obs["imu.quat.y"] = lowstate.imu_state.quaternion[2]
+            obs["imu.quat.z"] = lowstate.imu_state.quaternion[3]
+
+        # IMU - rpy
+        if lowstate.imu_state.rpy:
+            obs["imu.rpy.roll"] = lowstate.imu_state.rpy[0]
+            obs["imu.rpy.pitch"] = lowstate.imu_state.rpy[1]
+            obs["imu.rpy.yaw"] = lowstate.imu_state.rpy[2]
+
+        # Wireless remote (raw bytes for teleoperator)
+        if lowstate.wireless_remote:
+            obs["wireless_remote"] = lowstate.wireless_remote
 
         # Cameras - read images from ZMQ cameras
         for cam_name, cam in self._cameras.items():
@@ -759,22 +471,11 @@ class UnitreeG1(Robot):
         return obs
 
     def send_action(self, action: RobotAction) -> RobotAction:
-        if self._client:
-            return self._send_action_client(action)
-
         action_to_publish = action
         if self.controller is not None:
-            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)
-            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
@@ -802,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 (e.g. locomotion ``remote.*``
-        axes/buttons) 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:
@@ -820,8 +515,6 @@ class UnitreeG1(Robot):
 
     @property
     def is_connected(self) -> bool:
-        if self._client:
-            return self._client_action_sock is not None
         with self._lowstate_lock:
             return self._lowstate is not None
 
@@ -844,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 / OpenHLM) 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_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/utils/test_random_utils.py

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

uv.lock

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