add pico teleoperator, add sonic VR support

Co-authored-by: Cursor <cursoragent@cursor.com>

examples/unitree_g1/smpl_stream.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CLI/self-test shim for the live ``rt/smpl`` SONIC reference stream.
+
+The implementation now lives in the installed package so it can be shared with the
+``pico_headset`` teleoperator:
+``lerobot.robots.unitree_g1.smpl_stream``. This example re-exports it (so
+``from smpl_stream import SmplStream`` keeps working next to ``sonic.py``) and adds a
+standalone smoke-test entrypoint.
+
+Example:
+    python examples/unitree_g1/smpl_stream.py --smpl-host 127.0.0.1 --smpl-port 5560
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import numpy as np
+
+from lerobot.robots.unitree_g1.smpl_stream import (
+    DEFAULT_SMPL_HOST,
+    DEFAULT_SMPL_PORT,
+    SMPL_OBS_DIM,
+    SMPL_TOPIC,
+    SmplStream,
+)
+
+__all__ = [
+    "DEFAULT_SMPL_HOST",
+    "DEFAULT_SMPL_PORT",
+    "SMPL_OBS_DIM",
+    "SMPL_TOPIC",
+    "SmplStream",
+]
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--smpl-host", default=DEFAULT_SMPL_HOST)
+    parser.add_argument("--smpl-port", type=int, default=DEFAULT_SMPL_PORT)
+    parser.add_argument("--ticks", type=int, default=250, help="control ticks to sample")
+    args = parser.parse_args()
+
+    stream = SmplStream(host=args.smpl_host, port=args.smpl_port)
+    print(f"Subscribed to {SMPL_TOPIC} @ tcp://{args.smpl_host}:{args.smpl_port}")
+    print("Start pico_manager_thread_server.py --manager on the publisher host.")
+    try:
+        for t in range(args.ticks):
+            w = stream.step()
+            assert w.shape == (SMPL_OBS_DIM,), w.shape
+            if t % 25 == 0:
+                print(
+                    f"  t={t} idx={stream._last_index} window_norm={np.linalg.norm(w):.3f} "
+                    f"first={stream.has_data}"
+                )
+            time.sleep(1.0 / 50.0)
+    finally:
+        stream.close()
+    print("OK: rt/smpl stream yields SONIC-format 720-vec windows.")
+
+
+if __name__ == "__main__":
+    main()

src/lerobot/robots/unitree_g1/smpl_stream.py

@@ -0,0 +1,185 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Live SMPL stream as a SONIC reference motion (drop-in for ``SmplMotion``).
+
+Instead of reading an ``.npz`` clip, this pulls per-frame SMPL joints live off the
+``rt/smpl`` ZMQ channel published by the GEAR PICO teleop
+(``gear_sonic/scripts/pico_manager_thread_server.py``). Each message carries one
+frame of **canonical** (root-orientation-removed) SMPL local joints ``(24, 3)`` --
+the exact per-frame format ``SmplMotion`` emits -- so this class exposes the same
+``step() -> (720,)`` window interface and can be handed to ``sonic.py`` (or the
+``pico_headset`` teleoperator) wherever a reference motion is expected
+(``encode_mode == 2``).
+
+Transport mirrors the Unitree SDK socket bridge (``unitree_sdk2_socket.py``): a ZMQ
+``SUB`` socket with ``CONFLATE`` (keep only the latest frame) subscribed to the
+``rt/smpl`` topic, JSON payloads.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import time
+from collections import deque
+
+import numpy as np
+import zmq
+
+SMPL_TOPIC = "rt/smpl"
+DEFAULT_SMPL_HOST = "127.0.0.1"
+DEFAULT_SMPL_PORT = 5560
+
+WINDOW = 10  # frames per encoder window (smpl_joints_10frame_step1)
+N_JOINTS = 24
+JOINT_DIM = 3
+SMPL_OBS_DIM = WINDOW * N_JOINTS * JOINT_DIM  # 720
+
+
+class SmplStream:
+    """Live ``rt/smpl`` consumer with the ``SmplMotion`` interface.
+
+    Args:
+        host: publisher host (the laptop running pico_manager_thread_server.py).
+        port: publisher port for the ``rt/smpl`` channel.
+        fps: nominal source rate, only used for status/reporting.
+        stale_after_s: log a warning if no fresh frame arrives within this window.
+        loop: accepted for API parity with ``SmplMotion`` (ignored; a stream never ends).
+    """
+
+    def __init__(
+        self,
+        host: str = DEFAULT_SMPL_HOST,
+        port: int = DEFAULT_SMPL_PORT,
+        fps: float = 50.0,
+        stale_after_s: float = 0.5,
+        loop: bool = True,
+    ):
+        self.host = host
+        self.port = port
+        self.fps = float(fps)
+        self.loop = loop
+        self.stale_after_s = stale_after_s
+
+        self._ctx = zmq.Context.instance()
+        self._sock = self._ctx.socket(zmq.SUB)
+        self._sock.setsockopt(zmq.CONFLATE, 1)  # keep only the most recent frame
+        self._sock.connect(f"tcp://{host}:{port}")
+        # Single-frame JSON messages (topic embedded in payload); CONFLATE does not
+        # support multipart, so subscribe to everything on this dedicated port.
+        self._sock.setsockopt_string(zmq.SUBSCRIBE, "")
+        self._poller = zmq.Poller()
+        self._poller.register(self._sock, zmq.POLLIN)
+
+        # Rolling window, oldest -> newest (matches SmplMotion.window layout).
+        self._buf: deque[np.ndarray] = deque(maxlen=WINDOW)
+        self._last_frame = np.zeros((N_JOINTS, JOINT_DIM), np.float32)
+        # Latest root/torso pose (updated every received frame).
+        self.root_quat = np.array([1.0, 0.0, 0.0, 0.0], np.float32)  # (w, x, y, z)
+        self.root_transl = np.zeros(3, np.float32)
+        self._last_index = -1
+        self._last_recv_t = 0.0
+        self._warned_stale = False
+        self._got_first = False
+
+    # -- SmplMotion-compatible attributes ------------------------------------
+    @property
+    def num_frames(self) -> int:
+        """Streams are unbounded; report 0 (kept for API parity)."""
+        return 0
+
+    @property
+    def done(self) -> bool:
+        """A live stream never finishes."""
+        return False
+
+    @property
+    def has_data(self) -> bool:
+        """True once at least one real frame has been received."""
+        return self._got_first
+
+    def reset(self):
+        self._buf.clear()
+        self._got_first = False
+
+    # -- core ----------------------------------------------------------------
+    def _drain_latest(self) -> np.ndarray | None:
+        """Return the newest available (24, 3) frame, or None if nothing new.
+
+        CONFLATE already keeps only the last message, but we poll non-blocking so
+        the 50 Hz control loop never stalls waiting on the headset.
+        """
+        frame = None
+        while dict(self._poller.poll(0)).get(self._sock) == zmq.POLLIN:
+            payload = self._sock.recv()
+            data = json.loads(payload.decode("utf-8")).get("data", {})
+            joints = np.asarray(data.get("smpl_joints_local", []), np.float32)
+            if joints.size != N_JOINTS * JOINT_DIM:
+                continue
+            frame = joints.reshape(N_JOINTS, JOINT_DIM)
+            self._last_index = int(data.get("frame_index", self._last_index + 1))
+            rq = data.get("root_quat")
+            if rq is not None and len(rq) == 4:
+                self.root_quat = np.asarray(rq, np.float32)
+            rt = data.get("root_transl")
+            if rt is not None and len(rt) == 3:
+                self.root_transl = np.asarray(rt, np.float32)
+        return frame
+
+    def step(self) -> np.ndarray:
+        """Advance one control tick, returning the current 720-vec window.
+
+        If no new headset frame arrived this tick we hold the last one, so the
+        policy keeps tracking the latest pose rather than snapping to zero.
+        """
+        frame = self._drain_latest()
+        now = time.time()
+
+        if frame is not None:
+            self._last_frame = frame
+            self._last_recv_t = now
+            self._warned_stale = False
+            if not self._got_first:
+                # Pre-fill the window so the first send is a full, coherent clip.
+                self._buf.extend([frame.copy() for _ in range(WINDOW)])
+                self._got_first = True
+            else:
+                self._buf.append(frame)
+        elif self._got_first:
+            # No fresh frame: repeat the most recent to keep the window moving.
+            self._buf.append(self._last_frame.copy())
+            if (
+                self.stale_after_s
+                and not self._warned_stale
+                and (now - self._last_recv_t) > self.stale_after_s
+            ):
+                print(
+                    f"[SmplStream] no {SMPL_TOPIC} frame for "
+                    f"{now - self._last_recv_t:.2f}s (holding last pose)"
+                )
+                self._warned_stale = True
+
+        if not self._got_first:
+            return np.zeros(SMPL_OBS_DIM, np.float32)
+        # Flatten to (720,): frames oldest->newest, joint-major within a frame
+        # [f0_j0_xyz, f0_j1_xyz, ..., f9_j23_xyz] — matches SmplMotion.window.
+        return np.concatenate(list(self._buf), dtype=np.float32).reshape(-1)
+
+    def close(self):
+        with contextlib.suppress(Exception):
+            self._poller.unregister(self._sock)
+        self._sock.close(0)

src/lerobot/teleoperators/pico_headset/__init__.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .config_pico_headset import PicoHeadsetConfig
+from .pico_headset import PicoHeadset
+
+__all__ = ["PicoHeadset", "PicoHeadsetConfig"]

src/lerobot/teleoperators/pico_headset/config_pico_headset.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+from ..config import TeleoperatorConfig
+
+
+@TeleoperatorConfig.register_subclass("pico_headset")
+@dataclass
+class PicoHeadsetConfig(TeleoperatorConfig):
+    """PICO full-body headset teleop: live SMPL over the rt/smpl ZMQ stream.
+
+    Consumes the ``rt/smpl`` channel published by the GEAR PICO manager
+    (``gear_sonic/scripts/pico_manager_thread_server.py``) and emits the whole-body
+    SONIC reference window (``encode_mode == 2``) for SonicWholeBodyController.
+    """
+
+    smpl_host: str = "127.0.0.1"
+    """Host of the pico_manager rt/smpl publisher (the laptop bridging the PICO)."""
+    smpl_port: int = 5560
+    """Port of the rt/smpl publisher."""
+    stale_after_s: float = 0.5
+    """Warn if no fresh headset frame arrives within this many seconds."""

src/lerobot/teleoperators/pico_headset/pico_headset.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""PICO full-body headset teleoperator (live SMPL -> SONIC whole-body)."""
+
+import logging
+from typing import Any
+
+from lerobot.robots.unitree_g1.smpl_stream import SMPL_OBS_DIM, SmplStream
+from lerobot.types import RobotAction
+
+from ..teleoperator import Teleoperator
+from .config_pico_headset import PicoHeadsetConfig
+
+logger = logging.getLogger(__name__)
+
+# Flat action keys carrying the 720-vec SONIC whole-body reference window. Kept as
+# 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."
+
+
+class PicoHeadset(Teleoperator):
+    """Streams full-body SMPL from a PICO headset as a SONIC whole-body reference.
+
+    Subscribes to the ``rt/smpl`` ZMQ channel and, once real frames are flowing,
+    emits the 720-element encoder window as ``smpl.{i}`` floats. Before the first
+    frame arrives it emits no SMPL keys, so the robot stays in safe locomotion mode
+    rather than tracking a zero pose.
+    """
+
+    config_class = PicoHeadsetConfig
+    name = "pico_headset"
+
+    def __init__(self, config: PicoHeadsetConfig):
+        super().__init__(config)
+        self.config = config
+        self._stream: SmplStream | None = None
+
+    @property
+    def action_features(self) -> dict:
+        return {f"{SMPL_ACTION_PREFIX}{i}": float for i in range(SMPL_OBS_DIM)}
+
+    @property
+    def feedback_features(self) -> dict:
+        return {}
+
+    @property
+    def is_connected(self) -> bool:
+        return self._stream is not None
+
+    def connect(self, calibrate: bool = True) -> None:
+        if self._stream is not None:
+            raise RuntimeError(f"{self} already connected")
+        self._stream = SmplStream(
+            host=self.config.smpl_host,
+            port=self.config.smpl_port,
+            stale_after_s=self.config.stale_after_s,
+        )
+        logger.info(
+            "PicoHeadset subscribed to rt/smpl @ tcp://%s:%d",
+            self.config.smpl_host,
+            self.config.smpl_port,
+        )
+
+    @property
+    def is_calibrated(self) -> bool:
+        return True
+
+    def calibrate(self) -> None:
+        # Calibration happens on the headset / pico_manager side, not here.
+        pass
+
+    def configure(self) -> None:
+        pass
+
+    def get_action(self) -> RobotAction:
+        if self._stream is None:
+            raise RuntimeError(f"{self} is not connected")
+        window = self._stream.step()
+        # Hold back until the headset is actually streaming, so the controller does
+        # not switch to whole-body tracking of an all-zero (collapsed) pose.
+        if not self._stream.has_data:
+            return {}
+        return {f"{SMPL_ACTION_PREFIX}{i}": float(v) for i, v in enumerate(window)}
+
+    def send_feedback(self, feedback: dict[str, Any]) -> None:
+        pass
+
+    def disconnect(self) -> None:
+        if self._stream is not None:
+            self._stream.close()
+            self._stream = None