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).
This commit is contained in:
Martino Russi
2026-07-30 11:40:01 +02:00
parent bbfc4ff443
commit 95d9029039
6 changed files with 105 additions and 748 deletions
+1 -56
View File
@@ -8,15 +8,6 @@
The Unitree G1 humanoid is now supported in LeRobot! You can teleoperate, train locomanipulation policies, test in sim, and more. Both 29 and 23 DoF variants are supported.
<Tip>
**New: SONIC whole-body control.** The `SonicWholeBodyController` runs NVIDIA's
[GEAR-SONIC](https://huggingface.co/nvidia/GEAR-SONIC) decoder on the G1, turning a
64-D latent motion token into full-body joint targets at 50 Hz. This lets you drive
the robot from a VLA policy trained on SONIC motion tokens (token in → whole-body
motion out) with `lerobot-rollout`, in sim or on the physical robot. See
[Whole-body control with SONIC](#whole-body-control-with-sonic) below.
</Tip>
---
## Part 1: Getting Started
@@ -68,7 +59,7 @@ lerobot-teleoperate \
--robot.controller=GrootLocomotionController
```
This will launch a [MuJoCo sim instance](https://huggingface.co/lerobot/unitree-g1-mujoco/tree/main) for the G1. You can connect a gamepad to your machine before launching in order to control the robot's locomotion in sim. We support [HolosomaLocomotionController](https://github.com/amazon-far/holosoma), [GrootLocomotionController](https://github.com/NVlabs/GR00T-WholeBodyControl), and [SonicWholeBodyController](https://huggingface.co/nvidia/GEAR-SONIC) via `--robot.controller`.
This will launch a [MuJoCo sim instance](https://huggingface.co/lerobot/unitree-g1-mujoco/tree/main) for the G1. You can connect a gamepad to your machine before launching in order to control the robot's locomotion in sim. We support both [HolosomaLocomotionController](https://github.com/amazon-far/holosoma) and [GrootLocomotionController](https://github.com/NVlabs/GR00T-WholeBodyControl) via `--robot.controller`.
- Press `9` to release the robot
- Press `7` / `8` to increase / decrease waist height
@@ -299,52 +290,6 @@ lerobot-rollout \
---
## Whole-body control with SONIC
The `SonicWholeBodyController` runs NVIDIA's [GEAR-SONIC](https://huggingface.co/nvidia/GEAR-SONIC)
decoder on the G1. Each 50 Hz tick it consumes a **64-D latent motion token** and emits
full-body joint targets — the encoder is bypassed, so a policy feeds tokens in and the
decoder turns them into motion. Before the first token arrives the controller holds a
neutral (idle) pose.
This makes the G1 drivable by a VLA policy trained to output SONIC motion tokens (token
as both `observation.state` and `action`, e.g. [`nepyope/sonic_walk`](https://huggingface.co/nepyope/sonic_walk))
using the standard `lerobot-rollout`. The controller always runs **onboard** the robot;
the laptop is a thin client that streams tokens and receives camera frames over ZMQ.
**On the robot** — start the server in handshake mode so it instantiates and runs the
controller onboard against local DDS at full rate:
```bash
cd ~/lerobot
python src/lerobot/robots/unitree_g1/run_g1_server.py --handshake --camera
```
**From your laptop** — run the token policy; selecting `--robot.controller=SonicWholeBodyController`
implicitly switches the robot to the 64-D latent-token action/observation interface:
```bash
lerobot-rollout \
--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}}' \
--task="walk back and forth" \
--duration=1000 \
--fps=30
```
<Tip>
SONIC is a token-only decoder in LeRobot: the only input path is the 64-D latent
vector. To train your own token policy, expose the 64-D token as the action (a config
choice, e.g. `pi05` with a 64-D action dim) — no policy code changes are needed.
</Tip>
---
## Additional Resources
- [Unitree SDK Documentation](https://github.com/unitreerobotics/unitree_sdk2_python)
@@ -62,24 +62,6 @@ 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)
@@ -14,8 +14,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import logging
from collections import deque
@@ -70,15 +68,9 @@ def load_groot_policies(
filename="GR00T-WholeBodyControl-Walk.onnx",
)
# Load ONNX policies with a capped thread pool. GR00T runs at 50 Hz in a
# background thread alongside the (torch) upper-body policy, IK and sim; letting
# ORT grab every core starves those and makes the whole rollout stutter. These
# are small MLPs, so 1 thread is both enough and lowest-latency.
from ..g1_utils import make_ort_session_options
so = make_ort_session_options(intra_op_num_threads=1, inter_op_num_threads=1)
policy_balance = ort.InferenceSession(balance_path, sess_options=so)
policy_walk = ort.InferenceSession(walk_path, sess_options=so)
# Load ONNX policies
policy_balance = ort.InferenceSession(balance_path)
policy_walk = ort.InferenceSession(walk_path)
logger.info("GR00T policies loaded successfully")
@@ -204,16 +196,6 @@ class GrootLocomotionController:
# Transform action back to target joint positions
target_dof_pos_15 = GROOT_DEFAULT_ANGLES[:15] + self.groot_action * ACTION_SCALE
# Waist override: an external upper-body IK can command the 3 waist joints
# (indices 12/13/14) via ``kWaist{Yaw,Roll,Pitch}.q`` in the action dict. When
# present, we substitute the balance policy's waist target so the torso tracks
# the IK while the policy keeps only the legs balanced. Single-publisher stays
# intact (this thread still owns joints 0-14).
for idx in (G1_29_JointIndex.kWaistYaw, G1_29_JointIndex.kWaistRoll, G1_29_JointIndex.kWaistPitch):
key = f"{idx.name}.q"
if key in action and action[key] is not None:
target_dof_pos_15[idx.value] = float(action[key])
# Build action dict
action_dict = {}
for i in range(15):
@@ -14,11 +14,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
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
@@ -26,7 +26,6 @@ from ..g1_utils import (
REMOTE_AXES,
G1_29_JointArmIndex,
G1_29_JointIndex,
compute_pd_gains,
get_gravity_orientation,
)
@@ -58,23 +57,12 @@ POLICY_FILES = {
"ppo": "ppo_g1_29dof.onnx",
}
# Per-joint motor model in Holosoma's joint order, plus the joints that get a x2
# stiffness/damping factor. These reproduce the kp/kd that used to be read from the
# policy's ONNX metadata exactly (both fastsac and ppo), so gains are now derived from
# the shared motor model (see g1_utils.compute_pd_gains) instead.
HOLOSOMA_MOTOR_MODELS = (
["7520_14", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2
+ ["7520_14", "5020", "5020"]
+ ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2
)
HOLOSOMA_DOUBLE = {4, 5, 10, 11, 13, 14}
def load_policy(
repo_id: str = DEFAULT_HOLOSOMA_REPO_ID,
policy_type: str = "fastsac",
) -> tuple[ort.InferenceSession, np.ndarray, np.ndarray]:
"""Load the Holosoma locomotion policy and its motor-model-derived PD gains.
"""Load Holosoma locomotion policy and extract KP/KD from metadata.
Args:
repo_id: Hugging Face Hub repo ID
@@ -93,7 +81,16 @@ def load_policy(
policy = ort.InferenceSession(policy_path)
logger.info(f"Policy loaded: {policy.get_inputs()[0].shape}{policy.get_outputs()[0].shape}")
kp, kd = compute_pd_gains(HOLOSOMA_MOTOR_MODELS, HOLOSOMA_DOUBLE)
# Extract KP/KD from ONNX metadata
model = onnx.load(policy_path, load_external_data=False)
metadata = {prop.key: prop.value for prop in model.metadata_props}
if "kp" not in metadata or "kd" not in metadata:
raise ValueError("ONNX model must contain 'kp' and 'kd' in metadata")
kp = np.array(json.loads(metadata["kp"]), dtype=np.float32)
kd = np.array(json.loads(metadata["kd"]), dtype=np.float32)
logger.info(f"Loaded KP/KD from ONNX ({len(kp)} joints)")
return policy, kp, kd
+1 -343
View File
@@ -22,33 +22,16 @@ 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
@@ -67,253 +50,6 @@ 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."""
@@ -424,86 +160,8 @@ 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:
@@ -547,7 +205,6 @@ def main() -> None:
# initialize ZMQ
ctx = zmq.Context.instance()
shutdown_event = threading.Event()
# receive commands from remote client
lowcmd_sock = ctx.socket(zmq.PULL)
@@ -558,6 +215,7 @@ 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(
+86 -293
View File
@@ -16,8 +16,6 @@
from __future__ import annotations
import contextlib
import json
import logging
import threading
import time
@@ -28,7 +26,6 @@ from typing import TYPE_CHECKING, Protocol, runtime_checkable
import numpy as np
from lerobot.cameras import make_cameras_from_configs
from lerobot.utils.errors import DeviceNotConnectedError
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.import_utils import _unitree_sdk_available, require_package
@@ -40,7 +37,6 @@ from .g1_utils import (
G1_29_JointArmIndex,
G1_29_JointIndex,
default_remote_input,
lowstate_to_obs,
make_locomotion_controller,
)
@@ -50,9 +46,7 @@ if TYPE_CHECKING or _unitree_sdk_available:
ChannelPublisher as _SDKChannelPublisher,
ChannelSubscriber as _SDKChannelSubscriber,
)
from unitree_sdk2py.idl.default import (
unitree_hg_msg_dds__LowCmd_,
)
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
LowCmd_ as hg_LowCmd,
LowState_ as hg_LowState,
@@ -84,14 +78,6 @@ 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:
@@ -132,37 +118,24 @@ 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)
# DDS channel classes are only needed by the in-process control roles. The thin
# client never touches DDS, so we don't import the socket shim at all.
if config.is_simulation or config.onboard:
# Import channel classes based on mode
if config.is_simulation:
self._ChannelFactoryInitialize = _SDKChannelFactoryInitialize
self._ChannelPublisher = _SDKChannelPublisher
self._ChannelSubscriber = _SDKChannelSubscriber
else:
self._ChannelFactoryInitialize = None
self._ChannelPublisher = None
self._ChannelSubscriber = None
from .unitree_sdk2_socket import (
ChannelFactoryInitialize,
ChannelPublisher,
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
# Optional arm gravity compensation (feed-forward torque via the arm IK solver).
self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
self._ChannelFactoryInitialize = ChannelFactoryInitialize
self._ChannelPublisher = ChannelPublisher
self._ChannelSubscriber = ChannelSubscriber
# Initialize state variables
self.sim_env = None
@@ -172,17 +145,14 @@ class UnitreeG1(Robot):
self._shutdown_event = threading.Event()
self.subscribe_thread = 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)
self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
# Token-driven deploy: a SONIC whole-body controller always runs in token
# mode -- it holds a neutral token until the first real one arrives, then
# holds the last token between control ticks.
if hasattr(self.controller, "token_mode"):
# Lower-body / whole-body controller loaded dynamically
self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
# A SONIC whole-body controller always runs in token mode: it holds a neutral
# token until the first real one arrives, then holds the last token between ticks.
if self.controller is not None and hasattr(self.controller, "token_mode"):
self.controller.token_mode = True
# Controller thread state
@@ -194,10 +164,6 @@ class UnitreeG1(Robot):
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
# Token-mode state: last 64-D SONIC latent token commanded by the policy,
# echoed back as ``observation.state`` so a token-output VLA closes the loop
# on its own previous token. Implicit whenever the SONIC whole-body controller
@@ -317,20 +283,12 @@ class UnitreeG1(Robot):
@cached_property
def observation_features(self) -> dict[str, type | tuple]:
return {
**self._motors_ft,
**self._token_state_ft,
**self._cameras_ft,
}
return {**self._motors_ft, **self._token_state_ft, **self._cameras_ft}
@cached_property
def action_features(self) -> dict[str, type]:
# Role-agnostic: the schema is a pure function of the controller name. 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:
if self.controller is None:
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
@@ -378,13 +336,6 @@ 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)
@@ -407,163 +358,7 @@ 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
@@ -572,28 +367,6 @@ 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]
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
else:
self._ChannelFactoryInitialize(0, config=self.config)
@@ -645,8 +418,7 @@ 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
# Start the 50 Hz controller thread (runs the locomotion/whole-body policy and
# publishes low commands to DDS).
# Start controller thread if enabled
if self.controller is not None:
self._controller_thread = threading.Thread(target=self._controller_loop, daemon=True)
self._controller_thread.start()
@@ -668,34 +440,12 @@ class UnitreeG1(Robot):
logger.warning(f"Failed to send zero-torque on disconnect: {e}")
def disconnect(self):
if self._client:
self._disconnect_client()
return
# Stop the controller loop first so it isn't fighting the shutdown ramp.
self._shutdown_event.set()
controller_stopped = True
if self._controller_thread is not None:
# Wait long enough for any in-flight inference tick to finish and the loop
# to observe the shutdown flag, so no stray low command is published while
# the ramp runs (the shutdown routine must be the single publisher).
self._controller_thread.join(timeout=5.0)
if self._controller_thread.is_alive():
controller_stopped = False
logger.error(
"Controller thread did not stop; skipping graceful ramp to avoid "
"concurrent low commands (fail-safe: joints keep last command until exit)"
)
# Put the robot in passive mode (zero-torque) before stopping the rest (real
# robot only; the subscribe thread is still alive here to supply the current
# pose). Only publish once the controller thread has definitely exited so the
# two aren't publishing at once.
if not self.config.is_simulation and controller_stopped:
# Put robot in passive mode before stopping threads
if not self.config.is_simulation:
self._send_zero_torque()
if self.controller is not None and hasattr(self.controller, "shutdown"):
self.controller.shutdown()
# Signal thread to stop and unblock any waits
self._shutdown_event.set()
# Wait for subscribe thread to finish
if self.subscribe_thread is not None:
@@ -703,6 +453,16 @@ class UnitreeG1(Robot):
if self.subscribe_thread.is_alive():
logger.warning("Subscribe thread did not stop cleanly")
# Wait for controller thread to finish
if self._controller_thread is not None:
self._controller_thread.join(timeout=2.0)
if self._controller_thread.is_alive():
logger.warning("Controller thread did not stop cleanly")
# Release controller resources (e.g. SONIC decoder sessions).
if self.controller is not None and hasattr(self.controller, "shutdown"):
self.controller.shutdown()
# Close simulation environment
if self.config.is_simulation and self.sim_env is not None:
try:
@@ -729,16 +489,49 @@ 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:
return {}
# Motors + IMU + wireless remote (shared lowstate -> obs mapping)
obs = lowstate_to_obs(lowstate)
obs = {}
# Motors - q, dq, tau for all joints
for motor in G1_29_JointIndex:
name = motor.name
idx = motor.value
obs[f"{name}.q"] = lowstate.motor_state[idx].q
obs[f"{name}.dq"] = lowstate.motor_state[idx].dq
obs[f"{name}.tau"] = lowstate.motor_state[idx].tau_est
# IMU - gyroscope
if lowstate.imu_state.gyroscope:
obs["imu.gyro.x"] = lowstate.imu_state.gyroscope[0]
obs["imu.gyro.y"] = lowstate.imu_state.gyroscope[1]
obs["imu.gyro.z"] = lowstate.imu_state.gyroscope[2]
# IMU - accelerometer
if lowstate.imu_state.accelerometer:
obs["imu.accel.x"] = lowstate.imu_state.accelerometer[0]
obs["imu.accel.y"] = lowstate.imu_state.accelerometer[1]
obs["imu.accel.z"] = lowstate.imu_state.accelerometer[2]
# IMU - quaternion
if lowstate.imu_state.quaternion:
obs["imu.quat.w"] = lowstate.imu_state.quaternion[0]
obs["imu.quat.x"] = lowstate.imu_state.quaternion[1]
obs["imu.quat.y"] = lowstate.imu_state.quaternion[2]
obs["imu.quat.z"] = lowstate.imu_state.quaternion[3]
# IMU - rpy
if lowstate.imu_state.rpy:
obs["imu.rpy.roll"] = lowstate.imu_state.rpy[0]
obs["imu.rpy.pitch"] = lowstate.imu_state.rpy[1]
obs["imu.rpy.yaw"] = lowstate.imu_state.rpy[2]
# Wireless remote (raw bytes for teleoperator)
if lowstate.wireless_remote:
obs["wireless_remote"] = lowstate.wireless_remote
# Token mode: echo the last commanded latent token as observation.state so a
# token-output VLA closes the loop on its own previous token.
@@ -759,11 +552,11 @@ 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:
# SONIC decoder: pull the 64-D latent token out of the action and remember it
# for the observation.state echo. The controller thread reads it back from
# controller_input (populated below) and decodes it into a 29-DoF command.
if self._sonic_token:
from .controllers.sonic_whole_body import _extract_token_from_action
@@ -771,6 +564,8 @@ class UnitreeG1(Robot):
if token is not None:
self._last_token = token
self._update_controller_action(action)
# Full-body controllers (SONIC) own the whole 29-DoF command; nothing to
# publish here (the controller thread is the sole publisher).
if getattr(self.controller, "full_body", False):
return action
# Controller thread owns legs/waist. Here we only update joystick inputs
@@ -804,10 +599,10 @@ class UnitreeG1(Robot):
def _update_controller_action(self, action: RobotAction) -> None:
"""Update controller input state from an incoming teleop action.
Controller-agnostic: every value-carrying key (e.g. locomotion ``remote.*``
axes/buttons) is forwarded verbatim into ``controller_input`` and each
controller extracts only the keys it understands. The robot deliberately does
not enumerate any controller's key schema here.
Controller-agnostic: every value-carrying key (locomotion ``remote.*`` axes or
SONIC ``motion_token.*`` values) is forwarded verbatim into ``controller_input``
and each controller extracts only the keys it understands. The robot deliberately
does not enumerate any controller's key schema here.
"""
with self._controller_action_lock:
for key, value in action.items():
@@ -820,8 +615,6 @@ 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
@@ -844,10 +637,10 @@ class UnitreeG1(Robot):
if default_positions is None:
default_positions = np.array(self.config.default_positions, dtype=np.float32)
# Full-body controllers (SONIC / OpenHLM) own the whole 29-DoF command and
# ignore ``<joint>.q`` in send_action(), so reset() must publish the default
# pose directly. Pause the background controller first so the two aren't both
# writing low commands while the robot moves to the default pose.
# Full-body controllers (SONIC) own the whole 29-DoF command and ignore
# ``<joint>.q`` in send_action(), so reset() must publish the default pose
# directly. Pause the background controller first so the two aren't both writing
# low commands while the robot moves to the default pose.
full_body = getattr(self.controller, "full_body", False)
paused = False
if full_body and self._controller_thread is not None: