diff --git a/src/lerobot/robots/unitree_g1/config_unitree_g1.py b/src/lerobot/robots/unitree_g1/config_unitree_g1.py index e768ed94b..476412930 100644 --- a/src/lerobot/robots/unitree_g1/config_unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/config_unitree_g1.py @@ -65,6 +65,35 @@ class UnitreeG1Config(RobotConfig): # Cameras (ZMQ-based remote cameras) 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 gravity_compensation: bool = False diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py index 63ac5bcf6..120550504 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -20,6 +20,7 @@ from __future__ import annotations import logging import math +from collections import deque from typing import TYPE_CHECKING import numpy as np @@ -38,10 +39,12 @@ from lerobot.teleoperators.pico_headset.smpl_constants import ( VR3_ORN_PREFIX, VR3_POS_DIM, VR3_POS_PREFIX, + WB_ACTION_DIM, + wb_action_key, ) 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 ( CONTROL_DT, DEBUG_PRINT_EVERY, @@ -132,6 +135,51 @@ def _extract_vr3_from_action(action: dict | None) -> tuple[np.ndarray, np.ndarra 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: """Reassemble controller-stick locomotion from ``loco_axes.{i}`` / ``loco_btn.{i}``. @@ -234,6 +282,10 @@ class SonicWholeBodyController: control_dt = CONTROL_DT 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__( self, @@ -277,6 +329,13 @@ class SonicWholeBodyController: self._init_step = 0 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 # teleop works with ANY teleoperator (e.g. --teleop.type=unitree_g1 for the # 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) 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: """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 # stream (headset silent past its timeout) is treated as "no SMPL" so the # 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) root_quat = _extract_root_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._start_pose = {} self._smoothed_root_quat = None + self._wb_step = 0 + self._wb_traj.clear() + self._wb_quat_traj.clear() def shutdown(self): if self._smpl_stream is not None: diff --git a/src/lerobot/robots/unitree_g1/g1_utils.py b/src/lerobot/robots/unitree_g1/g1_utils.py index 0449a5a4f..723bf94aa 100644 --- a/src/lerobot/robots/unitree_g1/g1_utils.py +++ b/src/lerobot/robots/unitree_g1/g1_utils.py @@ -104,6 +104,21 @@ REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS # whole-body controller (see SonicWholeBodyController._process_keyboard). 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]: """Return a zeroed-out remote input dict (axes + buttons).""" @@ -181,6 +196,43 @@ def lowstate_to_obs(lowstate) -> dict: 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 ``.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): """Instantiate a locomotion controller by class name. Returns None if name is None.""" if name is None: diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 2146e293b..5349aec23 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -40,6 +40,7 @@ from .g1_utils import ( default_remote_input, lowstate_to_obs, make_locomotion_controller, + obs_to_wb34_state, ) if TYPE_CHECKING or _unitree_sdk_available: @@ -48,8 +49,12 @@ 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__HandCmd_ as hg_HandCmd_default, + unitree_hg_msg_dds__LowCmd_, + ) from unitree_sdk2py.idl.unitree_hg.msg.dds_ import ( + HandCmd_ as hg_HandCmd, LowCmd_ as hg_LowCmd, LowState_ as hg_LowState, ) @@ -59,6 +64,8 @@ else: _SDKChannelPublisher = None _SDKChannelSubscriber = None unitree_hg_msg_dds__LowCmd_ = None + hg_HandCmd_default = None + hg_HandCmd = None hg_LowCmd = None hg_LowState = None CRC = None @@ -158,6 +165,37 @@ class UnitreeG1(Robot): self.controller_input = default_remote_input() 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 while not self._shutdown_event.is_set(): start_time = time.time() @@ -232,15 +270,54 @@ class UnitreeG1(Robot): features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) 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 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 def action_features(self) -> dict[str, type]: if self.controller is None: 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} remote_features = dict.fromkeys(REMOTE_AXES, float) return {**arm_features, **remote_features} @@ -312,6 +389,17 @@ class UnitreeG1(Robot): self.lowstate_subscriber = self._ChannelSubscriber(kTopicLowState, hg_LowState) 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 self.subscribe_thread = threading.Thread(target=self._subscribe_lowstate) self.subscribe_thread.start() @@ -460,6 +548,31 @@ class UnitreeG1(Robot): # Motors + IMU + wireless remote (shared lowstate -> obs mapping) 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 for cam_name, cam in self._cameras.items(): if getattr(cam, "use_rgb", True): @@ -473,6 +586,8 @@ class UnitreeG1(Robot): action_to_publish = action if self.controller is not None: 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): return action # 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. Controller-agnostic: every value-carrying key is forwarded verbatim into - ``controller_input`` (joystick ``remote.*``, whole-body ``smpl.*``/``root.*`` - from pico_headset, or whatever a future controller expects), and each - controller extracts only the keys it understands. The robot deliberately does - not enumerate any controller's key schema here. + ``controller_input`` (whole-body ``wb.{i}.pos`` from a 34-D VLA, or whatever a + future controller expects), and each controller extracts only the keys it + understands. The robot deliberately does not enumerate any controller's key + schema here. 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 @@ -528,6 +643,44 @@ class UnitreeG1(Robot): if isinstance(key, str) and value is not None: 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 def is_calibrated(self) -> bool: return True diff --git a/src/lerobot/teleoperators/pico_headset/pico_publisher.py b/src/lerobot/teleoperators/pico_headset/pico_publisher.py index d918ac552..e4d249ff6 100644 --- a/src/lerobot/teleoperators/pico_headset/pico_publisher.py +++ b/src/lerobot/teleoperators/pico_headset/pico_publisher.py @@ -195,7 +195,13 @@ def main() -> None: # 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 # 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 if calibrator is not None: print( diff --git a/src/lerobot/teleoperators/pico_headset/smpl_constants.py b/src/lerobot/teleoperators/pico_headset/smpl_constants.py index c19b1295d..a59b320df 100644 --- a/src/lerobot/teleoperators/pico_headset/smpl_constants.py +++ b/src/lerobot/teleoperators/pico_headset/smpl_constants.py @@ -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_AXES_PREFIX = "loco_axes." 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" diff --git a/src/lerobot/teleoperators/pico_headset/smpl_fk.py b/src/lerobot/teleoperators/pico_headset/smpl_fk.py index 235b19932..a0791919b 100644 --- a/src/lerobot/teleoperators/pico_headset/smpl_fk.py +++ b/src/lerobot/teleoperators/pico_headset/smpl_fk.py @@ -394,8 +394,20 @@ class ThreePointCalibrator: neck orientation via the torso->neck kinematic chain. 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) _wrist_pos_offset: np.ndarray | None = field(default=None, 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) if self._neck_quat_inv is None: 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_rot_offset = [] for k in range(2): - corrected_pos = neck_inv.apply(pos[k]) - corrected_rot = neck_inv * R.from_quat(orn[k], scalar_first=True) + corrected_pos = wrist_inv.apply(pos[k]) + corrected_rot = wrist_inv * R.from_quat(orn[k], scalar_first=True) self._wrist_pos_offset[k] = corrected_pos - _G1_NEUTRAL_WRIST_POS[k] # rot_offset maps the corrected rest orientation onto the G1 neutral wrist # orientation: calibrated = rot_offset * (neck_inv * current). @@ -461,12 +475,13 @@ class ThreePointCalibrator: pos = np.asarray(pos, np.float64).reshape(3, 3) orn = np.asarray(orn, np.float64).reshape(3, 4) 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_orn = np.zeros((3, 4), np.float64) for k in range(2): # wrists - out_pos[k] = neck_inv.apply(pos[k]) - self._wrist_pos_offset[k] - corrected_rot = neck_inv * R.from_quat(orn[k], scalar_first=True) + out_pos[k] = wrist_inv.apply(pos[k]) - self._wrist_pos_offset[k] + 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) # Head/neck: orientation de-tilted, position from the torso->neck chain.