mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 13:09:40 +00:00
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:
@@ -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
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,18 +145,15 @@ 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"):
|
||||
self.controller.token_mode = True
|
||||
# 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
|
||||
self._controller_thread = None
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user