Compare commits

...

83 Commits

Author SHA1 Message Date
Martino Russi af7711c9f8 feat(unitree_g1): run controller onboard with thin-client laptop role
Add the onboard-controller deploy path on top of the SONIC-only branch: the
locomotion / whole-body controller runs ON the robot (local DDS, full rate)
while the laptop becomes a thin client that only relays high-level actions and
reads back the state echo + cameras over ZMQ.

- config: onboard / dds_interface / release_motion_control / physical_remote flags.
- unitree_g1.py: three mutually-exclusive roles (simulation / onboard / client).
  Onboard uses the real Unitree SDK channels, releases the built-in motion
  services, and reads the physical wireless remote from lowstate (priority over
  laptop axes). The client has no DDS/controller: it handshakes a controller with
  run_g1_server, PUSHes actions (arm targets or 64-D SONIC tokens) and SUBs the
  observation.state echo, advertising the same action/observation schema by
  controller name so the exact same policy output routes straight through.
- run_g1_server.py: serve_onboard_controller + request_controller handshake so
  the server instantiates and runs the negotiated controller onboard.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-31 14:45:28 +02:00
Martino Russi 09fb2b4580 update docs 2026-07-31 12:55:20 +02:00
Martino Russi 90699fdece inline token helpers in wbc, move properties to init 2026-07-31 11:50:39 +02:00
Martino Russi d1eed119a6 refactor 2026-07-31 11:43:17 +02:00
Martino Russi 87b0dc470c dedup reset 2026-07-31 11:25:08 +02:00
Martino Russi ec6da5c92e reduce verbosity 2026-07-31 11:09:59 +02:00
Martino Russi d29fecfa2c refactor sonic_wbc
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-31 10:48:37 +02:00
Martino Russi 1799eceb22 make every controller consume lowstate for obs 2026-07-31 10:31:11 +02:00
Martino Russi 89ac5b4408 remove comments 2026-07-31 10:00:22 +02:00
Martino Russi a8ea5e6ea0 (refactor) deduplicate _clear_history(), remove unused debug 2026-07-31 09:53:36 +02:00
Martino Russi 6ea5dd95e9 merge sonic runtime into sonicwbc 2026-07-30 16:22:42 +02:00
Martino Russi f8359c19d6 remove verbose 2026-07-30 16:13:58 +02:00
Martino Russi f8521958d8 (refactor) move sonic logic to controller 2026-07-30 16:00:35 +02:00
Martino Russi 35341dc603 restore REMOTE_KEYS 2026-07-30 15:53:33 +02:00
Martino Russi 686ac2cc21 move controller logic out of unitree_g1, reduce verbosity 2026-07-30 15:45:17 +02:00
Martino Russi 48e2530455 (refactor) match main sim logic, move extract_token_from_action to controller 2026-07-30 15:38:29 +02:00
Martino Russi b6a283193c (refactor) restore remote_keys 2026-07-30 15:22:38 +02:00
Martino Russi 44edbd704f refactor(unitree_g1): revert reset() to simple version
Drop the full-body pause/direct-publish reset special-casing (and the
_controller_paused plumbing) so SONIC uses the same reset path as the other
controllers. The removed logic is preserved on the better-reset branch.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-30 15:05:51 +02:00
Martino Russi 4cd9fc73f8 (refactor) token mode is default 2026-07-30 14:50:44 +02:00
Martino Russi 57d33c295e delete lowstate_to_obs 2026-07-30 14:48:09 +02:00
Martino Russi 6e23429105 inline mujoco indexing 2026-07-30 14:37:38 +02:00
Martino Russi fffa42cc5e refactor(unitree_g1): load SONIC constants from onnx 2026-07-30 13:44:24 +02:00
Martino Russi f33089b027 style(unitree_g1): ruff-format blank line in g1_utils
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-30 12:07:02 +02:00
Martino Russi 9f32f57b59 chore(unitree_g1): silence bandit B105 on SONIC token prefixes
The motion_token / motion_token_state constants are feature-key prefixes, not
secrets; mark them #nosec so the pre-commit bandit hook passes.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-30 12:03:37 +02:00
Martino Russi 0a9b9d9a93 refactor(unitree_g1): move lowstate_to_obs out of g1_utils into unitree_g1
Keep g1_utils close to main: the lowstate -> obs mapping lives in unitree_g1
(where get_observation builds it, left unchanged from main) and the SONIC
controller imports it from there.
2026-07-30 11:56:57 +02:00
Martino Russi 16b915ede5 refactor(unitree_g1): drop make_ort_session_options from SONIC branch
The ORT SessionOptions helper (thread-pool cap) is being split into its own
"ORT fix" PR. SONIC only needs quiet logs, so inline a minimal SessionOptions in
the decoder instead of depending on the shared helper.
2026-07-30 11:49:48 +02:00
Martino Russi 95d9029039 refactor(unitree_g1): reduce branch to SONIC-only diff vs main
Split the onboard-controller server/handshake/thin-client work out to
feat/g1_onboard_controller and revert it here:

- unitree_g1.py: drop client + onboard roles, wireless-remote parsing and
  motion-service release; restore main's sim/socket-bridge transport. Keep only
  the SONIC integration (implicit token action/state, full-body reset/pause,
  controller kp/kd + shutdown).
- config: drop onboard/dds_interface/release_motion_control/physical_remote.
- run_g1_server.py + unitree_g1.mdx: reverted to main.
- gr00t/holosoma: keep the move into controllers/ but revert their content to
  main (only the package-relative import changes).
2026-07-30 11:40:01 +02:00
Martino Russi bbfc4ff443 feat(unitree_g1): make SONIC token interface implicit
Drop the ``sonic_token_action`` config flag; the 64-D latent-token
action/observation interface now switches on automatically whenever the
SONIC whole-body controller is selected (``controller == "SonicWholeBodyController"``).
Keyed via a ``_sonic_token`` property so client, onboard and sim roles agree.
2026-07-30 11:30:07 +02:00
Martino Russi 0c57cd03f2 Merge branch 'main' into feat/sonic_encoder_decoder 2026-07-30 11:15:18 +02:00
Martino Russi 962ed236af config cleanup 2026-07-29 18:45:49 +02:00
Martino Russi 3ae036ee40 restore gravity_compensation 2026-07-29 18:12:19 +02:00
Martino Russi cdf5141688 remove ort_providers 2026-07-29 17:01:17 +02:00
Martino Russi 9f0663e9e3 clean init and utils 2026-07-29 16:43:36 +02:00
Martino Russi 503f3e57ae refactor(unitree_g1): minimize diff vs main (drop onnx guards, e-stop)
Revert pyproject.toml and import_utils.py to main: the onnxruntime-gpu
detection fallback and the _onnxruntime_available/_onnx_available flags
are unnecessary. Controllers now import onnx/onnxruntime unconditionally
(matching main) which also works with onnxruntime-gpu since the import
name is stable. Drop the require_package() calls we had added.

Remove the stdin e-stop listener from serve_onboard_controller and the
now-unused os/sys imports; Ctrl-C still triggers graceful shutdown.
2026-07-29 16:11:55 +02:00
Martino Russi 6d24f20eb4 drop stray artifacts 2026-07-29 15:47:46 +02:00
Martino Russi af163fd032 refactor(unitree_g1): minimize diff w.r.t main 2026-07-29 15:46:17 +02:00
Martino Russi 77259f436e feat(unitree_g1): use captured neutral SONIC token instead of zeros
The all-zero token is off the encoder's learned FSQ manifold and decodes to a
slightly goofy stance. Replace it with a NEUTRAL_TOKEN captured from the encoder's
own idle output in sim (stored as integer FSQ codes, rescaled by the encoder's
1/16 quantization step to an exact on-grid token). token_mode now seeds this
neutral, and the onboard sender starts observation.state from it so the first
inference sees the token the decoder is actually holding.
2026-07-27 11:05:42 +02:00
Martino Russi 85f5c3606d feat(unitree_g1): hold neutral SONIC token until first command
Move the token-hold idle logic into SonicWholeBodyController via a
token_mode flag (set by UnitreeG1 when sonic_token_action is enabled):
before any real token arrives the decoder is fed the all-zero neutral
token (stable neutral stance), and afterwards the last received token is
held between control ticks (the ~50 Hz control loop outruns the ~30 Hz
token stream). Living in the controller, this applies uniformly to
run_g1_onboard, lerobot-rollout and the sim replays, so the explicit
neutral seeding in run_g1_onboard is removed.
2026-07-27 10:34:42 +02:00
Martino Russi b587e81587 feat(unitree_g1): onboard controller deployment for SONIC walk
Run the whole-body controller (SONIC decoder / GR00T) onboard the G1 against
local DDS at full rate, with the laptop shipping only high-level actions over
ZMQ instead of 50Hz lowcmd via the socket bridge.

- config: add onboard, dds_interface, release_motion_control, physical_remote
- unitree_g1: onboard connect() branch (local DDS + MotionSwitcher release +
  physical wireless remote), _release_motion_control, _wireless_remote_input,
  controller-loop wireless priority; SDK channels when sim OR onboard
- run_g1_server: port Gripper/build_gripper/parse_camera_specs; add --cameras
  spec supporting by-path device names (survive USB re-enumeration) + FOURCC
- run_g1_onboard: onboard entry point (ZMQ actions -> send_action), with a
  --sonic-token-action flag for the 64-D latent-token interface
- infer_sonic_g1_onboard: laptop-side sender that runs nepyope/sonic_walk
  (pi0.5) and PUSHes 64-D tokens to the onboard controller
2026-07-26 21:36:02 +02:00
Martino Russi 4658dada9b feat(unitree_g1): 64-D SONIC token interface for lerobot-rollout + GR00T waist override
Add a token-output VLA path (sonic_token_action) so a policy trained on 64-D SONIC
motion tokens (e.g. nepyope/sonic_walk) drives the decoder directly via lerobot-rollout:
the robot advertises a 64-D motion_token.{i}.pos action and echoes the last commanded
token as a 64-D observation.state (motion_token_state.{i}.pos), encoder bypassed.

Also:
- gr00t_locomotion: allow an external upper-body IK to override the 3 waist joints, and
  cap ORT to 1 intra/inter thread so the 50Hz loop doesn't stutter under contention.
- sonic_pipeline: make_ort_session_options takes optional thread caps; report the
  provider actually bound.
- unitree_g1: build the sim env with publish_images=False/cameras=[] to avoid the
  offscreen EGL context crash (we drive image policies from recorded/live frames), and
  guard the startup sim-step race (zero-norm pelvis quat) so the sim thread survives.
2026-07-26 20:49:32 +02:00
Martino Russi 57ea6f4106 feat(unitree_g1): episode reset, lazy replay decode, safe shutdown
- reset(): pause the background controller and, for full-body controllers,
  publish the default pose directly (new _controller_paused flag) so reset and
  the controller loop aren't both writing low commands.
- SONIC pipeline: add reset() to StandingEncoderDecoder and PlannerController
  (clear token/proprio history/heading, rewind motion buffer); SonicRuntime.reset()
  now calls controller.reset().
- sonic_whole_body: require the full dense 34-D command (no silent zero-fill of a
  partial action) and integrate yaw-rate (idx 33) into heading.
- controllers/__init__: import the controller classes referenced in __all__.
- unitree_g1: lazy replay-frame decode + small cache instead of decoding all
  frames up front; safer disconnect (longer controller-thread join + fail-safe
  that skips the graceful ramp if the thread won't stop).
- lint: ruff-format config_unitree_g1 hand_closed_pose; prettier README table.
2026-07-24 12:02:49 +02:00
Martino Russi 4209639f33 refactor(unitree_g1): isolate SONIC encoder/decoder whole-body path
Strip everything except the OpenHLM/pi0.5 -> SONIC encoder/decoder rollout
path so this branch does exactly that and nothing more:

- Remove the SONIC motion planner (planner ONNX + subprocess worker, PlannerMotion,
  replanning, MovementState/LocomotionMode, joystick) from sonic_pipeline; keep the
  encoder/decoder and the caller-fed reference buffer (PlannerController) intact.
- Slim SonicRuntime to load only the encoder/decoder; SonicWholeBodyController now
  runs solely the 34-D whole-body command path (drop SMPL/VR3/keyboard teleop).
- Delete the pico_headset teleoperator (SONIC's SMPL/VR3 teleop source).
- Move WB action constants into g1_utils; repoint imports.

GR00T/Holosoma locomotion controllers are left untouched.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-20 20:15:51 +02:00
Martino Russi fc7a0bc2fd feat(unitree_g1): drive SONIC whole-body from a 34-D OpenHLM/pi0.5 VLA
Add a dense 34-D whole-body command path so lerobot-rollout can drive the
G1 directly with an OpenHLM / pi0.5 policy through the SONIC encoder/decoder:

- SonicWholeBodyController: wb.{i}.pos action interface, mode-0 reference with
  a rolling 50-frame trajectory (finite-diff velocities) and first-tick anchor
  init; correct MuJoCo->IsaacLab joint reordering.
- unitree_g1: expose 34-D wb_state.{i}.pos proprio; empty/replay camera feeds
  for image-conditioned policies; Dex3 hand publishing from the grip scalars.
- g1_utils: obs_to_wb34_state + WB action constants.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-20 20:01:48 +02:00
Martino Russi 5f6513551c Merge branch 'main' into feat/unitree_g1_sonic_rebased 2026-07-18 13:23:23 +02:00
Martino Russi 70e157e00f fix ruff 2026-07-18 13:22:53 +02:00
Martino Russi 1837be51bf add 3 point calibration + waist coupling, remote controller and smoothed motion 2026-07-17 17:56:30 +02:00
Martino Russi bedd56eed9 Remove g1_sonic_slider, examples/onnx, and SONIC debugging docs 2026-07-16 14:40:32 +02:00
Martino Russi c165e4df68 Merge branch 'main' into feat/unitree_g1_sonic_rebased 2026-07-16 14:33:10 +02:00
Martino Russi 5e24da483a (add) sonic 3-point teleop, safe startup/shutdown, tested on real g1 2026-07-16 13:38:49 +02:00
Martino Russi 9c54665a76 test 3-point teleop
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-15 18:20:26 +02:00
Martino Russi f6a845c30c Merge branch 'main' into feat/unitree_g1_sonic_rebased 2026-07-15 17:13:50 +02:00
Martino Russi 45e8336854 replace quat operations with scipy 2026-07-15 17:07:09 +02:00
Martino Russi 5046e2df32 fix ruff 2026-07-15 16:42:46 +02:00
Martino Russi 1c88e26c6d clean up sonic-side 2026-07-15 16:40:56 +02:00
Martino Russi 69a3edfa33 fix lint 2026-07-15 16:00:42 +02:00
Martino Russi 2492ce2c29 switch to logging 2026-07-15 15:30:54 +02:00
Martino Russi c8e75da55f Merge remote-tracking branch 'origin/main' into feat/unitree_g1_sonic_rebased 2026-07-15 14:59:53 +02:00
Martino Russi 2eae31ea2b fix(unitree_g1): disable SMPL root-motion anchor to prevent sim instability
Feeding the per-frame SMPL root quaternion into the mode-2 anchor produced
root-acceleration spikes (NaN QACC at DOF 0) mid-episode during replay. Keep the
anchor self-driven until the reference root trajectory is smoothed/rate-matched
(30 Hz dataset -> 50 Hz control).

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-15 14:59:22 +02:00
Martino Russi c997abe739 (fix) keep num of ORTthreads under core count 2026-07-14 18:11:03 +02:00
Martino Russi c73579055e refactor(unitree_g1): drop duplicate keyboard code, clarify smpl sentinel
- Remove unused RawKeyboard/drain_keyboard/process_keyboard from sonic_pipeline
  (dead code duplicating lerobot.utils.keyboard_input); the G1 integration uses
  the joystick path. Drop now-unused sys/select/termios/tty imports.
- Add a comment explaining the smpl.0 presence check is a sentinel for a full
  SMPL window (review question).

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-14 15:09:03 +02:00
Martino Russi 4be438161b style: apply ruff format to sonic_pipeline and smpl_fk
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-14 14:46:05 +02:00
Martino Russi 806d28a883 docs(unitree_g1): add docstrings and comments to sonic_pipeline
Address review feedback that sonic_pipeline.py was dense and hard to read.
Adds a module-level architecture overview plus class and key-function
docstrings (planner subprocess, encoder/decoder, movement state, input
helpers). No behavior change.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-14 14:40:39 +02:00
Martino Russi 573b65ff6b (fix) hardcode smpl_skeleton, remove .npz 2026-07-14 13:07:30 +02:00
Martino Russi bc55713e7c fix relative imports 2026-07-13 18:50:00 +02:00
Martino Russi 4f53c42583 Apply ruff-format
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-10 16:03:28 +02:00
Martino Russi bfced3d149 Silence ruff N817 on scipy Rotation import 2026-07-10 15:59:55 +02:00
Martino Russi 4969813d4e Silence ruff N817 on scipy Rotation import 2026-07-10 15:49:09 +02:00
Martino Russi 1c87ca31a3 remove examples inlcuding npz motion files 2026-07-10 15:47:21 +02:00
Martino Russi 4bcde762cc add heading to SMPL, stream dataset 2026-07-10 15:44:48 +02:00
Martino Russi 943ae78cfe feat(unitree_g1): standalone PICO SMPL publisher + dedup/replay fixes
Add a self-contained rt/smpl publisher in the pico_headset teleoperator
(pico_publisher.py + numpy SMPL FK in smpl_fk.py + vendored skeleton table)
so headset whole-body teleop no longer depends on gear_sonic/torch; only
xrobotoolkit_sdk is needed at the headset.

Also: share lowstate_to_obs/get_gravity_orientation via g1_utils (dedup
sonic_pipeline and UnitreeG1.get_observation), and fix dataset-replay joint
ordering (Unitree -> IsaacLab) for sonic.py --replay-dataset.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-09 19:13:22 +02:00
Martino Russi 3363688f1e Merge branch 'main' into feat/unitree_g1_sonic_rebased 2026-07-09 18:02:53 +02:00
Martino Russi 0876629e72 Merge branch 'main' into feat/unitree_g1_sonic_rebased 2026-07-06 18:21:16 +02:00
Martino Russi 305614b8c6 add pico teleoperator, add sonic VR support
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-06 18:16:12 +02:00
Martino Russi 02d3202c4f add SMPL wiring into sonic controller 2026-07-06 18:13:46 +02:00
Martino Russi 3b6de2fdf8 fix(unitree_g1): fix typo flagged by spellchecker in motion_loader docstring 2026-06-26 13:46:33 +02:00
Martino Russi 744f3667c0 fix(unitree_g1): silence bandit findings in SONIC example/pipeline 2026-06-26 13:40:53 +02:00
Martino Russi fdde436776 Merge branch 'main' into feat/unitree_g1_sonic_rebased 2026-06-26 13:35:28 +02:00
Martino Russi 5c683c65c6 Merge branch 'main' into feat/unitree_g1_sonic_rebased 2026-06-25 14:38:48 +02:00
Martino Russi dfbc25c58f fix(unitree_g1): satisfy ruff lint/format and address review comments 2026-06-25 14:37:44 +02:00
Martino Russi 804c76bcc2 Merge branch 'main' into feat/unitree_g1_sonic_rebased 2026-06-25 13:41:04 +02:00
Martino Russi e6afa69be9 add motion loader 2026-06-17 12:31:08 +02:00
Martino Russi 31d1439e29 add custom motion loader 2026-06-17 12:29:36 +02:00
Martino Russi 1c118c6359 feat(unitree_g1): add SONIC whole-body controller
Move GrootLocomotionController and HolosomaLocomotionController into a new
controllers/ subpackage and add the SONIC whole-body controller
(sonic_pipeline.py, sonic_whole_body.py) plus the examples/unitree_g1/sonic.py
standalone script. UnitreeG1 now honors a controller's kp/kd, calls
controller.shutdown() on disconnect, and skips arm publishing for full_body
controllers.
2026-06-16 17:12:20 +02:00
9 changed files with 1021 additions and 33 deletions
+42 -1
View File
@@ -290,6 +290,47 @@ lerobot-rollout \
---
## Part 5: SONIC Whole-Body Control (Latent-Token Policies)
LeRobot now ports the **decoder** half of NVIDIA's [SONIC](https://github.com/NVIDIA/sonic) whole-body deploy stack as the `SonicWholeBodyController`. Instead of commanding joints, the policy outputs a 64-D **latent motion token** (`motion_token.{i}.pos`) each tick; the SONIC decoder maps that token plus recent proprioception history onto a residual action that is scaled and added to the standing pose, producing 50 Hz joint-position targets for all 29 DoF.
The controller loads its ONNX decoder and all deploy constants (PD gains, `default_angles`, `action_scale`, and the neutral token) from [`lerobot/sonic_decoder`](https://huggingface.co/lerobot/sonic_decoder). The repo also ships a distilled low-latency decoder, selectable with the controller's `policy_type="low_latency"` argument (defaults to `"default"`).
### Test in Simulation
Any token-output policy trained on SONIC motion tokens (e.g. [`nepyope/sonic_walk`](https://huggingface.co/nepyope/sonic_walk)) can drive it:
```bash
lerobot-rollout \
--strategy.type=base \
--policy.path=nepyope/sonic_walk \
--policy.device=cuda \
--robot.type=unitree_g1 \
--robot.is_simulation=true \
--robot.controller=SonicWholeBodyController \
--robot.cameras='{"ego_view": {"type": "zmq", "server_address": "localhost", "port": 5555, "camera_name": "head_camera", "width": 640, "height": 480, "fps": 30, "warmup_s": 15}}'
```
### Run on the Physical Robot
Start the robot server (see Part 2), then point the rollout at the robot instead of the sim:
```bash
lerobot-rollout \
--strategy.type=base \
--policy.path=nepyope/sonic_walk \
--policy.device=cuda \
--robot.type=unitree_g1 \
--robot.is_simulation=false \
--robot.robot_ip=<ROBOT_IP> \
--robot.controller=SonicWholeBodyController \
--robot.cameras='{"ego_view": {"type": "zmq", "server_address": "<ROBOT_IP>", "port": 5555, "camera_name": "head_camera", "width": 640, "height": 480, "fps": 30, "warmup_s": 5}}'
```
On connect, the robot eases every joint to the decoder's `default_angles` before the controller takes over, so the first policy commands don't snap from the connect-time pose.
---
## Additional Resources
- [Unitree SDK Documentation](https://github.com/unitreerobotics/unitree_sdk2_python)
@@ -300,4 +341,4 @@ lerobot-rollout \
---
_Last updated: March 2026_
_Last updated: July 2026_
@@ -62,12 +62,30 @@ class UnitreeG1Config(RobotConfig):
# Socket config for ZMQ bridge
robot_ip: str = "192.168.123.164" # default G1 IP
# Run the locomotion / whole-body controller ONBOARD the robot (policy on the G1
# itself, against local DDS at full rate) instead of on the laptop over the ZMQ
# socket bridge. In this mode the robot object uses the real Unitree SDK channels
# and expects high-level actions (arm targets + joystick axes, or 64-D SONIC
# tokens) fed via send_action -- e.g. by run_g1_server's serve_onboard_controller,
# which receives them from the laptop over ZMQ. Mutually exclusive with is_simulation.
onboard: bool = False
# DDS network interface for onboard mode (None = SDK default, matching
# run_g1_server.py's ChannelFactoryInitialize(0)).
dds_interface: str | None = None
# Onboard sub-flags. On a real G1 both are True: the built-in motion services
# must be released before we can write lowcmd, and locomotion axes are read from
# the physical wireless remote. Against a DDS sim neither applies (no
# MotionSwitcher, no physical remote), so set both False so the controller takes
# its locomotion axes purely from send_action (ZMQ) input.
release_motion_control: bool = True
physical_remote: bool = True
# Cameras (ZMQ-based remote cameras)
cameras: dict[str, CameraConfig] = field(default_factory=dict)
# Compensates for gravity on the unitree's arms using the arm ik solver
gravity_compensation: bool = False
# Lower-body controller class name, e.g. "GrootLocomotionController" or
# "HolosomaLocomotionController". None disables it.
# Controller class name, e.g. GrootLocomotionController / HolosomaLocomotionController /
# SonicWholeBodyController. None disables it.
controller: str | None = None
@@ -0,0 +1,27 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
from .gr00t_locomotion import GrootLocomotionController
from .holosoma_locomotion import HolosomaLocomotionController
from .sonic_whole_body import SonicWholeBodyController
__all__ = [
"GrootLocomotionController",
"HolosomaLocomotionController",
"SonicWholeBodyController",
]
@@ -21,7 +21,7 @@ import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from .g1_utils import (
from ..g1_utils import (
REMOTE_AXES,
REMOTE_BUTTONS,
G1_29_JointIndex,
@@ -86,6 +86,7 @@ class GrootLocomotionController:
# Load policies
self.policy_balance, self.policy_walk = load_groot_policies()
self.default_angles = GROOT_DEFAULT_ANGLES
self.cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32) # vx, vy, theta_dot
# Robot state
@@ -22,7 +22,7 @@ import onnx
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from .g1_utils import (
from ..g1_utils import (
REMOTE_AXES,
G1_29_JointArmIndex,
G1_29_JointIndex,
@@ -104,6 +104,7 @@ class HolosomaLocomotionController:
# Load policy and gains
self.policy, self.kp, self.kd = load_policy()
self.default_angles = DEFAULT_ANGLES
self.cmd = np.zeros(3, dtype=np.float32)
# Robot state
@@ -0,0 +1,195 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""SONIC decoder whole-body controller for the Unitree G1 (token-only).
Pure-Python/ONNX re-implementation of the *decode* half of NVIDIA's SONIC deploy stack.
The encoder is intentionally absent: a token-output VLA (e.g. ``nepyope/sonic_walk``)
supplies the 64-D latent ``motion_token`` directly each tick, and the SONIC **decoder**
maps ``token + recent proprioception history`` to a residual action that is scaled and
added onto the standing pose (``default_angles``) to produce 50 Hz joint-position targets
for the robot's PD controller.
Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order) and
**MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` (in g1_utils)
convert between them. Quaternions are scalar-first ``(w, x, y, z)``.
"""
from __future__ import annotations
import json
import logging
import numpy as np
import onnx
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from ..g1_utils import (
ISAACLAB_TO_MUJOCO,
MUJOCO_TO_ISAACLAB,
G1_29_JointIndex,
get_gravity_orientation,
)
logger = logging.getLogger(__name__)
CONTROL_DT = 0.02 # 50 Hz control period (s)
TOKEN_DIM = 64 # decoder latent size
# Latent-token feature-key prefixes: action carries the token, obs echoes it back.
TOKEN_ACTION_PREFIX = "motion_token" # nosec B105 - feature-key prefix, not a secret
TOKEN_STATE_PREFIX = "motion_token_state" # nosec B105 - feature-key prefix, not a secret
# SONIC decoder checkpoint. Deploy constants (kp/kd, default_angles, action_scale,
# neutral_token) are baked into the ONNX metadata; see upload_sonic_decoder.py.
DEFAULT_SONIC_REPO_ID = "lerobot/sonic_decoder"
DECODER_INPUT_DIM = 994 # token(64) + 10-frame proprio history + gravity
# Decoder filename mapping: the full decoder (default) or NVIDIA's distilled low-latency one.
POLICY_FILES = {
"default": "model_decoder.onnx",
"low_latency": "low_latency/model_decoder.onnx",
}
def load_policy(
repo_id: str = DEFAULT_SONIC_REPO_ID,
policy_type: str = "default",
) -> tuple[ort.InferenceSession, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Load the SONIC decoder and its baked-in deploy constants from ONNX metadata.
Args:
repo_id: Hugging Face Hub repo ID
policy_type: Either "default" (full decoder) or "low_latency" (distilled)
Returns:
(decoder, kp, kd, default_angles, action_scale, neutral_token) tuple. The gains/pose/
scale are (29,) float32 in IsaacLab joint order; neutral_token is the (64,) idle latent.
"""
if policy_type not in POLICY_FILES:
raise ValueError(f"Unknown policy type: {policy_type}. Choose from: {list(POLICY_FILES.keys())}")
filename = POLICY_FILES[policy_type]
logger.info(f"Loading {policy_type.upper()} SONIC decoder from: {repo_id}/{filename}")
decoder_path = hf_hub_download(repo_id=repo_id, filename=filename)
decoder = ort.InferenceSession(decoder_path)
logger.info(f"Decoder loaded: {decoder.get_inputs()[0].shape}{decoder.get_outputs()[0].shape}")
# Extract deploy constants from ONNX metadata
model = onnx.load(decoder_path, load_external_data=False)
metadata = {prop.key: prop.value for prop in model.metadata_props}
required = ("kp", "kd", "default_angles", "action_scale", "neutral_token")
missing = [k for k in required if k not in metadata]
if missing:
raise ValueError(f"ONNX model must contain {list(required)} in metadata (missing {missing})")
arr = {k: np.array(json.loads(metadata[k]), dtype=np.float32) for k in required}
logger.info(f"Loaded SONIC deploy constants from ONNX ({len(arr['kp'])} joints)")
return decoder, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"]
class SonicWholeBodyController:
"""Full-body SONIC decoder controller for UnitreeG1's background controller thread.
Token-only deploy (encoder bypassed): each tick it appends the latest robot state to
10-frame history buffers, then maps the policy-supplied 64-D token + that history to a
residual action added onto ``default_angles`` -> 50 Hz joint-position targets.
"""
control_dt = CONTROL_DT
def __init__(self, policy_type: str = "default"):
self.decoder, self.kp, self.kd, self.default_angles, self.action_scale, self.neutral_token = (
load_policy(policy_type=policy_type)
)
self.decoder_input = self.decoder.get_inputs()[0].name
self.default_angles_mj = self.default_angles[MUJOCO_TO_ISAACLAB]
# 64-D latent-token action space; rollout maps the policy's 64-D output onto these keys.
self.action_ft = {f"{TOKEN_ACTION_PREFIX}.{i}.pos": float for i in range(TOKEN_DIM)}
# 64-D token proprio state, aggregated by rollout into observation.state (last token).
self.observation_ft = {f"{TOKEN_STATE_PREFIX}.{i}.pos": float for i in range(TOKEN_DIM)}
self.reset()
logger.info("SonicWholeBodyController initialized")
def reset(self) -> None:
"""Reset internal state for a new episode: held token and 10-frame history buffers."""
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._last_token = None # neutral token is re-seeded on the first tick
def observation_state(self) -> dict[str, float]:
"""Echo the last decoded token as ``observation.state`` so a token-output VLA closes
the loop on its own previous token."""
token = self._last_token if self._last_token is not None else np.zeros(TOKEN_DIM, dtype=np.float32)
return {f"{TOKEN_STATE_PREFIX}.{i}.pos": float(v) for i, v in enumerate(token)}
def run_step(self, action: dict, lowstate) -> dict:
if lowstate is None:
return {}
# Token: reassemble the dense 64-D latent from motion_token.{i}.pos (all keys required);
# else hold the last one (neutral until the first real token, which decodes to a stand).
keys = [f"{TOKEN_ACTION_PREFIX}.{i}.pos" for i in range(TOKEN_DIM)]
if action and all(k in action for k in keys):
self._last_token = np.fromiter(
(float(action[k]) for k in keys), dtype=np.float32, count=TOKEN_DIM
)
elif self._last_token is None:
self._last_token = self.neutral_token.copy()
# Read proprioception from lowstate (IsaacLab joint order).
q = np.array([lowstate.motor_state[m.value].q for m in G1_29_JointIndex], np.float32)
dq = np.array([lowstate.motor_state[m.value].dq for m in G1_29_JointIndex], np.float32)
quat = np.array(lowstate.imu_state.quaternion, np.float32) # (w, x, y, z)
quat = quat / (np.linalg.norm(quat) + 1e-8)
ang = np.array(lowstate.imu_state.gyroscope, np.float32)
# Push into the 10-frame history (newest first). The decoder consumes MuJoCo joint
# order, so reorder q/dq via MUJOCO_TO_ISAACLAB (validated against the ONNX; don't flip).
self.h_q_mj = [q[MUJOCO_TO_ISAACLAB] - self.default_angles_mj] + self.h_q_mj[:-1]
self.h_dq_mj = [dq[MUJOCO_TO_ISAACLAB]] + self.h_dq_mj[:-1]
self.h_ang = [ang] + self.h_ang[:-1]
self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
self.h_quat = [quat] + self.h_quat[:-1]
# Assemble the 994-D decoder input: token + oldest->newest history + gravity.
obs = np.zeros(DECODER_INPUT_DIM, np.float32)
obs[:TOKEN_DIM] = self._last_token
off = TOKEN_DIM
for hist, sz in ((self.h_ang, 3), (self.h_q_mj, 29), (self.h_dq_mj, 29), (self.h_act_mj, 29)):
for frame in reversed(hist):
obs[off : off + sz] = frame
off += sz
for hquat in reversed(self.h_quat):
obs[off : off + 3] = get_gravity_orientation(hquat)
off += 3
# Decode -> residual action (MuJoCo order) added onto the standing pose.
action_mj = (
self.decoder.run(None, {self.decoder_input: obs.reshape(1, -1)})[0].squeeze().astype(np.float32)
)
self.last_action_mj = action_mj.copy()
target = self.default_angles + action_mj[ISAACLAB_TO_MUJOCO] * self.action_scale
return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
+40 -2
View File
@@ -23,6 +23,43 @@ import numpy as np
NUM_MOTORS = 29
# Joint-order permutation between IsaacLab and Mujoco convention
ISAACLAB_TO_MUJOCO = np.array(
[
0,
3,
6,
9,
13,
17,
1,
4,
7,
10,
14,
18,
2,
5,
8,
11,
15,
19,
21,
23,
25,
27,
12,
16,
20,
22,
24,
26,
28,
],
dtype=np.int32,
)
MUJOCO_TO_ISAACLAB = np.argsort(ISAACLAB_TO_MUJOCO).astype(np.int32)
REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
@@ -68,8 +105,9 @@ def make_locomotion_controller(name: str | None):
if name is None:
return None
controllers = {
"GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
"HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
"GrootLocomotionController": "lerobot.robots.unitree_g1.controllers.gr00t_locomotion",
"HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion",
"SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
}
module_path = controllers.get(name)
if module_path is None:
+372 -1
View File
@@ -22,16 +22,33 @@ This server runs on the robot and forwards:
- Robot commands (LowCmd) from ZMQ to DDS (from remote clients)
Uses JSON for secure serialization instead of pickle.
Controller-negotiation handshake
--------------------------------
The first message from a client agrees on which controller the server will run onboard
(``serve_onboard_controller``); the controller NEVER runs on the laptop client.
Test the handshake in isolation (no DDS, runs on a laptop) in two terminals::
# terminal A: handshake-only server
python -m lerobot.robots.unitree_g1.run_g1_server --handshake-only
# terminal B: client proposes a controller
python -m lerobot.robots.unitree_g1.run_g1_server \\
--handshake-client SonicWholeBodyController --sonic-token-action --server-ip 127.0.0.1
On the real robot, add ``--handshake`` to the normal bridge to require agreement first.
"""
import argparse
import base64
import contextlib
import json
import signal
import threading
import time
from typing import Any
import numpy as np
import zmq
from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
from unitree_sdk2py.core.channel import ChannelFactoryInitialize, ChannelPublisher, ChannelSubscriber
@@ -50,6 +67,257 @@ LOWCMD_PORT = 6000
LOWSTATE_PORT = 6001
NUM_MOTORS = 35
# Onboard high-level channels (serve_onboard_controller): compact actions in, state out.
ACTION_PORT = 6004
STATE_PORT = 6005
# Controller-negotiation handshake (REQ/REP). The client's first message agrees on
# which controller the server will run before any control data flows.
HANDSHAKE_PORT = 6002
PROTOCOL_VERSION = 1
# Controllers that can run ONBOARD (must match g1_utils.make_locomotion_controller).
# ``None`` (a.k.a. "bridge") means no onboard controller: the laptop owns control and
# streams raw lowcmd over the ZMQ DDS bridge (the legacy run_g1_server behavior).
VALID_CONTROLLERS = (
"GrootLocomotionController",
"HolosomaLocomotionController",
"SonicWholeBodyController",
)
# SONIC latent-token dimensionality (mirrors sonic_whole_body.TOKEN_DIM; kept local so
# the handshake can run without importing the heavy controller / onnxruntime).
TOKEN_DIM = 64
_BRIDGE_ALIASES = {"", "none", "null", "bridge", "raw"}
def _normalize_controller(name: str | None) -> str | None:
"""Map a requested controller name to a canonical value (or None for raw bridge)."""
if name is None:
return None
low = str(name).strip().lower()
if low in _BRIDGE_ALIASES:
return None
for c in VALID_CONTROLLERS:
if c.lower() == low:
return c
raise ValueError(f"Unknown controller {name!r}. Available: {list(VALID_CONTROLLERS)} or 'bridge'")
def _capabilities(controller: str | None, sonic_token_action: bool) -> dict[str, Any]:
"""The interface the server advertises for an agreed controller."""
caps: dict[str, Any] = {
"controller": controller,
"sonic_token_action": bool(sonic_token_action),
"protocol": PROTOCOL_VERSION,
}
if controller is None:
# Raw DDS bridge: the laptop runs the controller and streams lowcmd.
caps["mode"] = "bridge"
caps["lowcmd_port"] = LOWCMD_PORT
caps["lowstate_port"] = LOWSTATE_PORT
else:
# Onboard: the controller runs here; the laptop ships compact high-level actions.
caps["mode"] = "onboard"
caps["action_port"] = ACTION_PORT
caps["state_port"] = STATE_PORT
if sonic_token_action:
caps["action_space"] = "motion_token"
caps["action_dim"] = TOKEN_DIM
return caps
def negotiate_controller(sock: zmq.Socket, shutdown_event: threading.Event) -> dict[str, Any]:
"""Server side of the handshake: block on one REP socket until a client sends a
valid ``hello``, then reply with the negotiated capabilities and return them.
Rejects malformed / unknown-controller requests with an error reply and keeps
waiting (a rejected client can retry). Honors ``shutdown_event`` so Ctrl-C works.
"""
poller = zmq.Poller()
poller.register(sock, zmq.POLLIN)
while not shutdown_event.is_set():
if not dict(poller.poll(timeout=200)):
continue
raw = sock.recv()
try:
hello = json.loads(raw.decode("utf-8"))
except (ValueError, UnicodeDecodeError) as e:
sock.send_json({"type": "error", "ok": False, "error": f"bad hello: {e}"})
continue
try:
controller = _normalize_controller(hello.get("controller"))
except ValueError as e:
sock.send_json(
{"type": "error", "ok": False, "error": str(e), "available": list(VALID_CONTROLLERS)}
)
continue
reply = {
"type": "welcome",
"ok": True,
**_capabilities(controller, hello.get("sonic_token_action", False)),
}
sock.send_json(reply)
return reply
raise KeyboardInterrupt
def request_controller(
server_ip: str,
controller: str | None,
*,
sonic_token_action: bool = False,
port: int = HANDSHAKE_PORT,
timeout_s: float = 5.0,
) -> dict[str, Any]:
"""Client side of the handshake: propose a controller, return the server's agreed
capabilities (or raise on rejection / timeout)."""
ctx = zmq.Context.instance()
sock = ctx.socket(zmq.REQ)
sock.setsockopt(zmq.LINGER, 0)
sock.setsockopt(zmq.RCVTIMEO, int(timeout_s * 1000))
sock.setsockopt(zmq.SNDTIMEO, int(timeout_s * 1000))
sock.connect(f"tcp://{server_ip}:{port}")
hello = {
"type": "hello",
"controller": controller,
"sonic_token_action": bool(sonic_token_action),
"protocol": PROTOCOL_VERSION,
}
try:
sock.send_json(hello)
reply = sock.recv_json()
except zmq.Again as e:
raise TimeoutError(f"no handshake reply from {server_ip}:{port} within {timeout_s}s") from e
finally:
sock.close(linger=0)
if not reply.get("ok"):
raise RuntimeError(f"handshake rejected: {reply.get('error')} (available: {reply.get('available')})")
return reply
def serve_onboard_controller(
*,
controller: str,
sonic_token_action: bool,
dds_interface: str | None = None,
sim: bool = False,
cameras: dict | None = None,
camera_fps: int = 30,
camera_port: int = 5555,
action_port: int = ACTION_PORT,
state_port: int = STATE_PORT,
state_fps: float = 30.0,
stop: threading.Event | None = None,
) -> None:
"""Run the negotiated controller ONBOARD -- the single control path on the robot.
Builds ``UnitreeG1(onboard=True, controller=...)`` so the controller/balance loop runs
locally against DDS at full rate (the 50 Hz ``_controller_loop`` thread lives in
UnitreeG1), then receives compact high-level actions from the laptop over ZMQ
(:action_port), decodes them via the controller, publishes ``observation.state``
(:state_port), and optionally serves the ego camera. The controller NEVER runs on the
laptop; the laptop (lerobot-rollout thin-client) only ships tokens/axes and reads back
state + camera frames.
"""
# Imported lazily: UnitreeG1 imports request_controller from this module, so a
# top-level import here would be circular.
from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
if stop is None:
stop = threading.Event()
signal.signal(signal.SIGINT, lambda *_: stop.set())
signal.signal(signal.SIGTERM, lambda *_: stop.set())
cfg = UnitreeG1Config(
is_simulation=False,
onboard=True,
controller=controller,
dds_interface=dds_interface,
release_motion_control=not sim,
physical_remote=not sim,
cameras={},
)
# Optional camera server (background daemon thread; independent of DDS).
if cameras:
camera_server = ImageServer({"fps": camera_fps, "cameras": cameras}, port=camera_port)
threading.Thread(target=camera_server.run, daemon=True).start()
cam_summary = ", ".join(f"{name}(dev {c['device_id']})" for name, c in cameras.items())
print(f"Camera server started on :{camera_port}: {cam_summary}")
robot = UnitreeG1(cfg)
print(f"Connecting onboard robot (controller={controller}, token={sonic_token_action})...")
robot.connect()
ctx = zmq.Context.instance()
sock = ctx.socket(zmq.PULL)
sock.setsockopt(zmq.CONFLATE, 1) # only ever act on the freshest command
sock.setsockopt(zmq.RCVTIMEO, 200) # keeps the loop responsive to the stop event
sock.bind(f"tcp://0.0.0.0:{action_port}")
print(f"Onboard controller live. Waiting for laptop actions on :{action_port} ...")
print("Ctrl-C for graceful shutdown.")
state_sock = None
if state_fps > 0:
state_sock = ctx.socket(zmq.PUB)
state_sock.setsockopt(zmq.SNDHWM, 2)
state_sock.setsockopt(zmq.LINGER, 0)
state_sock.bind(f"tcp://0.0.0.0:{state_port}")
print(f"Publishing observation.state on :{state_port} at {state_fps:.0f} Hz")
def publish_state() -> None:
period = 1.0 / state_fps
while not stop.is_set():
t0 = time.time()
obs = robot.get_observation()
if obs:
# Forward every scalar proprio key the robot exposes (29 joint .q, IMU,
# and the SONIC token echo: 64-D motion_token_state.*). Camera arrays are
# streamed separately by the ImageServer, so drop ndarrays here. This
# makes the laptop thin-client a pure relay.
state = {
k: float(v)
for k, v in obs.items()
if isinstance(v, (bool, int, float, np.floating, np.integer))
}
with contextlib.suppress(zmq.Again):
state_sock.send_json(state, zmq.NOBLOCK)
time.sleep(max(0.0, period - (time.time() - t0)))
threading.Thread(target=publish_state, daemon=True).start()
else:
print("observation.state PUB disabled (state_fps<=0)")
n = 0
try:
while not stop.is_set():
try:
payload = sock.recv()
except zmq.Again:
continue
except zmq.ContextTerminated:
break
try:
action = json.loads(payload.decode("utf-8"))
except (ValueError, UnicodeDecodeError) as e:
print(f"Dropping malformed action: {e}")
continue
robot.send_action(action)
n += 1
if n % 60 == 0:
print(f"Applied {n} actions")
finally:
print("Shutting down onboard controller...")
stop.set()
if state_sock is not None:
with contextlib.suppress(Exception):
state_sock.close(linger=0)
robot.disconnect()
def lowstate_to_dict(msg: hg_LowState) -> dict[str, Any]:
"""Convert LowState SDK message to a JSON-serializable dictionary."""
@@ -160,8 +428,111 @@ def main() -> None:
parser.add_argument("--camera-width", type=int, default=640, help="Camera width (default: 640)")
parser.add_argument("--camera-height", type=int, default=480, help="Camera height (default: 480)")
parser.add_argument("--camera-port", type=int, default=5555, help="Camera ZMQ port (default: 5555)")
# Controller-negotiation handshake (first message agrees on the controller).
parser.add_argument(
"--handshake",
action="store_true",
help="Wait for a client to negotiate the controller before bridging",
)
parser.add_argument(
"--handshake-port",
type=int,
default=HANDSHAKE_PORT,
help=f"Handshake REQ/REP port (default: {HANDSHAKE_PORT})",
)
parser.add_argument(
"--handshake-only",
action="store_true",
help="Run ONLY the handshake server (no DDS/cameras) to test negotiation",
)
parser.add_argument(
"--handshake-client",
default=None,
metavar="CONTROLLER",
help="Act as a client: propose CONTROLLER (or 'bridge') to --server-ip and print the reply",
)
parser.add_argument("--server-ip", default="127.0.0.1", help="[--handshake-client] server IP")
parser.add_argument(
"--sonic-token-action",
action="store_true",
help="[handshake] negotiate the 64-D SONIC token action interface",
)
args = parser.parse_args()
# --- Isolated handshake test paths (no DDS, safe to run on a laptop) ---
if args.handshake_client is not None:
controller = (
None if args.handshake_client.strip().lower() in _BRIDGE_ALIASES else args.handshake_client
)
reply = request_controller(
args.server_ip,
controller,
sonic_token_action=args.sonic_token_action,
port=args.handshake_port,
)
print(json.dumps(reply, indent=2))
return
if args.handshake_only:
ctx = zmq.Context.instance()
rep = ctx.socket(zmq.REP)
rep.bind(f"tcp://0.0.0.0:{args.handshake_port}")
print(f"[handshake] server listening on :{args.handshake_port} (no DDS). Ctrl-C to stop.")
shutdown = threading.Event()
try:
while True:
reply = negotiate_controller(rep, shutdown)
print(
f"[handshake] agreed: controller={reply['controller']} mode={reply['mode']} "
f"sonic_token_action={reply['sonic_token_action']}"
)
except KeyboardInterrupt:
print("\n[handshake] stopping")
finally:
rep.close(linger=0)
ctx.term()
return
# Controller-negotiation handshake: the client's first message agrees on the
# controller, which we then run ONBOARD (the controller NEVER runs on the laptop).
# Bridge/None falls through to the legacy raw DDS forward (deprecated laptop control).
if args.handshake:
ctx = zmq.Context.instance()
hs = ctx.socket(zmq.REP)
hs.bind(f"tcp://0.0.0.0:{args.handshake_port}")
print(f"[handshake] waiting for client controller agreement on :{args.handshake_port} ...")
shutdown = threading.Event()
try:
agreed = negotiate_controller(hs, shutdown)
except KeyboardInterrupt:
print("[handshake] interrupted before agreement; exiting")
hs.close(linger=0)
ctx.term()
return
hs.close(linger=0)
if agreed["controller"] is not None:
print(
f"[handshake] running controller ONBOARD: {agreed['controller']} "
f"(sonic_token_action={agreed['sonic_token_action']})"
)
cameras = None
if args.camera:
cameras = {
"head_camera": {
"device_id": args.camera_device,
"shape": [args.camera_height, args.camera_width],
}
}
serve_onboard_controller(
controller=agreed["controller"],
sonic_token_action=bool(agreed["sonic_token_action"]),
cameras=cameras,
camera_fps=args.camera_fps,
camera_port=args.camera_port,
)
return
print("[handshake] client selected raw DDS bridge (laptop owns control) -> legacy forward.")
# Optionally start camera server in background thread
camera_thread = None
if args.camera:
@@ -205,6 +576,7 @@ def main() -> None:
# initialize ZMQ
ctx = zmq.Context.instance()
shutdown_event = threading.Event()
# receive commands from remote client
lowcmd_sock = ctx.socket(zmq.PULL)
@@ -215,7 +587,6 @@ def main() -> None:
lowstate_sock.bind(f"tcp://0.0.0.0:{LOWSTATE_PORT}")
state_period = 0.002 # ~500 hz
shutdown_event = threading.Event()
# start observation forwarding in background thread
t_state = threading.Thread(
+321 -25
View File
@@ -16,6 +16,8 @@
from __future__ import annotations
import contextlib
import json
import logging
import threading
import time
@@ -27,6 +29,7 @@ import numpy as np
from lerobot.cameras import make_cameras_from_configs
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.errors import DeviceNotConnectedError
from lerobot.utils.import_utils import _unitree_sdk_available, require_package
from ..robot import Robot
@@ -34,7 +37,6 @@ from .config_unitree_g1 import UnitreeG1Config
from .g1_kinematics import G1_29_ArmIK
from .g1_utils import (
REMOTE_AXES,
REMOTE_KEYS,
G1_29_JointArmIndex,
G1_29_JointIndex,
default_remote_input,
@@ -79,6 +81,28 @@ class LocomotionController(Protocol):
kTopicLowCommand_Debug = "rt/lowcmd"
kTopicLowState = "rt/lowstate"
# Wireless-remote button byte layout, mapped to the positional button indices the
# locomotion controllers expect. Used in onboard mode to read the physical Unitree
# remote from lowstate (mirrors the exo teleoperator's RemoteController).
_REMOTE_BUTTON_MAP: list[str] = [
"RB",
"LB",
"start",
"back",
"RT",
"LT",
"",
"",
"A",
"B",
"X",
"Y",
"up",
"right",
"down",
"left",
]
@dataclass
class MotorState:
@@ -119,24 +143,34 @@ class UnitreeG1(Robot):
self.config = config
self.control_dt = config.control_dt
# Three mutually-exclusive roles:
# * simulation : local DDS + controller run in-process against a MuJoCo world.
# * onboard : local DDS + controller run in-process on the robot NX.
# * client : thin laptop client. No DDS, no controller. It negotiates a
# controller with ``run_g1_server`` (which runs it onboard),
# PUSHes high-level actions and reads back state + cameras over
# ZMQ. The controller *always* runs on the robot, never here.
self._client = not config.is_simulation and not config.onboard
# Initialize cameras config (ZMQ-based) - actual connection in connect()
self._cameras = make_cameras_from_configs(config.cameras)
# Import channel classes based on mode
if config.is_simulation:
# DDS channels are only needed by the in-process control roles (sim / onboard),
# which both drive the real Unitree SDK. The thin client never touches DDS.
if config.is_simulation or config.onboard:
self._ChannelFactoryInitialize = _SDKChannelFactoryInitialize
self._ChannelPublisher = _SDKChannelPublisher
self._ChannelSubscriber = _SDKChannelSubscriber
else:
from .unitree_sdk2_socket import (
ChannelFactoryInitialize,
ChannelPublisher,
ChannelSubscriber,
)
self._ChannelFactoryInitialize = None
self._ChannelPublisher = None
self._ChannelSubscriber = None
self._ChannelFactoryInitialize = ChannelFactoryInitialize
self._ChannelPublisher = ChannelPublisher
self._ChannelSubscriber = ChannelSubscriber
# Client-side ZMQ handles / negotiated capabilities (populated in connect()).
self._client_action_sock = None
self._client_state_sock = None
self._client_state_latest: dict[str, float] = {}
self._client_caps: dict | None = None
# Initialize state variables
self.sim_env = None
@@ -148,15 +182,33 @@ class UnitreeG1(Robot):
self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
# Lower-body controller loaded dynamically
self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
# 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)
# Controller thread state
self._controller_thread = None
self._controller_action_lock = threading.Lock()
self.controller_input = default_remote_input()
self.controller_output = {}
# Onboard-only: parser for the physical Unitree wireless remote (read straight
# from local lowstate so joystick locomotion works without a laptop round-trip).
self._joystick = None
@property
def _sonic_token(self) -> bool:
"""Whether the SONIC whole-body decoder is active.
A SONIC controller consumes a 64-D latent motion token as its action and echoes
the last commanded token as ``observation.state``. Keyed purely off the selected
controller so the token interface is implicit -- no separate config flag, and the
thin client (which has no controller instance) can still advertise the schema.
"""
return self.config.controller == "SonicWholeBodyController"
def _subscribe_lowstate(self): # polls robot state @ 250Hz
while not self._shutdown_event.is_set():
start_time = time.time()
@@ -233,13 +285,38 @@ class UnitreeG1(Robot):
@cached_property
def observation_features(self) -> dict[str, type | tuple]:
return {**self._motors_ft, **self._cameras_ft}
# Controllers may contribute their own proprio features (e.g. SONIC's token state).
# The thin client has no controller instance, so mirror the onboard token schema
# by controller name (SONIC echoes its last token as observation.state).
controller_ft = getattr(self.controller, "observation_ft", {})
if self._client and self._sonic_token:
from .controllers.sonic_whole_body import TOKEN_DIM, TOKEN_STATE_PREFIX
controller_ft = {f"{TOKEN_STATE_PREFIX}.{i}.pos": float for i in range(TOKEN_DIM)}
return {**self._motors_ft, **controller_ft, **self._cameras_ft}
@cached_property
def action_features(self) -> dict[str, type]:
if self.controller is None:
# Role-agnostic: the schema is a pure function of the configured controller name,
# so the thin client advertises the same action space as the onboard robot.
# No controller configured at all: raw 29-DoF joint teleop.
if self.config.controller is None:
return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
# Whole-body controllers (SONIC): 64-D latent token. On the thin client there is
# no controller instance, so advertise the same token schema by controller name.
controller_ft = getattr(self.controller, "action_ft", None)
if controller_ft is not None:
return dict(controller_ft)
if self._client and self._sonic_token:
from .controllers.sonic_whole_body import TOKEN_ACTION_PREFIX, TOKEN_DIM
return {f"{TOKEN_ACTION_PREFIX}.{i}.pos": float for i in range(TOKEN_DIM)}
# Locomotion controllers (GR00T / Holosoma): arm joint targets + joystick axes.
# TODO: have GR00T/Holosoma advertise their own action_features too, so every
# controller declares its action space and this fallthrough can be dropped.
arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
remote_features = dict.fromkeys(REMOTE_AXES, float)
return {**arm_features, **remote_features}
@@ -271,6 +348,13 @@ class UnitreeG1(Robot):
with self._controller_action_lock:
controller_input = dict(self.controller_input)
# Onboard: the physical Unitree remote (in local lowstate) takes
# priority for locomotion when active; otherwise laptop/ZMQ axes stand.
if self.config.onboard:
wl = self._wireless_remote_input(lowstate)
if wl is not None:
controller_input.update(wl)
# Run controller step
controller_action = self.controller.run_step(controller_input, lowstate)
@@ -293,7 +377,170 @@ class UnitreeG1(Robot):
def configure(self) -> None:
pass
def _wireless_remote_input(self, lowstate) -> dict | None:
"""Parse the physical Unitree remote from lowstate into controller inputs.
Onboard only. Returns None when the remote is idle so the laptop-provided
(ZMQ) axes keep control; otherwise the physical remote takes priority.
"""
js = self._joystick
if js is None:
return None
wr = getattr(lowstate, "wireless_remote", None)
if not wr or len(wr) < 24:
return None
try:
js.extract(wr)
except Exception: # noqa: BLE001
return None
axes = {
"remote.lx": float(js.lx.data),
"remote.ly": float(js.ly.data),
"remote.rx": float(js.rx.data),
"remote.ry": float(js.ry.data),
}
active = any(abs(v) > 1e-2 for v in axes.values())
out = dict(axes)
for i, name in enumerate(_REMOTE_BUTTON_MAP):
if name:
val = float(getattr(js, name).data)
out[f"remote.button.{i}"] = val
if val:
active = True
return out if active else None
def _release_motion_control(self) -> None:
"""Release the robot's built-in motion services so we can send raw lowcmd.
Onboard-only. Mirrors run_g1_server.py: on the real robot the factory
locomotion/hand services must relinquish control before our controller can
write to ``rt/lowcmd``, otherwise commands are ignored or fought.
"""
from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
msc = MotionSwitcherClient()
msc.SetTimeout(5.0)
msc.Init()
_, result = msc.CheckMode()
while result is not None and "name" in result and result["name"]:
logger.info("[UnitreeG1] Releasing built-in mode '%s'...", result["name"])
msc.ReleaseMode()
_, result = msc.CheckMode()
time.sleep(1.0)
# ------------------------------------------------------------------ #
# Thin-client role (laptop): no DDS, no controller. Talks to run_g1_server
# over ZMQ. The controller ALWAYS runs onboard; we only relay high-level
# actions and read back the state echo + camera frames.
# ------------------------------------------------------------------ #
def _connect_client(self) -> None:
import zmq
from .run_g1_server import ACTION_PORT, HANDSHAKE_PORT, STATE_PORT, request_controller
server_ip = self.config.robot_ip
if not server_ip:
raise ValueError("client mode requires config.robot_ip (the G1 running run_g1_server)")
# 1) Handshake: agree with the server on which controller it will run onboard.
logger.info(
"[client] handshaking with %s:%d (controller=%s, token=%s)...",
server_ip,
HANDSHAKE_PORT,
self.config.controller,
self._sonic_token,
)
self._client_caps = request_controller(
server_ip,
self.config.controller,
sonic_token_action=self._sonic_token,
port=HANDSHAKE_PORT,
)
logger.info("[client] server agreed: %s", self._client_caps)
ctx = zmq.Context.instance()
# 2) Action PUSH: ship compact high-level actions to the onboard controller.
self._client_action_sock = ctx.socket(zmq.PUSH)
self._client_action_sock.setsockopt(zmq.SNDHWM, 2)
self._client_action_sock.setsockopt(zmq.LINGER, 0)
self._client_action_sock.connect(f"tcp://{server_ip}:{ACTION_PORT}")
# 3) State SUB: read the onboard observation.state echo (last token / joints).
self._client_state_sock = ctx.socket(zmq.SUB)
self._client_state_sock.setsockopt(zmq.CONFLATE, 1)
self._client_state_sock.setsockopt_string(zmq.SUBSCRIBE, "")
self._client_state_sock.connect(f"tcp://{server_ip}:{STATE_PORT}")
# 4) Cameras (ZMQ ImageServer served by run_g1_server) - same as any client.
for cam in self._cameras.values():
if not cam.is_connected:
cam.connect()
logger.info(
"[client] connected: actions ->:%d, state <-:%d, %d camera(s).",
ACTION_PORT,
STATE_PORT,
len(self._cameras),
)
def _recv_client_state(self) -> None:
"""Drain the state SUB (CONFLATE keeps only the freshest) into the latest cache."""
import zmq
if self._client_state_sock is None:
return
while True:
try:
state = self._client_state_sock.recv_json(flags=zmq.NOBLOCK)
except zmq.Again:
break
except (ValueError, zmq.ZMQError):
break
if isinstance(state, dict):
self._client_state_latest = {k: float(v) for k, v in state.items()}
def _get_observation_client(self) -> RobotObservation:
self._recv_client_state()
obs: dict = dict(self._client_state_latest)
for cam_name, cam in self._cameras.items():
if getattr(cam, "use_rgb", True):
obs[cam_name] = cam.read_latest()
if getattr(cam, "use_depth", False):
obs[f"{cam_name}_depth"] = cam.read_latest_depth()
return obs
def _send_action_client(self, action: RobotAction) -> RobotAction:
"""Relay the raw action straight to the onboard controller. NO processing here:
the controller negotiated in the handshake interprets it (token / wb / arm)."""
import zmq
if self._client_action_sock is None:
raise DeviceNotConnectedError("UnitreeG1 client is not connected")
payload = json.dumps({k: float(v) for k, v in action.items()}).encode("utf-8")
with contextlib.suppress(zmq.Again):
self._client_action_sock.send(payload, zmq.NOBLOCK)
return action
def _disconnect_client(self) -> None:
for sock in (self._client_action_sock, self._client_state_sock):
if sock is not None:
with contextlib.suppress(Exception):
sock.close(linger=0)
self._client_action_sock = None
self._client_state_sock = None
for cam in self._cameras.values():
with contextlib.suppress(Exception):
cam.disconnect()
def connect(self, calibrate: bool = True) -> None: # connect to DDS
# Thin-client role: no DDS, no controller. Negotiate the controller with
# run_g1_server (which runs it onboard), then open the high-level ZMQ links:
# PUSH actions on :ACTION_PORT, SUB state echo on :STATE_PORT, cameras via ZMQ.
if self._client:
self._connect_client()
return
# Initialize DDS channel and simulation environment
if self.config.is_simulation:
from lerobot.envs import make_env
@@ -302,8 +549,28 @@ class UnitreeG1(Robot):
self._env_wrapper = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True)
# Extract the actual gym env from the dict structure
self.sim_env = self._env_wrapper["hub_env"][0].envs[0]
else:
self._ChannelFactoryInitialize(0, config=self.config)
elif self.config.onboard:
# Real robot, controller running onboard against local DDS. Initialize the
# real SDK channel factory on the robot's DDS interface and take low-level
# control from the built-in services before we start writing lowcmd.
if self.config.dds_interface:
self._ChannelFactoryInitialize(0, self.config.dds_interface)
else:
self._ChannelFactoryInitialize(0)
# Real robot: hand low-level control over from the built-in services.
# A DDS sim has no MotionSwitcher, so this is skipped there.
if self.config.release_motion_control:
self._release_motion_control()
# Real robot: read the physical wireless remote from lowstate for
# locomotion. A sim has no physical remote, so leave _joystick=None and
# let send_action (ZMQ) drive the locomotion axes instead.
if self.config.physical_remote:
from unitree_sdk2py.utils.joystick import Joystick
self._joystick = Joystick()
for axis in (self._joystick.lx, self._joystick.ly, self._joystick.rx, self._joystick.ry):
axis.smooth = 1.0
axis.deadzone = 0.0
# Initialize direct motor control interface
self.lowcmd_publisher = self._ChannelPublisher(kTopicLowCommand_Debug, hg_LowCmd)
@@ -341,8 +608,14 @@ class UnitreeG1(Robot):
logger.info("[UnitreeG1] Connected to robot.")
self.msg.mode_machine = lowstate.mode_machine
self.kp = np.array(self.config.kp, dtype=np.float32)
self.kd = np.array(self.config.kd, dtype=np.float32)
# Prefer the active controller's gains (e.g. SONIC loads kp/kd from its ONNX);
# otherwise fall back to the config defaults.
if self.controller is not None and hasattr(self.controller, "kp"):
self.kp = np.array(self.controller.kp, dtype=np.float32)
self.kd = np.array(self.controller.kd, dtype=np.float32)
else:
self.kp = np.array(self.config.kp, dtype=np.float32)
self.kd = np.array(self.config.kd, dtype=np.float32)
for joint in G1_29_JointIndex:
self.msg.motor_cmd[joint].mode = 1
@@ -350,6 +623,11 @@ class UnitreeG1(Robot):
self.msg.motor_cmd[joint].kd = self.kd[joint.value]
self.msg.motor_cmd[joint].q = lowstate.motor_state[joint.value].q
# Ease into the controller's home pose before it takes over, so the first commands
# don't snap from the connect-time pose.
if self.controller is not None and hasattr(self.controller, "default_angles"):
self.reset(default_positions=self.controller.default_angles)
# Start controller thread if enabled
if self.controller is not None:
self._controller_thread = threading.Thread(target=self._controller_loop, daemon=True)
@@ -372,6 +650,10 @@ class UnitreeG1(Robot):
logger.warning(f"Failed to send zero-torque on disconnect: {e}")
def disconnect(self):
if self._client:
self._disconnect_client()
return
# Put robot in passive mode before stopping threads
if not self.config.is_simulation:
self._send_zero_torque()
@@ -417,6 +699,9 @@ class UnitreeG1(Robot):
cam.disconnect()
def get_observation(self) -> RobotObservation:
if self._client:
return self._get_observation_client()
with self._lowstate_lock:
lowstate = self._lowstate
if lowstate is None:
@@ -461,6 +746,11 @@ class UnitreeG1(Robot):
if lowstate.wireless_remote:
obs["wireless_remote"] = lowstate.wireless_remote
# Controller-contributed observation (e.g. SONIC echoes its last decoded token as
# observation.state so a token-output VLA closes the loop on its own previous token).
if self.controller is not None and hasattr(self.controller, "observation_state"):
obs.update(self.controller.observation_state())
# Cameras - read images from ZMQ cameras
for cam_name, cam in self._cameras.items():
if getattr(cam, "use_rgb", True):
@@ -471,6 +761,9 @@ class UnitreeG1(Robot):
return obs
def send_action(self, action: RobotAction) -> RobotAction:
if self._client:
return self._send_action_client(action)
action_to_publish = action
if self.controller is not None:
# Controller thread owns legs/waist. Here we only update joystick inputs
@@ -503,11 +796,12 @@ class UnitreeG1(Robot):
return action
def _update_controller_action(self, action: RobotAction) -> None:
"""Update controller input state from incoming teleop action."""
"""Forward incoming teleop action values into ``controller_input``; each controller
reads only the keys it understands."""
with self._controller_action_lock:
for key in REMOTE_KEYS:
if key in action:
self.controller_input[key] = action[key]
for key, value in action.items():
if isinstance(key, str) and value is not None:
self.controller_input[key] = value
@property
def is_calibrated(self) -> bool:
@@ -515,6 +809,8 @@ class UnitreeG1(Robot):
@property
def is_connected(self) -> bool:
if self._client:
return self._client_action_sock is not None
with self._lowstate_lock:
return self._lowstate is not None
@@ -565,7 +861,7 @@ class UnitreeG1(Robot):
interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
action_dict[f"{motor.name}.q"] = float(interp_pos)
self.send_action(action_dict)
self.publish_lowcmd(action_dict)
# Maintain constant control rate
elapsed = time.time() - start_time