mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
Compare commits
86 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 7ee4dc3bbb | |||
| d647bb9e19 | |||
| a3c90dd2c0 | |||
| ae174cb220 | |||
| 09fb2b4580 | |||
| 90699fdece | |||
| d1eed119a6 | |||
| 87b0dc470c | |||
| ec6da5c92e | |||
| d29fecfa2c | |||
| 1799eceb22 | |||
| 89ac5b4408 | |||
| a8ea5e6ea0 | |||
| 6ea5dd95e9 | |||
| f8359c19d6 | |||
| f8521958d8 | |||
| 35341dc603 | |||
| 686ac2cc21 | |||
| 48e2530455 | |||
| b6a283193c | |||
| 44edbd704f | |||
| 4cd9fc73f8 | |||
| 57d33c295e | |||
| 6e23429105 | |||
| fffa42cc5e | |||
| f33089b027 | |||
| 9f32f57b59 | |||
| 0a9b9d9a93 | |||
| 16b915ede5 | |||
| 95d9029039 | |||
| bbfc4ff443 | |||
| 0c57cd03f2 | |||
| 962ed236af | |||
| 3ae036ee40 | |||
| cdf5141688 | |||
| 9f0663e9e3 | |||
| 503f3e57ae | |||
| 6d24f20eb4 | |||
| af163fd032 | |||
| 77259f436e | |||
| 85f5c3606d | |||
| b587e81587 | |||
| 4658dada9b | |||
| 57ea6f4106 | |||
| 4209639f33 | |||
| fc7a0bc2fd | |||
| 5f6513551c | |||
| 70e157e00f | |||
| 1837be51bf | |||
| bedd56eed9 | |||
| c165e4df68 | |||
| 5e24da483a | |||
| 9c54665a76 | |||
| f6a845c30c | |||
| 45e8336854 | |||
| 5046e2df32 | |||
| 1c88e26c6d | |||
| 69a3edfa33 | |||
| 2492ce2c29 | |||
| c8e75da55f | |||
| 2eae31ea2b | |||
| c997abe739 | |||
| c73579055e | |||
| 4be438161b | |||
| 806d28a883 | |||
| 573b65ff6b | |||
| bc55713e7c | |||
| 4f53c42583 | |||
| bfced3d149 | |||
| 4969813d4e | |||
| 1c87ca31a3 | |||
| 4bcde762cc | |||
| 943ae78cfe | |||
| 3363688f1e | |||
| 0876629e72 | |||
| 305614b8c6 | |||
| 02d3202c4f | |||
| 3b6de2fdf8 | |||
| 744f3667c0 | |||
| fdde436776 | |||
| 5c683c65c6 | |||
| dfbc25c58f | |||
| 804c76bcc2 | |||
| e6afa69be9 | |||
| 31d1439e29 | |||
| 1c118c6359 |
@@ -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_
|
||||
|
||||
@@ -31,7 +31,7 @@ import cv2
|
||||
import numpy as np
|
||||
import zmq
|
||||
|
||||
from ..configs import ColorMode
|
||||
from ..configs import ColorMode, Cv2Backends
|
||||
from ..opencv import OpenCVCamera, OpenCVCameraConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -103,6 +103,9 @@ class ImageServer:
|
||||
width=shape[1],
|
||||
height=shape[0],
|
||||
color_mode=ColorMode.RGB,
|
||||
# Some UVC devices are only settable through V4L2: the auto-selected
|
||||
# backend can open them but refuses `set()`, which reads as a failure.
|
||||
backend=cfg.get("backend", Cv2Backends.ANY),
|
||||
)
|
||||
camera = OpenCVCamera(cam_config)
|
||||
camera.connect()
|
||||
|
||||
@@ -62,12 +62,21 @@ 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,
|
||||
# releases the built-in motion services, 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
|
||||
|
||||
# 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",
|
||||
]
|
||||
+2
-1
@@ -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
|
||||
+2
-1
@@ -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,201 @@
|
||||
#!/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,
|
||||
make_ort_session_options,
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
# Cap the thread pool: onboard, this decoder steps at 50 Hz in the same process as the
|
||||
# ZMQ camera server (capture + JPEG encode). Default session options let ORT grab every
|
||||
# core on the NX, which starves the camera thread (stale frames) and jitters the control
|
||||
# loop (limping gait). It is a small MLP, so 1 thread is enough and lowest-latency.
|
||||
session_options = make_ort_session_options(intra_op_num_threads=1, inter_op_num_threads=1)
|
||||
decoder = ort.InferenceSession(decoder_path, sess_options=session_options)
|
||||
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}
|
||||
@@ -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
|
||||
@@ -43,6 +80,26 @@ def get_gravity_orientation(quaternion: list[float] | np.ndarray) -> np.ndarray:
|
||||
return gravity_orientation
|
||||
|
||||
|
||||
def make_ort_session_options(
|
||||
intra_op_num_threads: int | None = None, inter_op_num_threads: int | None = None
|
||||
):
|
||||
"""Build quiet ONNX Runtime SessionOptions, optionally capping the CPU thread pool.
|
||||
|
||||
These tiny MLP policies are latency-bound, not throughput-bound, so letting ORT grab
|
||||
every core starves the real-time control loop / torch policy and causes stutter. Pass
|
||||
1 intra + 1 inter thread for lowest-latency per-step inference.
|
||||
"""
|
||||
import onnxruntime as ort
|
||||
|
||||
so = ort.SessionOptions()
|
||||
so.log_severity_level = 3
|
||||
if intra_op_num_threads is not None:
|
||||
so.intra_op_num_threads = intra_op_num_threads
|
||||
if inter_op_num_threads is not None:
|
||||
so.inter_op_num_threads = inter_op_num_threads
|
||||
return so
|
||||
|
||||
|
||||
class G1_29_JointArmIndex(IntEnum):
|
||||
# Left arm
|
||||
kLeftShoulderPitch = 15
|
||||
@@ -68,8 +125,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:
|
||||
|
||||
@@ -15,11 +15,15 @@
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
DDS-to-ZMQ bridge server for Unitree G1 robot.
|
||||
DDS-to-ZMQ server for the Unitree G1 robot. Two modes, both run ON the robot:
|
||||
|
||||
This server runs on the robot and forwards:
|
||||
- Robot state (LowState) from DDS to ZMQ (for remote clients)
|
||||
- Robot commands (LowCmd) from ZMQ to DDS (from remote clients)
|
||||
* bridge (default): forward raw robot state (LowState) DDS -> ZMQ and raw commands
|
||||
(LowCmd) ZMQ -> DDS. The controller runs on the laptop and streams lowcmd.
|
||||
|
||||
* onboard (``--onboard --controller NAME``): run the controller ONBOARD instead. Builds
|
||||
``UnitreeG1(onboard=True, controller=NAME)`` so its control loop runs locally against
|
||||
DDS at full rate; the laptop only PUSHes compact high-level actions (e.g. the 64-D
|
||||
SONIC token) on :ACTION_PORT and reads back ``observation.state`` on :STATE_PORT.
|
||||
|
||||
Uses JSON for secure serialization instead of pickle.
|
||||
"""
|
||||
@@ -28,10 +32,12 @@ 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
|
||||
@@ -39,6 +45,7 @@ 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
|
||||
from unitree_sdk2py.utils.crc import CRC
|
||||
|
||||
from lerobot.cameras.configs import Cv2Backends
|
||||
from lerobot.cameras.zmq.image_server import ImageServer
|
||||
|
||||
# DDS topic names follow Unitree SDK naming conventions
|
||||
@@ -50,6 +57,131 @@ 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
|
||||
|
||||
|
||||
def cameras_from_args(args: argparse.Namespace) -> dict:
|
||||
"""Build the ImageServer camera map from the CLI camera flags.
|
||||
|
||||
``--camera-device`` takes an index (``4``) or a V4L2 path (``/dev/video0``). Some
|
||||
devices only open by path, so the path form is not just a convenience.
|
||||
"""
|
||||
device = args.camera_device
|
||||
return {
|
||||
"head_camera": {
|
||||
"device_id": int(device) if str(device).isdigit() else device,
|
||||
"shape": [args.camera_height, args.camera_width],
|
||||
"backend": Cv2Backends[args.camera_backend.upper()],
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def start_camera_server(cameras: dict, *, fps: int, port: int) -> threading.Thread:
|
||||
"""Launch the ZMQ ImageServer in a background daemon thread (independent of DDS)."""
|
||||
server = ImageServer({"fps": fps, "cameras": cameras}, port=port)
|
||||
thread = threading.Thread(target=server.run, daemon=True)
|
||||
thread.start()
|
||||
summary = ", ".join(f"{name}(dev {c['device_id']})" for name, c in cameras.items())
|
||||
print(f"Camera server started on :{port}: {summary}")
|
||||
return thread
|
||||
|
||||
|
||||
def serve_onboard_controller(
|
||||
*,
|
||||
controller: str,
|
||||
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 controller ONBOARD -- the single control path on the robot.
|
||||
|
||||
Builds ``UnitreeG1(onboard=True, controller=...)`` so its control loop runs locally
|
||||
against DDS at full rate (the ``_controller_loop`` thread lives in UnitreeG1), then
|
||||
receives compact high-level actions from the laptop over ZMQ (:action_port), feeds
|
||||
them to the controller, and publishes ``observation.state`` (:state_port). Camera
|
||||
frames are streamed separately by the ImageServer. The controller NEVER runs on the
|
||||
laptop; the laptop thin-client only ships tokens/axes and reads back state + frames.
|
||||
"""
|
||||
# Imported lazily to avoid importing the heavy controller stack until we serve.
|
||||
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())
|
||||
|
||||
# Optional camera server (background daemon thread; independent of DDS).
|
||||
if cameras:
|
||||
start_camera_server(cameras, fps=camera_fps, port=camera_port)
|
||||
|
||||
cfg = UnitreeG1Config(is_simulation=False, onboard=True, controller=controller, cameras={})
|
||||
robot = UnitreeG1(cfg)
|
||||
print(f"Connecting onboard robot (controller={controller})...")
|
||||
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} ...")
|
||||
|
||||
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 (joint .q, IMU, SONIC token echo).
|
||||
# Camera arrays are streamed separately by the ImageServer.
|
||||
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()
|
||||
|
||||
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)
|
||||
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."""
|
||||
@@ -152,32 +284,58 @@ def cmd_forward_loop(
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Main entry point for the robot server bridge."""
|
||||
parser = argparse.ArgumentParser(description="DDS-to-ZMQ bridge server for Unitree G1")
|
||||
"""Main entry point for the robot server."""
|
||||
parser = argparse.ArgumentParser(description="DDS-to-ZMQ server for Unitree G1")
|
||||
parser.add_argument("--camera", action="store_true", help="Also launch camera server")
|
||||
parser.add_argument("--camera-device", type=int, default=4, help="Camera device ID (default: 4)")
|
||||
parser.add_argument(
|
||||
"--camera-device",
|
||||
default="4",
|
||||
help="Camera index or V4L2 path, e.g. 4 or /dev/video0 (default: 4)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--camera-backend",
|
||||
default="any",
|
||||
choices=[b.name.lower() for b in Cv2Backends],
|
||||
help="OpenCV capture backend; use v4l2 for UVC cams that reject set() under any (default: any)",
|
||||
)
|
||||
parser.add_argument("--camera-fps", type=int, default=30, help="Camera FPS (default: 30)")
|
||||
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)")
|
||||
# Onboard mode: run the controller on the robot instead of bridging raw lowcmd.
|
||||
parser.add_argument(
|
||||
"--onboard",
|
||||
action="store_true",
|
||||
help="Run the controller ONBOARD (requires --controller) instead of the raw DDS bridge",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--controller",
|
||||
default=None,
|
||||
metavar="NAME",
|
||||
help="[--onboard] controller to run onboard, e.g. SonicWholeBodyController",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# --- Onboard mode: controller runs on the robot; laptop ships high-level actions. ---
|
||||
if args.onboard:
|
||||
if not args.controller:
|
||||
parser.error("--onboard requires --controller (e.g. --controller SonicWholeBodyController)")
|
||||
cameras = cameras_from_args(args) if args.camera else None
|
||||
serve_onboard_controller(
|
||||
controller=args.controller,
|
||||
cameras=cameras,
|
||||
camera_fps=args.camera_fps,
|
||||
camera_port=args.camera_port,
|
||||
)
|
||||
return
|
||||
|
||||
# --- Bridge mode (default): forward raw lowstate/lowcmd; controller runs on laptop. ---
|
||||
# Optionally start camera server in background thread
|
||||
camera_thread = None
|
||||
if args.camera:
|
||||
camera_config = {
|
||||
"fps": args.camera_fps,
|
||||
"cameras": {
|
||||
"head_camera": {
|
||||
"device_id": args.camera_device,
|
||||
"shape": [args.camera_height, args.camera_width],
|
||||
}
|
||||
},
|
||||
}
|
||||
camera_server = ImageServer(camera_config, port=args.camera_port)
|
||||
camera_thread = threading.Thread(target=camera_server.run, daemon=True)
|
||||
camera_thread.start()
|
||||
print(f"Camera server started on port {args.camera_port} (device {args.camera_device})")
|
||||
camera_thread = start_camera_server(
|
||||
cameras_from_args(args), fps=args.camera_fps, port=args.camera_port
|
||||
)
|
||||
|
||||
# initialize DDS
|
||||
ChannelFactoryInitialize(0)
|
||||
|
||||
@@ -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,
|
||||
@@ -119,24 +121,32 @@ 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 PUSHes high-level
|
||||
# actions to ``run_g1_server --onboard`` and reads back state +
|
||||
# cameras over ZMQ. The controller *always* runs on the robot.
|
||||
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 (populated in connect()).
|
||||
self._client_action_sock = None
|
||||
self._client_state_sock = None
|
||||
self._client_state_latest: dict[str, float] = {}
|
||||
|
||||
# Initialize state variables
|
||||
self.sim_env = None
|
||||
@@ -148,15 +158,29 @@ 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 = {}
|
||||
|
||||
@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 +257,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}
|
||||
@@ -293,7 +342,124 @@ class UnitreeG1(Robot):
|
||||
def configure(self) -> None:
|
||||
pass
|
||||
|
||||
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, STATE_PORT
|
||||
|
||||
server_ip = self.config.robot_ip
|
||||
if not server_ip:
|
||||
raise ValueError("client mode requires config.robot_ip (the G1 running run_g1_server)")
|
||||
|
||||
ctx = zmq.Context.instance()
|
||||
|
||||
# Action PUSH: ship compact high-level actions to the onboard controller. The
|
||||
# server runs the controller selected by its own ``--controller`` flag; both
|
||||
# sides use ``config.controller`` to agree on the action schema (no handshake).
|
||||
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}")
|
||||
|
||||
# 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}")
|
||||
|
||||
# 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 to %s: actions ->:%d, state <-:%d, %d camera(s).",
|
||||
server_ip,
|
||||
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 onboard controller 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. Open the high-level ZMQ links to
|
||||
# run_g1_server --onboard (which runs the controller): 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 +468,12 @@ 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, then take low-level control from the built-in
|
||||
# motion services before we start writing lowcmd.
|
||||
self._ChannelFactoryInitialize(0)
|
||||
self._release_motion_control()
|
||||
|
||||
# Initialize direct motor control interface
|
||||
self.lowcmd_publisher = self._ChannelPublisher(kTopicLowCommand_Debug, hg_LowCmd)
|
||||
@@ -341,8 +511,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 +526,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 +553,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 +602,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 +649,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 +664,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 +699,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 +712,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 +764,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
|
||||
|
||||
Reference in New Issue
Block a user