mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-29 20:49:42 +00:00
refactor(unitree_g1): minimize diff w.r.t main
This commit is contained in:
Martino Russi
Binary file not shown.
|
After Width: | Height: | Size: 167 KiB |
@@ -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) |
|
||||||
@@ -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.
|
||||||
@@ -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
|
This package runs NVIDIA's **SONIC** decoder on the Unitree G1, in MuJoCo simulation or
|
||||||
simulation or on real hardware, driven by a dense **34-D whole-body command** (the
|
on real hardware, driven by a **64-D latent motion token**. It is a pure-Python/ONNX
|
||||||
OpenHLM / pi0.5 action layout). It is a pure-Python/ONNX reimplementation of the
|
reimplementation of the decode half of the SONIC deploy stack (no `gear_sonic`/torch
|
||||||
reference-tracking half of the SONIC deploy stack (no `gear_sonic`/torch dependency, and
|
dependency): the decoder maps a 64-D latent token + proprioception history into 50 Hz
|
||||||
no motion planner): the encoder compresses a reference motion window into a latent token
|
joint-position targets for the robot's PD controller. The encoder is bypassed — a policy
|
||||||
and the decoder maps that token + proprioception history into 50 Hz joint-position
|
(e.g. `nepyope/sonic_walk`) emits the token directly.
|
||||||
targets for the robot's PD controller.
|
|
||||||
|
|
||||||
## Controllers
|
## Controllers
|
||||||
|
|
||||||
Selected with `--robot.controller=<ClassName>`:
|
Selected with `--robot.controller=<ClassName>`:
|
||||||
|
|
||||||
| Controller | Purpose |
|
| Controller | Purpose |
|
||||||
| ------------------------------ | ------------------------------------------------------------ |
|
| ------------------------------ | --------------------------------------------------- |
|
||||||
| `SonicWholeBodyController` | SONIC encoder/decoder driven by a 34-D OpenHLM/pi0.5 command |
|
| `SonicWholeBodyController` | SONIC decoder driven by a 64-D latent motion token |
|
||||||
| `GrootLocomotionController` | GR00T locomotion policy |
|
| `GrootLocomotionController` | GR00T locomotion policy |
|
||||||
| `HolosomaLocomotionController` | Holosoma 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
|
Each tick the `SonicWholeBodyController` takes a 64-D latent token
|
||||||
layout:
|
(`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
|
||||||
[L-arm(7), L-grip(1), R-arm(7), R-grip(1), L-leg(6), R-leg(6), waist(3),
|
~50 Hz control loop). On startup the controller **interpolates** from the robot's measured
|
||||||
root roll/pitch + yaw-rate(3)]
|
pose into the policy's commanded target over ~3 s (no snap).
|
||||||
```
|
|
||||||
|
|
||||||
The 29 joint targets become the SONIC encode-mode-0 reference (accumulated into a rolling
|
|
||||||
50-frame trajectory with finite-difference velocities so the encoder's lookahead sees a
|
|
||||||
real motion sequence), the root roll/pitch set the anchor orientation, and the two grip
|
|
||||||
scalars can drive the Dex3 hands (see below). On startup the controller **interpolates**
|
|
||||||
from the robot's measured pose into the policy's commanded target over ~3 s (no snap).
|
|
||||||
|
|
||||||
## Requirements
|
## 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
|
- The SONIC encoder/decoder ONNX models download automatically from the
|
||||||
`nvidia/GEAR-SONIC` Hub repo.
|
`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
|
```bash
|
||||||
lerobot-rollout \
|
lerobot-rollout \
|
||||||
--strategy.type=base \
|
--policy.path=nepyope/sonic_walk \
|
||||||
--policy.path=<pi05_openhlm_dir> \
|
|
||||||
--robot.type=unitree_g1 \
|
--robot.type=unitree_g1 \
|
||||||
--robot.controller=SonicWholeBodyController \
|
--robot.is_simulation=false --robot.onboard=false \
|
||||||
--robot.is_simulation=true \
|
--robot.robot_ip=<ROBOT_IP> \
|
||||||
--robot.publish_hands=true \
|
--robot.controller=SonicWholeBodyController --robot.sonic_token_action=true \
|
||||||
--task="<language instruction>" \
|
--robot.cameras='{ego_view: {type: zmq, server_address: <ROBOT_IP>, port: 5555, camera_name: ego_view}}' \
|
||||||
--duration=45 --device=cuda
|
--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
|
The SONIC token interface needs **no modeling changes** to pi05. A 64-D token action is
|
||||||
cameras:
|
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]'`.
|
Requirements:
|
||||||
- **Replay a recorded episode** as the camera feed:
|
|
||||||
```bash
|
|
||||||
--robot.replay_camera_parquet=<episode.parquet> \
|
|
||||||
--robot.replay_camera_map='{base: head_image_left, left_wrist: left_wrist_image, right_wrist: right_wrist_image}'
|
|
||||||
```
|
|
||||||
|
|
||||||
### Hands (Dex3)
|
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
|
```bash
|
||||||
scalars (`wb.7.pos` left, `wb.15.pos` right). The scalar is interpolated between
|
lerobot-train \
|
||||||
`hand_open_grip_value` (default 1.0 = open) and `hand_closed_grip_value` (default 0.0 =
|
--dataset.repo_id=nepyope/walk_back_and_forth \
|
||||||
closed) and scaled onto `hand_closed_pose` (7 joints:
|
--policy.type=pi05 \
|
||||||
`thumb_0, thumb_1, thumb_2, middle_0, middle_1, index_0, index_1`). Flip the signs in
|
--policy.max_action_dim=64 \
|
||||||
`hand_closed_pose` if the fingers curl the wrong way, or raise `hand_kp` for a firmer
|
--policy.max_state_dim=64 \
|
||||||
grip.
|
--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
|
## Observation / action interface (token mode)
|
||||||
(`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.
|
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.
|
||||||
|
|||||||
@@ -66,8 +66,8 @@ class UnitreeG1Config(RobotConfig):
|
|||||||
# itself, against local DDS at full rate) instead of on the laptop over the ZMQ
|
# 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
|
# 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
|
# 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
|
# tokens) fed via send_action -- e.g. by run_g1_server's serve_onboard_controller,
|
||||||
# from the laptop over ZMQ. Mutually exclusive with is_simulation.
|
# which receives them from the laptop over ZMQ. Mutually exclusive with is_simulation.
|
||||||
onboard: bool = False
|
onboard: bool = False
|
||||||
# DDS network interface for onboard mode (None = SDK default, matching
|
# DDS network interface for onboard mode (None = SDK default, matching
|
||||||
# run_g1_server.py's ChannelFactoryInitialize(0)).
|
# run_g1_server.py's ChannelFactoryInitialize(0)).
|
||||||
@@ -83,50 +83,20 @@ class UnitreeG1Config(RobotConfig):
|
|||||||
# Cameras (ZMQ-based remote cameras)
|
# Cameras (ZMQ-based remote cameras)
|
||||||
cameras: dict[str, CameraConfig] = field(default_factory=dict)
|
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
|
# When False, connect() does not start the background controller thread, so a
|
||||||
# caller can drive the controller synchronously (one decode per fed action),
|
# caller can drive the controller synchronously (one decode per fed action),
|
||||||
# reproducing the deploy's single 50Hz control clock for faithful replay.
|
# reproducing the deploy's single 50Hz control clock for faithful replay.
|
||||||
run_controller_thread: bool = True
|
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
|
# Token-output VLA interface for the SONIC decoder. When True (and the controller
|
||||||
# is ``SonicWholeBodyController``), the robot advertises a 64-D latent-token action
|
# 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
|
# space (``motion_token.{i}.pos``) and exposes the last commanded token as a 64-D
|
||||||
# exposes the last commanded token as a 64-D ``observation.state``
|
# ``observation.state`` (``motion_token_state.{i}.pos``). This lets
|
||||||
# (``motion_token_state.{i}.pos``). This lets ``lerobot-rollout`` drive a policy
|
# ``lerobot-rollout`` drive a policy trained with 64-D SONIC motion tokens as both
|
||||||
# that was trained with 64-D SONIC motion tokens as both state and action
|
# state and action (e.g. nepyope/sonic_walk): the decoder consumes the token
|
||||||
# (e.g. nepyope/sonic_walk): the decoder consumes the token directly, encoder
|
# directly, encoder bypassed. Ignored unless a SONIC whole-body controller is active.
|
||||||
# bypassed. Ignored unless a SONIC whole-body controller is active.
|
|
||||||
sonic_token_action: bool = False
|
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",
|
# Locomotion controller class name, e.g. "GrootLocomotionController",
|
||||||
# "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
|
# "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
|
||||||
controller: str | None = None
|
controller: str | None = None
|
||||||
|
|||||||
@@ -81,7 +81,7 @@ def load_groot_policies(
|
|||||||
# background thread alongside the (torch) upper-body policy, IK and sim; letting
|
# 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
|
# 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.
|
# 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)
|
so = make_ort_session_options(intra_op_num_threads=1, inter_op_num_threads=1)
|
||||||
policy_balance = ort.InferenceSession(balance_path, sess_options=so)
|
policy_balance = ort.InferenceSession(balance_path, sess_options=so)
|
||||||
|
|||||||
@@ -1,670 +0,0 @@
|
|||||||
#!/usr/bin/env python
|
|
||||||
|
|
||||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
|
||||||
#
|
|
||||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
||||||
# you may not use this file except in compliance with the License.
|
|
||||||
# You may obtain a copy of the License at
|
|
||||||
#
|
|
||||||
# http://www.apache.org/licenses/LICENSE-2.0
|
|
||||||
#
|
|
||||||
# Unless required by applicable law or agreed to in writing, software
|
|
||||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
||||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
||||||
# See the License for the specific language governing permissions and
|
|
||||||
# limitations under the License.
|
|
||||||
|
|
||||||
"""SONIC encoder/decoder pipeline for the Unitree G1 whole-body controller.
|
|
||||||
|
|
||||||
Pure-Python/ONNX re-implementation of the reference-tracking half of NVIDIA's SONIC
|
|
||||||
deploy stack (mirrors ``g1_deploy_onnx_ref.cpp``). Given a reference motion buffer
|
|
||||||
(joint targets + body orientation per frame) it produces 50 Hz joint-position targets
|
|
||||||
for the robot's PD controller. The upstream *motion planner* is intentionally absent:
|
|
||||||
here the reference is supplied directly by the caller (e.g. a 34-D OpenHLM / pi0.5 VLA
|
|
||||||
command per tick, in ``sonic_whole_body.py``).
|
|
||||||
|
|
||||||
Two cooperating ONNX models:
|
|
||||||
* **encoder** – compresses the reference window into a 64-D latent ``token``
|
|
||||||
(refreshed every ``ENCODER_UPDATE_EVERY`` ticks).
|
|
||||||
* **decoder** – every tick, maps the token + recent proprioception history to a
|
|
||||||
residual action that is scaled and added to ``DEFAULT_ANGLES``.
|
|
||||||
|
|
||||||
Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order)
|
|
||||||
and **MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` convert
|
|
||||||
between them. Quaternions are scalar-first ``(w, x, y, z)``.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import logging
|
|
||||||
import threading
|
|
||||||
from typing import TYPE_CHECKING
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
|
|
||||||
from lerobot.utils.import_utils import _onnxruntime_available
|
|
||||||
|
|
||||||
from ..g1_utils import (
|
|
||||||
ISAACLAB_TO_MUJOCO,
|
|
||||||
MUJOCO_TO_ISAACLAB,
|
|
||||||
G1_29_JointIndex,
|
|
||||||
get_gravity_orientation,
|
|
||||||
)
|
|
||||||
|
|
||||||
if TYPE_CHECKING or _onnxruntime_available:
|
|
||||||
import onnxruntime as ort
|
|
||||||
else:
|
|
||||||
ort = None
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
|
||||||
# ── Constants ────────────────────────────────────────────────────────────────
|
|
||||||
# Robot/motor physical constants and the joint-order permutation tables. All
|
|
||||||
# 29-vectors are in IsaacLab joint order unless the name says ``_MUJOCO``.
|
|
||||||
|
|
||||||
# Nominal standing pose (rad), 29 joints in IsaacLab order. Actions are residuals
|
|
||||||
# added on top of this; also used as the planner/encoder standing reference.
|
|
||||||
DEFAULT_ANGLES = np.array(
|
|
||||||
[
|
|
||||||
-0.312,
|
|
||||||
0.0,
|
|
||||||
0.0,
|
|
||||||
0.669,
|
|
||||||
-0.363,
|
|
||||||
0.0,
|
|
||||||
-0.312,
|
|
||||||
0.0,
|
|
||||||
0.0,
|
|
||||||
0.669,
|
|
||||||
-0.363,
|
|
||||||
0.0,
|
|
||||||
0.0,
|
|
||||||
0.0,
|
|
||||||
0.0,
|
|
||||||
0.2,
|
|
||||||
0.2,
|
|
||||||
0.0,
|
|
||||||
0.6,
|
|
||||||
0.0,
|
|
||||||
0.0,
|
|
||||||
0.0,
|
|
||||||
0.2,
|
|
||||||
-0.2,
|
|
||||||
0.0,
|
|
||||||
0.6,
|
|
||||||
0.0,
|
|
||||||
0.0,
|
|
||||||
0.0,
|
|
||||||
],
|
|
||||||
dtype=np.float32,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Per-motor-type parameters used to derive action scaling and PD gains. Keys are
|
|
||||||
# Unitree motor model names; ARMATURE = rotor inertia, EFFORT = torque limit (N·m).
|
|
||||||
NATURAL_FREQ = 10.0 * 2.0 * np.pi # target closed-loop stiffness bandwidth (rad/s)
|
|
||||||
ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425}
|
|
||||||
EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0}
|
|
||||||
|
|
||||||
|
|
||||||
def _action_scale(k):
|
|
||||||
"""Per-motor residual-action scale (maps policy output to joint-angle delta)."""
|
|
||||||
return 0.25 * EFFORT[k] / (ARMATURE[k] * NATURAL_FREQ**2)
|
|
||||||
|
|
||||||
|
|
||||||
# Per-joint motor model (IsaacLab order): legs, waist, then arms. Single source of
|
|
||||||
# truth for both ACTION_SCALE and compute_kp_kd().
|
|
||||||
MOTOR_MODELS = (
|
|
||||||
["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2
|
|
||||||
+ ["7520_14", "5020", "5020"]
|
|
||||||
+ ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2
|
|
||||||
)
|
|
||||||
ACTION_SCALE = np.array([_action_scale(k) for k in MOTOR_MODELS], dtype=np.float32) # (29,) IsaacLab order
|
|
||||||
|
|
||||||
CONTROL_DT = 0.02 # 50 Hz control period (s)
|
|
||||||
DEFAULT_HEIGHT = 0.788740 # nominal pelvis height (m)
|
|
||||||
TOKEN_DIM = 64 # encoder latent size
|
|
||||||
ENCODER_UPDATE_EVERY = 5 # refresh the encoder token every N ticks (decoder runs every tick)
|
|
||||||
DEBUG_PRINT_EVERY = 100 # ticks between debug prints
|
|
||||||
|
|
||||||
|
|
||||||
def _to_mujoco(a):
|
|
||||||
"""Apply the ``MUJOCO_TO_ISAACLAB`` gather to a 29-vector (deploy-order reorder).
|
|
||||||
|
|
||||||
NOTE: this returns ``a[MUJOCO_TO_ISAACLAB]``. The ``_mj`` suffixes and the exact
|
|
||||||
permutation direction throughout this module are a fixed convention validated
|
|
||||||
against the deployed SONIC ONNX policy (the encoder/decoder consume vectors in
|
|
||||||
this order). Do not "correct" the table or rename toward the opposite direction
|
|
||||||
without re-validating on hardware — the labels are historical, the ordering is
|
|
||||||
load-bearing.
|
|
||||||
"""
|
|
||||||
return a[MUJOCO_TO_ISAACLAB]
|
|
||||||
|
|
||||||
|
|
||||||
DEFAULT_ANGLES_MUJOCO = _to_mujoco(DEFAULT_ANGLES)
|
|
||||||
ENCODER_STANDING_REF = DEFAULT_ANGLES.copy()
|
|
||||||
|
|
||||||
# Joint-index subsets (IsaacLab order) used to slice encoder observations.
|
|
||||||
LOWER_BODY_IL = np.array([0, 3, 6, 9, 13, 17, 1, 4, 7, 10, 14, 18], dtype=np.int32) # 12 leg joints
|
|
||||||
WRIST_IL = np.array([23, 24, 25, 26, 27, 28], dtype=np.int32) # 6 wrist joints
|
|
||||||
VR_TARGET_DEF = np.zeros(9, dtype=np.float32) # 3-point VR position targets (mode 1)
|
|
||||||
VR_ORN_DEF = np.array([1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0], dtype=np.float32) # VR orn targets (mode 1)
|
|
||||||
SMPL_DEF = np.zeros(720, dtype=np.float32) # SMPL whole-body window default (mode 2)
|
|
||||||
|
|
||||||
# ── PD gains ─────────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
|
|
||||||
def compute_kp_kd():
|
|
||||||
"""Derive per-joint PD gains (kp, kd) from motor armature and target bandwidth.
|
|
||||||
|
|
||||||
Ankle and waist joints get a x2 factor for extra stiffness. Returns two
|
|
||||||
(29,) float32 arrays in IsaacLab joint order.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def s(k):
|
|
||||||
return ARMATURE[k] * NATURAL_FREQ**2
|
|
||||||
|
|
||||||
def d(k):
|
|
||||||
return 2.0 * 2.0 * ARMATURE[k] * NATURAL_FREQ
|
|
||||||
|
|
||||||
_double = {4, 5, 10, 11, 13, 14} # ankle + waist indices with factor 2
|
|
||||||
kp = np.array([2 * s(k) if i in _double else s(k) for i, k in enumerate(MOTOR_MODELS)], dtype=np.float32)
|
|
||||||
kd = np.array([2 * d(k) if i in _double else d(k) for i, k in enumerate(MOTOR_MODELS)], dtype=np.float32)
|
|
||||||
return kp, kd
|
|
||||||
|
|
||||||
|
|
||||||
_kp_kd = compute_kp_kd # backward-compatible alias
|
|
||||||
|
|
||||||
|
|
||||||
# ── Quaternion helpers ────────────────────────────────────────────────────────
|
|
||||||
# All quaternions are scalar-first (w, x, y, z). "heading" = yaw-only quaternion.
|
|
||||||
|
|
||||||
|
|
||||||
def quat_conj(q):
|
|
||||||
"""Quaternion conjugate (inverse for unit quaternions)."""
|
|
||||||
return np.array([q[0], -q[1], -q[2], -q[3]], dtype=np.float32)
|
|
||||||
|
|
||||||
|
|
||||||
def quat_mul(q1, q2):
|
|
||||||
"""Hamilton product ``q1 ⊗ q2``."""
|
|
||||||
w1, x1, y1, z1 = q1
|
|
||||||
w2, x2, y2, z2 = q2
|
|
||||||
return np.array(
|
|
||||||
[
|
|
||||||
w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
|
|
||||||
w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
|
|
||||||
w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
|
|
||||||
w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
|
|
||||||
],
|
|
||||||
dtype=np.float32,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def quat_to_6d(q):
|
|
||||||
"""Quaternion → 6-D rotation representation (first two rotated basis rows)."""
|
|
||||||
w, x, y, z = q
|
|
||||||
return np.array(
|
|
||||||
[
|
|
||||||
1 - 2 * (y * y + z * z),
|
|
||||||
2 * (x * y - z * w),
|
|
||||||
2 * (x * y + z * w),
|
|
||||||
1 - 2 * (x * x + z * z),
|
|
||||||
2 * (x * z - y * w),
|
|
||||||
2 * (y * z + x * w),
|
|
||||||
],
|
|
||||||
dtype=np.float32,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def calc_heading(q):
|
|
||||||
"""Extract the yaw (heading) angle in radians from a quaternion."""
|
|
||||||
w, x, y, z = q
|
|
||||||
return float(np.arctan2(2 * (x * y + w * z), 1 - 2 * (y * y + z * z)))
|
|
||||||
|
|
||||||
|
|
||||||
def heading_quat(q, sign=1.0):
|
|
||||||
"""Yaw-only quaternion for ``q``'s heading (``sign=-1`` gives its inverse)."""
|
|
||||||
a = sign * calc_heading(q) / 2.0
|
|
||||||
return np.array([np.cos(a), 0, 0, np.sin(a)], dtype=np.float64)
|
|
||||||
|
|
||||||
|
|
||||||
def heading_quat_inv(q):
|
|
||||||
"""Inverse yaw-only quaternion for ``q``'s heading."""
|
|
||||||
return heading_quat(q, -1.0)
|
|
||||||
|
|
||||||
|
|
||||||
def quat_slerp(q0, q1, t):
|
|
||||||
"""Spherical linear interpolation between two quaternions (scalar ``t``)."""
|
|
||||||
q0 = q0 / (np.linalg.norm(q0) + 1e-12)
|
|
||||||
q1 = q1 / (np.linalg.norm(q1) + 1e-12)
|
|
||||||
dot = float(np.dot(q0, q1))
|
|
||||||
if dot < 0:
|
|
||||||
q1, dot = -q1, -dot
|
|
||||||
dot = min(dot, 1.0)
|
|
||||||
if dot > 0.9995:
|
|
||||||
r = q0 + t * (q1 - q0)
|
|
||||||
return r / (np.linalg.norm(r) + 1e-12)
|
|
||||||
th = np.arccos(dot)
|
|
||||||
st = np.sin(th)
|
|
||||||
return (np.sin((1 - t) * th) / st) * q0 + (np.sin(t * th) / st) * q1
|
|
||||||
|
|
||||||
|
|
||||||
def quat_slerp_batch(q0, q1, t):
|
|
||||||
"""Vectorized slerp over arrays of quaternions with a per-row parameter ``t``."""
|
|
||||||
q0 = q0 / (np.linalg.norm(q0, axis=1, keepdims=True) + 1e-12)
|
|
||||||
q1 = q1 / (np.linalg.norm(q1, axis=1, keepdims=True) + 1e-12)
|
|
||||||
dot = np.sum(q0 * q1, axis=1)
|
|
||||||
neg = dot < 0
|
|
||||||
q1 = q1.copy()
|
|
||||||
q1[neg] = -q1[neg]
|
|
||||||
dot[neg] = -dot[neg]
|
|
||||||
dot = np.clip(dot, -1, 1)
|
|
||||||
lin = dot > 0.9995
|
|
||||||
th = np.arccos(dot)
|
|
||||||
st = np.where(np.sin(th) == 0, 1, np.sin(th))
|
|
||||||
c0 = np.sin((1 - t) * th) / st
|
|
||||||
c1 = np.sin(t * th) / st
|
|
||||||
c0[lin] = 1 - t[lin]
|
|
||||||
c1[lin] = t[lin]
|
|
||||||
r = c0[:, None] * q0 + c1[:, None] * q1
|
|
||||||
return r / (np.linalg.norm(r, axis=1, keepdims=True) + 1e-12)
|
|
||||||
|
|
||||||
|
|
||||||
def ort_providers(force_cpu: bool = False) -> list[str]:
|
|
||||||
"""Prefer CUDA for enc/dec/planner (matches deploy when onnxruntime-gpu is installed)."""
|
|
||||||
avail = ort.get_available_providers()
|
|
||||||
if not force_cpu and "CUDAExecutionProvider" in avail:
|
|
||||||
return ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
|
||||||
return ["CPUExecutionProvider"]
|
|
||||||
|
|
||||||
|
|
||||||
def make_ort_session_options(intra_op_num_threads: int | None = None,
|
|
||||||
inter_op_num_threads: int | None = None):
|
|
||||||
"""Build ONNX Runtime SessionOptions (quiet logging).
|
|
||||||
|
|
||||||
Pass thread counts to cap ORT's CPU pool. These tiny MLP policies are latency-
|
|
||||||
bound, not throughput-bound, so letting ORT grab every core just starves the
|
|
||||||
real-time control loop / torch policy / IK solver and causes stutter. 1 intra +
|
|
||||||
1 inter thread is plenty and lowest-latency for a per-step MLP inference.
|
|
||||||
"""
|
|
||||||
so = ort.SessionOptions()
|
|
||||||
so.log_severity_level = 3
|
|
||||||
if intra_op_num_threads is not None:
|
|
||||||
so.intra_op_num_threads = intra_op_num_threads
|
|
||||||
if inter_op_num_threads is not None:
|
|
||||||
so.inter_op_num_threads = inter_op_num_threads
|
|
||||||
return so
|
|
||||||
|
|
||||||
|
|
||||||
# ── Encoder / Decoder ─────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
|
|
||||||
class StandingEncoderDecoder:
|
|
||||||
"""Runs the encoder + decoder ONNX models and owns the proprioception history.
|
|
||||||
|
|
||||||
Each tick it appends the latest robot state to 10-frame history buffers, builds
|
|
||||||
the encoder observation (1762-D, layout depends on ``encode_mode``) to refresh
|
|
||||||
the 64-D ``token``, then builds the decoder observation (994-D) and maps
|
|
||||||
``token + history`` to a residual action added onto ``DEFAULT_ANGLES``.
|
|
||||||
|
|
||||||
``PlannerController`` subclasses this to source the reference from a live,
|
|
||||||
planner-generated motion buffer instead of a fixed standing pose.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, encoder, decoder):
|
|
||||||
self.encoder, self.decoder = encoder, decoder
|
|
||||||
self.encoder_input = encoder.get_inputs()[0].name
|
|
||||||
self.decoder_input = decoder.get_inputs()[0].name
|
|
||||||
enc_dim = int(encoder.get_inputs()[0].shape[1])
|
|
||||||
dec_dim = int(decoder.get_inputs()[0].shape[1])
|
|
||||||
if enc_dim != 1762 or dec_dim != 994:
|
|
||||||
raise RuntimeError(f"Unexpected dims encoder={enc_dim}, decoder={dec_dim}")
|
|
||||||
self.token = np.zeros(TOKEN_DIM, np.float32)
|
|
||||||
self.last_action_mj = np.zeros(29, np.float32)
|
|
||||||
self.h_q_mj = [np.zeros(29, np.float32)] * 10
|
|
||||||
self.h_dq_mj = [np.zeros(29, np.float32)] * 10
|
|
||||||
self.h_ang = [np.zeros(3, np.float32)] * 10
|
|
||||||
self.h_act_mj = [np.zeros(29, np.float32)] * 10
|
|
||||||
self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
|
|
||||||
self.init_base_quat = np.array([1, 0, 0, 0], np.float32)
|
|
||||||
self.init_ref_quat = np.array([1, 0, 0, 0], np.float32)
|
|
||||||
self._heading_init = False
|
|
||||||
self.encode_mode = 0
|
|
||||||
self.vr_3point_local_target = VR_TARGET_DEF.copy()
|
|
||||||
self.vr_3point_local_orn_target = VR_ORN_DEF.copy()
|
|
||||||
self.smpl_joints_10frame_step1 = SMPL_DEF.copy()
|
|
||||||
# Optional per-frame SMPL root orientation (wxyz) for the mode-2 anchor.
|
|
||||||
# When None, the anchor falls back to the planner reference body quat.
|
|
||||||
self.smpl_root_quat = None
|
|
||||||
self.set_zero_reference()
|
|
||||||
|
|
||||||
def reset(self):
|
|
||||||
"""Clear the token, 10-frame proprioception history and heading init.
|
|
||||||
|
|
||||||
``UnitreeG1.reset()`` relies on this so the first decoder outputs of a new
|
|
||||||
episode are not contaminated by the previous episode's state.
|
|
||||||
"""
|
|
||||||
self.token = np.zeros(TOKEN_DIM, np.float32)
|
|
||||||
self.last_action_mj = np.zeros(29, np.float32)
|
|
||||||
self.h_q_mj = [np.zeros(29, np.float32)] * 10
|
|
||||||
self.h_dq_mj = [np.zeros(29, np.float32)] * 10
|
|
||||||
self.h_ang = [np.zeros(3, np.float32)] * 10
|
|
||||||
self.h_act_mj = [np.zeros(29, np.float32)] * 10
|
|
||||||
self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
|
|
||||||
self.init_base_quat = np.array([1, 0, 0, 0], np.float32)
|
|
||||||
self.init_ref_quat = np.array([1, 0, 0, 0], np.float32)
|
|
||||||
self._heading_init = False
|
|
||||||
|
|
||||||
def update_history(self, q, dq, ang, quat):
|
|
||||||
"""Push the latest proprioception (pos/vel/gyro/orientation) into the 10-frame buffers."""
|
|
||||||
quat = quat / (np.linalg.norm(quat) + 1e-8)
|
|
||||||
q_mj = _to_mujoco(q)
|
|
||||||
dq_mj = _to_mujoco(dq)
|
|
||||||
self.h_q_mj = [q_mj - DEFAULT_ANGLES_MUJOCO] + self.h_q_mj[:-1]
|
|
||||||
self.h_dq_mj = [dq_mj] + self.h_dq_mj[:-1]
|
|
||||||
self.h_ang = [ang.copy()] + self.h_ang[:-1]
|
|
||||||
self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
|
|
||||||
self.h_quat = [quat.copy()] + self.h_quat[:-1]
|
|
||||||
if not self._heading_init:
|
|
||||||
self.init_base_quat = quat.copy()
|
|
||||||
self._heading_init = True
|
|
||||||
|
|
||||||
def _heading_quat(self, q):
|
|
||||||
h = calc_heading(q) / 2.0
|
|
||||||
return np.array([np.cos(h), 0, 0, np.sin(h)], np.float32)
|
|
||||||
|
|
||||||
def _heading_quat_inv(self, q):
|
|
||||||
h = calc_heading(q) / 2.0
|
|
||||||
return np.array([np.cos(-h), 0, 0, np.sin(-h)], np.float32)
|
|
||||||
|
|
||||||
def _anchor_6d(self, base_quat, ref_quat=None):
|
|
||||||
"""6-D orientation error between the robot base and the (heading-aligned) reference."""
|
|
||||||
if ref_quat is None:
|
|
||||||
ref_quat = self.init_ref_quat
|
|
||||||
delta = quat_mul(self._heading_quat(self.init_base_quat), self._heading_quat_inv(self.init_ref_quat))
|
|
||||||
new_ref = quat_mul(delta, ref_quat)
|
|
||||||
return quat_to_6d(quat_mul(quat_conj(base_quat), new_ref))
|
|
||||||
|
|
||||||
def set_zero_reference(self):
|
|
||||||
"""Initialize the reference to a single standing frame (used before a plan exists)."""
|
|
||||||
self.motion_joint_positions = [ENCODER_STANDING_REF.copy()]
|
|
||||||
self.motion_joint_velocities = [np.zeros(29, np.float32)]
|
|
||||||
self.motion_body_quats = [np.array([1, 0, 0, 0], np.float32)]
|
|
||||||
self.motion_body_z = [DEFAULT_HEIGHT]
|
|
||||||
self.motion_timesteps = 1
|
|
||||||
self.freeze_ref_frame = 0
|
|
||||||
self.init_ref_quat = self.motion_body_quats[0].copy()
|
|
||||||
|
|
||||||
def build_encoder_obs(self):
|
|
||||||
"""Assemble the 1762-D encoder input; slot layout depends on ``encode_mode``.
|
|
||||||
|
|
||||||
mode 0 = locomotion (ref joint pos + anchor), 1 = 3-point VR teleop
|
|
||||||
(lower-body ref + VR targets), 2 = SMPL whole-body window + anchor/wrist.
|
|
||||||
"""
|
|
||||||
obs = np.zeros(1762, np.float32)
|
|
||||||
obs[0] = float(self.encode_mode)
|
|
||||||
rf = min(self.freeze_ref_frame, self.motion_timesteps - 1)
|
|
||||||
ref_pos, ref_quat = self.motion_joint_positions[rf], self.motion_body_quats[rf]
|
|
||||||
if self.encode_mode == 0:
|
|
||||||
for f in range(10):
|
|
||||||
obs[4 + 29 * f : 4 + 29 * (f + 1)] = ref_pos
|
|
||||||
obs[601 + 6 * f : 601 + 6 * (f + 1)] = self._anchor_6d(self.h_quat[0], ref_quat)
|
|
||||||
elif self.encode_mode == 1:
|
|
||||||
ref_lower = ref_pos[LOWER_BODY_IL]
|
|
||||||
for f in range(10):
|
|
||||||
obs[661 + 12 * f : 661 + 12 * (f + 1)] = ref_lower
|
|
||||||
obs[901:910] = self.vr_3point_local_target
|
|
||||||
obs[910:922] = self.vr_3point_local_orn_target
|
|
||||||
obs[595:601] = self._anchor_6d(self.h_quat[0], ref_quat)
|
|
||||||
elif self.encode_mode == 2:
|
|
||||||
# Prefer the SMPL clip/stream root orientation for the anchor; fall
|
|
||||||
# back to the planner reference body quat when no root is provided.
|
|
||||||
anchor_ref = self.smpl_root_quat if self.smpl_root_quat is not None else ref_quat
|
|
||||||
obs[922:1642] = self.smpl_joints_10frame_step1
|
|
||||||
for f in range(10):
|
|
||||||
obs[1642 + 6 * f : 1642 + 6 * (f + 1)] = self._anchor_6d(self.h_quat[0], anchor_ref)
|
|
||||||
obs[1702 + 6 * f : 1702 + 6 * (f + 1)] = ref_pos[WRIST_IL]
|
|
||||||
else:
|
|
||||||
raise RuntimeError(f"Unsupported encoder mode: {self.encode_mode}")
|
|
||||||
return obs
|
|
||||||
|
|
||||||
def build_decoder_obs(self):
|
|
||||||
"""Assemble the 994-D decoder input: token + 10-frame proprioception history + gravity."""
|
|
||||||
obs = np.zeros(994, np.float32)
|
|
||||||
off = 0
|
|
||||||
obs[off : off + 64] = self.token
|
|
||||||
off += 64
|
|
||||||
for h, sz in [
|
|
||||||
(list(reversed(self.h_ang)), 3),
|
|
||||||
(list(reversed(self.h_q_mj)), 29),
|
|
||||||
(list(reversed(self.h_dq_mj)), 29),
|
|
||||||
(list(reversed(self.h_act_mj)), 29),
|
|
||||||
]:
|
|
||||||
for f in range(10):
|
|
||||||
obs[off : off + sz] = h[f]
|
|
||||||
off += sz
|
|
||||||
for q in reversed(self.h_quat):
|
|
||||||
obs[off : off + 3] = get_gravity_orientation(q)
|
|
||||||
off += 3
|
|
||||||
assert off == 994, f"Decoder obs mismatch: {off}"
|
|
||||||
return obs
|
|
||||||
|
|
||||||
def run_encoder(self):
|
|
||||||
"""Run the encoder ONNX model and return the fresh 64-D token."""
|
|
||||||
return (
|
|
||||||
self.encoder.run(None, {self.encoder_input: self.build_encoder_obs().reshape(1, -1)})[0]
|
|
||||||
.squeeze()
|
|
||||||
.astype(np.float32)
|
|
||||||
)
|
|
||||||
|
|
||||||
def step(self, robot_obs, update_encoder, debug=False):
|
|
||||||
"""One control tick: read robot obs, (optionally) re-encode, decode → joint targets.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
robot_obs: dict with ``<joint>.q``/``.dq`` and ``imu.*`` fields.
|
|
||||||
update_encoder: refresh the token this tick (else reuse the cached one).
|
|
||||||
debug: print action/delta norms.
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
dict of ``<joint>.q`` target positions (rad) in IsaacLab joint order.
|
|
||||||
"""
|
|
||||||
jnames = [m.name for m in G1_29_JointIndex]
|
|
||||||
q = np.array(
|
|
||||||
[
|
|
||||||
robot_obs.get(f"{n}.q", DEFAULT_ANGLES[m.value])
|
|
||||||
for m, n in zip(G1_29_JointIndex, jnames, strict=False)
|
|
||||||
],
|
|
||||||
np.float32,
|
|
||||||
)
|
|
||||||
dq = np.array([robot_obs.get(f"{n}.dq", 0.0) for n in jnames], np.float32)
|
|
||||||
quat = np.array(
|
|
||||||
[
|
|
||||||
robot_obs.get("imu.quat.w", 1),
|
|
||||||
robot_obs.get("imu.quat.x", 0),
|
|
||||||
robot_obs.get("imu.quat.y", 0),
|
|
||||||
robot_obs.get("imu.quat.z", 0),
|
|
||||||
],
|
|
||||||
np.float32,
|
|
||||||
)
|
|
||||||
ang = np.array([robot_obs.get(f"imu.gyro.{a}", 0) for a in "xyz"], np.float32)
|
|
||||||
self.update_history(q, dq, ang, quat)
|
|
||||||
if update_encoder:
|
|
||||||
self.token = self.run_encoder()
|
|
||||||
action_mj = (
|
|
||||||
self.decoder.run(None, {self.decoder_input: self.build_decoder_obs().reshape(1, -1)})[0]
|
|
||||||
.squeeze()
|
|
||||||
.astype(np.float32)
|
|
||||||
)
|
|
||||||
self.last_action_mj = action_mj.copy()
|
|
||||||
target = DEFAULT_ANGLES + action_mj[ISAACLAB_TO_MUJOCO] * ACTION_SCALE
|
|
||||||
if debug:
|
|
||||||
delta = target - q
|
|
||||||
logger.debug(
|
|
||||||
"token_norm=%.4f action_norm=%.4f delta_max=%.4f delta_rms=%.4f",
|
|
||||||
np.linalg.norm(self.token),
|
|
||||||
np.linalg.norm(action_mj),
|
|
||||||
np.max(np.abs(delta)),
|
|
||||||
np.sqrt(np.mean(delta**2)),
|
|
||||||
)
|
|
||||||
return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
|
|
||||||
|
|
||||||
|
|
||||||
class PlannerController(StandingEncoderDecoder):
|
|
||||||
"""Encoder/decoder driven by a caller-supplied, rolling motion buffer.
|
|
||||||
|
|
||||||
Extends ``StandingEncoderDecoder`` so the reference comes from a motion buffer
|
|
||||||
(a lookahead window with per-frame velocities) instead of a single fixed pose,
|
|
||||||
and handles heading re-initialization on the first frame / after a reset.
|
|
||||||
``motion_lock`` guards the buffer, which the whole-body controller rewrites each
|
|
||||||
tick from the incoming command. The class name is retained for continuity with
|
|
||||||
the SONIC reference; no motion planner is involved.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, encoder, decoder):
|
|
||||||
super().__init__(encoder, decoder)
|
|
||||||
self.ref_cursor = 0
|
|
||||||
self.motion_timesteps = 0
|
|
||||||
self.motion_joint_positions = np.zeros((1500, 29), np.float64)
|
|
||||||
self.motion_joint_velocities = np.zeros((1500, 29), np.float64)
|
|
||||||
self.motion_body_quats = np.zeros((1500, 4), np.float64)
|
|
||||||
self.motion_body_quats[:, 0] = 1.0
|
|
||||||
self.motion_body_pos = np.zeros((1500, 3), np.float64)
|
|
||||||
self.init_ref_quat = np.array([1, 0, 0, 0], np.float64)
|
|
||||||
self.heading_init_base_quat = np.array([1, 0, 0, 0], np.float64)
|
|
||||||
self.delta_heading = 0.0
|
|
||||||
self.reinit_heading = False
|
|
||||||
self.playing = self.first_motion = False
|
|
||||||
self.motion_lock = threading.Lock()
|
|
||||||
|
|
||||||
def reset(self):
|
|
||||||
"""Full reset: clear enc/dec state (super) plus the motion buffer and heading.
|
|
||||||
|
|
||||||
Forces a heading re-init on the next ``step`` so the reference frame is
|
|
||||||
re-latched to the post-reset robot orientation.
|
|
||||||
"""
|
|
||||||
super().reset()
|
|
||||||
with self.motion_lock:
|
|
||||||
self.ref_cursor = 0
|
|
||||||
self.motion_timesteps = 0
|
|
||||||
self.motion_joint_positions[:] = 0.0
|
|
||||||
self.motion_joint_velocities[:] = 0.0
|
|
||||||
self.motion_body_quats[:] = 0.0
|
|
||||||
self.motion_body_quats[:, 0] = 1.0
|
|
||||||
self.motion_body_pos[:] = 0.0
|
|
||||||
self.init_ref_quat = np.array([1, 0, 0, 0], np.float64)
|
|
||||||
self.heading_init_base_quat = np.array([1, 0, 0, 0], np.float64)
|
|
||||||
self.delta_heading = 0.0
|
|
||||||
self.first_motion = False
|
|
||||||
self.playing = False
|
|
||||||
self.reinit_heading = True
|
|
||||||
|
|
||||||
def _heading_apply_delta(self):
|
|
||||||
"""Heading correction quaternion (init base-vs-ref heading + operator ``delta_heading``)."""
|
|
||||||
delta = quat_mul(
|
|
||||||
heading_quat(self.heading_init_base_quat).astype(np.float32),
|
|
||||||
heading_quat_inv(self.init_ref_quat).astype(np.float32),
|
|
||||||
)
|
|
||||||
if self.delta_heading:
|
|
||||||
h = self.delta_heading / 2.0
|
|
||||||
delta = quat_mul(np.array([np.cos(h), 0, 0, np.sin(h)], np.float32), delta)
|
|
||||||
return delta
|
|
||||||
|
|
||||||
def _anchor_6d(self, base_quat, ref_quat=None):
|
|
||||||
"""6-D base-vs-reference orientation error, including the operator heading delta."""
|
|
||||||
if ref_quat is None:
|
|
||||||
ref_quat = self.init_ref_quat
|
|
||||||
new_ref = quat_mul(self._heading_apply_delta(), ref_quat.astype(np.float32))
|
|
||||||
return quat_to_6d(quat_mul(quat_conj(base_quat.astype(np.float32)), new_ref))
|
|
||||||
|
|
||||||
def build_encoder_obs(self):
|
|
||||||
"""Encoder input sourced from the live motion buffer (mode 0/2), lock-protected."""
|
|
||||||
obs = np.zeros(1762, np.float32)
|
|
||||||
obs[0] = float(self.encode_mode)
|
|
||||||
with self.motion_lock:
|
|
||||||
if self.encode_mode == 2:
|
|
||||||
# SMPL whole-body imitation: the 720-dim SMPL window carries the
|
|
||||||
# target pose; the planner reference frame supplies anchor + wrist.
|
|
||||||
rf = min(self.ref_cursor, self.motion_timesteps - 1)
|
|
||||||
ref_pos = self.motion_joint_positions[rf].astype(np.float32)
|
|
||||||
ref_quat = self.motion_body_quats[rf].astype(np.float32)
|
|
||||||
# Prefer the SMPL clip/stream root orientation (if provided) so the
|
|
||||||
# anchor tracks the operator's/clip's heading; else planner ref.
|
|
||||||
if self.smpl_root_quat is not None:
|
|
||||||
ref_quat = np.asarray(self.smpl_root_quat, np.float32)
|
|
||||||
anchor = self._anchor_6d(self.h_quat[0], ref_quat)
|
|
||||||
wrist = ref_pos[WRIST_IL]
|
|
||||||
obs[922:1642] = self.smpl_joints_10frame_step1
|
|
||||||
for f in range(10):
|
|
||||||
obs[1642 + 6 * f : 1642 + 6 * (f + 1)] = anchor
|
|
||||||
obs[1702 + 6 * f : 1702 + 6 * (f + 1)] = wrist
|
|
||||||
return obs
|
|
||||||
if self.encode_mode == 1:
|
|
||||||
# 3-point VR teleop: the upper body tracks the VR wrist/neck targets
|
|
||||||
# while the planner reference supplies the lower body + anchor. Lower
|
|
||||||
# body is per-frame (step 5) like mode 0; the VR targets are current.
|
|
||||||
rf = min(self.ref_cursor, self.motion_timesteps - 1)
|
|
||||||
obs[595:601] = self._anchor_6d(self.h_quat[0], self.motion_body_quats[rf].astype(np.float32))
|
|
||||||
for f in range(10):
|
|
||||||
tf = min(
|
|
||||||
self.ref_cursor + f * 5 if self.playing else self.ref_cursor,
|
|
||||||
self.motion_timesteps - 1,
|
|
||||||
)
|
|
||||||
ref_lower = self.motion_joint_positions[tf].astype(np.float32)[LOWER_BODY_IL]
|
|
||||||
obs[661 + 12 * f : 661 + 12 * (f + 1)] = ref_lower
|
|
||||||
obs[901:910] = self.vr_3point_local_target
|
|
||||||
obs[910:922] = self.vr_3point_local_orn_target
|
|
||||||
return obs
|
|
||||||
for f in range(10):
|
|
||||||
tf = min(
|
|
||||||
self.ref_cursor + f * 5 if self.playing else self.ref_cursor, self.motion_timesteps - 1
|
|
||||||
)
|
|
||||||
obs[4 + 29 * f : 4 + 29 * (f + 1)] = self.motion_joint_positions[tf].astype(np.float32)
|
|
||||||
if self.playing:
|
|
||||||
obs[294 + 29 * f : 294 + 29 * (f + 1)] = self.motion_joint_velocities[tf].astype(
|
|
||||||
np.float32
|
|
||||||
)
|
|
||||||
obs[601 + 6 * f : 601 + 6 * (f + 1)] = self._anchor_6d(
|
|
||||||
self.h_quat[0], self.motion_body_quats[tf].astype(np.float32)
|
|
||||||
)
|
|
||||||
return obs
|
|
||||||
|
|
||||||
def step(self, robot_obs, update_encoder, debug=False):
|
|
||||||
"""Re-init the heading reference on first frame / after a reset, then run the base step."""
|
|
||||||
if robot_obs and (self.first_motion or self.reinit_heading):
|
|
||||||
q = None
|
|
||||||
if "imu.quat.w" in robot_obs:
|
|
||||||
q = np.array(
|
|
||||||
[
|
|
||||||
robot_obs["imu.quat.w"],
|
|
||||||
robot_obs["imu.quat.x"],
|
|
||||||
robot_obs["imu.quat.y"],
|
|
||||||
robot_obs["imu.quat.z"],
|
|
||||||
],
|
|
||||||
np.float64,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
q = robot_obs.get("imu.quaternion")
|
|
||||||
if q is not None:
|
|
||||||
q = np.array(q, np.float64)
|
|
||||||
if q is not None:
|
|
||||||
self.heading_init_base_quat = np.array(q, np.float64)
|
|
||||||
with self.motion_lock:
|
|
||||||
rf = min(self.ref_cursor, self.motion_timesteps - 1)
|
|
||||||
if self.encode_mode == 2 and self.smpl_root_quat is not None:
|
|
||||||
# Anchor the heading delta to the SMPL root at init so the
|
|
||||||
# robot turns *relative* to the clip/operator start heading.
|
|
||||||
self.init_ref_quat = np.asarray(self.smpl_root_quat, np.float64)
|
|
||||||
else:
|
|
||||||
self.init_ref_quat = self.motion_body_quats[rf].copy()
|
|
||||||
self.delta_heading = 0.0
|
|
||||||
self.first_motion = False
|
|
||||||
self.reinit_heading = False
|
|
||||||
logger.debug("[Heading] init quat: %s", self.heading_init_base_quat)
|
|
||||||
return super().step(robot_obs, update_encoder=update_encoder, debug=debug)
|
|
||||||
|
|
||||||
def advance_cursor(self):
|
|
||||||
"""Advance the reference cursor one frame per 50 Hz tick (no wall-clock catch-up)."""
|
|
||||||
if not self.playing:
|
|
||||||
return
|
|
||||||
with self.motion_lock:
|
|
||||||
if self.motion_timesteps > 0:
|
|
||||||
self.ref_cursor = min(self.ref_cursor + 1, self.motion_timesteps - 1)
|
|
||||||
@@ -14,57 +14,40 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# 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 __future__ import annotations
|
||||||
|
|
||||||
from collections import deque
|
|
||||||
import logging
|
import logging
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from huggingface_hub import hf_hub_download
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
from huggingface_hub import hf_hub_download
|
||||||
|
|
||||||
from lerobot.utils.import_utils import _onnxruntime_available, require_package
|
from lerobot.utils.import_utils import _onnxruntime_available, require_package
|
||||||
|
|
||||||
from ..g1_utils import (
|
from ..g1_utils import (
|
||||||
|
ISAACLAB_TO_MUJOCO,
|
||||||
MUJOCO_TO_ISAACLAB,
|
MUJOCO_TO_ISAACLAB,
|
||||||
WB_ACTION_DIM,
|
|
||||||
G1_29_JointIndex,
|
G1_29_JointIndex,
|
||||||
|
get_gravity_orientation,
|
||||||
lowstate_to_obs,
|
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,
|
make_ort_session_options,
|
||||||
ort_providers,
|
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:
|
if TYPE_CHECKING or _onnxruntime_available:
|
||||||
import onnxruntime as ort
|
import onnxruntime as ort
|
||||||
else:
|
else:
|
||||||
@@ -72,18 +55,110 @@ else:
|
|||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
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
|
# 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
|
# from the robot's initial measured pose into the policy's commanded target, so control
|
||||||
# eases in without a snap on the first command.
|
# eases in without a snap on the first command.
|
||||||
INIT_RAMP_S = 3.0
|
INIT_RAMP_S = 3.0
|
||||||
|
|
||||||
# Neutral ("zero pose") SONIC token, held by token_mode until the first real token
|
# Neutral ("zero pose") SONIC token, held by token_mode until the first real token arrives.
|
||||||
# arrives. Captured from the encoder's own output while the robot stood idle in sim
|
# Captured from the encoder's own output while the robot stood idle in sim: the encoder is
|
||||||
# (capture_neutral_token.py): the encoder is an FSQ bottleneck (~5 bit/dim, 15.5 half-
|
# an FSQ bottleneck (~5 bit/dim, Div(16)), so its tokens live on the 1/16 grid. We store the
|
||||||
# width, Div(16)), so its tokens live on the 1/16 grid. We store the integer FSQ codes
|
# integer FSQ codes and rescale by 1/16, giving an exact on-grid token -- unlike the literal
|
||||||
# and rescale by the same 1/16 step, 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.
|
||||||
# all-zero token, which is off the encoder's learned manifold and decodes to a slightly
|
# This one decodes to a stable, natural standing pose.
|
||||||
# goofy stance. This one decodes to a stable, natural standing pose.
|
|
||||||
_NEUTRAL_TOKEN_CODES = np.array(
|
_NEUTRAL_TOKEN_CODES = np.array(
|
||||||
[-1, 3, 1, -1, 1, -3, 6, 1, 1, 1, -2, -4, -2, 0, -3, -1,
|
[-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,
|
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
|
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:
|
def _extract_token_from_action(action: dict | None) -> np.ndarray | None:
|
||||||
"""Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or 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
|
The token-only interface: the caller supplies the 64-D encoder latent directly (e.g. a
|
||||||
encoder, the caller supplies the 64-D encoder latent directly (e.g. a recorded
|
token-output VLA's action), which the decoder consumes with the encoder bypassed.
|
||||||
``action.motion_token`` column), which the decoder consumes with the encoder
|
Requires the full dense token; a partial one is ignored (returns None).
|
||||||
bypassed. Requires the full dense token; a partial one is ignored (returns None).
|
|
||||||
"""
|
"""
|
||||||
if not action:
|
if not action:
|
||||||
return None
|
return None
|
||||||
keys = [token_action_key(i) for i in range(TOKEN_DIM)]
|
keys = [token_action_key(i) for i in range(TOKEN_DIM)]
|
||||||
if any(key not in action for key in keys):
|
if any(key not in action for key in keys):
|
||||||
return None
|
return None
|
||||||
return np.fromiter(
|
return np.fromiter((float(action[key]) for key in keys), dtype=np.float32, count=TOKEN_DIM)
|
||||||
(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]:
|
class SonicDecoder:
|
||||||
"""Map a 34-D OpenHLM whole-body command to a SONIC mode-0 reference.
|
"""Runs the SONIC decoder ONNX model and owns the proprioception history.
|
||||||
|
|
||||||
Returns ``(ref29, anchor_quat)`` where ``ref29`` is the 29 joint targets in
|
Each tick it appends the latest robot state to 10-frame history buffers, then maps the
|
||||||
IsaacLab order (what SONIC's ``motion_joint_positions`` expects) and
|
supplied 64-D ``token`` + that history to a residual action added onto
|
||||||
``anchor_quat`` (wxyz) encodes the root roll/pitch (yaw=0).
|
``DEFAULT_ANGLES``. The encoder is bypassed entirely (token supplied by the policy).
|
||||||
|
|
||||||
OpenHLM layout : [L-arm 0:7, L-grip 7, R-arm 8:15, R-grip 15,
|
|
||||||
L-leg 16:22, R-leg 22:28, waist 28:31, root rp+yaw 31:34]
|
|
||||||
The 29 joints are first assembled in MuJoCo / Unitree-SDK order
|
|
||||||
([L-leg 0:6, R-leg 6:12, waist 12:15, L-arm 15:22, R-arm 22:29] — the
|
|
||||||
``G1_29_JointIndex`` grouping OpenHLM uses), then permuted to IsaacLab order via
|
|
||||||
``MUJOCO_TO_ISAACLAB``. Grippers (7, 15) are not part of the 29-DoF SONIC
|
|
||||||
reference, and yaw-rate (33) is integrated into the heading by the caller (it
|
|
||||||
cannot be represented in this static per-tick anchor).
|
|
||||||
"""
|
"""
|
||||||
ref_mj = np.zeros(29, np.float32) # MuJoCo / Unitree-SDK grouped order
|
|
||||||
ref_mj[0:6] = wb[16:22] # left leg
|
def __init__(self, decoder):
|
||||||
ref_mj[6:12] = wb[22:28] # right leg
|
self.decoder = decoder
|
||||||
ref_mj[12:15] = wb[28:31] # waist
|
self.decoder_input = decoder.get_inputs()[0].name
|
||||||
ref_mj[15:22] = wb[0:7] # left arm
|
dec_dim = int(decoder.get_inputs()[0].shape[1])
|
||||||
ref_mj[22:29] = wb[8:15] # right arm
|
if dec_dim != 994:
|
||||||
ref = ref_mj[MUJOCO_TO_ISAACLAB].astype(np.float32) # -> IsaacLab order for SONIC
|
raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected 994)")
|
||||||
roll, pitch = float(wb[31]), float(wb[32])
|
self.token = np.zeros(TOKEN_DIM, np.float32)
|
||||||
cr, sr, cp, sp = np.cos(roll / 2), np.sin(roll / 2), np.cos(pitch / 2), np.sin(pitch / 2)
|
self.last_action_mj = np.zeros(29, np.float32)
|
||||||
anchor = np.array([cr * cp, sr * cp, cr * sp, sr * sp], np.float32) # Rx(roll)·Ry(pitch)
|
self.h_q_mj = [np.zeros(29, np.float32)] * 10
|
||||||
return ref, anchor
|
self.h_dq_mj = [np.zeros(29, np.float32)] * 10
|
||||||
|
self.h_ang = [np.zeros(3, np.float32)] * 10
|
||||||
|
self.h_act_mj = [np.zeros(29, np.float32)] * 10
|
||||||
|
self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
|
||||||
|
|
||||||
|
def reset(self):
|
||||||
|
"""Clear the token and 10-frame proprioception history.
|
||||||
|
|
||||||
|
``UnitreeG1.reset()`` relies on this so the first decoder outputs of a new episode
|
||||||
|
are not contaminated by the previous episode's state.
|
||||||
|
"""
|
||||||
|
self.token = np.zeros(TOKEN_DIM, np.float32)
|
||||||
|
self.last_action_mj = np.zeros(29, np.float32)
|
||||||
|
self.h_q_mj = [np.zeros(29, np.float32)] * 10
|
||||||
|
self.h_dq_mj = [np.zeros(29, np.float32)] * 10
|
||||||
|
self.h_ang = [np.zeros(3, np.float32)] * 10
|
||||||
|
self.h_act_mj = [np.zeros(29, np.float32)] * 10
|
||||||
|
self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
|
||||||
|
|
||||||
|
def update_history(self, q, dq, ang, quat):
|
||||||
|
"""Push the latest proprioception (pos/vel/gyro/orientation) into the 10-frame buffers."""
|
||||||
|
quat = quat / (np.linalg.norm(quat) + 1e-8)
|
||||||
|
q_mj = _to_mujoco(q)
|
||||||
|
dq_mj = _to_mujoco(dq)
|
||||||
|
self.h_q_mj = [q_mj - DEFAULT_ANGLES_MUJOCO] + self.h_q_mj[:-1]
|
||||||
|
self.h_dq_mj = [dq_mj] + self.h_dq_mj[:-1]
|
||||||
|
self.h_ang = [ang.copy()] + self.h_ang[:-1]
|
||||||
|
self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
|
||||||
|
self.h_quat = [quat.copy()] + self.h_quat[:-1]
|
||||||
|
|
||||||
|
def build_decoder_obs(self):
|
||||||
|
"""Assemble the 994-D decoder input: token + 10-frame proprioception history + gravity."""
|
||||||
|
obs = np.zeros(994, np.float32)
|
||||||
|
off = 0
|
||||||
|
obs[off : off + 64] = self.token
|
||||||
|
off += 64
|
||||||
|
for h, sz in [
|
||||||
|
(list(reversed(self.h_ang)), 3),
|
||||||
|
(list(reversed(self.h_q_mj)), 29),
|
||||||
|
(list(reversed(self.h_dq_mj)), 29),
|
||||||
|
(list(reversed(self.h_act_mj)), 29),
|
||||||
|
]:
|
||||||
|
for f in range(10):
|
||||||
|
obs[off : off + sz] = h[f]
|
||||||
|
off += sz
|
||||||
|
for q in reversed(self.h_quat):
|
||||||
|
obs[off : off + 3] = get_gravity_orientation(q)
|
||||||
|
off += 3
|
||||||
|
assert off == 994, f"Decoder obs mismatch: {off}"
|
||||||
|
return obs
|
||||||
|
|
||||||
|
def step(self, robot_obs, token, debug=False):
|
||||||
|
"""One control tick: read robot obs, decode the supplied token -> joint targets.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
robot_obs: dict with ``<joint>.q``/``.dq`` and ``imu.*`` fields.
|
||||||
|
token: 64-D latent supplied by the policy (encoder bypassed).
|
||||||
|
debug: log action/delta norms.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict of ``<joint>.q`` target positions (rad) in IsaacLab joint order.
|
||||||
|
"""
|
||||||
|
self.token = np.asarray(token, np.float32)
|
||||||
|
jnames = [m.name for m in G1_29_JointIndex]
|
||||||
|
q = np.array(
|
||||||
|
[
|
||||||
|
robot_obs.get(f"{n}.q", DEFAULT_ANGLES[m.value])
|
||||||
|
for m, n in zip(G1_29_JointIndex, jnames, strict=False)
|
||||||
|
],
|
||||||
|
np.float32,
|
||||||
|
)
|
||||||
|
dq = np.array([robot_obs.get(f"{n}.dq", 0.0) for n in jnames], np.float32)
|
||||||
|
quat = np.array(
|
||||||
|
[
|
||||||
|
robot_obs.get("imu.quat.w", 1),
|
||||||
|
robot_obs.get("imu.quat.x", 0),
|
||||||
|
robot_obs.get("imu.quat.y", 0),
|
||||||
|
robot_obs.get("imu.quat.z", 0),
|
||||||
|
],
|
||||||
|
np.float32,
|
||||||
|
)
|
||||||
|
ang = np.array([robot_obs.get(f"imu.gyro.{a}", 0) for a in "xyz"], np.float32)
|
||||||
|
self.update_history(q, dq, ang, quat)
|
||||||
|
action_mj = (
|
||||||
|
self.decoder.run(None, {self.decoder_input: self.build_decoder_obs().reshape(1, -1)})[0]
|
||||||
|
.squeeze()
|
||||||
|
.astype(np.float32)
|
||||||
|
)
|
||||||
|
self.last_action_mj = action_mj.copy()
|
||||||
|
target = DEFAULT_ANGLES + action_mj[ISAACLAB_TO_MUJOCO] * ACTION_SCALE
|
||||||
|
if debug:
|
||||||
|
delta = target - q
|
||||||
|
logger.debug(
|
||||||
|
"token_norm=%.4f action_norm=%.4f delta_max=%.4f delta_rms=%.4f",
|
||||||
|
np.linalg.norm(self.token),
|
||||||
|
np.linalg.norm(action_mj),
|
||||||
|
np.max(np.abs(delta)),
|
||||||
|
np.sqrt(np.mean(delta**2)),
|
||||||
|
)
|
||||||
|
return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
|
||||||
|
|
||||||
|
|
||||||
class SonicRuntime:
|
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
|
Token-only deploy: the encoder is bypassed; each tick the decoder consumes a 64-D
|
||||||
:class:`SonicWholeBodyController` from the incoming 34-D whole-body command.
|
latent token supplied directly by the policy.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, force_cpu: bool = False):
|
def __init__(self, force_cpu: bool = False):
|
||||||
require_package("onnxruntime", extra="unitree_g1")
|
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")
|
decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx")
|
||||||
|
|
||||||
providers = ort_providers(force_cpu=force_cpu)
|
providers = ort_providers(force_cpu=force_cpu)
|
||||||
so = make_ort_session_options()
|
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)
|
decoder_sess = ort.InferenceSession(decoder_path, sess_options=so, providers=providers)
|
||||||
|
|
||||||
# Report the provider actually bound, not the one requested: ORT silently falls
|
# Report the provider actually bound, not the one requested: ORT silently falls back
|
||||||
# back to CPU if CUDA can't load (e.g. libcudnn not on LD_LIBRARY_PATH), and a
|
# to CPU if CUDA can't load (e.g. libcudnn not on LD_LIBRARY_PATH), and a CPU decoder
|
||||||
# CPU decoder drifts the closed-loop heading. Warn loudly so it can't hide.
|
# drifts the closed-loop heading. Warn loudly so it can't hide.
|
||||||
self.use_gpu = decoder_sess.get_providers()[0] == "CUDAExecutionProvider"
|
self.use_gpu = decoder_sess.get_providers()[0] == "CUDAExecutionProvider"
|
||||||
if not force_cpu and not self.use_gpu:
|
if not force_cpu and not self.use_gpu:
|
||||||
print(
|
print(
|
||||||
@@ -196,16 +331,13 @@ class SonicRuntime:
|
|||||||
)
|
)
|
||||||
|
|
||||||
self.kp, self.kd = compute_kp_kd()
|
self.kp, self.kd = compute_kp_kd()
|
||||||
self.controller = PlannerController(encoder_sess, decoder_sess)
|
self.controller = SonicDecoder(decoder_sess)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def pipeline(self):
|
def pipeline(self):
|
||||||
return self.controller
|
return self.controller
|
||||||
|
|
||||||
def reset(self):
|
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()
|
self.controller.reset()
|
||||||
|
|
||||||
def shutdown(self):
|
def shutdown(self):
|
||||||
@@ -217,10 +349,6 @@ class SonicWholeBodyController:
|
|||||||
|
|
||||||
control_dt = CONTROL_DT
|
control_dt = CONTROL_DT
|
||||||
full_body = True
|
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):
|
def __init__(self, force_cpu: bool = False):
|
||||||
logger.info("Loading SONIC whole-body controller...")
|
logger.info("Loading SONIC whole-body controller...")
|
||||||
@@ -229,113 +357,22 @@ class SonicWholeBodyController:
|
|||||||
self.kd = self._runtime.kd
|
self.kd = self._runtime.kd
|
||||||
self.controller = self._runtime.controller
|
self.controller = self._runtime.controller
|
||||||
|
|
||||||
# Startup blend: ease from the robot's initial pose into the first commanded
|
# Startup blend: ease from the robot's initial pose into the first commanded policy
|
||||||
# policy targets over INIT_RAMP_S (captured on the first control tick).
|
# 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_ramp_steps = max(1, round(INIT_RAMP_S / CONTROL_DT))
|
||||||
self._init_step = 0
|
self._init_step = 0
|
||||||
self._start_pose: dict[str, float] = {}
|
self._start_pose: dict[str, float] = {}
|
||||||
|
|
||||||
# Tick counter for the dense whole-body (OpenHLM, mode-0) path's encoder cadence.
|
# Token-interface state. ``token_mode`` is set True by the robot when the deploy is
|
||||||
self._wb_step = 0
|
# token-driven (``UnitreeG1Config.sonic_token_action``): the controller then holds a
|
||||||
# Rolling 50-frame reference trajectory (ref29 + anchor quat) built from the
|
# stable *neutral* token until the first real token arrives, and afterwards holds the
|
||||||
# stream of per-tick whole-body commands, fed to the encoder as a batch.
|
# *last* token received between ticks (the async controller runs ~50 Hz while a token
|
||||||
self._wb_traj: deque[np.ndarray] = deque(maxlen=50)
|
# VLA streams ~30 Hz). This lives here (not in the entry-point script) so it applies
|
||||||
self._wb_quat_traj: deque[np.ndarray] = deque(maxlen=50)
|
# uniformly to run_g1_server, lerobot-rollout and the sim replays.
|
||||||
# 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.
|
|
||||||
self.token_mode = False
|
self.token_mode = False
|
||||||
self._last_token: np.ndarray | None = None
|
self._last_token: np.ndarray | None = None
|
||||||
|
|
||||||
logger.info("SONIC ready (encoder/decoder, 34-D whole-body command path)")
|
logger.info("SONIC ready (decoder, 64-D token command path)")
|
||||||
|
|
||||||
def _run_wholebody34(self, obs: dict, wb: np.ndarray) -> dict:
|
|
||||||
"""Feed a dense 34-D OpenHLM whole-body command as the mode-0 encoder reference.
|
|
||||||
|
|
||||||
The 29 joint targets are held across the encoder lookahead window (zero
|
|
||||||
velocity) and the root roll/pitch set the anchor orientation, then the
|
|
||||||
encoder/decoder run directly (planner bypassed). One command per tick, so the
|
|
||||||
VLA's commanded pose is what SONIC tracks.
|
|
||||||
"""
|
|
||||||
ref, anchor = _wb34_to_reference(wb)
|
|
||||||
c = self.controller
|
|
||||||
if c.encode_mode != 0:
|
|
||||||
c.encode_mode = 0
|
|
||||||
c.reinit_heading = True
|
|
||||||
# Index 33 is a yaw-rate (rad/s): integrate it into a heading offset and hand
|
|
||||||
# it to the pipeline as ``delta_heading`` so commanded turns are tracked rather
|
|
||||||
# than silently dropped (the anchor from _wb34_to_reference only carries r/p).
|
|
||||||
self._heading += float(wb[33]) * CONTROL_DT
|
|
||||||
c.delta_heading = self._heading
|
|
||||||
# Capture the heading/anchor reference on the first whole-body tick. The
|
|
||||||
# controller only latches ``init_ref_quat`` (and the base heading) inside
|
|
||||||
# ``step()`` when ``first_motion or reinit_heading`` — but it already boots in
|
|
||||||
# mode 0, so the mode-switch guard above misses the very first command and the
|
|
||||||
# anchor would stay identity. This mirrors the GEAR reference, which seeds
|
|
||||||
# ``init_ref_quat`` from the first anchor. Must run before the buffers below so
|
|
||||||
# ``step()`` latches ``motion_body_quats[0]`` = this tick's anchor.
|
|
||||||
if self._wb_step == 0:
|
|
||||||
c.reinit_heading = True
|
|
||||||
|
|
||||||
# Accumulate the per-tick commands into a rolling 50-frame reference
|
|
||||||
# trajectory so the encoder's 10-frame, step-5 lookahead sees an actual
|
|
||||||
# motion sequence (with velocities) instead of one repeated pose. 50 frames
|
|
||||||
# == chunk horizon == 10 lookahead frames × step 5.
|
|
||||||
self._wb_traj.append(ref)
|
|
||||||
self._wb_quat_traj.append(anchor)
|
|
||||||
traj = np.asarray(self._wb_traj, np.float32) # (L, 29), oldest -> newest
|
|
||||||
quats = np.asarray(self._wb_quat_traj, np.float32) # (L, 4)
|
|
||||||
n = len(traj)
|
|
||||||
# Per-frame velocities from finite differences (rad/s at the control rate).
|
|
||||||
vel = np.zeros_like(traj)
|
|
||||||
if n > 1:
|
|
||||||
vel[1:] = (traj[1:] - traj[:-1]) / CONTROL_DT
|
|
||||||
vel[0] = vel[1]
|
|
||||||
with c.motion_lock:
|
|
||||||
c.motion_joint_positions[:n] = traj
|
|
||||||
c.motion_joint_velocities[:n] = vel
|
|
||||||
c.motion_body_quats[:n] = quats
|
|
||||||
c.motion_body_pos[:n] = 0.0
|
|
||||||
c.motion_timesteps = n
|
|
||||||
c.ref_cursor = 0
|
|
||||||
c.playing = True
|
|
||||||
do_enc = self._wb_step % ENCODER_UPDATE_EVERY == 0
|
|
||||||
out = c.step(obs, update_encoder=do_enc, debug=False)
|
|
||||||
if self._wb_step % 25 == 0:
|
|
||||||
tgt = np.array([out[f"{m.name}.q"] for m in G1_29_JointIndex], np.float32)
|
|
||||||
logger.info(
|
|
||||||
"[WB34] step=%d |ref|mean=%.3f |target|mean=%.3f target_std=%.3f init_ref_quat=%s",
|
|
||||||
self._wb_step,
|
|
||||||
float(np.abs(ref).mean()),
|
|
||||||
float(np.abs(tgt).mean()),
|
|
||||||
float(tgt.std()),
|
|
||||||
np.round(c.init_ref_quat, 3).tolist(),
|
|
||||||
)
|
|
||||||
self._wb_step += 1
|
|
||||||
return out
|
|
||||||
|
|
||||||
def _run_token(self, obs: dict, token: np.ndarray) -> dict:
|
|
||||||
"""Decode a supplied 64-D latent token directly (encoder bypassed).
|
|
||||||
|
|
||||||
Token-only replay: set the pipeline's cached token to the supplied one and run
|
|
||||||
a decode-only step (``update_encoder=False``). The decoder still closes the loop
|
|
||||||
on live proprioception (history is refreshed inside ``step`` from ``obs``); only
|
|
||||||
the encoder — which would recompute the token from a motion reference — is
|
|
||||||
skipped. Returns the ``<joint>.q`` target dict.
|
|
||||||
"""
|
|
||||||
c = self.controller
|
|
||||||
c.token = np.asarray(token, np.float32)
|
|
||||||
self._wb_step += 1
|
|
||||||
return c.step(obs, update_encoder=False, debug=False)
|
|
||||||
|
|
||||||
def _startup_blend(self, obs: dict, out: dict) -> dict:
|
def _startup_blend(self, obs: dict, out: dict) -> dict:
|
||||||
"""Ease into policy control at startup: for the first ``INIT_RAMP_S`` seconds,
|
"""Ease into policy control at startup: for the first ``INIT_RAMP_S`` seconds,
|
||||||
@@ -368,46 +405,26 @@ class SonicWholeBodyController:
|
|||||||
return {}
|
return {}
|
||||||
obs = lowstate_to_obs(lowstate)
|
obs = lowstate_to_obs(lowstate)
|
||||||
|
|
||||||
# Token-only interface (latent replay / token-output VLA): a dense 64-D
|
# Token-only interface (token-output VLA): a dense 64-D ``motion_token.{i}`` command
|
||||||
# ``motion_token.{i}`` command is decoded directly, bypassing the encoder.
|
# is decoded directly, encoder bypassed.
|
||||||
# Checked before the joint path so a token action takes precedence.
|
|
||||||
token = _extract_token_from_action(action)
|
token = _extract_token_from_action(action)
|
||||||
if token is not None:
|
if token is not None:
|
||||||
self._last_token = token
|
self._last_token = token
|
||||||
elif self._last_token is None and self.token_mode:
|
elif self._last_token is None and self.token_mode:
|
||||||
# Token-driven deploy, but no token has arrived yet: hold the captured
|
# Token-driven deploy, but no token has arrived yet: hold the captured neutral
|
||||||
# neutral token (NEUTRAL_TOKEN), which the decoder maps to a stable, natural
|
# token (NEUTRAL_TOKEN), which the decoder maps to a stable, natural standing pose.
|
||||||
# standing pose (the encoder's own idle output; see NEUTRAL_TOKEN).
|
|
||||||
self._last_token = NEUTRAL_TOKEN.copy()
|
self._last_token = NEUTRAL_TOKEN.copy()
|
||||||
if self._last_token is not None:
|
if self._last_token is None:
|
||||||
# Either a fresh token this tick or the last one received (held between the
|
# No token yet and not in token_mode: hold (keep last target).
|
||||||
# ~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],
|
|
||||||
)
|
|
||||||
return {}
|
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):
|
def reset(self):
|
||||||
self._runtime.reset()
|
self._runtime.reset()
|
||||||
self._init_step = 0 # re-run the startup blend after a reset
|
self._init_step = 0 # re-run the startup blend after a reset
|
||||||
self._start_pose = {}
|
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.
|
# Drop the held token so token_mode re-seeds the neutral token after a reset.
|
||||||
self._last_token = None
|
self._last_token = None
|
||||||
|
|
||||||
|
|||||||
@@ -105,21 +105,6 @@ REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
|
|||||||
KEYBOARD_KEYS_FIELD = "keyboard.keys"
|
KEYBOARD_KEYS_FIELD = "keyboard.keys"
|
||||||
|
|
||||||
# ── Dense whole-body joint reference (SONIC encode_mode 0, OpenHLM / pi0.5) ──────
|
# ── 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]:
|
def default_remote_input() -> dict[str, float]:
|
||||||
"""Return a zeroed-out remote input dict (axes + buttons)."""
|
"""Return a zeroed-out remote input dict (axes + buttons)."""
|
||||||
return dict.fromkeys(REMOTE_KEYS, 0.0)
|
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
|
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):
|
class G1_29_JointArmIndex(IntEnum):
|
||||||
# Left arm
|
# Left arm
|
||||||
kLeftShoulderPitch = 15
|
kLeftShoulderPitch = 15
|
||||||
@@ -196,43 +210,6 @@ def lowstate_to_obs(lowstate) -> dict:
|
|||||||
return obs
|
return obs
|
||||||
|
|
||||||
|
|
||||||
def obs_to_wb34_state(obs: dict) -> np.ndarray:
|
|
||||||
"""Build the 34-D OpenHLM / pi0.5 proprio state from a G1 observation dict.
|
|
||||||
|
|
||||||
Mirrors the whole-body *action* layout so the policy sees state and action in
|
|
||||||
the same coordinates::
|
|
||||||
|
|
||||||
[L-arm(7), L-grip(1), R-arm(7), R-grip(1),
|
|
||||||
L-leg(6), R-leg(6), waist(3), root roll/pitch + yaw-rate(3)]
|
|
||||||
|
|
||||||
Joint positions come from the ``<joint>.q`` obs keys, which are already in
|
|
||||||
MuJoCo / Unitree-SDK order — the same body-part grouping OpenHLM uses
|
|
||||||
([L-leg 0:6, R-leg 6:12, waist 12:15, L-arm 15:22, R-arm 22:29]) — so they are
|
|
||||||
regrouped directly (no IsaacLab permutation). The G1 has no grippers in its
|
|
||||||
29-DoF body, so both gripper slots are 0. Root roll/pitch are the IMU RPY and
|
|
||||||
the last slot is the IMU yaw rate (gyro z).
|
|
||||||
"""
|
|
||||||
q_mj = np.array(
|
|
||||||
[float(obs.get(f"{m.name}.q", 0.0)) for m in G1_29_JointIndex],
|
|
||||||
dtype=np.float32,
|
|
||||||
)
|
|
||||||
lleg, rleg, waist = q_mj[0:6], q_mj[6:12], q_mj[12:15]
|
|
||||||
larm, rarm = q_mj[15:22], q_mj[22:29]
|
|
||||||
|
|
||||||
state = np.zeros(34, dtype=np.float32)
|
|
||||||
state[0:7] = larm
|
|
||||||
# state[7] left gripper — none on 29-DoF G1
|
|
||||||
state[8:15] = rarm
|
|
||||||
# state[15] right gripper — none on 29-DoF G1
|
|
||||||
state[16:22] = lleg
|
|
||||||
state[22:28] = rleg
|
|
||||||
state[28:31] = waist
|
|
||||||
state[31] = float(obs.get("imu.rpy.roll", 0.0))
|
|
||||||
state[32] = float(obs.get("imu.rpy.pitch", 0.0))
|
|
||||||
state[33] = float(obs.get("imu.gyro.z", 0.0))
|
|
||||||
return state
|
|
||||||
|
|
||||||
|
|
||||||
def make_locomotion_controller(name: str | None):
|
def make_locomotion_controller(name: str | None):
|
||||||
"""Instantiate a locomotion controller by class name. Returns None if name is None."""
|
"""Instantiate a locomotion controller by class name. Returns None if name is None."""
|
||||||
if name is None:
|
if name is None:
|
||||||
|
|||||||
@@ -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()
|
|
||||||
@@ -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()
|
|
||||||
@@ -22,18 +22,36 @@ This server runs on the robot and forwards:
|
|||||||
- Robot commands (LowCmd) from ZMQ to DDS (from remote clients)
|
- Robot commands (LowCmd) from ZMQ to DDS (from remote clients)
|
||||||
|
|
||||||
Uses JSON for secure serialization instead of pickle.
|
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 argparse
|
||||||
import base64
|
import base64
|
||||||
import contextlib
|
import contextlib
|
||||||
import json
|
import json
|
||||||
|
import os
|
||||||
import re
|
import re
|
||||||
|
import signal
|
||||||
|
import sys
|
||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
from dataclasses import dataclass
|
from typing import Any
|
||||||
from typing import TYPE_CHECKING, Any
|
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
import zmq
|
import zmq
|
||||||
from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
|
from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
|
||||||
from unitree_sdk2py.core.channel import ChannelFactoryInitialize, ChannelPublisher, ChannelSubscriber
|
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
|
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
|
# DDS topic names follow Unitree SDK naming conventions
|
||||||
# ruff: noqa: N816
|
# ruff: noqa: N816
|
||||||
kTopicLowCommand_Debug = "rt/lowcmd" # action to robot
|
kTopicLowCommand_Debug = "rt/lowcmd" # action to robot
|
||||||
@@ -55,59 +70,271 @@ LOWCMD_PORT = 6000
|
|||||||
LOWSTATE_PORT = 6001
|
LOWSTATE_PORT = 6001
|
||||||
NUM_MOTORS = 35
|
NUM_MOTORS = 35
|
||||||
|
|
||||||
|
# Onboard high-level channels (serve_onboard_controller): compact actions in, state out.
|
||||||
|
ACTION_PORT = 6004
|
||||||
|
STATE_PORT = 6005
|
||||||
|
|
||||||
@dataclass
|
# Controller-negotiation handshake (REQ/REP). The client's first message agrees on
|
||||||
class Gripper:
|
# which controller the server will run before any control data flows.
|
||||||
"""A single Damiao gripper that only writes to CAN when the open/close state changes."""
|
HANDSHAKE_PORT = 6002
|
||||||
|
PROTOCOL_VERSION = 1
|
||||||
|
|
||||||
name: str
|
# Controllers that can run ONBOARD (must match g1_utils.make_locomotion_controller).
|
||||||
bus: "DamiaoMotorsBus"
|
# ``None`` (a.k.a. "bridge") means no onboard controller: the laptop owns control and
|
||||||
open_deg: float
|
# streams raw lowcmd over the ZMQ DDS bridge (the legacy run_g1_server behavior).
|
||||||
close_deg: float
|
VALID_CONTROLLERS = (
|
||||||
_last_cmd: str | None = None # "open" | "close"
|
"GrootLocomotionController",
|
||||||
|
"HolosomaLocomotionController",
|
||||||
def apply(self, want_close: bool) -> None:
|
"SonicWholeBodyController",
|
||||||
want = "close" if want_close else "open"
|
)
|
||||||
if want == self._last_cmd:
|
# SONIC latent-token dimensionality (mirrors sonic_whole_body.TOKEN_DIM; kept local so
|
||||||
return
|
# the handshake can run without importing the heavy controller / onnxruntime).
|
||||||
target = self.close_deg if want_close else self.open_deg
|
TOKEN_DIM = 64
|
||||||
self.bus.write("Goal_Position", "gripper", target)
|
_BRIDGE_ALIASES = {"", "none", "null", "bridge", "raw"}
|
||||||
self._last_cmd = want
|
|
||||||
print(f"[gripper] {self.name} -> {want.upper()} ({target:.1f} deg)")
|
|
||||||
|
|
||||||
|
|
||||||
def build_gripper(
|
def _normalize_controller(name: str | None) -> str | None:
|
||||||
name: str,
|
"""Map a requested controller name to a canonical value (or None for raw bridge)."""
|
||||||
port: str,
|
if name is None:
|
||||||
send_id: int,
|
return None
|
||||||
recv_id: int,
|
low = str(name).strip().lower()
|
||||||
motor_type: str,
|
if low in _BRIDGE_ALIASES:
|
||||||
use_can_fd: bool,
|
return None
|
||||||
open_deg: float,
|
for c in VALID_CONTROLLERS:
|
||||||
close_deg: float,
|
if c.lower() == low:
|
||||||
kp: float,
|
return c
|
||||||
kd: float,
|
raise ValueError(f"Unknown controller {name!r}. Available: {list(VALID_CONTROLLERS)} or 'bridge'")
|
||||||
) -> Gripper:
|
|
||||||
from lerobot.motors.damiao.damiao import DamiaoMotorsBus
|
|
||||||
from lerobot.motors.motors_bus import Motor, MotorNormMode
|
|
||||||
|
|
||||||
motors = {
|
|
||||||
"gripper": Motor(
|
def _capabilities(controller: str | None, sonic_token_action: bool) -> dict[str, Any]:
|
||||||
id=send_id,
|
"""The interface the server advertises for an agreed controller."""
|
||||||
model=motor_type,
|
caps: dict[str, Any] = {
|
||||||
norm_mode=MotorNormMode.DEGREES,
|
"controller": controller,
|
||||||
motor_type_str=motor_type,
|
"sonic_token_action": bool(sonic_token_action),
|
||||||
recv_id=recv_id,
|
"protocol": PROTOCOL_VERSION,
|
||||||
)
|
|
||||||
}
|
}
|
||||||
bus = DamiaoMotorsBus(port=port, motors=motors, use_can_fd=use_can_fd)
|
if controller is None:
|
||||||
print(f"Connecting {name} gripper on {port} (fd={use_can_fd})...")
|
# Raw DDS bridge: the laptop runs the controller and streams lowcmd.
|
||||||
bus.connect(handshake=True)
|
caps["mode"] = "bridge"
|
||||||
bus.write("Kp", "gripper", kp)
|
caps["lowcmd_port"] = LOWCMD_PORT
|
||||||
bus.write("Kd", "gripper", kd)
|
caps["lowstate_port"] = LOWSTATE_PORT
|
||||||
bus.write("Goal_Position", "gripper", open_deg) # start open
|
else:
|
||||||
print(f" {name}: connected, torque enabled, opened.")
|
# Onboard: the controller runs here; the laptop ships compact high-level actions.
|
||||||
return Gripper(name, bus, open_deg, close_deg, _last_cmd="open")
|
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]:
|
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-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-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)")
|
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()
|
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
|
# Optionally start camera server in background thread
|
||||||
camera_thread = None
|
camera_thread = None
|
||||||
if args.camera or args.cameras:
|
if args.camera or args.cameras:
|
||||||
@@ -313,6 +618,7 @@ def main() -> None:
|
|||||||
|
|
||||||
# initialize ZMQ
|
# initialize ZMQ
|
||||||
ctx = zmq.Context.instance()
|
ctx = zmq.Context.instance()
|
||||||
|
shutdown_event = threading.Event()
|
||||||
|
|
||||||
# receive commands from remote client
|
# receive commands from remote client
|
||||||
lowcmd_sock = ctx.socket(zmq.PULL)
|
lowcmd_sock = ctx.socket(zmq.PULL)
|
||||||
@@ -323,7 +629,6 @@ def main() -> None:
|
|||||||
lowstate_sock.bind(f"tcp://0.0.0.0:{LOWSTATE_PORT}")
|
lowstate_sock.bind(f"tcp://0.0.0.0:{LOWSTATE_PORT}")
|
||||||
|
|
||||||
state_period = 0.002 # ~500 hz
|
state_period = 0.002 # ~500 hz
|
||||||
shutdown_event = threading.Event()
|
|
||||||
|
|
||||||
# start observation forwarding in background thread
|
# start observation forwarding in background thread
|
||||||
t_state = threading.Thread(
|
t_state = threading.Thread(
|
||||||
|
|||||||
@@ -16,6 +16,8 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import contextlib
|
||||||
|
import json
|
||||||
import logging
|
import logging
|
||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
@@ -26,12 +28,12 @@ from typing import TYPE_CHECKING, Protocol, runtime_checkable
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
from lerobot.cameras import make_cameras_from_configs
|
from lerobot.cameras import make_cameras_from_configs
|
||||||
|
from lerobot.utils.errors import DeviceNotConnectedError
|
||||||
from lerobot.types import RobotAction, RobotObservation
|
from lerobot.types import RobotAction, RobotObservation
|
||||||
from lerobot.utils.import_utils import _unitree_sdk_available, require_package
|
from lerobot.utils.import_utils import _unitree_sdk_available, require_package
|
||||||
|
|
||||||
from ..robot import Robot
|
from ..robot import Robot
|
||||||
from .config_unitree_g1 import UnitreeG1Config
|
from .config_unitree_g1 import UnitreeG1Config
|
||||||
from .g1_kinematics import G1_29_ArmIK
|
|
||||||
from .g1_utils import (
|
from .g1_utils import (
|
||||||
KEYBOARD_KEYS_FIELD,
|
KEYBOARD_KEYS_FIELD,
|
||||||
REMOTE_AXES,
|
REMOTE_AXES,
|
||||||
@@ -40,7 +42,6 @@ from .g1_utils import (
|
|||||||
default_remote_input,
|
default_remote_input,
|
||||||
lowstate_to_obs,
|
lowstate_to_obs,
|
||||||
make_locomotion_controller,
|
make_locomotion_controller,
|
||||||
obs_to_wb34_state,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
if TYPE_CHECKING or _unitree_sdk_available:
|
if TYPE_CHECKING or _unitree_sdk_available:
|
||||||
@@ -50,11 +51,9 @@ if TYPE_CHECKING or _unitree_sdk_available:
|
|||||||
ChannelSubscriber as _SDKChannelSubscriber,
|
ChannelSubscriber as _SDKChannelSubscriber,
|
||||||
)
|
)
|
||||||
from unitree_sdk2py.idl.default import (
|
from unitree_sdk2py.idl.default import (
|
||||||
unitree_hg_msg_dds__HandCmd_ as hg_HandCmd_default,
|
|
||||||
unitree_hg_msg_dds__LowCmd_,
|
unitree_hg_msg_dds__LowCmd_,
|
||||||
)
|
)
|
||||||
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
|
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
|
||||||
HandCmd_ as hg_HandCmd,
|
|
||||||
LowCmd_ as hg_LowCmd,
|
LowCmd_ as hg_LowCmd,
|
||||||
LowState_ as hg_LowState,
|
LowState_ as hg_LowState,
|
||||||
)
|
)
|
||||||
@@ -64,8 +63,6 @@ else:
|
|||||||
_SDKChannelPublisher = None
|
_SDKChannelPublisher = None
|
||||||
_SDKChannelSubscriber = None
|
_SDKChannelSubscriber = None
|
||||||
unitree_hg_msg_dds__LowCmd_ = None
|
unitree_hg_msg_dds__LowCmd_ = None
|
||||||
hg_HandCmd_default = None
|
|
||||||
hg_HandCmd = None
|
|
||||||
hg_LowCmd = None
|
hg_LowCmd = None
|
||||||
hg_LowState = None
|
hg_LowState = None
|
||||||
CRC = None
|
CRC = None
|
||||||
@@ -135,26 +132,34 @@ class UnitreeG1(Robot):
|
|||||||
self.config = config
|
self.config = config
|
||||||
self.control_dt = config.control_dt
|
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()
|
# Initialize cameras config (ZMQ-based) - actual connection in connect()
|
||||||
self._cameras = make_cameras_from_configs(config.cameras)
|
self._cameras = make_cameras_from_configs(config.cameras)
|
||||||
|
|
||||||
# Import channel classes based on mode. Simulation and onboard both talk to a
|
# DDS channel classes are only needed by the in-process control roles. The thin
|
||||||
# real (local) DDS via the Unitree SDK; only the laptop-side bridge client uses
|
# client never touches DDS, so we don't import the socket shim at all.
|
||||||
# the ZMQ socket shim.
|
|
||||||
if config.is_simulation or config.onboard:
|
if config.is_simulation or config.onboard:
|
||||||
self._ChannelFactoryInitialize = _SDKChannelFactoryInitialize
|
self._ChannelFactoryInitialize = _SDKChannelFactoryInitialize
|
||||||
self._ChannelPublisher = _SDKChannelPublisher
|
self._ChannelPublisher = _SDKChannelPublisher
|
||||||
self._ChannelSubscriber = _SDKChannelSubscriber
|
self._ChannelSubscriber = _SDKChannelSubscriber
|
||||||
else:
|
else:
|
||||||
from .unitree_sdk2_socket import (
|
self._ChannelFactoryInitialize = None
|
||||||
ChannelFactoryInitialize,
|
self._ChannelPublisher = None
|
||||||
ChannelPublisher,
|
self._ChannelSubscriber = None
|
||||||
ChannelSubscriber,
|
|
||||||
)
|
|
||||||
|
|
||||||
self._ChannelFactoryInitialize = ChannelFactoryInitialize
|
# Client-side ZMQ handles / negotiated capabilities (populated in connect()).
|
||||||
self._ChannelPublisher = ChannelPublisher
|
self._client_action_sock = None
|
||||||
self._ChannelSubscriber = ChannelSubscriber
|
self._client_state_sock = None
|
||||||
|
self._client_state_latest: dict[str, float] = {}
|
||||||
|
self._client_caps: dict | None = None
|
||||||
|
|
||||||
# Initialize state variables
|
# Initialize state variables
|
||||||
self.sim_env = None
|
self.sim_env = None
|
||||||
@@ -164,15 +169,17 @@ class UnitreeG1(Robot):
|
|||||||
self._shutdown_event = threading.Event()
|
self._shutdown_event = threading.Event()
|
||||||
self.subscribe_thread = None
|
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
|
# Token-driven deploy: let a SONIC controller hold a neutral token until the
|
||||||
self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
|
# first real one arrives, then hold the last token between control ticks.
|
||||||
|
if config.sonic_token_action and hasattr(self.controller, "token_mode"):
|
||||||
# Token-driven deploy: let a SONIC controller hold a neutral token until the
|
self.controller.token_mode = True
|
||||||
# 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
|
# Controller thread state
|
||||||
self._controller_thread = None
|
self._controller_thread = None
|
||||||
@@ -187,17 +194,6 @@ class UnitreeG1(Robot):
|
|||||||
# from local lowstate so joystick locomotion works without a laptop round-trip).
|
# from local lowstate so joystick locomotion works without a laptop round-trip).
|
||||||
self._joystick = None
|
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,
|
# 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
|
# 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;
|
# 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)
|
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
|
def _subscribe_lowstate(self): # polls robot state @ 250Hz
|
||||||
while not self._shutdown_event.is_set():
|
while not self._shutdown_event.is_set():
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
@@ -330,22 +288,6 @@ class UnitreeG1(Robot):
|
|||||||
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
|
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
|
||||||
return features
|
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
|
@property
|
||||||
def _token_state_ft(self) -> dict[str, type]:
|
def _token_state_ft(self) -> dict[str, type]:
|
||||||
"""64-D SONIC latent-token proprio state (``motion_token_state.{i}.pos``).
|
"""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)}
|
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
|
@cached_property
|
||||||
def observation_features(self) -> dict[str, type | tuple]:
|
def observation_features(self) -> dict[str, type | tuple]:
|
||||||
return {
|
return {
|
||||||
**self._motors_ft,
|
**self._motors_ft,
|
||||||
**self._wb_state_ft,
|
|
||||||
**self._token_state_ft,
|
**self._token_state_ft,
|
||||||
**self._empty_cameras_ft,
|
|
||||||
**self._replay_cameras_ft,
|
|
||||||
**self._cameras_ft,
|
**self._cameras_ft,
|
||||||
}
|
}
|
||||||
|
|
||||||
@cached_property
|
@cached_property
|
||||||
def action_features(self) -> dict[str, type]:
|
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}
|
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
|
# 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)}
|
return {token_action_key(i): float for i in range(TOKEN_DIM)}
|
||||||
|
|
||||||
# Dense whole-body controllers (SONIC / OpenHLM, pi0.5) consume a single
|
# Locomotion controllers (GR00T / Holosoma): arm joint targets + joystick axes.
|
||||||
# 34-D command per tick. Expose it as ``wb.{i}.pos`` joint-position features
|
|
||||||
# so ``lerobot-rollout`` maps a 34-D policy output straight onto the robot.
|
|
||||||
if getattr(self.controller, "wb_action", False):
|
|
||||||
from .g1_utils import WB_ACTION_DIM, wb_action_key
|
|
||||||
|
|
||||||
return {wb_action_key(i): float for i in range(WB_ACTION_DIM)}
|
|
||||||
|
|
||||||
arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
|
arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
|
||||||
remote_features = dict.fromkeys(REMOTE_AXES, float)
|
remote_features = dict.fromkeys(REMOTE_AXES, float)
|
||||||
return {**arm_features, **remote_features}
|
return {**arm_features, **remote_features}
|
||||||
@@ -521,7 +445,111 @@ class UnitreeG1(Robot):
|
|||||||
_, result = msc.CheckMode()
|
_, result = msc.CheckMode()
|
||||||
time.sleep(1.0)
|
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
|
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
|
# Initialize DDS channel and simulation environment
|
||||||
if self.config.is_simulation:
|
if self.config.is_simulation:
|
||||||
from lerobot.envs.utils import (
|
from lerobot.envs.utils import (
|
||||||
@@ -532,9 +560,9 @@ class UnitreeG1(Robot):
|
|||||||
|
|
||||||
self._ChannelFactoryInitialize(0, "lo")
|
self._ChannelFactoryInitialize(0, "lo")
|
||||||
# Call the hub env's make_env directly so we can disable the offscreen
|
# Call the hub env's make_env directly so we can disable the offscreen
|
||||||
# head_camera renderer. We drive image-conditioned policies from recorded
|
# head_camera renderer. We drive image-conditioned policies from external
|
||||||
# frames (see replay_camera_parquet / external obs), never the sim's own
|
# camera frames, never the sim's own camera, so building a MuJoCo offscreen
|
||||||
# camera, so building a MuJoCo offscreen GL context is pure liability: it
|
# GL context is pure liability: it
|
||||||
# crashes with "Failed to make the EGL context current" when GLFW/SDL
|
# crashes with "Failed to make the EGL context current" when GLFW/SDL
|
||||||
# already own a context, killing the sim thread and hanging on
|
# already own a context, killing the sim thread and hanging on
|
||||||
# "Waiting for robot state...". publish_images=False -> no renderer.
|
# "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 = self._ChannelSubscriber(kTopicLowState, hg_LowState)
|
||||||
self.lowstate_subscriber.Init()
|
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
|
# Start subscribe thread to read robot state
|
||||||
self.subscribe_thread = threading.Thread(target=self._subscribe_lowstate)
|
self.subscribe_thread = threading.Thread(target=self._subscribe_lowstate)
|
||||||
self.subscribe_thread.start()
|
self.subscribe_thread.start()
|
||||||
@@ -685,6 +702,10 @@ class UnitreeG1(Robot):
|
|||||||
self._send_zero_torque()
|
self._send_zero_torque()
|
||||||
|
|
||||||
def disconnect(self):
|
def disconnect(self):
|
||||||
|
if self._client:
|
||||||
|
self._disconnect_client()
|
||||||
|
return
|
||||||
|
|
||||||
# Stop the controller loop first so it isn't fighting the shutdown ramp.
|
# Stop the controller loop first so it isn't fighting the shutdown ramp.
|
||||||
self._shutdown_event.set()
|
self._shutdown_event.set()
|
||||||
controller_stopped = True
|
controller_stopped = True
|
||||||
@@ -741,6 +762,9 @@ class UnitreeG1(Robot):
|
|||||||
cam.disconnect()
|
cam.disconnect()
|
||||||
|
|
||||||
def get_observation(self) -> RobotObservation:
|
def get_observation(self) -> RobotObservation:
|
||||||
|
if self._client:
|
||||||
|
return self._get_observation_client()
|
||||||
|
|
||||||
with self._lowstate_lock:
|
with self._lowstate_lock:
|
||||||
lowstate = self._lowstate
|
lowstate = self._lowstate
|
||||||
if lowstate is None:
|
if lowstate is None:
|
||||||
@@ -749,38 +773,14 @@ class UnitreeG1(Robot):
|
|||||||
# Motors + IMU + wireless remote (shared lowstate -> obs mapping)
|
# Motors + IMU + wireless remote (shared lowstate -> obs mapping)
|
||||||
obs = lowstate_to_obs(lowstate)
|
obs = lowstate_to_obs(lowstate)
|
||||||
|
|
||||||
# Dense whole-body controllers (OpenHLM / pi0.5): expose the 34-D proprio
|
# Token mode: echo the last commanded latent token as observation.state so a
|
||||||
# state as ``wb_state.{i}.pos`` so the rollout aggregates it into
|
# token-output VLA closes the loop on its own previous token.
|
||||||
# ``observation.state`` for the policy.
|
|
||||||
if self.config.sonic_token_action:
|
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
|
from .controllers.sonic_whole_body import token_state_key
|
||||||
|
|
||||||
token = self._last_token if self._last_token is not None else []
|
token = self._last_token if self._last_token is not None else []
|
||||||
for i, v in enumerate(token):
|
for i, v in enumerate(token):
|
||||||
obs[token_state_key(i)] = float(v)
|
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
|
# Cameras - read images from ZMQ cameras
|
||||||
for cam_name, cam in self._cameras.items():
|
for cam_name, cam in self._cameras.items():
|
||||||
@@ -792,6 +792,9 @@ class UnitreeG1(Robot):
|
|||||||
return obs
|
return obs
|
||||||
|
|
||||||
def send_action(self, action: RobotAction) -> RobotAction:
|
def send_action(self, action: RobotAction) -> RobotAction:
|
||||||
|
if self._client:
|
||||||
|
return self._send_action_client(action)
|
||||||
|
|
||||||
action_to_publish = action
|
action_to_publish = action
|
||||||
if self.controller is not None:
|
if self.controller is not None:
|
||||||
if self.config.sonic_token_action:
|
if self.config.sonic_token_action:
|
||||||
@@ -801,8 +804,6 @@ class UnitreeG1(Robot):
|
|||||||
if token is not None:
|
if token is not None:
|
||||||
self._last_token = token
|
self._last_token = token
|
||||||
self._update_controller_action(action)
|
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):
|
if getattr(self.controller, "full_body", False):
|
||||||
return action
|
return action
|
||||||
# Controller thread owns legs/waist. Here we only update joystick inputs
|
# 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)
|
if key.endswith(".q") and key.startswith(arm_prefixes)
|
||||||
}
|
}
|
||||||
|
|
||||||
tau = None
|
self.publish_lowcmd(action_to_publish)
|
||||||
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)
|
|
||||||
return action
|
return action
|
||||||
|
|
||||||
def _update_controller_action(self, action: RobotAction) -> None:
|
def _update_controller_action(self, action: RobotAction) -> None:
|
||||||
@@ -858,50 +843,14 @@ class UnitreeG1(Robot):
|
|||||||
if isinstance(key, str) and value is not None:
|
if isinstance(key, str) and value is not None:
|
||||||
self.controller_input[key] = value
|
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
|
@property
|
||||||
def is_calibrated(self) -> bool:
|
def is_calibrated(self) -> bool:
|
||||||
return True
|
return True
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def is_connected(self) -> bool:
|
def is_connected(self) -> bool:
|
||||||
|
if self._client:
|
||||||
|
return self._client_action_sock is not None
|
||||||
with self._lowstate_lock:
|
with self._lowstate_lock:
|
||||||
return self._lowstate is not None
|
return self._lowstate is not None
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user