feat(unitree_g1): episode reset, lazy replay decode, safe shutdown

- reset(): pause the background controller and, for full-body controllers,
  publish the default pose directly (new _controller_paused flag) so reset and
  the controller loop aren't both writing low commands.
- SONIC pipeline: add reset() to StandingEncoderDecoder and PlannerController
  (clear token/proprio history/heading, rewind motion buffer); SonicRuntime.reset()
  now calls controller.reset().
- sonic_whole_body: require the full dense 34-D command (no silent zero-fill of a
  partial action) and integrate yaw-rate (idx 33) into heading.
- controllers/__init__: import the controller classes referenced in __all__.
- unitree_g1: lazy replay-frame decode + small cache instead of decoding all
  frames up front; safer disconnect (longer controller-thread join + fail-safe
  that skips the graceful ramp if the thread won't stop).
- lint: ruff-format config_unitree_g1 hand_closed_pose; prettier README table.

pyproject.toml

@@ -374,7 +374,11 @@ torch = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
 torchvision = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
 
 [tool.setuptools.package-data]
-lerobot = ["envs/*.json", "annotations/steerable_pipeline/prompts/*.txt"]
+lerobot = [
+    "envs/*.json",
+    "annotations/steerable_pipeline/prompts/*.txt",
+    "teleoperators/pico_headset/assets/*.npz",
+]
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/lerobot/robots/unitree_g1/README.md

@@ -0,0 +1,89 @@
+# Unitree G1 — SONIC encoder/decoder whole-body control
+
+This package runs NVIDIA's **SONIC** encoder/decoder on the Unitree G1, in MuJoCo
+simulation or on real hardware, driven by a dense **34-D whole-body command** (the
+OpenHLM / pi0.5 action layout). It is a pure-Python/ONNX reimplementation of the
+reference-tracking half of the SONIC deploy stack (no `gear_sonic`/torch dependency, and
+no motion planner): the encoder compresses a reference motion window into a latent token
+and the decoder maps that token + proprioception history into 50 Hz joint-position
+targets for the robot's PD controller.
+
+## Controllers
+
+Selected with `--robot.controller=<ClassName>`:
+
+| Controller                     | Purpose                                                      |
+| ------------------------------ | ------------------------------------------------------------ |
+| `SonicWholeBodyController`     | SONIC encoder/decoder driven by a 34-D OpenHLM/pi0.5 command |
+| `GrootLocomotionController`    | GR00T locomotion policy                                      |
+| `HolosomaLocomotionController` | Holosoma locomotion policy                                   |
+
+The rest of this document covers the SONIC whole-body path.
+
+Each tick the `SonicWholeBodyController` takes a 34-D command (`wb.0.pos … wb.33.pos`) in the OpenHLM
+layout:
+
+```
+[L-arm(7), L-grip(1), R-arm(7), R-grip(1), L-leg(6), R-leg(6), waist(3),
+ root roll/pitch + yaw-rate(3)]
+```
+
+The 29 joint targets become the SONIC encode-mode-0 reference (accumulated into a rolling
+50-frame trajectory with finite-difference velocities so the encoder's lookahead sees a
+real motion sequence), the root roll/pitch set the anchor orientation, and the two grip
+scalars can drive the Dex3 hands (see below). On startup the controller **interpolates**
+from the robot's measured pose into the policy's commanded target over ~3 s (no snap).
+
+## Requirements
+
+- `onnxruntime` (CPU) **or** `onnxruntime-gpu` (recommended). Verify with:
+  ```bash
+  python -c "import onnxruntime as ort; print(ort.get_available_providers())"
+  ```
+- `mujoco` for simulation (`is_simulation=True`).
+- The SONIC encoder/decoder ONNX models download automatically from the
+  `nvidia/GEAR-SONIC` Hub repo.
+
+## Running a rollout
+
+Drive the G1 with a 34-D VLA policy (OpenHLM / pi0.5) via `lerobot-rollout`:
+
+```bash
+lerobot-rollout \
+  --strategy.type=base \
+  --policy.path=<pi05_openhlm_dir> \
+  --robot.type=unitree_g1 \
+  --robot.controller=SonicWholeBodyController \
+  --robot.is_simulation=true \
+  --robot.publish_hands=true \
+  --task="<language instruction>" \
+  --duration=45 --device=cuda
+```
+
+### Cameras
+
+Image-conditioned policies need camera frames. Two options are available without live
+cameras:
+
+- **Black frames**: `--robot.empty_cameras='[base, left_wrist, right_wrist]'`.
+- **Replay a recorded episode** as the camera feed:
+  ```bash
+  --robot.replay_camera_parquet=<episode.parquet> \
+  --robot.replay_camera_map='{base: head_image_left, left_wrist: left_wrist_image, right_wrist: right_wrist_image}'
+  ```
+
+### Hands (Dex3)
+
+`--robot.publish_hands=true` publishes `rt/dex3/{left,right}/cmd` from the two grip
+scalars (`wb.7.pos` left, `wb.15.pos` right). The scalar is interpolated between
+`hand_open_grip_value` (default 1.0 = open) and `hand_closed_grip_value` (default 0.0 =
+closed) and scaled onto `hand_closed_pose` (7 joints:
+`thumb_0, thumb_1, thumb_2, middle_0, middle_1, index_0, index_1`). Flip the signs in
+`hand_closed_pose` if the fingers curl the wrong way, or raise `hand_kp` for a firmer
+grip.
+
+## Observation state
+
+When the whole-body controller is active the robot exposes a 34-D proprio state
+(`wb_state.0.pos … wb_state.33.pos`) in the same OpenHLM layout as the action, which the
+rollout aggregates into `observation.state` for the policy.

src/lerobot/robots/unitree_g1/config_unitree_g1.py

@@ -65,9 +65,41 @@ class UnitreeG1Config(RobotConfig):
     # Cameras (ZMQ-based remote cameras)
     cameras: dict[str, CameraConfig] = field(default_factory=dict)
 
+    # Synthetic zero-image cameras exposed as ``observation.images.{name}`` (H×W×3
+    # black frames). Lets image-conditioned policies (e.g. pi0.5 / OpenHLM) run in
+    # sim before real cameras are wired. Empty = disabled.
+    empty_cameras: list[str] = field(default_factory=list)
+    empty_camera_hw: tuple[int, int] = (224, 224)
+
+    # Publish Dex3 hand commands (``rt/dex3/{left,right}/cmd``) driven by the OpenHLM
+    # gripper scalars (``wb.7.pos`` left, ``wb.15.pos`` right). Lets the 43-DoF sim
+    # (or a real Dex3-equipped G1) show grasping. The scalar in [0, 1] is remapped to
+    # a curl amount (``hand_open_grip_value`` -> open) and scaled onto
+    # ``hand_closed_pose`` (7 joints: thumb_0/1/2, middle_0/1, index_0/1). Flip signs
+    # in ``hand_closed_pose`` if fingers curl the wrong way.
+    publish_hands: bool = False
+    hand_open_grip_value: float = 1.0
+    hand_closed_grip_value: float = 0.0
+    hand_closed_pose: list[float] = field(default_factory=lambda: [1.0, 0.9, 0.9, 1.3, 1.3, 1.3, 1.3])
+    hand_kp: float = 1.5
+    hand_kd: float = 0.1
+
+    # Replay recorded camera frames from a LeRobot parquet episode as the camera
+    # feed (e.g. OpenHLM-data episode). Maps a robot camera name to a parquet image
+    # column; frames advance one per observation and loop. Lets a VLA see the real
+    # task video in sim without live cameras. Empty map = disabled.
+    replay_camera_parquet: str | None = None
+    replay_camera_map: dict[str, str] = field(default_factory=dict)
+    replay_camera_loop: bool = True
+
     # Compensates for gravity on the unitree's arms using the arm ik solver
     gravity_compensation: bool = False
 
-    # Lower-body controller class name, e.g. "GrootLocomotionController" or
-    # "HolosomaLocomotionController". None disables it.
+    # Locomotion controller class name, e.g. "GrootLocomotionController",
+    # "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
     controller: str | None = None
+
+    # On disconnect (e.g. Ctrl-C), seconds to hold the current pose while ramping joint
+    # stiffness (kp) to zero — a soft, damped settle instead of an instant limp /
+    # free-fall. 0 disables it (immediate zero-torque). Real robot only.
+    graceful_stop_s: float = 1.5

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

@@ -0,0 +1,28 @@
+#!/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 SonicRuntime, SonicWholeBodyController
+
+__all__ = [
+    "GrootLocomotionController",
+    "HolosomaLocomotionController",
+    "SonicRuntime",
+    "SonicWholeBodyController",
+]

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

@@ -14,20 +14,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import logging
 from collections import deque
+from typing import TYPE_CHECKING
 
 import numpy as np
-import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from .g1_utils import (
+from lerobot.utils.import_utils import _onnxruntime_available, require_package
+
+from ..g1_utils import (
     REMOTE_AXES,
     REMOTE_BUTTONS,
     G1_29_JointIndex,
     get_gravity_orientation,
 )
 
+if TYPE_CHECKING or _onnxruntime_available:
+    import onnxruntime as ort
+else:
+    ort = None
+
 logger = logging.getLogger(__name__)
 
 
@@ -83,6 +92,7 @@ class GrootLocomotionController:
     control_dt = CONTROL_DT  # Expose for unitree_g1.py
 
     def __init__(self):
+        require_package("onnxruntime", extra="unitree_g1")
         # Load policies
         self.policy_balance, self.policy_walk = load_groot_policies()
 

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

@@ -14,21 +14,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import json
 import logging
+from typing import TYPE_CHECKING
 
 import numpy as np
-import onnx
-import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from .g1_utils import (
+from lerobot.utils.import_utils import _onnx_available, _onnxruntime_available, require_package
+
+from ..g1_utils import (
     REMOTE_AXES,
     G1_29_JointArmIndex,
     G1_29_JointIndex,
     get_gravity_orientation,
 )
 
+if TYPE_CHECKING or _onnxruntime_available:
+    import onnxruntime as ort
+else:
+    ort = None
+
+if TYPE_CHECKING or _onnx_available:
+    import onnx
+else:
+    onnx = None
+
 logger = logging.getLogger(__name__)
 
 DEFAULT_ANGLES = np.zeros(29, dtype=np.float32)
@@ -101,6 +114,8 @@ class HolosomaLocomotionController:
     control_dt = CONTROL_DT  # Expose for unitree_g1.py
 
     def __init__(self):
+        require_package("onnxruntime", extra="unitree_g1")
+        require_package("onnx", extra="unitree_g1")
         # Load policy and gains
         self.policy, self.kp, self.kd = load_policy()
 

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

@@ -0,0 +1,659 @@
+#!/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 encoder/decoder pipeline for the Unitree G1 whole-body controller.
+
+Pure-Python/ONNX re-implementation of the reference-tracking half of NVIDIA's SONIC
+deploy stack (mirrors ``g1_deploy_onnx_ref.cpp``). Given a reference motion buffer
+(joint targets + body orientation per frame) it produces 50 Hz joint-position targets
+for the robot's PD controller. The upstream *motion planner* is intentionally absent:
+here the reference is supplied directly by the caller (e.g. a 34-D OpenHLM / pi0.5 VLA
+command per tick, in ``sonic_whole_body.py``).
+
+Two cooperating ONNX models:
+  * **encoder** – compresses the reference window into a 64-D latent ``token``
+    (refreshed every ``ENCODER_UPDATE_EVERY`` ticks).
+  * **decoder** – every tick, maps the token + recent proprioception history to a
+    residual action that is scaled and added to ``DEFAULT_ANGLES``.
+
+Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order)
+and **MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` convert
+between them. Quaternions are scalar-first ``(w, x, y, z)``.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from lerobot.utils.import_utils import _onnxruntime_available
+
+from ..g1_utils import (
+    ISAACLAB_TO_MUJOCO,
+    MUJOCO_TO_ISAACLAB,
+    G1_29_JointIndex,
+    get_gravity_orientation,
+)
+
+if TYPE_CHECKING or _onnxruntime_available:
+    import onnxruntime as ort
+else:
+    ort = None
+
+logger = logging.getLogger(__name__)
+
+# ── Constants ────────────────────────────────────────────────────────────────
+# Robot/motor physical constants and the joint-order permutation tables. All
+# 29-vectors are in IsaacLab joint order unless the name says ``_MUJOCO``.
+
+# Nominal standing pose (rad), 29 joints in IsaacLab order. Actions are residuals
+# added on top of this; also used as the planner/encoder standing reference.
+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-type parameters used to derive action scaling and PD gains. Keys are
+# Unitree motor model names; ARMATURE = rotor inertia, EFFORT = torque limit (N·m).
+NATURAL_FREQ = 10.0 * 2.0 * np.pi  # target closed-loop stiffness bandwidth (rad/s)
+ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425}
+EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0}
+
+
+def _action_scale(k):
+    """Per-motor residual-action scale (maps policy output to joint-angle delta)."""
+    return 0.25 * EFFORT[k] / (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 order
+
+CONTROL_DT = 0.02  # 50 Hz control period (s)
+DEFAULT_HEIGHT = 0.788740  # nominal pelvis height (m)
+TOKEN_DIM = 64  # encoder latent size
+ENCODER_UPDATE_EVERY = 5  # refresh the encoder token every N ticks (decoder runs every tick)
+DEBUG_PRINT_EVERY = 100  # ticks between debug prints
+
+
+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 throughout this module are a fixed convention validated
+    against the deployed SONIC ONNX policy (the encoder/decoder consume vectors in
+    this order). Do not "correct" the table or rename toward the opposite direction
+    without re-validating on hardware — the labels are historical, the ordering is
+    load-bearing.
+    """
+    return a[MUJOCO_TO_ISAACLAB]
+
+
+DEFAULT_ANGLES_MUJOCO = _to_mujoco(DEFAULT_ANGLES)
+ENCODER_STANDING_REF = DEFAULT_ANGLES.copy()
+
+# Joint-index subsets (IsaacLab order) used to slice encoder observations.
+LOWER_BODY_IL = np.array([0, 3, 6, 9, 13, 17, 1, 4, 7, 10, 14, 18], dtype=np.int32)  # 12 leg joints
+WRIST_IL = np.array([23, 24, 25, 26, 27, 28], dtype=np.int32)  # 6 wrist joints
+VR_TARGET_DEF = np.zeros(9, dtype=np.float32)  # 3-point VR position targets (mode 1)
+VR_ORN_DEF = np.array([1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0], dtype=np.float32)  # VR orn targets (mode 1)
+SMPL_DEF = np.zeros(720, dtype=np.float32)  # SMPL whole-body window default (mode 2)
+
+# ── PD gains ─────────────────────────────────────────────────────────────────
+
+
+def compute_kp_kd():
+    """Derive per-joint PD gains (kp, kd) from motor armature and target bandwidth.
+
+    Ankle and waist joints get a x2 factor for extra stiffness. Returns two
+    (29,) float32 arrays in IsaacLab joint order.
+    """
+
+    def s(k):
+        return ARMATURE[k] * NATURAL_FREQ**2
+
+    def d(k):
+        return 2.0 * 2.0 * ARMATURE[k] * NATURAL_FREQ
+
+    _double = {4, 5, 10, 11, 13, 14}  # ankle + waist indices with factor 2
+    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
+
+
+_kp_kd = compute_kp_kd  # backward-compatible alias
+
+
+# ── Quaternion helpers ────────────────────────────────────────────────────────
+# All quaternions are scalar-first (w, x, y, z). "heading" = yaw-only quaternion.
+
+
+def quat_conj(q):
+    """Quaternion conjugate (inverse for unit quaternions)."""
+    return np.array([q[0], -q[1], -q[2], -q[3]], dtype=np.float32)
+
+
+def quat_mul(q1, q2):
+    """Hamilton product ``q1 ⊗ q2``."""
+    w1, x1, y1, z1 = q1
+    w2, x2, y2, z2 = q2
+    return np.array(
+        [
+            w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
+            w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
+            w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
+            w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
+        ],
+        dtype=np.float32,
+    )
+
+
+def quat_to_6d(q):
+    """Quaternion → 6-D rotation representation (first two rotated basis rows)."""
+    w, x, y, z = q
+    return np.array(
+        [
+            1 - 2 * (y * y + z * z),
+            2 * (x * y - z * w),
+            2 * (x * y + z * w),
+            1 - 2 * (x * x + z * z),
+            2 * (x * z - y * w),
+            2 * (y * z + x * w),
+        ],
+        dtype=np.float32,
+    )
+
+
+def calc_heading(q):
+    """Extract the yaw (heading) angle in radians from a quaternion."""
+    w, x, y, z = q
+    return float(np.arctan2(2 * (x * y + w * z), 1 - 2 * (y * y + z * z)))
+
+
+def heading_quat(q, sign=1.0):
+    """Yaw-only quaternion for ``q``'s heading (``sign=-1`` gives its inverse)."""
+    a = sign * calc_heading(q) / 2.0
+    return np.array([np.cos(a), 0, 0, np.sin(a)], dtype=np.float64)
+
+
+def heading_quat_inv(q):
+    """Inverse yaw-only quaternion for ``q``'s heading."""
+    return heading_quat(q, -1.0)
+
+
+def quat_slerp(q0, q1, t):
+    """Spherical linear interpolation between two quaternions (scalar ``t``)."""
+    q0 = q0 / (np.linalg.norm(q0) + 1e-12)
+    q1 = q1 / (np.linalg.norm(q1) + 1e-12)
+    dot = float(np.dot(q0, q1))
+    if dot < 0:
+        q1, dot = -q1, -dot
+    dot = min(dot, 1.0)
+    if dot > 0.9995:
+        r = q0 + t * (q1 - q0)
+        return r / (np.linalg.norm(r) + 1e-12)
+    th = np.arccos(dot)
+    st = np.sin(th)
+    return (np.sin((1 - t) * th) / st) * q0 + (np.sin(t * th) / st) * q1
+
+
+def quat_slerp_batch(q0, q1, t):
+    """Vectorized slerp over arrays of quaternions with a per-row parameter ``t``."""
+    q0 = q0 / (np.linalg.norm(q0, axis=1, keepdims=True) + 1e-12)
+    q1 = q1 / (np.linalg.norm(q1, axis=1, keepdims=True) + 1e-12)
+    dot = np.sum(q0 * q1, axis=1)
+    neg = dot < 0
+    q1 = q1.copy()
+    q1[neg] = -q1[neg]
+    dot[neg] = -dot[neg]
+    dot = np.clip(dot, -1, 1)
+    lin = dot > 0.9995
+    th = np.arccos(dot)
+    st = np.where(np.sin(th) == 0, 1, np.sin(th))
+    c0 = np.sin((1 - t) * th) / st
+    c1 = np.sin(t * th) / st
+    c0[lin] = 1 - t[lin]
+    c1[lin] = t[lin]
+    r = c0[:, None] * q0 + c1[:, None] * q1
+    return r / (np.linalg.norm(r, axis=1, keepdims=True) + 1e-12)
+
+
+def ort_providers(force_cpu: bool = False) -> list[str]:
+    """Prefer CUDA for enc/dec/planner (matches deploy when onnxruntime-gpu is installed)."""
+    avail = ort.get_available_providers()
+    if not force_cpu and "CUDAExecutionProvider" in avail:
+        return ["CUDAExecutionProvider", "CPUExecutionProvider"]
+    return ["CPUExecutionProvider"]
+
+
+def make_ort_session_options():
+    """Build ONNX Runtime SessionOptions (quiet logging, default threading)."""
+    so = ort.SessionOptions()
+    so.log_severity_level = 3
+    return so
+
+
+# ── Encoder / Decoder ─────────────────────────────────────────────────────────
+
+
+class StandingEncoderDecoder:
+    """Runs the encoder + decoder ONNX models and owns the proprioception history.
+
+    Each tick it appends the latest robot state to 10-frame history buffers, builds
+    the encoder observation (1762-D, layout depends on ``encode_mode``) to refresh
+    the 64-D ``token``, then builds the decoder observation (994-D) and maps
+    ``token + history`` to a residual action added onto ``DEFAULT_ANGLES``.
+
+    ``PlannerController`` subclasses this to source the reference from a live,
+    planner-generated motion buffer instead of a fixed standing pose.
+    """
+
+    def __init__(self, encoder, decoder):
+        self.encoder, self.decoder = encoder, decoder
+        self.encoder_input = encoder.get_inputs()[0].name
+        self.decoder_input = decoder.get_inputs()[0].name
+        enc_dim = int(encoder.get_inputs()[0].shape[1])
+        dec_dim = int(decoder.get_inputs()[0].shape[1])
+        if enc_dim != 1762 or dec_dim != 994:
+            raise RuntimeError(f"Unexpected dims encoder={enc_dim}, decoder={dec_dim}")
+        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
+        self.init_base_quat = np.array([1, 0, 0, 0], np.float32)
+        self.init_ref_quat = np.array([1, 0, 0, 0], np.float32)
+        self._heading_init = False
+        self.encode_mode = 0
+        self.vr_3point_local_target = VR_TARGET_DEF.copy()
+        self.vr_3point_local_orn_target = VR_ORN_DEF.copy()
+        self.smpl_joints_10frame_step1 = SMPL_DEF.copy()
+        # Optional per-frame SMPL root orientation (wxyz) for the mode-2 anchor.
+        # When None, the anchor falls back to the planner reference body quat.
+        self.smpl_root_quat = None
+        self.set_zero_reference()
+
+    def reset(self):
+        """Clear the token, 10-frame proprioception history and heading init.
+
+        ``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
+        self.init_base_quat = np.array([1, 0, 0, 0], np.float32)
+        self.init_ref_quat = np.array([1, 0, 0, 0], np.float32)
+        self._heading_init = False
+
+    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]
+        if not self._heading_init:
+            self.init_base_quat = quat.copy()
+            self._heading_init = True
+
+    def _heading_quat(self, q):
+        h = calc_heading(q) / 2.0
+        return np.array([np.cos(h), 0, 0, np.sin(h)], np.float32)
+
+    def _heading_quat_inv(self, q):
+        h = calc_heading(q) / 2.0
+        return np.array([np.cos(-h), 0, 0, np.sin(-h)], np.float32)
+
+    def _anchor_6d(self, base_quat, ref_quat=None):
+        """6-D orientation error between the robot base and the (heading-aligned) reference."""
+        if ref_quat is None:
+            ref_quat = self.init_ref_quat
+        delta = quat_mul(self._heading_quat(self.init_base_quat), self._heading_quat_inv(self.init_ref_quat))
+        new_ref = quat_mul(delta, ref_quat)
+        return quat_to_6d(quat_mul(quat_conj(base_quat), new_ref))
+
+    def set_zero_reference(self):
+        """Initialize the reference to a single standing frame (used before a plan exists)."""
+        self.motion_joint_positions = [ENCODER_STANDING_REF.copy()]
+        self.motion_joint_velocities = [np.zeros(29, np.float32)]
+        self.motion_body_quats = [np.array([1, 0, 0, 0], np.float32)]
+        self.motion_body_z = [DEFAULT_HEIGHT]
+        self.motion_timesteps = 1
+        self.freeze_ref_frame = 0
+        self.init_ref_quat = self.motion_body_quats[0].copy()
+
+    def build_encoder_obs(self):
+        """Assemble the 1762-D encoder input; slot layout depends on ``encode_mode``.
+
+        mode 0 = locomotion (ref joint pos + anchor), 1 = 3-point VR teleop
+        (lower-body ref + VR targets), 2 = SMPL whole-body window + anchor/wrist.
+        """
+        obs = np.zeros(1762, np.float32)
+        obs[0] = float(self.encode_mode)
+        rf = min(self.freeze_ref_frame, self.motion_timesteps - 1)
+        ref_pos, ref_quat = self.motion_joint_positions[rf], self.motion_body_quats[rf]
+        if self.encode_mode == 0:
+            for f in range(10):
+                obs[4 + 29 * f : 4 + 29 * (f + 1)] = ref_pos
+                obs[601 + 6 * f : 601 + 6 * (f + 1)] = self._anchor_6d(self.h_quat[0], ref_quat)
+        elif self.encode_mode == 1:
+            ref_lower = ref_pos[LOWER_BODY_IL]
+            for f in range(10):
+                obs[661 + 12 * f : 661 + 12 * (f + 1)] = ref_lower
+            obs[901:910] = self.vr_3point_local_target
+            obs[910:922] = self.vr_3point_local_orn_target
+            obs[595:601] = self._anchor_6d(self.h_quat[0], ref_quat)
+        elif self.encode_mode == 2:
+            # Prefer the SMPL clip/stream root orientation for the anchor; fall
+            # back to the planner reference body quat when no root is provided.
+            anchor_ref = self.smpl_root_quat if self.smpl_root_quat is not None else ref_quat
+            obs[922:1642] = self.smpl_joints_10frame_step1
+            for f in range(10):
+                obs[1642 + 6 * f : 1642 + 6 * (f + 1)] = self._anchor_6d(self.h_quat[0], anchor_ref)
+                obs[1702 + 6 * f : 1702 + 6 * (f + 1)] = ref_pos[WRIST_IL]
+        else:
+            raise RuntimeError(f"Unsupported encoder mode: {self.encode_mode}")
+        return obs
+
+    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 run_encoder(self):
+        """Run the encoder ONNX model and return the fresh 64-D token."""
+        return (
+            self.encoder.run(None, {self.encoder_input: self.build_encoder_obs().reshape(1, -1)})[0]
+            .squeeze()
+            .astype(np.float32)
+        )
+
+    def step(self, robot_obs, update_encoder, debug=False):
+        """One control tick: read robot obs, (optionally) re-encode, decode → joint targets.
+
+        Args:
+            robot_obs: dict with ``<joint>.q``/``.dq`` and ``imu.*`` fields.
+            update_encoder: refresh the token this tick (else reuse the cached one).
+            debug: print action/delta norms.
+
+        Returns:
+            dict of ``<joint>.q`` target positions (rad) in IsaacLab joint order.
+        """
+        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)
+        if update_encoder:
+            self.token = self.run_encoder()
+        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 PlannerController(StandingEncoderDecoder):
+    """Encoder/decoder driven by a caller-supplied, rolling motion buffer.
+
+    Extends ``StandingEncoderDecoder`` so the reference comes from a motion buffer
+    (a lookahead window with per-frame velocities) instead of a single fixed pose,
+    and handles heading re-initialization on the first frame / after a reset.
+    ``motion_lock`` guards the buffer, which the whole-body controller rewrites each
+    tick from the incoming command. The class name is retained for continuity with
+    the SONIC reference; no motion planner is involved.
+    """
+
+    def __init__(self, encoder, decoder):
+        super().__init__(encoder, decoder)
+        self.ref_cursor = 0
+        self.motion_timesteps = 0
+        self.motion_joint_positions = np.zeros((1500, 29), np.float64)
+        self.motion_joint_velocities = np.zeros((1500, 29), np.float64)
+        self.motion_body_quats = np.zeros((1500, 4), np.float64)
+        self.motion_body_quats[:, 0] = 1.0
+        self.motion_body_pos = np.zeros((1500, 3), np.float64)
+        self.init_ref_quat = np.array([1, 0, 0, 0], np.float64)
+        self.heading_init_base_quat = np.array([1, 0, 0, 0], np.float64)
+        self.delta_heading = 0.0
+        self.reinit_heading = False
+        self.playing = self.first_motion = False
+        self.motion_lock = threading.Lock()
+
+    def reset(self):
+        """Full reset: clear enc/dec state (super) plus the motion buffer and heading.
+
+        Forces a heading re-init on the next ``step`` so the reference frame is
+        re-latched to the post-reset robot orientation.
+        """
+        super().reset()
+        with self.motion_lock:
+            self.ref_cursor = 0
+            self.motion_timesteps = 0
+            self.motion_joint_positions[:] = 0.0
+            self.motion_joint_velocities[:] = 0.0
+            self.motion_body_quats[:] = 0.0
+            self.motion_body_quats[:, 0] = 1.0
+            self.motion_body_pos[:] = 0.0
+        self.init_ref_quat = np.array([1, 0, 0, 0], np.float64)
+        self.heading_init_base_quat = np.array([1, 0, 0, 0], np.float64)
+        self.delta_heading = 0.0
+        self.first_motion = False
+        self.playing = False
+        self.reinit_heading = True
+
+    def _heading_apply_delta(self):
+        """Heading correction quaternion (init base-vs-ref heading + operator ``delta_heading``)."""
+        delta = quat_mul(
+            heading_quat(self.heading_init_base_quat).astype(np.float32),
+            heading_quat_inv(self.init_ref_quat).astype(np.float32),
+        )
+        if self.delta_heading:
+            h = self.delta_heading / 2.0
+            delta = quat_mul(np.array([np.cos(h), 0, 0, np.sin(h)], np.float32), delta)
+        return delta
+
+    def _anchor_6d(self, base_quat, ref_quat=None):
+        """6-D base-vs-reference orientation error, including the operator heading delta."""
+        if ref_quat is None:
+            ref_quat = self.init_ref_quat
+        new_ref = quat_mul(self._heading_apply_delta(), ref_quat.astype(np.float32))
+        return quat_to_6d(quat_mul(quat_conj(base_quat.astype(np.float32)), new_ref))
+
+    def build_encoder_obs(self):
+        """Encoder input sourced from the live motion buffer (mode 0/2), lock-protected."""
+        obs = np.zeros(1762, np.float32)
+        obs[0] = float(self.encode_mode)
+        with self.motion_lock:
+            if self.encode_mode == 2:
+                # SMPL whole-body imitation: the 720-dim SMPL window carries the
+                # target pose; the planner reference frame supplies anchor + wrist.
+                rf = min(self.ref_cursor, self.motion_timesteps - 1)
+                ref_pos = self.motion_joint_positions[rf].astype(np.float32)
+                ref_quat = self.motion_body_quats[rf].astype(np.float32)
+                # Prefer the SMPL clip/stream root orientation (if provided) so the
+                # anchor tracks the operator's/clip's heading; else planner ref.
+                if self.smpl_root_quat is not None:
+                    ref_quat = np.asarray(self.smpl_root_quat, np.float32)
+                anchor = self._anchor_6d(self.h_quat[0], ref_quat)
+                wrist = ref_pos[WRIST_IL]
+                obs[922:1642] = self.smpl_joints_10frame_step1
+                for f in range(10):
+                    obs[1642 + 6 * f : 1642 + 6 * (f + 1)] = anchor
+                    obs[1702 + 6 * f : 1702 + 6 * (f + 1)] = wrist
+                return obs
+            if self.encode_mode == 1:
+                # 3-point VR teleop: the upper body tracks the VR wrist/neck targets
+                # while the planner reference supplies the lower body + anchor. Lower
+                # body is per-frame (step 5) like mode 0; the VR targets are current.
+                rf = min(self.ref_cursor, self.motion_timesteps - 1)
+                obs[595:601] = self._anchor_6d(self.h_quat[0], self.motion_body_quats[rf].astype(np.float32))
+                for f in range(10):
+                    tf = min(
+                        self.ref_cursor + f * 5 if self.playing else self.ref_cursor,
+                        self.motion_timesteps - 1,
+                    )
+                    ref_lower = self.motion_joint_positions[tf].astype(np.float32)[LOWER_BODY_IL]
+                    obs[661 + 12 * f : 661 + 12 * (f + 1)] = ref_lower
+                obs[901:910] = self.vr_3point_local_target
+                obs[910:922] = self.vr_3point_local_orn_target
+                return obs
+            for f in range(10):
+                tf = min(
+                    self.ref_cursor + f * 5 if self.playing else self.ref_cursor, self.motion_timesteps - 1
+                )
+                obs[4 + 29 * f : 4 + 29 * (f + 1)] = self.motion_joint_positions[tf].astype(np.float32)
+                if self.playing:
+                    obs[294 + 29 * f : 294 + 29 * (f + 1)] = self.motion_joint_velocities[tf].astype(
+                        np.float32
+                    )
+                obs[601 + 6 * f : 601 + 6 * (f + 1)] = self._anchor_6d(
+                    self.h_quat[0], self.motion_body_quats[tf].astype(np.float32)
+                )
+        return obs
+
+    def step(self, robot_obs, update_encoder, debug=False):
+        """Re-init the heading reference on first frame / after a reset, then run the base step."""
+        if robot_obs and (self.first_motion or self.reinit_heading):
+            q = None
+            if "imu.quat.w" in robot_obs:
+                q = np.array(
+                    [
+                        robot_obs["imu.quat.w"],
+                        robot_obs["imu.quat.x"],
+                        robot_obs["imu.quat.y"],
+                        robot_obs["imu.quat.z"],
+                    ],
+                    np.float64,
+                )
+            else:
+                q = robot_obs.get("imu.quaternion")
+                if q is not None:
+                    q = np.array(q, np.float64)
+            if q is not None:
+                self.heading_init_base_quat = np.array(q, np.float64)
+                with self.motion_lock:
+                    rf = min(self.ref_cursor, self.motion_timesteps - 1)
+                    if self.encode_mode == 2 and self.smpl_root_quat is not None:
+                        # Anchor the heading delta to the SMPL root at init so the
+                        # robot turns *relative* to the clip/operator start heading.
+                        self.init_ref_quat = np.asarray(self.smpl_root_quat, np.float64)
+                    else:
+                        self.init_ref_quat = self.motion_body_quats[rf].copy()
+                self.delta_heading = 0.0
+                self.first_motion = False
+                self.reinit_heading = False
+                logger.debug("[Heading] init quat: %s", self.heading_init_base_quat)
+        return super().step(robot_obs, update_encoder=update_encoder, debug=debug)
+
+    def advance_cursor(self):
+        """Advance the reference cursor one frame per 50 Hz tick (no wall-clock catch-up)."""
+        if not self.playing:
+            return
+        with self.motion_lock:
+            if self.motion_timesteps > 0:
+                self.ref_cursor = min(self.ref_cursor + 1, self.motion_timesteps - 1)

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

@@ -0,0 +1,304 @@
+#!/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 full-body controller for Unitree G1."""
+
+from __future__ import annotations
+
+import logging
+from collections import deque
+from typing import TYPE_CHECKING
+
+import numpy as np
+from huggingface_hub import hf_hub_download
+
+from lerobot.utils.import_utils import _onnxruntime_available, require_package
+
+from ..g1_utils import (
+    MUJOCO_TO_ISAACLAB,
+    WB_ACTION_DIM,
+    G1_29_JointIndex,
+    lowstate_to_obs,
+    wb_action_key,
+)
+from .sonic_pipeline import (
+    CONTROL_DT,
+    DEFAULT_ANGLES,
+    ENCODER_UPDATE_EVERY,
+    PlannerController,
+    compute_kp_kd,
+    make_ort_session_options,
+    ort_providers,
+)
+
+if TYPE_CHECKING or _onnxruntime_available:
+    import onnxruntime as ort
+else:
+    ort = None
+
+logger = logging.getLogger(__name__)
+
+# Startup blend duration: over the first control ticks, linearly interpolate every joint
+# from the robot's initial measured pose into the policy's commanded target, so control
+# eases in without a snap on the first command.
+INIT_RAMP_S = 3.0
+
+
+def _extract_wb34_from_action(action: dict | None) -> np.ndarray | None:
+    """Reassemble a dense (34,) whole-body command from ``wb.{i}.pos`` keys, or None.
+
+    This is the OpenHLM / pi0.5 joint-based interface: one 34-D vector per tick
+    (sentinel: presence of ``wb.0.pos``) carrying absolute joint targets in real
+    units. The ``.pos`` suffix lets these flow through ``lerobot-rollout`` as normal
+    joint-position action features.
+    """
+    if not action:
+        return None
+    keys = [wb_action_key(i) for i in range(WB_ACTION_DIM)]
+    # Require the full dense command: a partial action (e.g. only ``wb.0.pos``)
+    # must not be silently zero-filled, which would drive most joints toward 0.
+    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=WB_ACTION_DIM,
+    )
+
+
+def _wb34_to_reference(wb: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Map a 34-D OpenHLM whole-body command to a SONIC mode-0 reference.
+
+    Returns ``(ref29, anchor_quat)`` where ``ref29`` is the 29 joint targets in
+    IsaacLab order (what SONIC's ``motion_joint_positions`` expects) and
+    ``anchor_quat`` (wxyz) encodes the root roll/pitch (yaw=0).
+
+    OpenHLM layout : [L-arm 0:7, L-grip 7, R-arm 8:15, R-grip 15,
+                      L-leg 16:22, R-leg 22:28, waist 28:31, root rp+yaw 31:34]
+    The 29 joints are first assembled in MuJoCo / Unitree-SDK order
+    ([L-leg 0:6, R-leg 6:12, waist 12:15, L-arm 15:22, R-arm 22:29] — the
+    ``G1_29_JointIndex`` grouping OpenHLM uses), then permuted to IsaacLab order via
+    ``MUJOCO_TO_ISAACLAB``. Grippers (7, 15) are not part of the 29-DoF SONIC
+    reference, and yaw-rate (33) is integrated into the heading by the caller (it
+    cannot be represented in this static per-tick anchor).
+    """
+    ref_mj = np.zeros(29, np.float32)  # MuJoCo / Unitree-SDK grouped order
+    ref_mj[0:6] = wb[16:22]  # left leg
+    ref_mj[6:12] = wb[22:28]  # right leg
+    ref_mj[12:15] = wb[28:31]  # waist
+    ref_mj[15:22] = wb[0:7]  # left arm
+    ref_mj[22:29] = wb[8:15]  # right arm
+    ref = ref_mj[MUJOCO_TO_ISAACLAB].astype(np.float32)  # -> IsaacLab order for SONIC
+    roll, pitch = float(wb[31]), float(wb[32])
+    cr, sr, cp, sp = np.cos(roll / 2), np.sin(roll / 2), np.cos(pitch / 2), np.sin(pitch / 2)
+    anchor = np.array([cr * cp, sr * cp, cr * sp, sr * sp], np.float32)  # Rx(roll)·Ry(pitch)
+    return ref, anchor
+
+
+class SonicRuntime:
+    """Loads the SONIC encoder/decoder ONNX models and owns the controller.
+
+    No motion planner: the reference motion buffer is written directly each tick by
+    :class:`SonicWholeBodyController` from the incoming 34-D whole-body command.
+    """
+
+    def __init__(self, force_cpu: bool = False):
+        require_package("onnxruntime", extra="unitree_g1")
+        encoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_encoder.onnx")
+        decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx")
+
+        providers = ort_providers(force_cpu=force_cpu)
+        self.use_gpu = providers[0] == "CUDAExecutionProvider"
+        so = make_ort_session_options()
+
+        encoder_sess = ort.InferenceSession(encoder_path, sess_options=so, providers=providers)
+        decoder_sess = ort.InferenceSession(decoder_path, sess_options=so, providers=providers)
+
+        self.kp, self.kd = compute_kp_kd()
+        self.controller = PlannerController(encoder_sess, decoder_sess)
+
+    @property
+    def pipeline(self):
+        return self.controller
+
+    def reset(self):
+        # Full pipeline reset: clears the encoder token, proprioception history and
+        # heading, and rewinds the motion buffer. reinit_heading is set so the next
+        # step re-latches the reference frame to the current robot orientation.
+        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
+    # Advertise a dense 34-D whole-body action space (OpenHLM / pi0.5) so the robot
+    # exposes ``wb.{i}.pos`` action features and ``lerobot-rollout`` can drive it
+    # directly with a 34-D VLA policy.
+    wb_action = True
+
+    def __init__(self, force_cpu: bool = False):
+        logger.info("Loading SONIC whole-body controller...")
+        self._runtime = SonicRuntime(force_cpu=force_cpu)
+        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] = {}
+
+        # Tick counter for the dense whole-body (OpenHLM, mode-0) path's encoder cadence.
+        self._wb_step = 0
+        # Rolling 50-frame reference trajectory (ref29 + anchor quat) built from the
+        # stream of per-tick whole-body commands, fed to the encoder as a batch.
+        self._wb_traj: deque[np.ndarray] = deque(maxlen=50)
+        self._wb_quat_traj: deque[np.ndarray] = deque(maxlen=50)
+        # Integrated heading (rad) from the whole-body command's yaw-rate (index 33),
+        # forwarded to the pipeline as ``delta_heading`` so turn commands take effect.
+        self._heading = 0.0
+
+        logger.info("SONIC ready (encoder/decoder, 34-D whole-body command path)")
+
+    def _run_wholebody34(self, obs: dict, wb: np.ndarray) -> dict:
+        """Feed a dense 34-D OpenHLM whole-body command as the mode-0 encoder reference.
+
+        The 29 joint targets are held across the encoder lookahead window (zero
+        velocity) and the root roll/pitch set the anchor orientation, then the
+        encoder/decoder run directly (planner bypassed). One command per tick, so the
+        VLA's commanded pose is what SONIC tracks.
+        """
+        ref, anchor = _wb34_to_reference(wb)
+        c = self.controller
+        if c.encode_mode != 0:
+            c.encode_mode = 0
+            c.reinit_heading = True
+        # Index 33 is a yaw-rate (rad/s): integrate it into a heading offset and hand
+        # it to the pipeline as ``delta_heading`` so commanded turns are tracked rather
+        # than silently dropped (the anchor from _wb34_to_reference only carries r/p).
+        self._heading += float(wb[33]) * CONTROL_DT
+        c.delta_heading = self._heading
+        # Capture the heading/anchor reference on the first whole-body tick. The
+        # controller only latches ``init_ref_quat`` (and the base heading) inside
+        # ``step()`` when ``first_motion or reinit_heading`` — but it already boots in
+        # mode 0, so the mode-switch guard above misses the very first command and the
+        # anchor would stay identity. This mirrors the GEAR reference, which seeds
+        # ``init_ref_quat`` from the first anchor. Must run before the buffers below so
+        # ``step()`` latches ``motion_body_quats[0]`` = this tick's anchor.
+        if self._wb_step == 0:
+            c.reinit_heading = True
+
+        # Accumulate the per-tick commands into a rolling 50-frame reference
+        # trajectory so the encoder's 10-frame, step-5 lookahead sees an actual
+        # motion sequence (with velocities) instead of one repeated pose. 50 frames
+        # == chunk horizon == 10 lookahead frames × step 5.
+        self._wb_traj.append(ref)
+        self._wb_quat_traj.append(anchor)
+        traj = np.asarray(self._wb_traj, np.float32)  # (L, 29), oldest -> newest
+        quats = np.asarray(self._wb_quat_traj, np.float32)  # (L, 4)
+        n = len(traj)
+        # Per-frame velocities from finite differences (rad/s at the control rate).
+        vel = np.zeros_like(traj)
+        if n > 1:
+            vel[1:] = (traj[1:] - traj[:-1]) / CONTROL_DT
+            vel[0] = vel[1]
+        with c.motion_lock:
+            c.motion_joint_positions[:n] = traj
+            c.motion_joint_velocities[:n] = vel
+            c.motion_body_quats[:n] = quats
+            c.motion_body_pos[:n] = 0.0
+            c.motion_timesteps = n
+            c.ref_cursor = 0
+        c.playing = True
+        do_enc = self._wb_step % ENCODER_UPDATE_EVERY == 0
+        out = c.step(obs, update_encoder=do_enc, debug=False)
+        if self._wb_step % 25 == 0:
+            tgt = np.array([out[f"{m.name}.q"] for m in G1_29_JointIndex], np.float32)
+            logger.info(
+                "[WB34] step=%d |ref|mean=%.3f |target|mean=%.3f target_std=%.3f init_ref_quat=%s",
+                self._wb_step,
+                float(np.abs(ref).mean()),
+                float(np.abs(tgt).mean()),
+                float(tgt.std()),
+                np.round(c.init_ref_quat, 3).tolist(),
+            )
+        self._wb_step += 1
+        return out
+
+    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)
+
+        # Dense 34-D whole-body command (OpenHLM / pi0.5 joint interface): a single
+        # vector per tick drives the mode-0 encoder reference directly. Until the
+        # policy produces one, hold (no command) so the robot keeps its last target.
+        wb = _extract_wb34_from_action(action)
+        if wb is None:
+            self._wb_miss = getattr(self, "_wb_miss", 0) + 1
+            if self._wb_miss % 50 == 1:
+                akeys = [k for k in action if isinstance(k, str)]
+                logger.info(
+                    "[WB34] no wb.*.pos in action this tick (miss=%d). action keys sample: %s",
+                    self._wb_miss,
+                    akeys[:8],
+                )
+            return {}
+        return self._startup_blend(obs, self._run_wholebody34(obs, wb))
+
+    def reset(self):
+        self._runtime.reset()
+        self._init_step = 0  # re-run the startup blend after a reset
+        self._start_pose = {}
+        self._wb_step = 0
+        self._wb_traj.clear()
+        self._wb_quat_traj.clear()
+        self._heading = 0.0
+
+    def shutdown(self):
+        self._runtime.shutdown()

src/lerobot/robots/unitree_g1/g1_utils.py

@@ -23,10 +23,102 @@ 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
 
+# Reserved action-dict field used to forward the set of currently-pressed keyboard
+# keys from a KeyboardTeleop through the standard action pipeline to the SONIC
+# whole-body controller (see SonicWholeBodyController._process_keyboard).
+KEYBOARD_KEYS_FIELD = "keyboard.keys"
+
+# ── Dense whole-body joint reference (SONIC encode_mode 0, OpenHLM / pi0.5) ──────
+# A single 34-D whole-body command per tick, in the OpenHLM action layout:
+#   [L-arm(7), L-grip(1), R-arm(7), R-grip(1), L-leg(6), R-leg(6), waist(3),
+#    root roll/pitch + yaw-rate(3)]
+# Fed as flat scalars ``wb.0.pos .. wb.33.pos``. The ``.pos`` suffix makes these
+# behave like ordinary joint-position action features so ``lerobot-rollout`` routes
+# them straight from a 34-D VLA (OpenHLM / pi0.5) onto the robot.
+WB_ACTION_PREFIX = "wb."
+WB_ACTION_DIM = 34
+
+
+def wb_action_key(i: int) -> str:
+    """Action-dict key for the ``i``-th whole-body command scalar (``wb.{i}.pos``)."""
+    return f"{WB_ACTION_PREFIX}{i}.pos"
+
 
 def default_remote_input() -> dict[str, float]:
     """Return a zeroed-out remote input dict (axes + buttons)."""
@@ -63,13 +155,92 @@ 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 obs_to_wb34_state(obs: dict) -> np.ndarray:
+    """Build the 34-D OpenHLM / pi0.5 proprio state from a G1 observation dict.
+
+    Mirrors the whole-body *action* layout so the policy sees state and action in
+    the same coordinates::
+
+        [L-arm(7), L-grip(1), R-arm(7), R-grip(1),
+         L-leg(6), R-leg(6), waist(3), root roll/pitch + yaw-rate(3)]
+
+    Joint positions come from the ``<joint>.q`` obs keys, which are already in
+    MuJoCo / Unitree-SDK order — the same body-part grouping OpenHLM uses
+    ([L-leg 0:6, R-leg 6:12, waist 12:15, L-arm 15:22, R-arm 22:29]) — so they are
+    regrouped directly (no IsaacLab permutation). The G1 has no grippers in its
+    29-DoF body, so both gripper slots are 0. Root roll/pitch are the IMU RPY and
+    the last slot is the IMU yaw rate (gyro z).
+    """
+    q_mj = np.array(
+        [float(obs.get(f"{m.name}.q", 0.0)) for m in G1_29_JointIndex],
+        dtype=np.float32,
+    )
+    lleg, rleg, waist = q_mj[0:6], q_mj[6:12], q_mj[12:15]
+    larm, rarm = q_mj[15:22], q_mj[22:29]
+
+    state = np.zeros(34, dtype=np.float32)
+    state[0:7] = larm
+    # state[7] left gripper — none on 29-DoF G1
+    state[8:15] = rarm
+    # state[15] right gripper — none on 29-DoF G1
+    state[16:22] = lleg
+    state[22:28] = rleg
+    state[28:31] = waist
+    state[31] = float(obs.get("imu.rpy.roll", 0.0))
+    state[32] = float(obs.get("imu.rpy.pitch", 0.0))
+    state[33] = float(obs.get("imu.gyro.z", 0.0))
+    return state
+
+
 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.gr00t_locomotion",
-        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
+        "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",
     }
     module_path = controllers.get(name)
     if module_path is None:

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -33,12 +33,14 @@ from ..robot import Robot
 from .config_unitree_g1 import UnitreeG1Config
 from .g1_kinematics import G1_29_ArmIK
 from .g1_utils import (
+    KEYBOARD_KEYS_FIELD,
     REMOTE_AXES,
-    REMOTE_KEYS,
     G1_29_JointArmIndex,
     G1_29_JointIndex,
     default_remote_input,
+    lowstate_to_obs,
     make_locomotion_controller,
+    obs_to_wb34_state,
 )
 
 if TYPE_CHECKING or _unitree_sdk_available:
@@ -47,8 +49,12 @@ 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__HandCmd_ as hg_HandCmd_default,
+        unitree_hg_msg_dds__LowCmd_,
+    )
     from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
+        HandCmd_ as hg_HandCmd,
         LowCmd_ as hg_LowCmd,
         LowState_ as hg_LowState,
     )
@@ -58,6 +64,8 @@ else:
     _SDKChannelPublisher = None
     _SDKChannelSubscriber = None
     unitree_hg_msg_dds__LowCmd_ = None
+    hg_HandCmd_default = None
+    hg_HandCmd = None
     hg_LowCmd = None
     hg_LowState = None
     CRC = None
@@ -153,10 +161,62 @@ class UnitreeG1(Robot):
 
         # 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 = {}
 
+        # Replay-camera state: keep the encoded (raw) cells per camera and decode
+        # frames lazily as the play cursor advances, with a small frame cache, so we
+        # don't materialize gigabytes of decoded RGB at construction time.
+        self._replay_raw: dict[str, list] = {}
+        self._replay_cache: dict[tuple[str, int], np.ndarray] = {}
+        self._replay_cache_cap = 8
+        self._replay_len = 0
+        self._replay_idx = 0
+        if config.replay_camera_parquet and config.replay_camera_map:
+            self._load_replay_frames()
+
+    def _load_replay_frames(self) -> None:
+        """Load only the mapped parquet columns (encoded frames); decode on demand."""
+        import pyarrow.parquet as pq
+
+        cols_needed = list(dict.fromkeys(self.config.replay_camera_map.values()))
+        table = pq.read_table(self.config.replay_camera_parquet, columns=cols_needed)
+        self._replay_len = table.num_rows
+        self._replay_raw = {
+            cam_name: table.column(column).to_pylist()
+            for cam_name, column in self.config.replay_camera_map.items()
+        }
+        logger.info(
+            "Loaded %d replay frames (lazy-decode) for cameras %s from %s",
+            self._replay_len,
+            list(self.config.replay_camera_map),
+            self.config.replay_camera_parquet,
+        )
+
+    def _decode_replay_cell(self, cell) -> np.ndarray:
+        import io
+
+        from PIL import Image
+
+        data = cell["bytes"] if isinstance(cell, dict) else cell
+        return np.asarray(Image.open(io.BytesIO(data)).convert("RGB"), dtype=np.uint8)
+
+    def _replay_frame(self, cam_name: str, idx: int) -> np.ndarray:
+        """Decode (and briefly cache) a single replay frame for a camera."""
+        key = (cam_name, idx)
+        cached = self._replay_cache.get(key)
+        if cached is not None:
+            return cached
+        frame = self._decode_replay_cell(self._replay_raw[cam_name][idx])
+        if len(self._replay_cache) >= self._replay_cache_cap:
+            self._replay_cache.pop(next(iter(self._replay_cache)))
+        self._replay_cache[key] = frame
+        return frame
+
     def _subscribe_lowstate(self):  # polls robot state @ 250Hz
         while not self._shutdown_event.is_set():
             start_time = time.time()
@@ -231,15 +291,56 @@ class UnitreeG1(Robot):
                 features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
         return features
 
+    @property
+    def _wb_state_ft(self) -> dict[str, type]:
+        """34-D whole-body proprio state (``wb_state.{i}.pos``) for dense controllers.
+
+        Exposed only when the controller consumes a dense whole-body command
+        (OpenHLM / pi0.5). These ``.pos`` scalars are aggregated by the rollout
+        pipeline into a single 34-D ``observation.state`` for the policy.
+        """
+        if not getattr(self.controller, "wb_action", False):
+            return {}
+        from .g1_utils import WB_ACTION_DIM
+
+        return {f"wb_state.{i}.pos": float for i in range(WB_ACTION_DIM)}
+
+    @property
+    def _empty_cameras_ft(self) -> dict[str, tuple]:
+        """Synthetic zero-image cameras (see ``UnitreeG1Config.empty_cameras``)."""
+        h, w = self.config.empty_camera_hw
+        return dict.fromkeys(self.config.empty_cameras, (h, w, 3))
+
+    @property
+    def _replay_cameras_ft(self) -> dict[str, tuple]:
+        """Replay cameras, shaped from their first (lazily decoded) frame."""
+        if not self._replay_len:
+            return {}
+        return {name: self._replay_frame(name, 0).shape for name in self._replay_raw}
+
     @cached_property
     def observation_features(self) -> dict[str, type | tuple]:
-        return {**self._motors_ft, **self._cameras_ft}
+        return {
+            **self._motors_ft,
+            **self._wb_state_ft,
+            **self._empty_cameras_ft,
+            **self._replay_cameras_ft,
+            **self._cameras_ft,
+        }
 
     @cached_property
     def action_features(self) -> dict[str, type]:
         if self.controller is None:
             return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
 
+        # Dense whole-body controllers (SONIC / OpenHLM, pi0.5) consume a single
+        # 34-D command per tick. Expose it as ``wb.{i}.pos`` joint-position features
+        # so ``lerobot-rollout`` maps a 34-D policy output straight onto the robot.
+        if getattr(self.controller, "wb_action", False):
+            from .g1_utils import WB_ACTION_DIM, wb_action_key
+
+            return {wb_action_key(i): float for i in range(WB_ACTION_DIM)}
+
         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}
@@ -255,6 +356,11 @@ 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
 
@@ -311,6 +417,17 @@ class UnitreeG1(Robot):
         self.lowstate_subscriber = self._ChannelSubscriber(kTopicLowState, hg_LowState)
         self.lowstate_subscriber.Init()
 
+        # Dex3 hand command publishers (grasping). Driven by the OpenHLM grip scalars.
+        self._hand_publishers = {}
+        if self.config.publish_hands:
+            self._left_hand_cmd = hg_HandCmd_default()
+            self._right_hand_cmd = hg_HandCmd_default()
+            self._hand_publishers["left"] = self._ChannelPublisher("rt/dex3/left/cmd", hg_HandCmd)
+            self._hand_publishers["right"] = self._ChannelPublisher("rt/dex3/right/cmd", hg_HandCmd)
+            for pub in self._hand_publishers.values():
+                pub.Init()
+            logger.info("Dex3 hand command publishers initialized (rt/dex3/{left,right}/cmd)")
+
         # Start subscribe thread to read robot state
         self.subscribe_thread = threading.Thread(target=self._subscribe_lowstate)
         self.subscribe_thread.start()
@@ -343,6 +460,9 @@ 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
@@ -371,13 +491,59 @@ class UnitreeG1(Robot):
         except Exception as e:
             logger.warning(f"Failed to send zero-torque on disconnect: {e}")
 
-    def disconnect(self):
-        # Put robot in passive mode before stopping threads
-        if not self.config.is_simulation:
-            self._send_zero_torque()
+    def _graceful_stop(self) -> None:
+        """Soft shutdown: hold the current pose and ramp joint stiffness (kp) to zero
+        over ``graceful_stop_s`` while keeping damping (kd), then go passive.
 
-        # Signal thread to stop and unblock any waits
+        Prevents the robot from collapsing the instant control ends (a bare
+        zero-torque command is kp=kd=0 ≈ free-fall). Must run after the controller
+        loop has stopped so the two aren't publishing at once.
+        """
+        if self.config.graceful_stop_s <= 0:
+            self._send_zero_torque()
+            return
+        with self._lowstate_lock:
+            lowstate = self._lowstate
+        if lowstate is None:
+            self._send_zero_torque()
+            return
+        q_hold = {f"{motor.name}.q": lowstate.motor_state[motor.value].q for motor in G1_29_JointIndex}
+        kp = np.array(self.kp, dtype=np.float32)
+        kd = np.array(self.kd, dtype=np.float32)
+        zeros = np.zeros(29, dtype=np.float32)
+        dt = self.controller.control_dt if self.controller is not None else self.config.control_dt
+        steps = max(1, int(self.config.graceful_stop_s / dt))
+        logger.info("Graceful stop: damping down over %.1fs", self.config.graceful_stop_s)
+        for i in range(steps):
+            ratio = (i + 1) / steps
+            self.publish_lowcmd(q_hold, kp=kp * (1.0 - ratio), kd=kd, tau=zeros)
+            time.sleep(dt)
+        self._send_zero_torque()
+
+    def disconnect(self):
+        # 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)"
+                )
+
+        # Soft, damped settle instead of an instant limp (real robot only; the
+        # subscribe thread is still alive here to supply the current pose). Only ramp
+        # once the controller thread has definitely exited.
+        if not self.config.is_simulation and controller_stopped:
+            self._graceful_stop()
+
+        if self.controller is not None and hasattr(self.controller, "shutdown"):
+            self.controller.shutdown()
 
         # Wait for subscribe thread to finish
         if self.subscribe_thread is not None:
@@ -385,12 +551,6 @@ 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:
@@ -422,44 +582,33 @@ class UnitreeG1(Robot):
         if lowstate is None:
             return {}
 
-        obs = {}
+        # Motors + IMU + wireless remote (shared lowstate -> obs mapping)
+        obs = lowstate_to_obs(lowstate)
 
-        # 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
+        # Dense whole-body controllers (OpenHLM / pi0.5): expose the 34-D proprio
+        # state as ``wb_state.{i}.pos`` so the rollout aggregates it into
+        # ``observation.state`` for the policy.
+        if getattr(self.controller, "wb_action", False):
+            wb_state = obs_to_wb34_state(obs)
+            for i, v in enumerate(wb_state):
+                obs[f"wb_state.{i}.pos"] = 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]
+        # Synthetic empty cameras: black frames so image-conditioned policies run
+        # before real cameras are wired.
+        if self.config.empty_cameras:
+            h, w = self.config.empty_camera_hw
+            black = np.zeros((h, w, 3), dtype=np.uint8)
+            for name in self.config.empty_cameras:
+                obs[name] = black
 
-        # 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
+        # Replay cameras: serve the current recorded frame per camera, then advance.
+        if self._replay_len:
+            idx = self._replay_idx
+            if idx >= self._replay_len:
+                idx = self._replay_len - 1 if not self.config.replay_camera_loop else idx % self._replay_len
+            for name in self._replay_raw:
+                obs[name] = self._replay_frame(name, idx)
+            self._replay_idx += 1
 
         # Cameras - read images from ZMQ cameras
         for cam_name, cam in self._cameras.items():
@@ -473,9 +622,13 @@ class UnitreeG1(Robot):
     def send_action(self, action: RobotAction) -> RobotAction:
         action_to_publish = action
         if self.controller is not None:
+            self._update_controller_action(action)
+            if self.config.publish_hands and getattr(self.controller, "wb_action", False):
+                self._publish_hand_cmds(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
@@ -503,11 +656,67 @@ class UnitreeG1(Robot):
         return action
 
     def _update_controller_action(self, action: RobotAction) -> None:
-        """Update controller input state from incoming teleop action."""
+        """Update controller input state from an incoming teleop action.
+
+        Controller-agnostic: every value-carrying key is forwarded verbatim into
+        ``controller_input`` (whole-body ``wb.{i}.pos`` from a 34-D VLA, or whatever a
+        future controller expects), and each controller extracts only the keys it
+        understands. The robot deliberately does not enumerate any controller's key
+        schema here.
+
+        KeyboardTeleop is the one special case: it emits the currently-pressed keys as
+        bare action keys with a ``None`` value (``dict.fromkeys(pressed, None)``), so
+        those are collected into a single held-key set under ``KEYBOARD_KEYS_FIELD``,
+        rebuilt each tick so releases clear. Special keys arrive as pynput objects and
+        are normalised to their name ("space", ...).
+        """
         with self._controller_action_lock:
-            for key in REMOTE_KEYS:
-                if key in action:
-                    self.controller_input[key] = action[key]
+            self.controller_input[KEYBOARD_KEYS_FIELD] = {
+                (k if isinstance(k, str) else getattr(k, "name", str(k)))
+                for k, value in action.items()
+                if value is None
+            }
+            for key, value in action.items():
+                if isinstance(key, str) and value is not None:
+                    self.controller_input[key] = value
+
+    def _publish_hand_cmds(self, action: RobotAction) -> None:
+        """Drive the Dex3 hands from the OpenHLM grip scalars in a 34-D wb action.
+
+        ``wb.7.pos`` is the left grip and ``wb.15.pos`` the right grip. Each scalar in
+        [0, 1] (``hand_open_grip_value`` == fully open) is turned into a curl amount and
+        scaled onto ``hand_closed_pose`` (7 joints), then published as a PD target on
+        ``rt/dex3/{left,right}/cmd`` so the fingers close when the policy grips.
+        """
+        if not self._hand_publishers:
+            return
+        from .g1_utils import wb_action_key
+
+        open_val = float(self.config.hand_open_grip_value)
+        closed_val = float(self.config.hand_closed_grip_value)
+        closed_pose = self.config.hand_closed_pose
+        kp, kd = float(self.config.hand_kp), float(self.config.hand_kd)
+        span = (closed_val - open_val) or 1.0
+
+        def curl_amount(grip: float) -> float:
+            # Fraction of the way from the open scalar to the closed scalar, in [0, 1].
+            return float(min(max((grip - open_val) / span, 0.0), 1.0))
+
+        for side, grip_idx, cmd in (
+            ("left", 7, self._left_hand_cmd),
+            ("right", 15, self._right_hand_cmd),
+        ):
+            grip = action.get(wb_action_key(grip_idx))
+            if grip is None:
+                continue
+            amount = curl_amount(float(grip))
+            for i, closed_q in enumerate(closed_pose):
+                cmd.motor_cmd[i].q = float(closed_q) * amount
+                cmd.motor_cmd[i].dq = 0.0
+                cmd.motor_cmd[i].kp = kp
+                cmd.motor_cmd[i].kd = kd
+                cmd.motor_cmd[i].tau = 0.0
+            self._hand_publishers[side].Write(cmd)
 
     @property
     def is_calibrated(self) -> bool:
@@ -537,43 +746,64 @@ class UnitreeG1(Robot):
         if default_positions is None:
             default_positions = np.array(self.config.default_positions, dtype=np.float32)
 
-        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)
+        # 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
 
-            # get current state
-            obs = self.get_observation()
+        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)
 
-            # 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"]
+                # get current state
+                obs = self.get_observation()
 
-            # Interpolate to default position
-            for step in range(num_steps):
-                start_time = time.time()
-
-                alpha = step / num_steps
-                action_dict = {}
+                # record current positions
+                init_dof_pos = np.zeros(29, dtype=np.float32)
                 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)
+                    init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
 
-                self.send_action(action_dict)
+                # Interpolate to default position
+                for step in range(num_steps):
+                    start_time = time.time()
 
-                # Maintain constant control rate
-                elapsed = time.time() - start_time
-                sleep_time = max(0, control_dt - elapsed)
-                time.sleep(sleep_time)
+                    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)
 
-        # Reset controller internal state (gait phase, obs history, etc.)
-        if self.controller is not None and hasattr(self.controller, "reset"):
-            self.controller.reset()
+                    # 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()
 
         logger.info("Reached default position")

src/lerobot/utils/import_utils.py

@@ -60,8 +60,18 @@ def is_package_available(
                     # If the package can't be imported, it's not available
                     package_exists = False
             else:
-                # For packages other than "torch", don't attempt the fallback and set as not available
-                package_exists = False
+                # The distribution may be published under a name that differs from the
+                # import name (e.g. ``onnxruntime`` imports from ``onnxruntime-gpu`` /
+                # ``onnxruntime-silicon``). Resolve the import name to its actual
+                # distribution(s) and read the version from there before giving up.
+                try:
+                    dists = importlib.metadata.packages_distributions().get(import_name, [])
+                    if dists:
+                        package_version = importlib.metadata.version(dists[0])
+                    else:
+                        package_exists = False
+                except importlib.metadata.PackageNotFoundError:
+                    package_exists = False
         logging.debug(f"Detected {pkg_name} version: {package_version}")
     if return_version:
         return package_exists, package_version
@@ -123,6 +133,8 @@ _pyrealsense2_available = is_package_available("pyrealsense2") or is_package_ava
     "pyrealsense2-macosx", import_name="pyrealsense2"
 )
 _zmq_available = is_package_available("pyzmq", import_name="zmq")
+_onnxruntime_available = is_package_available("onnxruntime")
+_onnx_available = is_package_available("onnx")
 _hebi_available = is_package_available("hebi-py", import_name="hebi")
 _teleop_available = is_package_available("teleop")
 _placo_available = is_package_available("placo")