admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-29 20:49:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(unitree_g1): minimize diff w.r.t main

Browse Source
This commit is contained in:
Martino Russi
2026-07-29 15:46:17 +02:00
parent 77259f436e
commit af163fd032
13 changed files with 1073 additions and 1770 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cam_test.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 167 KiB

docs/PR3827_REVIEW_COMMENTS.md
+133
View File
@@ -0,0 +1,133 @@
# PR #3827 — `feat(unitree_g1): pySONIC wbc` — Review Comments
Reviewer: **CarolinePascal** (Collaborator)
Branch: `feat/unitree_g1_sonic_rebased` → `main`
Status legend: [ ] open · [x] done · [~] partially done / needs reply
---
## 1. `controllers/gr00t_locomotion.py` & `holosoma_locomotion.py` — import style
> "It's actually better to have local imports."
- **Ask:** Use relative (local) imports for sibling/parent-package modules instead of the
full absolute `from lerobot.robots.unitree_g1.g1_utils import ...` path.
- **Resolution:** [x] Done — switched to `from ..g1_utils import (...)` (files live in
the `controllers/` subpackage). Pushed in commit "fix relative imports".
## 2. `controllers/sonic_whole_body.py` — private imports across files
> "These are private, they should technically not be imported outside their original file."
> (re: `_ort_providers`, `_snapshot_ms` imported from `sonic_pipeline`)
- **Ask:** Don't import underscore-prefixed (private) symbols outside their defining module.
- **Resolution:** [x] Done — renamed `_ort_providers` → `ort_providers` and
`_snapshot_ms` → `snapshot_ms` in `sonic_pipeline.py`; updated all call sites.
## 3. `controllers/sonic_whole_body.py` / `sonic_pipeline.py` — safe optional imports
> "This should be a safe import with require_package." (re: `import onnxruntime as ort`)
- **Ask:** Guard optional heavy deps (`onnxruntime`, and `onnx` in holosoma) the way the
rest of the repo does, with `require_package` + an availability flag.
- **Resolution:** [x] Done — added `_onnxruntime_available` / `_onnx_available` flags in
`import_utils.py`, guarded the imports (`if TYPE_CHECKING or _..._available: import ...`),
and call `require_package("onnxruntime", extra="unitree_g1")` (and `onnx`) in the
controllers' `__init__`.
## 4. `controllers/__init__.py` (line 17)
> (comment on the module docstring / `__all__` block)
- **Status:** [x] Marked **Resolved** on GitHub.
---
## 5. `controllers/sonic_pipeline.py` — whole-file readability
> "This file is dense and very technical, which makes it hard to read. Could you add some
> comments and docstrings so we can clearly identify what is each object's purpose?"
- **Ask:** Add module-level + class/function docstrings across the 1,297-line pipeline.
- **Resolution:** [x] Done — added an architecture module docstring (planner subprocess ↔
encoder/decoder ↔ controller data flow, `encode_mode` meanings, IsaacLab↔MuJoCo joint
ordering) plus docstrings on every class (`StandingEncoderDecoder`, `SonicPlanner`,
`PlannerController`, `MovementState`, etc.) and the key functions. No behavior change.
Committed as `806d28a8`.
---
## 6. `controllers/sonic_whole_body.py` (line 72) — `f"{SMPL_ACTION_PREFIX}0"`
> "Is the 0 intended?"
- **Ask:** Clarify whether checking the literal key `smpl.0` is intentional.
- **Answer:** Yes — it's a sentinel: the presence of the first element (`smpl.0`) signals
that a full 720-element SMPL window was sent in this action; if absent, there's no SMPL
reference this tick and the function returns `None`.
- **Action:** [ ] Reply to the thread and/or add a clarifying inline comment.
## 7. `smpl_stream.py` (lines 43–45) — `DEFAULT_SMPL_HOST` / `DEFAULT_SMPL_PORT`
> "This should live in the Pico config, right?"
- **Ask:** Move the default host/port into the Pico headset config.
- **Notes:** They already exist in `config_pico_headset.py` (`smpl_host`, `smpl_port`).
The module constants remain as fallbacks for the non-teleop paths
(`sonic_whole_body`'s `SONIC_SMPL_STREAM` direct subscription, and `pico_publisher`).
- **Action:** [ ] Either reply clarifying the config already owns these, or fold the
defaults into the shared constants module (see #8).
## 8. `smpl_stream.py` (lines 47–51) — `WINDOW`, `N_JOINTS`, `JOINT_DIM`, `SMPL_OBS_DIM`
> "These variables are used in other files, consider putting them in a shared 'constant'
> file so that we have a single source of truth (probably in the unitree folder)."
- **Ask:** Deduplicate the SMPL geometry constants (also re-derived as `SMPL_ACTION_DIM = 720`
in `sonic_whole_body.py`, and referenced in `pico_headset.py` / `smpl_to_dataset.py`).
- **Action:** [ ] Create a shared constants module (e.g.
`robots/unitree_g1/smpl_constants.py` or add to `g1_utils.py`) and import everywhere.
Can also absorb the host/port defaults from #7.
## 9. `smpl_stream.py` (line 53) — `class SmplStream`
> "Could be a Dataclass"
- **Ask:** Consider making `SmplStream` a `@dataclass`.
- **Notes:** Only the ctor args (`host/port/fps/stale_after_s/loop`) are dataclass-friendly;
most of `__init__` sets up a live ZMQ socket + rolling buffers, which suits a plain class.
A dataclass with `__post_init__` is possible but low value.
- **Action:** [ ] Reply explaining the socket setup makes a plain class cleaner (or convert
if she feels strongly). Non-blocking "consider".
## 10. `unitree_g1.py` (lines 488–490) — forwarding `smpl.` / `root.` keys
> "Won't these keys be included in the REMOTE_KEYS?"
- **Ask:** Is the `smpl.*`/`root.*` forwarding loop redundant with the `REMOTE_KEYS` loop?
- **Answer:** No — `REMOTE_KEYS` is the fixed joystick/button set; the `smpl.*`/`root.*`
keys are the ~724 whole-body reference floats, which are not in `REMOTE_KEYS`, so the
separate loop is required.
- **Action:** [ ] Reply confirming (and/or add a clarifying comment).
## 11. `sonic_pipeline.py` — custom keyboard handling duplicates existing utility
> (paraphrased) "We already have a keyboard utility — don't reimplement keyboard control."
- **Ask:** Reuse lerobot's shared keyboard infrastructure instead of a bespoke reader.
- **Notes:** `lerobot/utils/keyboard_input.py` already provides `TerminalKeyListener` /
`create_key_listener` (cbreak terminal reader + pynput backend, cross-platform). The
`RawKeyboard`/`drain_keyboard`/`process_keyboard` in `sonic_pipeline.py` reimplemented a
subset of this — and were **dead code**: nothing in the lerobot package imported them
(only the standalone `sonic_python/test_sonic_planner.py` has its own separate copies).
- **Resolution:** [x] Done — removed the unused `RawKeyboard`, `drain_keyboard`,
`process_keyboard` and their now-unused `sys`/`select`/`termios`/`tty` imports. The G1
integration drives movement via the joystick path (`process_joystick`), which is kept.
---
## Summary of remaining work
| # | File | Type | Remaining |
|---|------|------|-----------|
| 1 | gr00t/holosoma | relative imports | done |
| 2 | sonic_whole_body | private imports | done |
| 3 | sonic_whole_body/pipeline | safe onnx import | done |
| 4 | controllers/__init__ | — | resolved |
| 5 | sonic_pipeline | docstrings | done |
| 6 | sonic_whole_body:72 | `smpl.0` sentinel | reply / comment |
| 7 | smpl_stream:43–45 | host/port in config | reply / move to shared |
| 8 | smpl_stream:47–51 | shared constants | refactor |
| 9 | smpl_stream:53 | dataclass? | reply (push back) |
| 10 | unitree_g1:488–490 | REMOTE_KEYS overlap | reply / comment |
| 11 | sonic_pipeline | reuse keyboard utility | done (removed dead code) |
docs/SONIC_FIDELITY_TODO.md
+54
View File
@@ -0,0 +1,54 @@
# SONIC fidelity TODO — gaps vs gear_sonic / C++ deploy
Remaining differences between this lerobot SONIC port and the original gear_sonic +
`gear_sonic_deploy/g1_deploy_onnx_ref` reference. All three deployed encoder modes exist
(0 locomotion, 1 three-point, 2 full-body SMPL); these are faithfulness gaps in *how*
some pieces behave.
## 1. 3-point calibration & source (mode 1)
The original derives the 3 points from **full-body joints** and calibrates with
`ThreePointPose` against **live measured robot joints** (`reset_with_measured_q`), plus
neck→waist coupling. Our device path (head + controllers, fixed zero-q neutral) is an
approximation.
- Missing: measured-q recalibration + waist coupling.
- Note: the body-source `compute_3point` path (needs PICO Motion Trackers) is already faithful.
## 2. Locomotion stick/speed mapping (#6)
gear_sonic's `PlannerLoop` uses a yaw-accumulator + **mode-dependent speed curves**
(slow `0.1 + 0.5·mag`, run `1.5 + 3·mag`) with a facing-rotated movement vector. We reuse
the keyboard-parity `apply_joystick_axes` (no speed-by-mode curve).
- Also: PICO stick sign conventions untested (may need axis flips in `pico_publisher`).
- Self-contained and sim-testable.
## 3. Mode-2 root motion
`enable_smpl_root=False` by default; the original feeds the SMPL root orientation into the
anchor. Faithful only when enabled **and** the 30→50 Hz root trajectory is smoothed/
rate-matched (currently causes base-acceleration spikes → NaN QACC).
## 4. Start/stop handshake
C++ has `WAIT_FOR_CONTROL` (operator "start" gate) and a "stop" halt. We start policy
immediately after the startup ramp — no explicit arm/e-stop state machine.
## 5. Startup ramp target
C++ `InitControl` ramps to `default_angles`; we ramp to the policy's live commanded pose.
Deliberate (per request) — only a gap if exact C++ parity is wanted.
## 6. Hands / grippers
Original maps trigger/grip → Dex3 hand joints in modes 1 & 2 (simple gripping: thumb open,
trigger>0.5 closes middle finger, run through a gripper IK solver). Our 29-DOF path has no
hand control.
- Hardware here: **OpenArm grippers** (not Dex3) — could be routed through the same
trigger→open/close path, but needs the robot interface to expose/command the gripper
actuators (they're separate from the 29-DOF body).
## Not gaps (settled)
- `vr_5point_index` — vestigial constant in `policy_parameters.hpp`, not a deployed encoder
mode. Nothing to implement.
- Action scaling, kp/kd, IsaacLab↔MuJoCo remaps, encoder cadence (encode ⅕ / decode 1×),
per-mode replan intervals (0.1/0.2/1.0), 8-frame cross-fade blend — all faithful.
- Graceful damped shutdown — intentional safety addition, keep.
## Suggested priority
- High (change how the robot moves under teleop): #1, #2.
- Medium: #3, #4, #6 (grippers).
- Low / intentional: #5.
src/lerobot/robots/unitree_g1/README.md
+71 -57
View File
@@ -1,38 +1,30 @@
# Unitree G1 — SONIC encoder/decoder whole-body control
# Unitree G1 — SONIC 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.
This package runs NVIDIA's **SONIC** decoder on the Unitree G1, in MuJoCo simulation or
on real hardware, driven by a **64-D latent motion token**. It is a pure-Python/ONNX
reimplementation of the decode half of the SONIC deploy stack (no `gear_sonic`/torch
dependency): the decoder maps a 64-D latent token + proprioception history into 50 Hz
joint-position targets for the robot's PD controller. The encoder is bypassed — a policy
(e.g. `nepyope/sonic_walk`) emits the token directly.
## 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 |
| Controller | Purpose |
| ------------------------------ | --------------------------------------------------- |
| `SonicWholeBodyController` | SONIC decoder driven by a 64-D latent motion token |
| `GrootLocomotionController` | GR00T locomotion policy |
| `HolosomaLocomotionController` | Holosoma locomotion policy |
The rest of this document covers the SONIC whole-body path.
The rest of this document covers the SONIC token 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).
Each tick the `SonicWholeBodyController` takes a 64-D latent token
(`motion_token.0.pos … motion_token.63.pos`) and decodes it directly (encoder bypassed).
Before the first token arrives it holds a captured **neutral token** (a stable standing
pose), then holds the last token received between ticks (the ~30 Hz token stream vs. the
~50 Hz control loop). On startup the controller **interpolates** from the robot's measured
pose into the policy's commanded target over ~3 s (no snap).
## Requirements
@@ -44,46 +36,68 @@ from the robot's measured pose into the policy's commanded target over ~3 s (no
- The SONIC encoder/decoder ONNX models download automatically from the
`nvidia/GEAR-SONIC` Hub repo.
## Running a rollout
## Architecture: controller always runs onboard
Drive the G1 with a 34-D VLA policy (OpenHLM / pi0.5) via `lerobot-rollout`:
The controller runs **on the robot**, never on the laptop. The laptop is a thin client:
it negotiates the controller with `run_g1_server` (handshake), then PUSHes the 64-D token
and reads back the `observation.state` echo + camera frames over ZMQ.
## Running a rollout (real robot)
On the robot — host the SONIC decoder + camera onboard:
```bash
python -m lerobot.robots.unitree_g1.run_g1_server --handshake \
--cameras "ego_view:/dev/v4l/by-path/platform-3610000.usb-usb-0:2.1:1.3-video-index0:640x480"
```
On the laptop — `lerobot-rollout` drives the thin client:
```bash
lerobot-rollout \
--strategy.type=base \
--policy.path=<pi05_openhlm_dir> \
--policy.path=nepyope/sonic_walk \
--robot.type=unitree_g1 \
--robot.controller=SonicWholeBodyController \
--robot.is_simulation=true \
--robot.publish_hands=true \
--task="<language instruction>" \
--duration=45 --device=cuda
--robot.is_simulation=false --robot.onboard=false \
--robot.robot_ip=<ROBOT_IP> \
--robot.controller=SonicWholeBodyController --robot.sonic_token_action=true \
--robot.cameras='{ego_view: {type: zmq, server_address: <ROBOT_IP>, port: 5555, camera_name: ego_view}}' \
--task="walk back and forth" --device=cuda
```
### Cameras
## Training a token policy (no pi05 code patch)
Image-conditioned policies need camera frames. Two options are available without live
cameras:
The SONIC token interface needs **no modeling changes** to pi05. A 64-D token action is
handled entirely by config: pi05 builds its action projections straight from config
(`action_in_proj = nn.Linear(max_action_dim, …)`, `action_out_proj = nn.Linear(…,
max_action_dim)`), pads the action to `max_action_dim`, then slices back to the dataset's
action dim. Set both dims to 64 and the pad/slice is a no-op, so the full 64-D token is
supervised.
- **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}'
```
Requirements:
### Hands (Dex3)
1. The dataset carries a 64-D `action` and 64-D `observation.state` (the motion tokens).
2. Pass the dims to `lerobot-train`:
`--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.
```bash
lerobot-train \
--dataset.repo_id=nepyope/walk_back_and_forth \
--policy.type=pi05 \
--policy.max_action_dim=64 \
--policy.max_state_dim=64 \
--policy.chunk_size=50 --policy.n_action_steps=50
```
## Observation state
`nepyope/sonic_walk` was trained exactly this way (`config.json`: `max_action_dim=64`,
`max_state_dim=64`, `output_features.action.shape=[64]`). Same stock code path for train
and inference — the checkpoint's 64-wide `Linear`s load with unmodified pi05.
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.
## Observation / action interface (token mode)
With `--robot.sonic_token_action=true` the robot advertises:
- action: 64-D `motion_token.{i}.pos` (the decoder consumes it directly),
- `observation.state`: 64-D `motion_token_state.{i}.pos` (the last commanded token,
echoed so a token-output VLA closes the loop on its own previous token),
plus the ego camera image. The controller always runs onboard (or in sim); it is never
built on the laptop client.
src/lerobot/robots/unitree_g1/config_unitree_g1.py
+7 -37
View File
@@ -66,8 +66,8 @@ class UnitreeG1Config(RobotConfig):
# 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_onboard.py, which receives them
# from the laptop over ZMQ. Mutually exclusive with is_simulation.
# 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)).
@@ -83,50 +83,20 @@ 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
# When False, connect() does not start the background controller thread, so a
# caller can drive the controller synchronously (one decode per fed action),
# reproducing the deploy's single 50Hz control clock for faithful replay.
run_controller_thread: bool = True
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
# Token-output VLA interface for the SONIC decoder. When True (and the controller
# is ``SonicWholeBodyController``), the robot advertises a 64-D latent-token action
# space (``motion_token.{i}.pos``) instead of the 34-D whole-body command, and
# exposes the last commanded token as a 64-D ``observation.state``
# (``motion_token_state.{i}.pos``). This lets ``lerobot-rollout`` drive a policy
# that was trained with 64-D SONIC motion tokens as both state and action
# (e.g. nepyope/sonic_walk): the decoder consumes the token directly, encoder
# bypassed. Ignored unless a SONIC whole-body controller is active.
# space (``motion_token.{i}.pos``) and exposes the last commanded token as a 64-D
# ``observation.state`` (``motion_token_state.{i}.pos``). This lets
# ``lerobot-rollout`` drive a policy trained with 64-D SONIC motion tokens as both
# state and action (e.g. nepyope/sonic_walk): the decoder consumes the token
# directly, encoder bypassed. Ignored unless a SONIC whole-body controller is active.
sonic_token_action: bool = False
# 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.
controller: str | None = None
src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py
+1 -1
View File
@@ -81,7 +81,7 @@ def load_groot_policies(
# 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 .sonic_pipeline import make_ort_session_options
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)
src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py
-670
View File
@@ -1,670 +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 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(intra_op_num_threads: int | None = None,
inter_op_num_threads: int | None = None):
"""Build ONNX Runtime SessionOptions (quiet logging).
Pass thread counts to cap ORT's CPU pool. These tiny MLP policies are latency-
bound, not throughput-bound, so letting ORT grab every core just starves the
real-time control loop / torch policy / IK solver and causes stutter. 1 intra +
1 inter thread is plenty and lowest-latency for a per-step MLP inference.
"""
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
# ── 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
+259 -242
View File
@@ -14,57 +14,40 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""SONIC full-body controller for Unitree G1."""
"""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
from collections import deque
import logging
from typing import TYPE_CHECKING
from huggingface_hub import hf_hub_download
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 (
ISAACLAB_TO_MUJOCO,
MUJOCO_TO_ISAACLAB,
WB_ACTION_DIM,
G1_29_JointIndex,
get_gravity_orientation,
lowstate_to_obs,
wb_action_key,
)
from .sonic_pipeline import (
CONTROL_DT,
DEFAULT_ANGLES,
ENCODER_UPDATE_EVERY,
TOKEN_DIM,
PlannerController,
compute_kp_kd,
make_ort_session_options,
ort_providers,
)
# 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"
if TYPE_CHECKING or _onnxruntime_available:
import onnxruntime as ort
else:
@@ -72,18 +55,110 @@ else:
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-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
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)
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
# 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
# (capture_neutral_token.py): the encoder is an FSQ bottleneck (~5 bit/dim, 15.5 half-
# width, Div(16)), so its tokens live on the 1/16 grid. We store the integer FSQ codes
# and rescale by the same 1/16 step, giving an exact on-grid token -- unlike the literal
# all-zero token, which is off the encoder's learned manifold and decodes to a slightly
# goofy stance. This one decodes to a stable, natural standing pose.
# 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,
@@ -94,98 +169,158 @@ _NEUTRAL_TOKEN_CODES = np.array(
NEUTRAL_TOKEN = _NEUTRAL_TOKEN_CODES / 16.0 # FSQ Div(16): integer codes -> on-grid token
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 _extract_token_from_action(action: dict | None) -> np.ndarray | None:
"""Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or None.
This is the token-only replay interface: instead of a joint reference driving the
encoder, the caller supplies the 64-D encoder latent directly (e.g. a recorded
``action.motion_token`` column), which the decoder consumes with the encoder
bypassed. Requires the full dense token; a partial one is ignored (returns 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,
)
return np.fromiter((float(action[key]) for key in keys), dtype=np.float32, count=TOKEN_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.
class SonicDecoder:
"""Runs the SONIC decoder ONNX model and owns the proprioception history.
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).
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).
"""
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
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 encoder/decoder ONNX models and owns the controller.
"""Loads the SONIC decoder ONNX model and owns the decode controller.
No motion planner: the reference motion buffer is written directly each tick by
:class:`SonicWholeBodyController` from the incoming 34-D whole-body command.
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, 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)
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)
# Report the provider actually bound, not the one requested: ORT silently falls
# back to CPU if CUDA can't load (e.g. libcudnn not on LD_LIBRARY_PATH), and a
# CPU decoder drifts the closed-loop heading. Warn loudly so it can't hide.
# Report the provider actually bound, not the one requested: ORT silently falls back
# to CPU if CUDA can't load (e.g. libcudnn not on LD_LIBRARY_PATH), and a CPU decoder
# drifts the closed-loop heading. Warn loudly so it can't hide.
self.use_gpu = decoder_sess.get_providers()[0] == "CUDAExecutionProvider"
if not force_cpu and not self.use_gpu:
print(
@@ -196,16 +331,13 @@ class SonicRuntime:
)
self.kp, self.kd = compute_kp_kd()
self.controller = PlannerController(encoder_sess, decoder_sess)
self.controller = SonicDecoder(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):
@@ -217,10 +349,6 @@ class SonicWholeBodyController:
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...")
@@ -229,113 +357,22 @@ class SonicWholeBodyController:
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).
# 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
# Token-interface state. ``token_mode`` is set True by the robot when the deploy
# is token-driven (``UnitreeG1Config.sonic_token_action``): the controller then
# holds a stable *neutral* (all-zero) 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_onboard,
# lerobot-rollout and the sim replays. ``token_mode`` stays False for the dense
# 34-D whole-body / OpenHLM path, which keeps its own "hold last target" idle.
# Token-interface state. ``token_mode`` is set True by the robot when the deploy is
# token-driven (``UnitreeG1Config.sonic_token_action``): 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 (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 _run_token(self, obs: dict, token: np.ndarray) -> dict:
"""Decode a supplied 64-D latent token directly (encoder bypassed).
Token-only replay: set the pipeline's cached token to the supplied one and run
a decode-only step (``update_encoder=False``). The decoder still closes the loop
on live proprioception (history is refreshed inside ``step`` from ``obs``); only
the encoder — which would recompute the token from a motion reference — is
skipped. Returns the ``<joint>.q`` target dict.
"""
c = self.controller
c.token = np.asarray(token, np.float32)
self._wb_step += 1
return c.step(obs, update_encoder=False, debug=False)
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,
@@ -368,46 +405,26 @@ class SonicWholeBodyController:
return {}
obs = lowstate_to_obs(lowstate)
# Token-only interface (latent replay / token-output VLA): a dense 64-D
# ``motion_token.{i}`` command is decoded directly, bypassing the encoder.
# Checked before the joint path so a token action takes precedence.
# 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 (the encoder's own idle output; see NEUTRAL_TOKEN).
# 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 not None:
# 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._run_token(obs, self._last_token))
# 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],
)
if self._last_token is None:
# No token yet and not in token_mode: hold (keep last target).
return {}
return self._startup_blend(obs, self._run_wholebody34(obs, wb))
# 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 = {}
self._wb_step = 0
self._wb_traj.clear()
self._wb_quat_traj.clear()
self._heading = 0.0
# Drop the held token so token_mode re-seeds the neutral token after a reset.
self._last_token = None
src/lerobot/robots/unitree_g1/g1_utils.py
+29 -52
View File
@@ -105,21 +105,6 @@ REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
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)."""
return dict.fromkeys(REMOTE_KEYS, 0.0)
@@ -135,6 +120,35 @@ def get_gravity_orientation(quaternion: list[float] | np.ndarray) -> np.ndarray:
return gravity_orientation
def ort_providers(force_cpu: bool = False) -> list[str]:
"""ONNX Runtime providers, preferring CUDA when available (shared by the ONNX
controllers: SONIC decoder, GR00T). Falls back to CPU."""
import onnxruntime as ort
avail = ort.get_available_providers()
if not force_cpu and "CUDAExecutionProvider" in avail:
return ["CUDAExecutionProvider", "CPUExecutionProvider"]
return ["CPUExecutionProvider"]
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
@@ -196,43 +210,6 @@ def lowstate_to_obs(lowstate) -> dict:
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:
src/lerobot/robots/unitree_g1/infer_sonic_g1_onboard.py
-192
View File
@@ -1,192 +0,0 @@
#!/usr/bin/env python
# Copyright 2026 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.
"""Laptop-side sender for the SONIC whole-body walk policy, onboard deployment.
This is the counterpart to ``run_g1_onboard.py`` (which runs the SONIC decoder on the
robot). The heavy VLA (``nepyope/sonic_walk``, a pi0.5 token policy) runs here on the
laptop GPU; only the resulting 64-D latent token is shipped to the robot over ZMQ:
laptop: camera frame (ZMQ from robot :5555) + previous token
-> pi0.5 -> next 64-D token
-> PUSH JSON {motion_token.i.pos: ...} to robot :6004
robot: run_g1_onboard receives the token, SonicWholeBodyController decodes it
into whole-body joint commands against local DDS at full rate.
The policy's ``observation.state`` is the token currently being executed, so we close
the loop by feeding back the *last token we sent* (the decoder holds it until a new one
arrives). This mirrors what ``lerobot-rollout`` does via the robot's token echo, but
without a controller / DDS on the laptop.
The policy is pi0.5 with chunk_size=50, so a full diffusion inference runs only about
once every 50 ticks; ``select_action`` pops one queued token per tick in between.
Run ``run_g1_onboard.py --controller SonicWholeBodyController --sonic-token-action
--cameras ...`` on the robot first, then this on the laptop:
python -m lerobot.robots.unitree_g1.infer_sonic_g1_onboard \
--policy-path nepyope/sonic_walk --robot-ip 192.168.123.164 \
--task "walk back and forth"
"""
import argparse
import contextlib
import json
import logging
import signal
import time
import numpy as np
import torch
from lerobot.cameras.zmq import ZMQCamera, ZMQCameraConfig
from lerobot.configs.policies import PreTrainedConfig
from lerobot.policies.factory import get_policy_class, make_pre_post_processors
from lerobot.policies.utils import prepare_observation_for_inference
from lerobot.robots.unitree_g1.controllers.sonic_whole_body import (
NEUTRAL_TOKEN,
TOKEN_DIM,
token_action_key,
)
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", force=True)
logger = logging.getLogger("sonic_sender")
ACTION_PORT = 6004 # matches run_g1_onboard.py --action-port
IMAGE_KEY = "observation.images.ego_view" # pi05 sonic_walk VISUAL input
STATE_KEY = "observation.state"
def main() -> None:
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--policy-path", default="nepyope/sonic_walk", help="Policy repo id or local path")
p.add_argument("--robot-ip", default="192.168.123.164", help="Robot IP (camera + action ports)")
p.add_argument("--action-port", type=int, default=ACTION_PORT, help="Onboard ZMQ PULL port for actions")
p.add_argument("--camera-port", type=int, default=5555, help="Onboard ZMQ camera PUB port")
p.add_argument("--camera-name", default="head_camera", help="Camera name served by run_g1_onboard")
p.add_argument("--camera-width", type=int, default=640, help="Camera width")
p.add_argument("--camera-height", type=int, default=480, help="Camera height")
p.add_argument("--task", default="walk back and forth", help="Language prompt for the VLA")
p.add_argument("--fps", type=float, default=30.0, help="Token send rate (matches training inference)")
p.add_argument("--device", default="cuda", help="Torch device")
p.add_argument("--max-ticks", type=int, default=0, help="Stop after N ticks (0 = run forever)")
p.add_argument("--dry-run", action="store_true", help="Run inference but do not PUSH tokens to the robot")
args = p.parse_args()
device = torch.device(args.device)
# --- Policy + processors (normalization stats baked into the checkpoint) ---
logger.info("Loading policy from '%s'...", args.policy_path)
policy_cfg = PreTrainedConfig.from_pretrained(args.policy_path)
policy_cfg.pretrained_path = args.policy_path
policy = get_policy_class(policy_cfg.type).from_pretrained(args.policy_path, config=policy_cfg)
policy = policy.to(device)
policy.eval()
policy.reset()
preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=policy_cfg,
pretrained_path=args.policy_path,
preprocessor_overrides={"device_processor": {"device": str(device)}},
)
logger.info("Policy loaded (type=%s, device=%s, chunk=%s)", policy_cfg.type, device,
getattr(policy_cfg, "chunk_size", "?"))
# --- Camera (ZMQ from the robot's onboard image server) ---
cam = ZMQCamera(
ZMQCameraConfig(
server_address=args.robot_ip,
port=args.camera_port,
camera_name=args.camera_name,
width=args.camera_width,
height=args.camera_height,
fps=int(args.fps),
)
)
logger.info("Connecting camera %s@%s:%d ...", args.camera_name, args.robot_ip, args.camera_port)
cam.connect()
# --- Action PUSH socket to the onboard controller ---
import zmq
ctx = zmq.Context.instance()
sock = ctx.socket(zmq.PUSH)
sock.setsockopt(zmq.SNDHWM, 2)
sock.setsockopt(zmq.LINGER, 0)
sock.connect(f"tcp://{args.robot_ip}:{args.action_port}")
logger.info("Sending tokens to tcp://%s:%d (dry_run=%s)", args.robot_ip, args.action_port, args.dry_run)
stop = {"flag": False}
signal.signal(signal.SIGINT, lambda *_: stop.__setitem__("flag", True))
signal.signal(signal.SIGTERM, lambda *_: stop.__setitem__("flag", True))
# observation.state = the token currently executing on the robot (last one we sent);
# start at the neutral token the decoder holds before the first send, so the very
# first inference sees the true executing token (not zeros).
prev_token = NEUTRAL_TOKEN.copy()
period = 1.0 / args.fps
n = 0
t_infer_total = 0.0
logger.info("Streaming tokens at %.0f Hz. Ctrl-C to stop.", args.fps)
try:
while not stop["flag"]:
t0 = time.time()
try:
frame = cam.read() # HxWxC uint8 RGB
except Exception as e: # noqa: BLE001
logger.warning("Camera read failed: %s", e)
time.sleep(period)
continue
raw_obs = {
IMAGE_KEY: np.ascontiguousarray(frame),
STATE_KEY: prev_token.copy(),
}
with torch.inference_mode():
obs = prepare_observation_for_inference(raw_obs, device, args.task, "unitree_g1")
obs = preprocessor(obs)
action = policy.select_action(obs)
action = postprocessor(action)
token = action.squeeze(0).to("cpu").numpy().astype(np.float32)
prev_token = token
if not args.dry_run:
msg = {token_action_key(i): float(token[i]) for i in range(TOKEN_DIM)}
with contextlib.suppress(zmq.Again):
sock.send_string(json.dumps(msg), zmq.NOBLOCK)
n += 1
t_infer_total += time.time() - t0
if n % 30 == 0:
logger.info(
"tick %d | avg %.1f ms/tick | token[:3]=%s",
n, 1000.0 * t_infer_total / 30.0, np.round(token[:3], 3).tolist(),
)
t_infer_total = 0.0
if args.max_ticks and n >= args.max_ticks:
break
time.sleep(max(0.0, period - (time.time() - t0)))
finally:
logger.info("Stopping sender after %d ticks.", n)
with contextlib.suppress(Exception):
cam.disconnect()
with contextlib.suppress(Exception):
sock.close(linger=0)
if __name__ == "__main__":
main()
src/lerobot/robots/unitree_g1/run_g1_onboard.py
-254
View File
@@ -1,254 +0,0 @@
#!/usr/bin/env python
# Copyright 2026 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.
"""Run the G1 locomotion / whole-body controller ONBOARD, driven by high-level actions
from a laptop.
The controller (GR00T / Holosoma / SONIC whole-body) runs on the robot itself against
local DDS, at full control rate. The laptop ships only the resulting high-level action
(arm joint targets + joystick axes + gripper flags, or a 64-D SONIC motion token) as
JSON over ZMQ. This process applies each action via ``UnitreeG1.send_action`` while the
onboard controller thread keeps the legs balanced / decodes the token.
This is the real-deploy counterpart to running ``lerobot-rollout`` on the laptop with
``--robot.is_simulation=false`` (the ZMQ *socket bridge*): there the 50 Hz lowcmd
crosses the network; here only compact high-level actions do, and the control loop stays
local to the robot. Pair with a laptop client that produces actions (exo teleop, or a
policy such as ``nepyope/sonic_walk`` emitting ``motion_token.{i}.pos``).
Besides receiving actions, this process publishes ``observation.state`` (29 joint ``.q``)
on a ZMQ PUB port so a laptop policy client has proprioception.
Safety: type ``e`` then Enter in this terminal to stop immediately (zero-torque + exit).
Ctrl-C does the normal graceful shutdown (kp ramp).
Examples (on the robot):
# GR00T locomotion, arm targets from the laptop:
python -m lerobot.robots.unitree_g1.run_g1_onboard --controller GrootLocomotionController
# SONIC whole-body walk policy: laptop ships 64-D tokens, decoder runs here:
python -m lerobot.robots.unitree_g1.run_g1_onboard \
--controller SonicWholeBodyController --sonic-token-action \
--cameras "head_camera:/dev/v4l/by-path/platform-3610000.usb-usb-0:2.1:1.3-video-index0:640x480"
"""
import argparse
import contextlib
import json
import logging
import os
import signal
import sys
import threading
import time
import numpy as np
import zmq
from lerobot.cameras.zmq.image_server import ImageServer
from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
from lerobot.robots.unitree_g1.run_g1_server import Gripper, build_gripper, parse_camera_specs
from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", force=True)
logger = logging.getLogger("g1_onboard")
ACTION_PORT = 6004
STATE_PORT = 6005
def main() -> None:
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--controller", default="GrootLocomotionController", help="Controller class name")
p.add_argument("--dds-interface", default=None, help="DDS network interface (default: SDK default)")
p.add_argument(
"--sim",
action="store_true",
help="Attach to a DDS MuJoCo sim: skip MotionSwitcher + physical remote, default dds-interface 'lo'.",
)
p.add_argument(
"--sonic-token-action",
action="store_true",
help="SONIC token interface: actions carry a 64-D motion_token.{i}.pos that the decoder consumes.",
)
p.add_argument("--action-port", type=int, default=ACTION_PORT, help="ZMQ PULL port for laptop actions")
p.add_argument("--state-port", type=int, default=STATE_PORT, help="ZMQ PUB port for observation.state")
p.add_argument("--state-fps", type=float, default=30.0, help="observation.state publish rate; <=0 disables")
p.add_argument("--gravity-compensation", action="store_true", help="Enable arm gravity compensation")
# Gripper control (Damiao over CAN).
p.add_argument("--grippers", action="store_true", help="Drive Damiao grippers from action L3/R3 flags")
p.add_argument("--gripper-port-left", default="can1", help="CAN interface for LEFT gripper")
p.add_argument("--gripper-port-right", default="can0", help="CAN interface for RIGHT gripper")
p.add_argument("--gripper-send-id", type=lambda x: int(x, 0), default=0x08, help="Motor send CAN id")
p.add_argument("--gripper-recv-id", type=lambda x: int(x, 0), default=0x18, help="Motor recv CAN id")
p.add_argument("--gripper-motor-type", default="dm4310", help="Damiao motor type")
p.add_argument("--gripper-open-deg", type=float, default=-65.0, help="Gripper OPEN position (deg)")
p.add_argument("--gripper-close-deg", type=float, default=0.0, help="Gripper CLOSE position (deg)")
p.add_argument("--gripper-kp", type=float, default=15.0, help="MIT position gain (stiffness)")
p.add_argument("--gripper-kd", type=float, default=0.5, help="MIT damping gain")
p.add_argument("--gripper-no-fd", dest="gripper_fd", action="store_false", help="Classic CAN (non-FD)")
p.set_defaults(gripper_fd=True)
# Optional camera streaming (ZMQ) so the laptop policy client / viewer can connect.
p.add_argument("--cameras", default=None, help="Camera spec 'name:device[:WxH[:FOURCC]]', comma-sep")
p.add_argument("--camera-fps", type=int, default=30, help="Camera FPS")
p.add_argument("--camera-port", type=int, default=5555, help="Camera ZMQ port")
p.add_argument("--camera-width", type=int, default=640, help="Default camera width")
p.add_argument("--camera-height", type=int, default=480, help="Default camera height")
args = p.parse_args()
dds_interface = args.dds_interface
if args.sim and dds_interface is None:
dds_interface = "lo"
cfg = UnitreeG1Config(
is_simulation=False,
onboard=True,
controller=args.controller,
dds_interface=dds_interface,
gravity_compensation=args.gravity_compensation,
release_motion_control=not args.sim,
physical_remote=not args.sim,
sonic_token_action=args.sonic_token_action,
cameras={},
)
# Optional camera server (background thread; independent of DDS/CAN).
camera_server = None
if args.cameras:
cameras = parse_camera_specs(args.cameras, args.camera_width, args.camera_height)
camera_server = ImageServer({"fps": args.camera_fps, "cameras": cameras}, port=args.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())
logger.info("Camera server started on :%d: %s", args.camera_port, cam_summary)
robot = UnitreeG1(cfg)
logger.info("Connecting onboard robot (controller=%s, token=%s)...", args.controller, args.sonic_token_action)
robot.connect()
# Note: with --sonic-token-action the SonicWholeBodyController holds a neutral
# (all-zero) token until the first laptop token arrives, then holds the last token
# between ticks -- see SonicWholeBodyController.token_mode (set from config).
grippers: dict[str, Gripper] = {}
if args.grippers:
for side, port in (("L", args.gripper_port_left), ("R", args.gripper_port_right)):
grippers[side] = build_gripper(
side, port, args.gripper_send_id, args.gripper_recv_id, args.gripper_motor_type,
args.gripper_fd, args.gripper_open_deg, args.gripper_close_deg, args.gripper_kp, args.gripper_kd,
)
logger.info("Grippers enabled: L3 -> left, R3 -> right")
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:{args.action_port}")
logger.info("Onboard controller live. Waiting for laptop actions on :%d ...", args.action_port)
logger.info("Type 'e' then Enter to STOP immediately (or Ctrl-C for graceful shutdown).")
stop = threading.Event()
signal.signal(signal.SIGINT, lambda *_: stop.set())
signal.signal(signal.SIGTERM, lambda *_: stop.set())
def estop_listener() -> None:
for line in sys.stdin:
if line.strip().lower() == "e":
logger.warning("E-STOP ('e'): going passive NOW.")
try:
robot._shutdown_event.set() # stop the controller loop publishing
time.sleep(0.05)
robot._send_zero_torque() # motors limp; nothing overwrites it now
except Exception as e: # noqa: BLE001
logger.warning("E-stop zero-torque failed: %s", e)
os._exit(0) # immediate hard exit, no slow cleanup
threading.Thread(target=estop_listener, daemon=True).start()
# Proprioception feedback: publish observation.state (29 joint .q) so a laptop
# inference client can feed it to a policy. DDS stays local; only compact JSON
# state crosses the network. (For a token policy the laptop closes the loop on the
# token instead, but publishing joint state is harmless and useful for logging.)
state_sock = None
if args.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:{args.state_port}")
logger.info("Publishing observation.state on :%d at %.0f Hz", args.state_port, args.state_fps)
def publish_state() -> None:
period = 1.0 / args.state_fps
joint_names = [j.name for j in G1_29_JointIndex]
while not stop.is_set():
t0 = time.time()
obs = robot.get_observation()
if obs:
state = {f"{name}.q": float(obs.get(f"{name}.q", 0.0)) for name in joint_names}
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:
logger.info("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:
logger.warning("Dropping malformed action: %s", e)
continue
robot.send_action(action)
if grippers:
# L3 = remote.button.4 -> left, R3 = remote.button.0 -> right.
if "L" in grippers and "remote.button.4" in action:
grippers["L"].apply(bool(action["remote.button.4"]))
if "R" in grippers and "remote.button.0" in action:
grippers["R"].apply(bool(action["remote.button.0"]))
n += 1
if n % 60 == 0:
axes = {k: round(float(action.get(k, 0.0)), 3) for k in ("remote.lx", "remote.ly", "remote.rx", "remote.ry")}
logger.info("Applied %d actions | axes=%s", n, axes)
finally:
logger.info("Shutting down onboard controller...")
stop.set()
if state_sock is not None:
with contextlib.suppress(Exception):
state_sock.close(linger=0)
if camera_server is not None:
with contextlib.suppress(Exception):
camera_server.stop()
for g in grippers.values():
with contextlib.suppress(Exception):
g.bus.disconnect()
robot.disconnect()
if __name__ == "__main__":
main()
src/lerobot/robots/unitree_g1/run_g1_server.py
+358 -53
View File
@@ -22,18 +22,36 @@ 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 os
import re
import signal
import sys
import threading
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
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
@@ -43,9 +61,6 @@ from unitree_sdk2py.utils.crc import CRC
from lerobot.cameras.zmq.image_server import ImageServer
if TYPE_CHECKING:
from lerobot.motors.damiao.damiao import DamiaoMotorsBus
# DDS topic names follow Unitree SDK naming conventions
# ruff: noqa: N816
kTopicLowCommand_Debug = "rt/lowcmd" # action to robot
@@ -55,59 +70,271 @@ 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
@dataclass
class Gripper:
"""A single Damiao gripper that only writes to CAN when the open/close state changes."""
# 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
name: str
bus: "DamiaoMotorsBus"
open_deg: float
close_deg: float
_last_cmd: str | None = None # "open" | "close"
def apply(self, want_close: bool) -> None:
want = "close" if want_close else "open"
if want == self._last_cmd:
return
target = self.close_deg if want_close else self.open_deg
self.bus.write("Goal_Position", "gripper", target)
self._last_cmd = want
print(f"[gripper] {self.name} -> {want.upper()} ({target:.1f} deg)")
# 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 build_gripper(
name: str,
port: str,
send_id: int,
recv_id: int,
motor_type: str,
use_can_fd: bool,
open_deg: float,
close_deg: float,
kp: float,
kd: float,
) -> Gripper:
from lerobot.motors.damiao.damiao import DamiaoMotorsBus
from lerobot.motors.motors_bus import Motor, MotorNormMode
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'")
motors = {
"gripper": Motor(
id=send_id,
model=motor_type,
norm_mode=MotorNormMode.DEGREES,
motor_type_str=motor_type,
recv_id=recv_id,
)
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,
}
bus = DamiaoMotorsBus(port=port, motors=motors, use_can_fd=use_can_fd)
print(f"Connecting {name} gripper on {port} (fd={use_can_fd})...")
bus.connect(handshake=True)
bus.write("Kp", "gripper", kp)
bus.write("Kd", "gripper", kd)
bus.write("Goal_Position", "gripper", open_deg) # start open
print(f" {name}: connected, torque enabled, opened.")
return Gripper(name, bus, open_deg, close_deg, _last_cmd="open")
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,
sonic_token_action=sonic_token_action,
cameras={},
)
# Optional camera server (background thread; independent of DDS).
camera_server = None
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("Type 'e' then Enter to STOP immediately (or Ctrl-C for graceful shutdown).")
def estop_listener() -> None:
for line in sys.stdin:
if line.strip().lower() == "e":
print("E-STOP ('e'): going passive NOW.")
try:
robot._shutdown_event.set() # stop the controller loop publishing
time.sleep(0.05)
robot._send_zero_torque() # motors limp; nothing overwrites it now
except Exception as e: # noqa: BLE001
print(f"E-stop zero-torque failed: {e}")
os._exit(0) # immediate hard exit, no slow cleanup
threading.Thread(target=estop_listener, daemon=True).start()
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)
if camera_server is not None:
with contextlib.suppress(Exception):
camera_server.stop()
robot.disconnect()
def parse_camera_specs(spec: str, default_width: int, default_height: int) -> dict[str, dict]:
@@ -268,8 +495,86 @@ 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 or args.cameras:
if args.cameras:
cameras = parse_camera_specs(args.cameras, args.camera_width, args.camera_height)
else:
dev = args.camera_device
dev = int(dev) if str(dev).lstrip("-").isdigit() else dev
cameras = {"head_camera": {"device_id": dev, "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 or args.cameras:
@@ -313,6 +618,7 @@ def main() -> None:
# initialize ZMQ
ctx = zmq.Context.instance()
shutdown_event = threading.Event()
# receive commands from remote client
lowcmd_sock = ctx.socket(zmq.PULL)
@@ -323,7 +629,6 @@ 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
+161 -212
View File
@@ -16,6 +16,8 @@
from __future__ import annotations
import contextlib
import json
import logging
import threading
import time
@@ -26,12 +28,12 @@ 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
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,
@@ -40,7 +42,6 @@ from .g1_utils import (
default_remote_input,
lowstate_to_obs,
make_locomotion_controller,
obs_to_wb34_state,
)
if TYPE_CHECKING or _unitree_sdk_available:
@@ -50,11 +51,9 @@ if TYPE_CHECKING or _unitree_sdk_available:
ChannelSubscriber as _SDKChannelSubscriber,
)
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,
)
@@ -64,8 +63,6 @@ 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
@@ -135,26 +132,34 @@ 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)
# Import channel classes based on mode. Simulation and onboard both talk to a
# real (local) DDS via the Unitree SDK; only the laptop-side bridge client uses
# the ZMQ socket shim.
# 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:
self._ChannelFactoryInitialize = _SDKChannelFactoryInitialize
self._ChannelPublisher = _SDKChannelPublisher
self._ChannelSubscriber = _SDKChannelSubscriber
else:
from .unitree_sdk2_socket import (
ChannelFactoryInitialize,
ChannelPublisher,
ChannelSubscriber,
)
self._ChannelFactoryInitialize = None
self._ChannelPublisher = None
self._ChannelSubscriber = None
self._ChannelFactoryInitialize = ChannelFactoryInitialize
self._ChannelPublisher = ChannelPublisher
self._ChannelSubscriber = 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
# Initialize state variables
self.sim_env = None
@@ -164,15 +169,17 @@ class UnitreeG1(Robot):
self._shutdown_event = threading.Event()
self.subscribe_thread = None
self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else 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)
# Lower-body controller loaded dynamically
self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
# Token-driven deploy: let a SONIC controller hold a neutral token until the
# first real one arrives, then hold the last token between control ticks.
if config.sonic_token_action and hasattr(self.controller, "token_mode"):
self.controller.token_mode = True
# Token-driven deploy: let a SONIC controller hold a neutral token until the
# first real one arrives, then hold the last token between control ticks.
if config.sonic_token_action and hasattr(self.controller, "token_mode"):
self.controller.token_mode = True
# Controller thread state
self._controller_thread = None
@@ -187,17 +194,6 @@ class UnitreeG1(Robot):
# from local lowstate so joystick locomotion works without a laptop round-trip).
self._joystick = None
# 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()
# 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 (see ``sonic_token_action``). Seeded to zeros;
@@ -208,44 +204,6 @@ class UnitreeG1(Robot):
self._last_token = np.zeros(TOKEN_DIM, dtype=np.float32)
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()
@@ -330,22 +288,6 @@ 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 self.config.sonic_token_action:
return {}
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 _token_state_ft(self) -> dict[str, type]:
"""64-D SONIC latent-token proprio state (``motion_token_state.{i}.pos``).
@@ -359,33 +301,22 @@ class UnitreeG1(Robot):
return {token_state_key(i): float for i in range(TOKEN_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._wb_state_ft,
**self._token_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:
# Role-agnostic: the schema is a pure function of (controller name,
# sonic_token_action). 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 and not self.config.sonic_token_action:
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
@@ -396,14 +327,7 @@ class UnitreeG1(Robot):
return {token_action_key(i): float for i in range(TOKEN_DIM)}
# 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)}
# 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}
@@ -521,7 +445,111 @@ class UnitreeG1(Robot):
_, 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.config.sonic_token_action,
)
self._client_caps = request_controller(
server_ip,
self.config.controller,
sonic_token_action=self.config.sonic_token_action,
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.utils import (
@@ -532,9 +560,9 @@ class UnitreeG1(Robot):
self._ChannelFactoryInitialize(0, "lo")
# Call the hub env's make_env directly so we can disable the offscreen
# head_camera renderer. We drive image-conditioned policies from recorded
# frames (see replay_camera_parquet / external obs), never the sim's own
# camera, so building a MuJoCo offscreen GL context is pure liability: it
# head_camera renderer. We drive image-conditioned policies from external
# camera frames, never the sim's own camera, so building a MuJoCo offscreen
# GL context is pure liability: it
# crashes with "Failed to make the EGL context current" when GLFW/SDL
# already own a context, killing the sim thread and hanging on
# "Waiting for robot state...". publish_images=False -> no renderer.
@@ -577,17 +605,6 @@ 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()
@@ -685,6 +702,10 @@ class UnitreeG1(Robot):
self._send_zero_torque()
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
@@ -741,6 +762,9 @@ 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:
@@ -749,38 +773,14 @@ class UnitreeG1(Robot):
# Motors + IMU + wireless remote (shared lowstate -> obs mapping)
obs = lowstate_to_obs(lowstate)
# 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.
# 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.config.sonic_token_action:
# Token mode: echo the last commanded latent token as observation.state
# so a token-output VLA closes the loop on its own previous token.
from .controllers.sonic_whole_body import token_state_key
token = self._last_token if self._last_token is not None else []
for i, v in enumerate(token):
obs[token_state_key(i)] = float(v)
elif 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)
# 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
# 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():
@@ -792,6 +792,9 @@ 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.config.sonic_token_action:
@@ -801,8 +804,6 @@ class UnitreeG1(Robot):
if token is not None:
self._last_token = token
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
@@ -814,23 +815,7 @@ class UnitreeG1(Robot):
if key.endswith(".q") and key.startswith(arm_prefixes)
}
tau = None
if self.config.gravity_compensation and self.arm_ik is not None:
tau = np.zeros(29, dtype=np.float32)
action_np = np.array(
[
action_to_publish.get(f"{joint.name}.q", self.msg.motor_cmd[joint.value].q)
for joint in G1_29_JointArmIndex
],
dtype=np.float32,
)
arm_tau = self.arm_ik.solve_tau(action_np)
arm_start_idx = G1_29_JointArmIndex.kLeftShoulderPitch.value
for joint in G1_29_JointArmIndex:
local_idx = joint.value - arm_start_idx
tau[joint.value] = arm_tau[local_idx]
self.publish_lowcmd(action_to_publish, tau=tau)
self.publish_lowcmd(action_to_publish)
return action
def _update_controller_action(self, action: RobotAction) -> None:
@@ -858,50 +843,14 @@ class UnitreeG1(Robot):
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:
return True
@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