feat(unitree_g1): drive SONIC whole-body from a 34-D OpenHLM/pi0.5 VLA

Add a dense 34-D whole-body command path so lerobot-rollout can drive the
G1 directly with an OpenHLM / pi0.5 policy through the SONIC encoder/decoder:

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

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Martino Russi
2026-07-20 19:58:47 +02:00
parent 5f6513551c
commit fc7a0bc2fd
7 changed files with 423 additions and 13 deletions
@@ -65,6 +65,35 @@ class UnitreeG1Config(RobotConfig):
# Cameras (ZMQ-based remote cameras) # Cameras (ZMQ-based remote cameras)
cameras: dict[str, CameraConfig] = field(default_factory=dict) cameras: dict[str, CameraConfig] = field(default_factory=dict)
# Synthetic zero-image cameras exposed as ``observation.images.{name}`` (H×W×3
# black frames). Lets image-conditioned policies (e.g. pi0.5 / OpenHLM) run in
# sim before real cameras are wired. Empty = disabled.
empty_cameras: list[str] = field(default_factory=list)
empty_camera_hw: tuple[int, int] = (224, 224)
# Publish Dex3 hand commands (``rt/dex3/{left,right}/cmd``) driven by the OpenHLM
# gripper scalars (``wb.7.pos`` left, ``wb.15.pos`` right). Lets the 43-DoF sim
# (or a real Dex3-equipped G1) show grasping. The scalar in [0, 1] is remapped to
# a curl amount (``hand_open_grip_value`` -> open) and scaled onto
# ``hand_closed_pose`` (7 joints: thumb_0/1/2, middle_0/1, index_0/1). Flip signs
# in ``hand_closed_pose`` if fingers curl the wrong way.
publish_hands: bool = False
hand_open_grip_value: float = 1.0
hand_closed_grip_value: float = 0.0
hand_closed_pose: list[float] = field(
default_factory=lambda: [1.0, 0.9, 0.9, 1.3, 1.3, 1.3, 1.3]
)
hand_kp: float = 1.5
hand_kd: float = 0.1
# Replay recorded camera frames from a LeRobot parquet episode as the camera
# feed (e.g. OpenHLM-data episode). Maps a robot camera name to a parquet image
# column; frames advance one per observation and loop. Lets a VLA see the real
# task video in sim without live cameras. Empty map = disabled.
replay_camera_parquet: str | None = None
replay_camera_map: dict[str, str] = field(default_factory=dict)
replay_camera_loop: bool = True
# Compensates for gravity on the unitree's arms using the arm ik solver # Compensates for gravity on the unitree's arms using the arm ik solver
gravity_compensation: bool = False gravity_compensation: bool = False
@@ -20,6 +20,7 @@ from __future__ import annotations
import logging import logging
import math import math
from collections import deque
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
import numpy as np import numpy as np
@@ -38,10 +39,12 @@ from lerobot.teleoperators.pico_headset.smpl_constants import (
VR3_ORN_PREFIX, VR3_ORN_PREFIX,
VR3_POS_DIM, VR3_POS_DIM,
VR3_POS_PREFIX, VR3_POS_PREFIX,
WB_ACTION_DIM,
wb_action_key,
) )
from lerobot.utils.import_utils import _onnxruntime_available, require_package from lerobot.utils.import_utils import _onnxruntime_available, require_package
from ..g1_utils import KEYBOARD_KEYS_FIELD, G1_29_JointIndex, lowstate_to_obs from ..g1_utils import MUJOCO_TO_ISAACLAB, KEYBOARD_KEYS_FIELD, G1_29_JointIndex, lowstate_to_obs
from .sonic_pipeline import ( from .sonic_pipeline import (
CONTROL_DT, CONTROL_DT,
DEBUG_PRINT_EVERY, DEBUG_PRINT_EVERY,
@@ -132,6 +135,51 @@ def _extract_vr3_from_action(action: dict | None) -> tuple[np.ndarray, np.ndarra
return pos, orn return pos, orn
def _extract_wb34_from_action(action: dict | None) -> np.ndarray | None:
"""Reassemble a dense (34,) whole-body command from ``wb.{i}.pos`` keys, or None.
This is the OpenHLM / pi0.5 joint-based interface: one 34-D vector per tick
(sentinel: presence of ``wb.0.pos``) carrying absolute joint targets in real
units. The ``.pos`` suffix lets these flow through ``lerobot-rollout`` as normal
joint-position action features.
"""
if not action or wb_action_key(0) not in action:
return None
return np.fromiter(
(float(action.get(wb_action_key(i), 0.0)) for i in range(WB_ACTION_DIM)),
dtype=np.float32,
count=WB_ACTION_DIM,
)
def _wb34_to_reference(wb: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
"""Map a 34-D OpenHLM whole-body command to a SONIC mode-0 reference.
Returns ``(ref29, anchor_quat)`` where ``ref29`` is the 29 joint targets in
IsaacLab order (what SONIC's ``motion_joint_positions`` expects) and
``anchor_quat`` (wxyz) encodes the root roll/pitch (yaw=0).
OpenHLM layout : [L-arm 0:7, L-grip 7, R-arm 8:15, R-grip 15,
L-leg 16:22, R-leg 22:28, waist 28:31, root rp+yaw 31:34]
The 29 joints are first assembled in MuJoCo / Unitree-SDK order
([L-leg 0:6, R-leg 6:12, waist 12:15, L-arm 15:22, R-arm 22:29] — the
``G1_29_JointIndex`` grouping OpenHLM uses), then permuted to IsaacLab order via
``MUJOCO_TO_ISAACLAB``. Grippers (7, 15) and yaw-rate (33) are not part of the
29-DoF SONIC reference.
"""
ref_mj = np.zeros(29, np.float32) # MuJoCo / Unitree-SDK grouped order
ref_mj[0:6] = wb[16:22] # left leg
ref_mj[6:12] = wb[22:28] # right leg
ref_mj[12:15] = wb[28:31] # waist
ref_mj[15:22] = wb[0:7] # left arm
ref_mj[22:29] = wb[8:15] # right arm
ref = ref_mj[MUJOCO_TO_ISAACLAB].astype(np.float32) # -> IsaacLab order for SONIC
roll, pitch = float(wb[31]), float(wb[32])
cr, sr, cp, sp = np.cos(roll / 2), np.sin(roll / 2), np.cos(pitch / 2), np.sin(pitch / 2)
anchor = np.array([cr * cp, sr * cp, cr * sp, sr * sp], np.float32) # Rx(roll)·Ry(pitch)
return ref, anchor
def _extract_loco_from_action(action: dict | None) -> tuple[np.ndarray, np.ndarray] | None: def _extract_loco_from_action(action: dict | None) -> tuple[np.ndarray, np.ndarray] | None:
"""Reassemble controller-stick locomotion from ``loco_axes.{i}`` / ``loco_btn.{i}``. """Reassemble controller-stick locomotion from ``loco_axes.{i}`` / ``loco_btn.{i}``.
@@ -234,6 +282,10 @@ class SonicWholeBodyController:
control_dt = CONTROL_DT control_dt = CONTROL_DT
full_body = True full_body = True
# Advertise a dense 34-D whole-body action space (OpenHLM / pi0.5) so the robot
# exposes ``wb.{i}.pos`` action features and ``lerobot-rollout`` can drive it
# directly with a 34-D VLA policy.
wb_action = True
def __init__( def __init__(
self, self,
@@ -277,6 +329,13 @@ class SonicWholeBodyController:
self._init_step = 0 self._init_step = 0
self._start_pose: dict[str, float] = {} self._start_pose: dict[str, float] = {}
# Tick counter for the dense whole-body (OpenHLM, mode-0) path's encoder cadence.
self._wb_step = 0
# Rolling 50-frame reference trajectory (ref29 + anchor quat) built from the
# stream of per-tick whole-body commands, fed to the encoder as a batch.
self._wb_traj: deque[np.ndarray] = deque(maxlen=50)
self._wb_quat_traj: deque[np.ndarray] = deque(maxlen=50)
# Optional: subscribe directly to the rt/smpl headset stream so full-body # Optional: subscribe directly to the rt/smpl headset stream so full-body
# teleop works with ANY teleoperator (e.g. --teleop.type=unitree_g1 for the # teleop works with ANY teleoperator (e.g. --teleop.type=unitree_g1 for the
# estop/joystick) before the dedicated pico_headset teleop exists. # estop/joystick) before the dedicated pico_headset teleop exists.
@@ -445,6 +504,66 @@ class SonicWholeBodyController:
logger.info("SONIC 3-point: locomotion mode -> %s", LM(self.ms.mode).name) logger.info("SONIC 3-point: locomotion mode -> %s", LM(self.ms.mode).name)
self._prev_loco_mode_pair = (ab_now, xy_now) self._prev_loco_mode_pair = (ab_now, xy_now)
def _run_wholebody34(self, obs: dict, wb: np.ndarray) -> dict:
"""Feed a dense 34-D OpenHLM whole-body command as the mode-0 encoder reference.
The 29 joint targets are held across the encoder lookahead window (zero
velocity) and the root roll/pitch set the anchor orientation, then the
encoder/decoder run directly (planner bypassed). One command per tick, so the
VLA's commanded pose is what SONIC tracks.
"""
ref, anchor = _wb34_to_reference(wb)
c = self.controller
if c.encode_mode != 0:
c.encode_mode = 0
c.reinit_heading = True
# Capture the heading/anchor reference on the first whole-body tick. The
# controller only latches ``init_ref_quat`` (and the base heading) inside
# ``step()`` when ``first_motion or reinit_heading`` — but it already boots in
# mode 0, so the mode-switch guard above misses the very first command and the
# anchor would stay identity. This mirrors the GEAR reference, which seeds
# ``init_ref_quat`` from the first anchor. Must run before the buffers below so
# ``step()`` latches ``motion_body_quats[0]`` = this tick's anchor.
if self._wb_step == 0:
c.reinit_heading = True
# Accumulate the per-tick commands into a rolling 50-frame reference
# trajectory so the encoder's 10-frame, step-5 lookahead sees an actual
# motion sequence (with velocities) instead of one repeated pose. 50 frames
# == chunk horizon == 10 lookahead frames × step 5.
self._wb_traj.append(ref)
self._wb_quat_traj.append(anchor)
traj = np.asarray(self._wb_traj, np.float32) # (L, 29), oldest -> newest
quats = np.asarray(self._wb_quat_traj, np.float32) # (L, 4)
n = len(traj)
# Per-frame velocities from finite differences (rad/s at the control rate).
vel = np.zeros_like(traj)
if n > 1:
vel[1:] = (traj[1:] - traj[:-1]) / CONTROL_DT
vel[0] = vel[1]
with c.motion_lock:
c.motion_joint_positions[:n] = traj
c.motion_joint_velocities[:n] = vel
c.motion_body_quats[:n] = quats
c.motion_body_pos[:n] = 0.0
c.motion_timesteps = n
c.ref_cursor = 0
c.playing = True
do_enc = self._wb_step % ENCODER_UPDATE_EVERY == 0
out = c.step(obs, update_encoder=do_enc, debug=False)
if self._wb_step % 25 == 0:
tgt = np.array([out[f"{m.name}.q"] for m in G1_29_JointIndex], np.float32)
logger.info(
"[WB34] step=%d |ref|mean=%.3f |target|mean=%.3f target_std=%.3f init_ref_quat=%s",
self._wb_step,
float(np.abs(ref).mean()),
float(np.abs(tgt).mean()),
float(tgt.std()),
np.round(c.init_ref_quat, 3).tolist(),
)
self._wb_step += 1
return out
def _smooth_root_quat(self, root_quat: np.ndarray | None) -> np.ndarray | None: def _smooth_root_quat(self, root_quat: np.ndarray | None) -> np.ndarray | None:
"""Spherically smooth the per-frame SMPL root quaternion (mode-2 anchor). """Spherically smooth the per-frame SMPL root quaternion (mode-2 anchor).
@@ -513,6 +632,20 @@ class SonicWholeBodyController:
# direct rt/smpl subscription when enabled (enable_smpl_stream). A stale # direct rt/smpl subscription when enabled (enable_smpl_stream). A stale
# stream (headset silent past its timeout) is treated as "no SMPL" so the # stream (headset silent past its timeout) is treated as "no SMPL" so the
# robot doesn't stay frozen tracking the last pose. # robot doesn't stay frozen tracking the last pose.
# Dense whole-body command (OpenHLM / pi0.5 joint interface) takes priority:
# a single 34-D vector drives the mode-0 joint reference directly.
wb = _extract_wb34_from_action(action)
if wb is not None:
return self._startup_blend(obs, self._run_wholebody34(obs, wb))
self._wb_miss = getattr(self, "_wb_miss", 0) + 1
if self._wb_miss % 50 == 1:
akeys = [k for k in action if isinstance(k, str)]
logger.info(
"[WB34] no wb.*.pos in action this tick (miss=%d). action keys sample: %s",
self._wb_miss,
akeys[:8],
)
smpl = _extract_smpl_from_action(action) smpl = _extract_smpl_from_action(action)
root_quat = _extract_root_from_action(action) root_quat = _extract_root_from_action(action)
vr3 = _extract_vr3_from_action(action) vr3 = _extract_vr3_from_action(action)
@@ -575,6 +708,9 @@ class SonicWholeBodyController:
self._init_step = 0 # re-run the startup blend after a reset self._init_step = 0 # re-run the startup blend after a reset
self._start_pose = {} self._start_pose = {}
self._smoothed_root_quat = None self._smoothed_root_quat = None
self._wb_step = 0
self._wb_traj.clear()
self._wb_quat_traj.clear()
def shutdown(self): def shutdown(self):
if self._smpl_stream is not None: if self._smpl_stream is not None:
+52
View File
@@ -104,6 +104,21 @@ REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
# whole-body controller (see SonicWholeBodyController._process_keyboard). # whole-body controller (see SonicWholeBodyController._process_keyboard).
KEYBOARD_KEYS_FIELD = "keyboard.keys" KEYBOARD_KEYS_FIELD = "keyboard.keys"
# ── Dense whole-body joint reference (SONIC encode_mode 0, OpenHLM / pi0.5) ──────
# A single 34-D whole-body command per tick, in the OpenHLM action layout:
# [L-arm(7), L-grip(1), R-arm(7), R-grip(1), L-leg(6), R-leg(6), waist(3),
# root roll/pitch + yaw-rate(3)]
# Fed as flat scalars ``wb.0.pos .. wb.33.pos``. The ``.pos`` suffix makes these
# behave like ordinary joint-position action features so ``lerobot-rollout`` routes
# them straight from a 34-D VLA (OpenHLM / pi0.5) onto the robot.
WB_ACTION_PREFIX = "wb."
WB_ACTION_DIM = 34
def wb_action_key(i: int) -> str:
"""Action-dict key for the ``i``-th whole-body command scalar (``wb.{i}.pos``)."""
return f"{WB_ACTION_PREFIX}{i}.pos"
def default_remote_input() -> dict[str, float]: def default_remote_input() -> dict[str, float]:
"""Return a zeroed-out remote input dict (axes + buttons).""" """Return a zeroed-out remote input dict (axes + buttons)."""
@@ -181,6 +196,43 @@ def lowstate_to_obs(lowstate) -> dict:
return obs return obs
def obs_to_wb34_state(obs: dict) -> np.ndarray:
"""Build the 34-D OpenHLM / pi0.5 proprio state from a G1 observation dict.
Mirrors the whole-body *action* layout so the policy sees state and action in
the same coordinates::
[L-arm(7), L-grip(1), R-arm(7), R-grip(1),
L-leg(6), R-leg(6), waist(3), root roll/pitch + yaw-rate(3)]
Joint positions come from the ``<joint>.q`` obs keys, which are already in
MuJoCo / Unitree-SDK order — the same body-part grouping OpenHLM uses
([L-leg 0:6, R-leg 6:12, waist 12:15, L-arm 15:22, R-arm 22:29]) — so they are
regrouped directly (no IsaacLab permutation). The G1 has no grippers in its
29-DoF body, so both gripper slots are 0. Root roll/pitch are the IMU RPY and
the last slot is the IMU yaw rate (gyro z).
"""
q_mj = np.array(
[float(obs.get(f"{m.name}.q", 0.0)) for m in G1_29_JointIndex],
dtype=np.float32,
)
lleg, rleg, waist = q_mj[0:6], q_mj[6:12], q_mj[12:15]
larm, rarm = q_mj[15:22], q_mj[22:29]
state = np.zeros(34, dtype=np.float32)
state[0:7] = larm
# state[7] left gripper — none on 29-DoF G1
state[8:15] = rarm
# state[15] right gripper — none on 29-DoF G1
state[16:22] = lleg
state[22:28] = rleg
state[28:31] = waist
state[31] = float(obs.get("imu.rpy.roll", 0.0))
state[32] = float(obs.get("imu.rpy.pitch", 0.0))
state[33] = float(obs.get("imu.gyro.z", 0.0))
return state
def make_locomotion_controller(name: str | None): def make_locomotion_controller(name: str | None):
"""Instantiate a locomotion controller by class name. Returns None if name is None.""" """Instantiate a locomotion controller by class name. Returns None if name is None."""
if name is None: if name is None:
+159 -6
View File
@@ -40,6 +40,7 @@ from .g1_utils import (
default_remote_input, default_remote_input,
lowstate_to_obs, lowstate_to_obs,
make_locomotion_controller, make_locomotion_controller,
obs_to_wb34_state,
) )
if TYPE_CHECKING or _unitree_sdk_available: if TYPE_CHECKING or _unitree_sdk_available:
@@ -48,8 +49,12 @@ if TYPE_CHECKING or _unitree_sdk_available:
ChannelPublisher as _SDKChannelPublisher, ChannelPublisher as _SDKChannelPublisher,
ChannelSubscriber as _SDKChannelSubscriber, ChannelSubscriber as _SDKChannelSubscriber,
) )
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_ from unitree_sdk2py.idl.default import (
unitree_hg_msg_dds__HandCmd_ as hg_HandCmd_default,
unitree_hg_msg_dds__LowCmd_,
)
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import ( from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
HandCmd_ as hg_HandCmd,
LowCmd_ as hg_LowCmd, LowCmd_ as hg_LowCmd,
LowState_ as hg_LowState, LowState_ as hg_LowState,
) )
@@ -59,6 +64,8 @@ else:
_SDKChannelPublisher = None _SDKChannelPublisher = None
_SDKChannelSubscriber = None _SDKChannelSubscriber = None
unitree_hg_msg_dds__LowCmd_ = None unitree_hg_msg_dds__LowCmd_ = None
hg_HandCmd_default = None
hg_HandCmd = None
hg_LowCmd = None hg_LowCmd = None
hg_LowState = None hg_LowState = None
CRC = None CRC = None
@@ -158,6 +165,37 @@ class UnitreeG1(Robot):
self.controller_input = default_remote_input() self.controller_input = default_remote_input()
self.controller_output = {} self.controller_output = {}
# Replay-camera state (decoded frames per robot camera name + play cursor).
self._replay_frames: dict[str, list[np.ndarray]] = {}
self._replay_len = 0
self._replay_idx = 0
if config.replay_camera_parquet and config.replay_camera_map:
self._load_replay_frames()
def _load_replay_frames(self) -> None:
"""Decode recorded episode frames from a parquet into per-camera image lists."""
import io
import pyarrow.parquet as pq
from PIL import Image
table = pq.read_table(self.config.replay_camera_parquet)
cols = {col: table.column(col).to_pylist() for col in self.config.replay_camera_map.values()}
self._replay_len = table.num_rows
def decode(cell) -> np.ndarray:
data = cell["bytes"] if isinstance(cell, dict) else cell
return np.asarray(Image.open(io.BytesIO(data)).convert("RGB"), dtype=np.uint8)
for cam_name, column in self.config.replay_camera_map.items():
self._replay_frames[cam_name] = [decode(c) for c in cols[column]]
logger.info(
"Loaded %d replay frames for cameras %s from %s",
self._replay_len,
list(self.config.replay_camera_map),
self.config.replay_camera_parquet,
)
def _subscribe_lowstate(self): # polls robot state @ 250Hz def _subscribe_lowstate(self): # polls robot state @ 250Hz
while not self._shutdown_event.is_set(): while not self._shutdown_event.is_set():
start_time = time.time() start_time = time.time()
@@ -232,15 +270,54 @@ class UnitreeG1(Robot):
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
return features return features
@property
def _wb_state_ft(self) -> dict[str, type]:
"""34-D whole-body proprio state (``wb_state.{i}.pos``) for dense controllers.
Exposed only when the controller consumes a dense whole-body command
(OpenHLM / pi0.5). These ``.pos`` scalars are aggregated by the rollout
pipeline into a single 34-D ``observation.state`` for the policy.
"""
if not getattr(self.controller, "wb_action", False):
return {}
from .g1_utils import WB_ACTION_DIM
return {f"wb_state.{i}.pos": float for i in range(WB_ACTION_DIM)}
@property
def _empty_cameras_ft(self) -> dict[str, tuple]:
"""Synthetic zero-image cameras (see ``UnitreeG1Config.empty_cameras``)."""
h, w = self.config.empty_camera_hw
return {name: (h, w, 3) for name in self.config.empty_cameras}
@property
def _replay_cameras_ft(self) -> dict[str, tuple]:
"""Replay cameras, shaped from their first decoded frame."""
return {name: frames[0].shape for name, frames in self._replay_frames.items() if frames}
@cached_property @cached_property
def observation_features(self) -> dict[str, type | tuple]: def observation_features(self) -> dict[str, type | tuple]:
return {**self._motors_ft, **self._cameras_ft} return {
**self._motors_ft,
**self._wb_state_ft,
**self._empty_cameras_ft,
**self._replay_cameras_ft,
**self._cameras_ft,
}
@cached_property @cached_property
def action_features(self) -> dict[str, type]: def action_features(self) -> dict[str, type]:
if self.controller is None: if self.controller is None:
return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex} return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
# Dense whole-body controllers (SONIC / OpenHLM, pi0.5) consume a single
# 34-D command per tick. Expose it as ``wb.{i}.pos`` joint-position features
# so ``lerobot-rollout`` maps a 34-D policy output straight onto the robot.
if getattr(self.controller, "wb_action", False):
from .g1_utils import WB_ACTION_DIM, wb_action_key
return {wb_action_key(i): float for i in range(WB_ACTION_DIM)}
arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex} arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
remote_features = dict.fromkeys(REMOTE_AXES, float) remote_features = dict.fromkeys(REMOTE_AXES, float)
return {**arm_features, **remote_features} return {**arm_features, **remote_features}
@@ -312,6 +389,17 @@ class UnitreeG1(Robot):
self.lowstate_subscriber = self._ChannelSubscriber(kTopicLowState, hg_LowState) self.lowstate_subscriber = self._ChannelSubscriber(kTopicLowState, hg_LowState)
self.lowstate_subscriber.Init() self.lowstate_subscriber.Init()
# Dex3 hand command publishers (grasping). Driven by the OpenHLM grip scalars.
self._hand_publishers = {}
if self.config.publish_hands:
self._left_hand_cmd = hg_HandCmd_default()
self._right_hand_cmd = hg_HandCmd_default()
self._hand_publishers["left"] = self._ChannelPublisher("rt/dex3/left/cmd", hg_HandCmd)
self._hand_publishers["right"] = self._ChannelPublisher("rt/dex3/right/cmd", hg_HandCmd)
for pub in self._hand_publishers.values():
pub.Init()
logger.info("Dex3 hand command publishers initialized (rt/dex3/{left,right}/cmd)")
# Start subscribe thread to read robot state # Start subscribe thread to read robot state
self.subscribe_thread = threading.Thread(target=self._subscribe_lowstate) self.subscribe_thread = threading.Thread(target=self._subscribe_lowstate)
self.subscribe_thread.start() self.subscribe_thread.start()
@@ -460,6 +548,31 @@ class UnitreeG1(Robot):
# Motors + IMU + wireless remote (shared lowstate -> obs mapping) # Motors + IMU + wireless remote (shared lowstate -> obs mapping)
obs = lowstate_to_obs(lowstate) obs = lowstate_to_obs(lowstate)
# Dense whole-body controllers (OpenHLM / pi0.5): expose the 34-D proprio
# state as ``wb_state.{i}.pos`` so the rollout aggregates it into
# ``observation.state`` for the policy.
if getattr(self.controller, "wb_action", False):
wb_state = obs_to_wb34_state(obs)
for i, v in enumerate(wb_state):
obs[f"wb_state.{i}.pos"] = float(v)
# Synthetic empty cameras: black frames so image-conditioned policies run
# before real cameras are wired.
if self.config.empty_cameras:
h, w = self.config.empty_camera_hw
black = np.zeros((h, w, 3), dtype=np.uint8)
for name in self.config.empty_cameras:
obs[name] = black
# Replay cameras: serve the current recorded frame per camera, then advance.
if self._replay_len:
idx = self._replay_idx
if idx >= self._replay_len:
idx = self._replay_len - 1 if not self.config.replay_camera_loop else idx % self._replay_len
for name, frames in self._replay_frames.items():
obs[name] = frames[idx]
self._replay_idx += 1
# Cameras - read images from ZMQ cameras # Cameras - read images from ZMQ cameras
for cam_name, cam in self._cameras.items(): for cam_name, cam in self._cameras.items():
if getattr(cam, "use_rgb", True): if getattr(cam, "use_rgb", True):
@@ -473,6 +586,8 @@ class UnitreeG1(Robot):
action_to_publish = action action_to_publish = action
if self.controller is not None: if self.controller is not None:
self._update_controller_action(action) self._update_controller_action(action)
if self.config.publish_hands and getattr(self.controller, "wb_action", False):
self._publish_hand_cmds(action)
if getattr(self.controller, "full_body", False): if getattr(self.controller, "full_body", False):
return action return action
# Controller thread owns legs/waist. Here we only update joystick inputs # Controller thread owns legs/waist. Here we only update joystick inputs
@@ -507,10 +622,10 @@ class UnitreeG1(Robot):
"""Update controller input state from an incoming teleop action. """Update controller input state from an incoming teleop action.
Controller-agnostic: every value-carrying key is forwarded verbatim into Controller-agnostic: every value-carrying key is forwarded verbatim into
``controller_input`` (joystick ``remote.*``, whole-body ``smpl.*``/``root.*`` ``controller_input`` (whole-body ``wb.{i}.pos`` from a 34-D VLA, or whatever a
from pico_headset, or whatever a future controller expects), and each future controller expects), and each controller extracts only the keys it
controller extracts only the keys it understands. The robot deliberately does understands. The robot deliberately does not enumerate any controller's key
not enumerate any controller's key schema here. schema here.
KeyboardTeleop is the one special case: it emits the currently-pressed keys as KeyboardTeleop is the one special case: it emits the currently-pressed keys as
bare action keys with a ``None`` value (``dict.fromkeys(pressed, None)``), so bare action keys with a ``None`` value (``dict.fromkeys(pressed, None)``), so
@@ -528,6 +643,44 @@ class UnitreeG1(Robot):
if isinstance(key, str) and value is not None: if isinstance(key, str) and value is not None:
self.controller_input[key] = value self.controller_input[key] = value
def _publish_hand_cmds(self, action: RobotAction) -> None:
"""Drive the Dex3 hands from the OpenHLM grip scalars in a 34-D wb action.
``wb.7.pos`` is the left grip and ``wb.15.pos`` the right grip. Each scalar in
[0, 1] (``hand_open_grip_value`` == fully open) is turned into a curl amount and
scaled onto ``hand_closed_pose`` (7 joints), then published as a PD target on
``rt/dex3/{left,right}/cmd`` so the fingers close when the policy grips.
"""
if not self._hand_publishers:
return
from .g1_utils import wb_action_key
open_val = float(self.config.hand_open_grip_value)
closed_val = float(self.config.hand_closed_grip_value)
closed_pose = self.config.hand_closed_pose
kp, kd = float(self.config.hand_kp), float(self.config.hand_kd)
span = (closed_val - open_val) or 1.0
def curl_amount(grip: float) -> float:
# Fraction of the way from the open scalar to the closed scalar, in [0, 1].
return float(min(max((grip - open_val) / span, 0.0), 1.0))
for side, grip_idx, cmd in (
("left", 7, self._left_hand_cmd),
("right", 15, self._right_hand_cmd),
):
grip = action.get(wb_action_key(grip_idx))
if grip is None:
continue
amount = curl_amount(float(grip))
for i, closed_q in enumerate(closed_pose):
cmd.motor_cmd[i].q = float(closed_q) * amount
cmd.motor_cmd[i].dq = 0.0
cmd.motor_cmd[i].kp = kp
cmd.motor_cmd[i].kd = kd
cmd.motor_cmd[i].tau = 0.0
self._hand_publishers[side].Write(cmd)
@property @property
def is_calibrated(self) -> bool: def is_calibrated(self) -> bool:
return True return True
@@ -195,7 +195,13 @@ def main() -> None:
# 3-point operator calibration (device source only): map the operator's neutral # 3-point operator calibration (device source only): map the operator's neutral
# rest pose onto the G1's neutral stance. Trigger a (re)capture with the A+B+X+Y # rest pose onto the G1's neutral stance. Trigger a (re)capture with the A+B+X+Y
# controller combo, mirroring gear_sonic's ThreePointPose.calibrate_now. # controller combo, mirroring gear_sonic's ThreePointPose.calibrate_now.
calibrator = ThreePointCalibrator() if (xrt is not None and args.headset_source == "devices") else None # Device source: the "neck" is the headset (pitches when looking down), so keep the
# wrist targets in the yaw-local world frame rather than de-rotating by head tilt.
calibrator = (
ThreePointCalibrator(neck_relative_wrists=False)
if (xrt is not None and args.headset_source == "devices")
else None
)
calib_combo_last = False calib_combo_last = False
if calibrator is not None: if calibrator is not None:
print( print(
@@ -59,3 +59,22 @@ LOCO_N_AXES = 4 # [left_x, left_y, right_x, right_y]
LOCO_N_BTN = 4 # [A, B, X, Y] LOCO_N_BTN = 4 # [A, B, X, Y]
LOCO_AXES_PREFIX = "loco_axes." LOCO_AXES_PREFIX = "loco_axes."
LOCO_BTN_PREFIX = "loco_btn." LOCO_BTN_PREFIX = "loco_btn."
# ── Dense whole-body joint reference (SONIC encode_mode 0, OpenHLM-style) ─────
# A single 34-D whole-body command per tick, in the pi0.5 / OpenHLM action layout:
# [L-arm(7), L-grip(1), R-arm(7), R-grip(1), L-leg(6), R-leg(6), waist(3),
# root roll/pitch + yaw-rate(3)]
# The 29 joint targets (arms/legs/waist, grippers excluded) become the mode-0
# encoder joint reference and root roll/pitch become the anchor orientation.
#
# Fed as flat scalars ``wb.0.pos .. wb.33.pos``. The ``.pos`` suffix is required so
# these behave like ordinary joint-position action features: ``lerobot-rollout``
# only routes ``*.pos`` keys from ``robot.action_features`` into the policy<->robot
# action mapping, letting a 34-D VLA (OpenHLM / pi0.5) drive the robot directly.
WB_ACTION_PREFIX = "wb."
WB_ACTION_DIM = 34
def wb_action_key(i: int) -> str:
"""Action-dict key for the ``i``-th whole-body command scalar (``wb.{i}.pos``)."""
return f"{WB_ACTION_PREFIX}{i}.pos"
@@ -394,8 +394,20 @@ class ThreePointCalibrator:
neck orientation via the torso->neck kinematic chain. neck orientation via the torso->neck kinematic chain.
All quaternions are scalar-first (w, x, y, z), matching :func:`compute_3point`. All quaternions are scalar-first (w, x, y, z), matching :func:`compute_3point`.
``neck_relative_wrists`` selects how the wrist targets are framed:
- ``True`` (body-source, :func:`compute_3point`): wrists are pelvis-relative and
the neck sits roughly upright over the pelvis, so de-rotating them by ``neck_inv``
correctly expresses them in the neck/torso frame (gear_sonic's behaviour).
- ``False`` (device-source, :func:`compute_3point_from_devices`): the "neck" is the
**headset**, which pitches down when the operator looks at their hands; the wrists
are already yaw-stabilised, so applying the head pitch/roll would rotate "up" hand
motion into "forward". The wrist frame is left in the yaw-local world frame (only
the neck point itself is still de-tilted).
""" """
neck_relative_wrists: bool = True
_neck_quat_inv: R | None = field(default=None, init=False) _neck_quat_inv: R | None = field(default=None, init=False)
_wrist_pos_offset: np.ndarray | None = field(default=None, init=False) _wrist_pos_offset: np.ndarray | None = field(default=None, init=False)
_wrist_rot_offset: list[R] = field(default_factory=list, init=False) _wrist_rot_offset: list[R] = field(default_factory=list, init=False)
@@ -436,13 +448,15 @@ class ThreePointCalibrator:
orn = np.asarray(orn, np.float64).reshape(3, 4) orn = np.asarray(orn, np.float64).reshape(3, 4)
if self._neck_quat_inv is None: if self._neck_quat_inv is None:
self._neck_quat_inv = R.from_quat(orn[2], scalar_first=True).inv() self._neck_quat_inv = R.from_quat(orn[2], scalar_first=True).inv()
neck_inv = self._neck_quat_inv # Wrists use the neck frame only in body-source mode; device-source keeps them
# in the already yaw-stabilised world frame (see class docstring).
wrist_inv = self._neck_quat_inv if self.neck_relative_wrists else R.identity()
self._wrist_pos_offset = np.zeros((2, 3), np.float64) self._wrist_pos_offset = np.zeros((2, 3), np.float64)
self._wrist_rot_offset = [] self._wrist_rot_offset = []
for k in range(2): for k in range(2):
corrected_pos = neck_inv.apply(pos[k]) corrected_pos = wrist_inv.apply(pos[k])
corrected_rot = neck_inv * R.from_quat(orn[k], scalar_first=True) corrected_rot = wrist_inv * R.from_quat(orn[k], scalar_first=True)
self._wrist_pos_offset[k] = corrected_pos - _G1_NEUTRAL_WRIST_POS[k] self._wrist_pos_offset[k] = corrected_pos - _G1_NEUTRAL_WRIST_POS[k]
# rot_offset maps the corrected rest orientation onto the G1 neutral wrist # rot_offset maps the corrected rest orientation onto the G1 neutral wrist
# orientation: calibrated = rot_offset * (neck_inv * current). # orientation: calibrated = rot_offset * (neck_inv * current).
@@ -461,12 +475,13 @@ class ThreePointCalibrator:
pos = np.asarray(pos, np.float64).reshape(3, 3) pos = np.asarray(pos, np.float64).reshape(3, 3)
orn = np.asarray(orn, np.float64).reshape(3, 4) orn = np.asarray(orn, np.float64).reshape(3, 4)
neck_inv = self._neck_quat_inv neck_inv = self._neck_quat_inv
wrist_inv = neck_inv if self.neck_relative_wrists else R.identity()
out_pos = np.zeros((3, 3), np.float64) out_pos = np.zeros((3, 3), np.float64)
out_orn = np.zeros((3, 4), np.float64) out_orn = np.zeros((3, 4), np.float64)
for k in range(2): # wrists for k in range(2): # wrists
out_pos[k] = neck_inv.apply(pos[k]) - self._wrist_pos_offset[k] out_pos[k] = wrist_inv.apply(pos[k]) - self._wrist_pos_offset[k]
corrected_rot = neck_inv * R.from_quat(orn[k], scalar_first=True) corrected_rot = wrist_inv * R.from_quat(orn[k], scalar_first=True)
out_orn[k] = (self._wrist_rot_offset[k] * corrected_rot).as_quat(scalar_first=True) out_orn[k] = (self._wrist_rot_offset[k] * corrected_rot).as_quat(scalar_first=True)
# Head/neck: orientation de-tilted, position from the torso->neck chain. # Head/neck: orientation de-tilted, position from the torso->neck chain.