add heading to SMPL, stream dataset

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

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)
 

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

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

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:

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)."""