From 305614b8c66931cbbec4e5eab7c8b71fb3b6518c Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Mon, 6 Jul 2026 18:15:06 +0200 Subject: [PATCH] add pico teleoperator, add sonic VR support Co-authored-by: Cursor --- examples/unitree_g1/smpl_stream.py | 79 ++++++++ src/lerobot/robots/unitree_g1/smpl_stream.py | 185 ++++++++++++++++++ .../teleoperators/pico_headset/__init__.py | 20 ++ .../pico_headset/config_pico_headset.py | 37 ++++ .../pico_headset/pico_headset.py | 106 ++++++++++ 5 files changed, 427 insertions(+) create mode 100644 examples/unitree_g1/smpl_stream.py create mode 100644 src/lerobot/robots/unitree_g1/smpl_stream.py create mode 100644 src/lerobot/teleoperators/pico_headset/__init__.py create mode 100644 src/lerobot/teleoperators/pico_headset/config_pico_headset.py create mode 100644 src/lerobot/teleoperators/pico_headset/pico_headset.py diff --git a/examples/unitree_g1/smpl_stream.py b/examples/unitree_g1/smpl_stream.py new file mode 100644 index 000000000..4e3d637b7 --- /dev/null +++ b/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() diff --git a/src/lerobot/robots/unitree_g1/smpl_stream.py b/src/lerobot/robots/unitree_g1/smpl_stream.py new file mode 100644 index 000000000..33fe16c7d --- /dev/null +++ b/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) diff --git a/src/lerobot/teleoperators/pico_headset/__init__.py b/src/lerobot/teleoperators/pico_headset/__init__.py new file mode 100644 index 000000000..b800c4879 --- /dev/null +++ b/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"] diff --git a/src/lerobot/teleoperators/pico_headset/config_pico_headset.py b/src/lerobot/teleoperators/pico_headset/config_pico_headset.py new file mode 100644 index 000000000..c452dd9c1 --- /dev/null +++ b/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.""" diff --git a/src/lerobot/teleoperators/pico_headset/pico_headset.py b/src/lerobot/teleoperators/pico_headset/pico_headset.py new file mode 100644 index 000000000..68eed4a7d --- /dev/null +++ b/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