diff --git a/examples/unitree_g1/sonic.py b/examples/unitree_g1/sonic.py index efd5c4446..bd00a1c41 100644 --- a/examples/unitree_g1/sonic.py +++ b/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: 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 56958c742..9b813a486 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/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() diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 782f3519c..283137d37 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/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: diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index 5d806200d..c5f9d5bb2 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/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,