diff --git a/cam_test.png b/cam_test.png new file mode 100644 index 000000000..a4be93ad5 Binary files /dev/null and b/cam_test.png differ diff --git a/docs/PR3827_REVIEW_COMMENTS.md b/docs/PR3827_REVIEW_COMMENTS.md new file mode 100644 index 000000000..45cfb6de0 --- /dev/null +++ b/docs/PR3827_REVIEW_COMMENTS.md @@ -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) | diff --git a/docs/SONIC_FIDELITY_TODO.md b/docs/SONIC_FIDELITY_TODO.md new file mode 100644 index 000000000..8fcb149c8 --- /dev/null +++ b/docs/SONIC_FIDELITY_TODO.md @@ -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. diff --git a/src/lerobot/robots/unitree_g1/README.md b/src/lerobot/robots/unitree_g1/README.md index 45a912d1d..971684b0c 100644 --- a/src/lerobot/robots/unitree_g1/README.md +++ b/src/lerobot/robots/unitree_g1/README.md @@ -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=`: -| 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= \ + --policy.path=nepyope/sonic_walk \ --robot.type=unitree_g1 \ - --robot.controller=SonicWholeBodyController \ - --robot.is_simulation=true \ - --robot.publish_hands=true \ - --task="" \ - --duration=45 --device=cuda + --robot.is_simulation=false --robot.onboard=false \ + --robot.robot_ip= \ + --robot.controller=SonicWholeBodyController --robot.sonic_token_action=true \ + --robot.cameras='{ego_view: {type: zmq, server_address: , 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= \ - --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. diff --git a/src/lerobot/robots/unitree_g1/config_unitree_g1.py b/src/lerobot/robots/unitree_g1/config_unitree_g1.py index 818cb68b7..0ae028bf4 100644 --- a/src/lerobot/robots/unitree_g1/config_unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/config_unitree_g1.py @@ -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 diff --git a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py b/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py index b25198b42..1313ad325 100644 --- a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py +++ b/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py @@ -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) diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py deleted file mode 100644 index 2cda20a6d..000000000 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py +++ /dev/null @@ -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 ``.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 ``.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) diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py index 8320f199a..e1a187c5c 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -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 ``.q``/``.dq`` and ``imu.*`` fields. + token: 64-D latent supplied by the policy (encoder bypassed). + debug: log action/delta norms. + + Returns: + dict of ``.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 ``.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 diff --git a/src/lerobot/robots/unitree_g1/g1_utils.py b/src/lerobot/robots/unitree_g1/g1_utils.py index 723bf94aa..ac8efb0b1 100644 --- a/src/lerobot/robots/unitree_g1/g1_utils.py +++ b/src/lerobot/robots/unitree_g1/g1_utils.py @@ -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 ``.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: diff --git a/src/lerobot/robots/unitree_g1/infer_sonic_g1_onboard.py b/src/lerobot/robots/unitree_g1/infer_sonic_g1_onboard.py deleted file mode 100644 index 2fbc8ebdb..000000000 --- a/src/lerobot/robots/unitree_g1/infer_sonic_g1_onboard.py +++ /dev/null @@ -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() diff --git a/src/lerobot/robots/unitree_g1/run_g1_onboard.py b/src/lerobot/robots/unitree_g1/run_g1_onboard.py deleted file mode 100644 index 905766350..000000000 --- a/src/lerobot/robots/unitree_g1/run_g1_onboard.py +++ /dev/null @@ -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() diff --git a/src/lerobot/robots/unitree_g1/run_g1_server.py b/src/lerobot/robots/unitree_g1/run_g1_server.py index f769f5d5d..aadf81e11 100644 --- a/src/lerobot/robots/unitree_g1/run_g1_server.py +++ b/src/lerobot/robots/unitree_g1/run_g1_server.py @@ -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( diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 823d7e33d..5b9169b82 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -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