add SMPL wiring into sonic controller

examples/unitree_g1/sonic.py

@@ -35,6 +35,7 @@ import time
 
 import numpy as np
 from motion_loader import SmplMotion
+from smpl_stream import DEFAULT_SMPL_HOST, DEFAULT_SMPL_PORT, SmplStream
 
 from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
 from lerobot.robots.unitree_g1.controllers.sonic_pipeline import (
@@ -90,8 +91,29 @@ def main():
     parser.add_argument(
         "--no-loop", action="store_true", help="With --motion-file, play once instead of looping"
     )
+    parser.add_argument(
+        "--smpl-stream",
+        action="store_true",
+        help="Use the live rt/smpl headset stream as the reference motion "
+        "(SONIC whole-body, encode_mode=2), instead of a --motion-file clip.",
+    )
+    parser.add_argument(
+        "--smpl-host",
+        type=str,
+        default=DEFAULT_SMPL_HOST,
+        help=f"Host publishing rt/smpl (default: {DEFAULT_SMPL_HOST})",
+    )
+    parser.add_argument(
+        "--smpl-port",
+        type=int,
+        default=DEFAULT_SMPL_PORT,
+        help=f"Port for the rt/smpl stream (default: {DEFAULT_SMPL_PORT})",
+    )
     args = parser.parse_args()
 
+    if args.smpl_stream and args.motion_file:
+        parser.error("--smpl-stream and --motion-file are mutually exclusive")
+
     # Surface native crashes (onnxruntime / mujoco) with a real traceback, and
     # avoid losing buffered diagnostics if the process dies mid-loop.
     faulthandler.enable()
@@ -128,7 +150,14 @@ def main():
     ms = runtime.ms
 
     motion = None
-    if args.motion_file:
+    if args.smpl_stream:
+        motion = SmplStream(host=args.smpl_host, port=args.smpl_port)
+        controller.smpl_motion = motion  # lets 'M' key toggle streaming
+        controller.encode_mode = 2  # start in SONIC whole-body SMPL imitation
+        print(f"\n[Motion] live SMPL stream (rt/smpl @ {args.smpl_host}:{args.smpl_port})")
+        print("  Source: pico_manager_thread_server.py --manager on the publisher host.")
+        print("  encode_mode=2. Press 'M' to toggle SMPL stream <-> locomotion at runtime.")
+    elif args.motion_file:
         motion = SmplMotion(args.motion_file, loop=not args.no_loop)
         controller.smpl_motion = motion  # lets 'M' key toggle playback
         controller.encode_mode = 2  # start in SONIC whole-body SMPL imitation
@@ -283,6 +312,8 @@ def main():
             gc.enable()
             if args.log_csv:
                 print(f"\n[Log] Saved to {log_path}")
+            if motion is not None and hasattr(motion, "close"):
+                motion.close()
             runtime.shutdown()
             print("\nStopping...")
             if robot.is_connected:

src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py

@@ -17,7 +17,9 @@
 """SONIC full-body controller for Unitree G1."""
 
 import logging
+import os
 
+import numpy as np
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
@@ -42,6 +44,28 @@ from lerobot.robots.unitree_g1.controllers.sonic_pipeline import (
 
 logger = logging.getLogger(__name__)
 
+# Length of the flattened SONIC whole-body reference window
+# (10 frames x 24 SMPL joints x 3 coords). Matches smpl_joints_10frame_step1.
+SMPL_ACTION_DIM = 720
+# Prefix for per-element SMPL floats carried on the teleop action dict.
+SMPL_ACTION_PREFIX = "smpl."
+
+
+def _extract_smpl_from_action(action: dict | None) -> np.ndarray | None:
+    """Reassemble a (720,) SMPL window from ``smpl.{i}`` action keys, or None.
+
+    The pico_headset teleoperator emits the whole-body reference as flat floats so
+    it flows unchanged through the standard lerobot action pipeline.
+    """
+    if not action or f"{SMPL_ACTION_PREFIX}0" not in action:
+        return None
+    arr = np.fromiter(
+        (float(action.get(f"{SMPL_ACTION_PREFIX}{i}", 0.0)) for i in range(SMPL_ACTION_DIM)),
+        dtype=np.float32,
+        count=SMPL_ACTION_DIM,
+    )
+    return arr
+
 
 class SonicRuntime:
     """Shared SONIC control loop state (standalone demo + locomotion controller)."""
@@ -132,20 +156,60 @@ class SonicWholeBodyController:
         self.kd = self._runtime.kd
         self.controller = self._runtime.controller
         self.ms = self._runtime.ms
+
+        # 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. Enable
+        # with SONIC_SMPL_STREAM=1; override endpoint via SONIC_SMPL_HOST/PORT.
+        self._smpl_stream = None
+        if os.environ.get("SONIC_SMPL_STREAM", "0") not in ("0", "", "false", "False"):
+            self._init_smpl_stream()
+
         logger.info(
-            "SONIC ready: %s (default mode: %s)",
+            "SONIC ready: %s (default mode: %s, smpl_stream=%s)",
             MOTION_SETS[0][0],
             LM(self.ms.mode).name,
+            self._smpl_stream is not None,
         )
 
+    def _init_smpl_stream(self) -> None:
+        # Lazy import so the zmq dependency is only required when streaming is on.
+        from lerobot.robots.unitree_g1.smpl_stream import (
+            DEFAULT_SMPL_HOST,
+            DEFAULT_SMPL_PORT,
+            SmplStream,
+        )
+
+        host = os.environ.get("SONIC_SMPL_HOST", DEFAULT_SMPL_HOST)
+        port = int(os.environ.get("SONIC_SMPL_PORT", DEFAULT_SMPL_PORT))
+        self._smpl_stream = SmplStream(host=host, port=port)
+        logger.info("SONIC subscribed to rt/smpl @ tcp://%s:%d", host, port)
+
     def run_step(self, action: dict, lowstate) -> dict:
         if lowstate is None:
             return {}
         obs = lowstate_to_obs(lowstate)
+
+        # Prefer SMPL delivered via the teleop action (pico_headset). Fall back to a
+        # direct rt/smpl subscription when SONIC_SMPL_STREAM is enabled.
+        smpl = _extract_smpl_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:
+                smpl = window
+
+        if smpl is not None:
+            # Full-body whole-body tracking: SMPL drives the reference, not joystick.
+            self.controller.encode_mode = 2
+            self.controller.smpl_joints_10frame_step1 = smpl
+            return self._runtime.tick(obs, debug=False, use_joystick=False)
+
         return self._runtime.tick(obs, debug=False)
 
     def reset(self):
         self._runtime.reset()
 
     def shutdown(self):
+        if self._smpl_stream is not None:
+            self._smpl_stream.close()
         self._runtime.shutdown()

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -508,6 +508,13 @@ class UnitreeG1(Robot):
             for key in REMOTE_KEYS:
                 if key in action:
                     self.controller_input[key] = action[key]
+            # Forward the whole-body SMPL reference (pico_headset teleop) to the
+            # controller. Carried as flat smpl.{i} floats so it passes untouched
+            # through the standard action pipeline; SonicWholeBodyController
+            # reassembles it into the 720-vec encoder window.
+            for key in action:
+                if key.startswith("smpl."):
+                    self.controller_input[key] = action[key]
 
     @property
     def is_calibrated(self) -> bool:

src/lerobot/scripts/lerobot_teleoperate.py

@@ -100,6 +100,7 @@ from lerobot.teleoperators import (  # noqa: F401
     omx_leader,
     openarm_leader,
     openarm_mini,
+    pico_headset,
     reachy2_teleoperator,
     rebot_102_leader,
     so_leader,