From 4bcde762ccf7baaf9d6654e3db27629b759d555a Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Fri, 10 Jul 2026 15:44:48 +0200 Subject: [PATCH] add heading to SMPL, stream dataset --- .../unitree_g1/controllers/sonic_pipeline.py | 19 +++- .../controllers/sonic_whole_body.py | 23 +++++ src/lerobot/robots/unitree_g1/unitree_g1.py | 2 +- .../pico_headset/pico_headset.py | 13 ++- .../pico_headset/pico_publisher.py | 98 ++++++++++++++----- .../teleoperators/pico_headset/smpl_fk.py | 33 +++++++ 6 files changed, 161 insertions(+), 27 deletions(-) diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py index 0340f043b..18e1fdac1 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py @@ -521,6 +521,9 @@ class StandingEncoderDecoder: self.vr_3point_local_target = VR_TARGET_DEF.copy() self.vr_3point_local_orn_target = VR_ORN_DEF.copy() self.smpl_joints_10frame_step1 = SMPL_DEF.copy() + # Optional per-frame SMPL root orientation (wxyz) for the mode-2 anchor. + # When None, the anchor falls back to the planner reference body quat. + self.smpl_root_quat = None self.set_zero_reference() def update_history(self, q, dq, ang, quat): @@ -577,9 +580,12 @@ class StandingEncoderDecoder: obs[910:922] = self.vr_3point_local_orn_target obs[595:601] = self._anchor_6d(self.h_quat[0], ref_quat) elif self.encode_mode == 2: + # Prefer the SMPL clip/stream root orientation for the anchor; fall + # back to the planner reference body quat when no root is provided. + anchor_ref = self.smpl_root_quat if self.smpl_root_quat is not None else ref_quat obs[922:1642] = self.smpl_joints_10frame_step1 for f in range(10): - obs[1642 + 6 * f : 1642 + 6 * (f + 1)] = self._anchor_6d(self.h_quat[0], ref_quat) + obs[1642 + 6 * f : 1642 + 6 * (f + 1)] = self._anchor_6d(self.h_quat[0], anchor_ref) obs[1702 + 6 * f : 1702 + 6 * (f + 1)] = ref_pos[WRIST_IL] else: raise RuntimeError(f"Unsupported encoder mode: {self.encode_mode}") @@ -1033,6 +1039,10 @@ class PlannerController(StandingEncoderDecoder): rf = min(self.ref_cursor, self.motion_timesteps - 1) ref_pos = self.motion_joint_positions[rf].astype(np.float32) ref_quat = self.motion_body_quats[rf].astype(np.float32) + # Prefer the SMPL clip/stream root orientation (if provided) so the + # anchor tracks the operator's/clip's heading; else planner ref. + if self.smpl_root_quat is not None: + ref_quat = np.asarray(self.smpl_root_quat, np.float32) anchor = self._anchor_6d(self.h_quat[0], ref_quat) wrist = ref_pos[WRIST_IL] obs[922:1642] = self.smpl_joints_10frame_step1 @@ -1075,7 +1085,12 @@ class PlannerController(StandingEncoderDecoder): self.heading_init_base_quat = np.array(q, np.float64) with self.motion_lock: rf = min(self.ref_cursor, self.motion_timesteps - 1) - self.init_ref_quat = self.motion_body_quats[rf].copy() + if self.encode_mode == 2 and self.smpl_root_quat is not None: + # Anchor the heading delta to the SMPL root at init so the + # robot turns *relative* to the clip/operator start heading. + self.init_ref_quat = np.asarray(self.smpl_root_quat, np.float64) + else: + self.init_ref_quat = self.motion_body_quats[rf].copy() self.delta_heading = 0.0 self.first_motion = False self.reinit_heading = 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 0b4597eaf..d0b91b006 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -49,6 +49,9 @@ logger = logging.getLogger(__name__) SMPL_ACTION_DIM = 720 # Prefix for per-element SMPL floats carried on the teleop action dict. SMPL_ACTION_PREFIX = "smpl." +# Optional per-frame SMPL root orientation (wxyz) for the mode-2 anchor. +ROOT_ACTION_DIM = 4 +ROOT_ACTION_PREFIX = "root." def _extract_smpl_from_action(action: dict | None) -> np.ndarray | None: @@ -67,6 +70,21 @@ def _extract_smpl_from_action(action: dict | None) -> np.ndarray | None: return arr +def _extract_root_from_action(action: dict | None) -> np.ndarray | None: + """Reassemble a (4,) SMPL root quaternion (wxyz) from ``root.{i}`` keys, or None.""" + if not action or f"{ROOT_ACTION_PREFIX}0" not in action: + return None + q = np.fromiter( + (float(action.get(f"{ROOT_ACTION_PREFIX}{i}", 0.0)) for i in range(ROOT_ACTION_DIM)), + dtype=np.float32, + count=ROOT_ACTION_DIM, + ) + n = float(np.linalg.norm(q)) + if n < 1e-6: + return None + return q / n + + class SonicRuntime: """Shared SONIC control loop state (standalone demo + locomotion controller).""" @@ -213,20 +231,25 @@ class SonicWholeBodyController: # stream (headset silent past its timeout) is treated as "no SMPL" so the # robot doesn't stay frozen tracking the last pose. smpl = _extract_smpl_from_action(action) + root_quat = _extract_root_from_action(action) if smpl is None and self._smpl_stream is not None: window = self._smpl_stream.step() if self._smpl_stream.has_data and not self._smpl_stream.is_stale: smpl = window + root_quat = np.asarray(self._smpl_stream.root_quat, np.float32) if smpl is not None: # Full-body whole-body tracking: SMPL drives the reference, not joystick. if self.controller.encode_mode != 2: self._enter_wholebody() self.controller.smpl_joints_10frame_step1 = smpl + # Root orientation (if provided) steers the mode-2 anchor/heading. + self.controller.smpl_root_quat = root_quat return self._runtime.tick(obs, debug=False, use_joystick=False) # No (or stale) SMPL: fall back to locomotion so the robot stays balanced. if self.controller.encode_mode == 2: + self.controller.smpl_root_quat = None self._exit_wholebody() return self._runtime.tick(obs, debug=False) diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index a26c5c732..1ad5818ff 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -486,7 +486,7 @@ class UnitreeG1(Robot): # through the standard action pipeline; SonicWholeBodyController # reassembles it into the 720-vec encoder window. for key in action: - if key.startswith("smpl."): + if key.startswith("smpl.") or key.startswith("root."): self.controller_input[key] = action[key] @property diff --git a/src/lerobot/teleoperators/pico_headset/pico_headset.py b/src/lerobot/teleoperators/pico_headset/pico_headset.py index 108b29c03..4804daafc 100644 --- a/src/lerobot/teleoperators/pico_headset/pico_headset.py +++ b/src/lerobot/teleoperators/pico_headset/pico_headset.py @@ -31,6 +31,9 @@ logger = logging.getLogger(__name__) # scalar floats so the reference flows unchanged through the standard lerobot action # pipeline; SonicWholeBodyController reassembles them into smpl_joints_10frame_step1. SMPL_ACTION_PREFIX = "smpl." +# Per-frame SMPL root orientation (wxyz) that steers the SONIC mode-2 anchor/heading. +ROOT_ACTION_PREFIX = "root." +ROOT_ACTION_DIM = 4 class PicoHeadset(Teleoperator): @@ -52,7 +55,9 @@ class PicoHeadset(Teleoperator): @property def action_features(self) -> dict: - return {f"{SMPL_ACTION_PREFIX}{i}": float for i in range(SMPL_OBS_DIM)} + feats = {f"{SMPL_ACTION_PREFIX}{i}": float for i in range(SMPL_OBS_DIM)} + feats.update({f"{ROOT_ACTION_PREFIX}{i}": float for i in range(ROOT_ACTION_DIM)}) + return feats @property def feedback_features(self) -> dict: @@ -97,7 +102,11 @@ class PicoHeadset(Teleoperator): # safe standing/locomotion mode instead of freezing on the last pose). if not self._stream.has_data or self._stream.is_stale: return {} - return {f"{SMPL_ACTION_PREFIX}{i}": float(v) for i, v in enumerate(window)} + action = {f"{SMPL_ACTION_PREFIX}{i}": float(v) for i, v in enumerate(window)} + action.update( + {f"{ROOT_ACTION_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.root_quat)} + ) + return action def send_feedback(self, feedback: dict[str, Any]) -> None: pass diff --git a/src/lerobot/teleoperators/pico_headset/pico_publisher.py b/src/lerobot/teleoperators/pico_headset/pico_publisher.py index e9472e779..88a6aaf25 100644 --- a/src/lerobot/teleoperators/pico_headset/pico_publisher.py +++ b/src/lerobot/teleoperators/pico_headset/pico_publisher.py @@ -31,6 +31,10 @@ Usage: # No hardware — emit a synthetic waving motion to test the consumer end-to-end: python -m lerobot.teleoperators.pico_headset.pico_publisher --fake + + # Replay a canned SMPL clip to the robot through the same rt/smpl -> SONIC path: + python -m lerobot.teleoperators.pico_headset.pico_publisher \ + --motion-file examples/unitree_g1/motions/walk_forward.npz """ from __future__ import annotations @@ -43,7 +47,11 @@ import time import numpy as np import zmq -from lerobot.teleoperators.pico_headset.smpl_fk import SmplForwardKinematics +from lerobot.teleoperators.pico_headset.smpl_fk import ( + SmplForwardKinematics, + canonicalize_smpl_joints, + root_quats_from_aa, +) SMPL_TOPIC = "rt/smpl" DEFAULT_SMPL_PORT = 5560 @@ -80,24 +88,56 @@ def _fake_body_poses(t: float) -> np.ndarray: return poses +def _load_motion_clip(path: str) -> dict: + """Load an SMPL ``.npz`` clip and canonicalize it for rt/smpl streaming. + + Expects the same keys as ``motion_loader.SmplMotion``: + smpl_joints (T, 24, 3), pose_aa (T, 72) optional, transl (T, 3) optional. + Returns per-frame joints already in the encoder's root-removed convention, + plus optional per-frame root quat/translation. + """ + data = np.load(path) + joints = data["smpl_joints"].astype(np.float32) + if joints.ndim != 3 or joints.shape[1:] != (24, 3): + raise ValueError(f"Expected smpl_joints (T, 24, 3), got {joints.shape}") + + pose_aa = data["pose_aa"].astype(np.float32) if "pose_aa" in data.files else None + root_quat = None + if pose_aa is not None: + joints = canonicalize_smpl_joints(joints, pose_aa[:, :3]) + root_quat = root_quats_from_aa(pose_aa[:, :3]) + transl = data["transl"].astype(np.float32) if "transl" in data.files else None + return {"joints": joints, "root_quat": root_quat, "transl": transl} + + def main() -> None: p = argparse.ArgumentParser(description=__doc__) p.add_argument("--port", type=int, default=DEFAULT_SMPL_PORT, help="ZMQ PUB port for rt/smpl") p.add_argument("--fps", type=float, default=50.0, help="Target publish rate (Hz)") p.add_argument("--skeleton", type=str, default=None, help="Path to smpl_skeleton.npz") - p.add_argument("--fake", action="store_true", help="Publish synthetic motion (no headset)") + src = p.add_mutually_exclusive_group() + src.add_argument("--fake", action="store_true", help="Publish synthetic motion (no headset)") + src.add_argument( + "--motion-file", type=str, default=None, help="Replay an SMPL .npz clip over rt/smpl" + ) + p.add_argument("--no-loop", action="store_true", help="Play a --motion-file once, then stop") args = p.parse_args() - fk = SmplForwardKinematics(args.skeleton) if args.skeleton else SmplForwardKinematics() + clip = _load_motion_clip(args.motion_file) if args.motion_file else None + + # FK is only needed for live/synthetic (24,7) body poses; clips are pre-canonical. + fk = None + if clip is None: + fk = SmplForwardKinematics(args.skeleton) if args.skeleton else SmplForwardKinematics() xrt = None - if not args.fake: + if clip is None and not args.fake: try: import xrobotoolkit_sdk as xrt # noqa: PLC0415 except ImportError as e: raise SystemExit( - "xrobotoolkit_sdk not available. Install it, or run with --fake to test " - "the pipeline without a headset." + "xrobotoolkit_sdk not available. Install it, or run with --fake / --motion-file " + "to test the pipeline without a headset." ) from e xrt.init() print("[pico_publisher] XRoboToolkit initialized") @@ -105,7 +145,13 @@ def main() -> None: ctx = zmq.Context.instance() sock = ctx.socket(zmq.PUB) sock.bind(f"tcp://*:{args.port}") - print(f"[pico_publisher] '{SMPL_TOPIC}' bound to tcp://*:{args.port} @ {args.fps:.0f} Hz") + src_desc = f"motion-file {args.motion_file}" if clip else ("fake" if args.fake else "headset") + print( + f"[pico_publisher] '{SMPL_TOPIC}' bound to tcp://*:{args.port} @ {args.fps:.0f} Hz " + f"[source: {src_desc}]" + ) + if clip is not None: + print(f"[pico_publisher] clip frames={clip['joints'].shape[0]} loop={not args.no_loop}") period = 1.0 / max(1.0, args.fps) frame_index = 0 @@ -113,25 +159,33 @@ def main() -> None: try: while True: loop_start = time.time() - if args.fake: - body_poses = _fake_body_poses(loop_start - t0) + if clip is not None: + n = clip["joints"].shape[0] + if args.no_loop and frame_index >= n: + print("\n[pico_publisher] clip finished") + break + i = frame_index % n + joints = clip["joints"][i] + root_quat = None if clip["root_quat"] is None else clip["root_quat"][i] + root_transl = None if clip["transl"] is None else clip["transl"][i] stamp_ns = time.time_ns() else: - body_poses = np.asarray(xrt.get_body_joints_pose(), np.float32) - stamp_ns = int(xrt.get_time_stamp_ns()) - if body_poses.shape != (24, 7): - time.sleep(0.005) - continue + if args.fake: + body_poses = _fake_body_poses(loop_start - t0) + stamp_ns = time.time_ns() + else: + body_poses = np.asarray(xrt.get_body_joints_pose(), np.float32) + stamp_ns = int(xrt.get_time_stamp_ns()) + if body_poses.shape != (24, 7): + time.sleep(0.005) + continue + out = fk.compute(body_poses) + joints = out["smpl_joints_local"] + root_quat = out["root_quat"] + root_transl = out["root_transl"] - out = fk.compute(body_poses) sock.send( - pack_message( - out["smpl_joints_local"], - frame_index, - stamp_ns, - root_quat=out["root_quat"], - root_transl=out["root_transl"], - ) + pack_message(joints, frame_index, stamp_ns, root_quat=root_quat, root_transl=root_transl) ) frame_index += 1 if frame_index % int(max(1, args.fps)) == 0: diff --git a/src/lerobot/teleoperators/pico_headset/smpl_fk.py b/src/lerobot/teleoperators/pico_headset/smpl_fk.py index 7e665e414..a54e3529f 100644 --- a/src/lerobot/teleoperators/pico_headset/smpl_fk.py +++ b/src/lerobot/teleoperators/pico_headset/smpl_fk.py @@ -117,6 +117,39 @@ def remove_smpl_base_rot(root_quat: np.ndarray) -> np.ndarray: # ── forward kinematics ─────────────────────────────────────────────────────── +def canonicalize_smpl_joints(smpl_joints: np.ndarray, root_aa: np.ndarray) -> np.ndarray: + """Remove per-frame root orientation -> SONIC ``smpl_joints_local`` format. + + Mirrors the deploy transform (and ``motion_loader.canonicalize_smpl_joints``): + reference clips store world-frame joints, but the encoder wants each frame's + joints with the body root orientation removed. + + Args: + smpl_joints: (T, 24, 3) world-frame (z-up) SMPL joint positions. + root_aa: (T, 3) SMPL global-orient axis-angle (y-up convention). + + Returns: + (T, 24, 3) per-frame root-orientation-removed joints. + """ + rx90 = R.from_euler("x", 90, degrees=True) # smpl_root_ytoz_up + base120 = R.from_quat([0.5, 0.5, 0.5, 0.5]) # remove_smpl_base_rot + a = rx90 * R.from_rotvec(root_aa) + b_inv = base120 * a.inv() + return np.einsum("tij,tkj->tki", b_inv.as_matrix(), smpl_joints).astype(np.float32) + + +def root_quats_from_aa(root_aa: np.ndarray) -> np.ndarray: + """Per-frame root orientation as (T, 4) wxyz, matching the live ``root_quat``. + + Same convention as the headset stream: ytoz-up then base-rotation removed. + """ + rx90 = R.from_euler("x", 90, degrees=True) + base120 = R.from_quat([0.5, 0.5, 0.5, 0.5]) + r = (rx90 * R.from_rotvec(root_aa)) * base120.inv() + q = r.as_quat() # xyzw + return np.concatenate([q[:, 3:4], q[:, :3]], axis=1).astype(np.float32) # -> wxyz + + class SmplForwardKinematics: """Rest-skeleton SMPL forward kinematics (no mesh, no torch)."""