(add) sonic 3-point teleop, safe startup/shutdown, tested on real g1

src/lerobot/robots/unitree_g1/config_unitree_g1.py

@@ -71,3 +71,8 @@ class UnitreeG1Config(RobotConfig):
     # Locomotion controller class name, e.g. "GrootLocomotionController",
     # "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
     controller: str | None = None
+
+    # On disconnect (e.g. Ctrl-C), seconds to hold the current pose while ramping joint
+    # stiffness (kp) to zero — a soft, damped settle instead of an instant limp /
+    # free-fall. 0 disables it (immediate zero-torque). Real robot only.
+    graceful_stop_s: float = 1.5

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

@@ -1244,16 +1244,15 @@ def _parse_wireless(wr):
     return lx, ly, rx, ry
 
 
-def process_joystick(obs, ms, controller=None):
-    """Joystick mirrors keyboard: left stick=WASD, right stick X=Q/E, right stick Y=height."""
-    wr = obs.get("wireless_remote")
-    if wr is None:
-        return
-    parsed = _parse_wireless(wr)
-    if parsed is None:
-        return
-    lx, ly, rx, ry = parsed
+def apply_joystick_axes(lx, ly, rx, ry, ms, controller=None):
+    """Map raw stick axes onto ``MovementState`` (left stick=WASD, right stick X=Q/E,
+    right stick Y=height).
 
+    Shared by the G1 wireless remote (:func:`process_joystick`) and the PICO
+    controller sticks (3-point teleop), so both drive the planner identically. Axes
+    are expected pre-negated to the same convention as the parsed wireless remote:
+    ``ly`` and ``ry`` already flipped, dead zone not yet applied.
+    """
     # Dead zone + negate both Y axes (bridge already flips them once)
     lx = 0.0 if abs(lx) < DEADZONE else lx
     ly = 0.0 if abs(ly) < DEADZONE else -ly
@@ -1284,3 +1283,15 @@ def process_joystick(obs, ms, controller=None):
     if abs(ry) > 0:
         step = -0.005 * ry
         ms.height = max(0.1, min(1.0, (ms.height if ms.height >= 0 else DEFAULT_HEIGHT) + step))
+
+
+def process_joystick(obs, ms, controller=None):
+    """Drive ``MovementState`` from the G1 wireless remote in ``obs``."""
+    wr = obs.get("wireless_remote")
+    if wr is None:
+        return
+    parsed = _parse_wireless(wr)
+    if parsed is None:
+        return
+    lx, ly, rx, ry = parsed
+    apply_joystick_axes(lx, ly, rx, ry, ms, controller)

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

@@ -26,6 +26,10 @@ import numpy as np
 from huggingface_hub import hf_hub_download
 
 from lerobot.teleoperators.pico_headset.smpl_constants import (
+    LOCO_AXES_PREFIX,
+    LOCO_BTN_PREFIX,
+    LOCO_N_AXES,
+    LOCO_N_BTN,
     ROOT_ACTION_DIM,
     ROOT_ACTION_PREFIX,
     SMPL_ACTION_PREFIX,
@@ -37,7 +41,7 @@ from lerobot.teleoperators.pico_headset.smpl_constants import (
 )
 from lerobot.utils.import_utils import _onnxruntime_available, require_package
 
-from ..g1_utils import KEYBOARD_KEYS_FIELD, lowstate_to_obs
+from ..g1_utils import KEYBOARD_KEYS_FIELD, G1_29_JointIndex, lowstate_to_obs
 from .sonic_pipeline import (
     CONTROL_DT,
     DEBUG_PRINT_EVERY,
@@ -49,6 +53,7 @@ from .sonic_pipeline import (
     MovementState,
     PlannerController,
     SonicPlanner,
+    apply_joystick_axes,
     clamp_mode_params,
     compute_kp_kd,
     make_ort_session_options,
@@ -65,6 +70,11 @@ else:
 
 logger = logging.getLogger(__name__)
 
+# Startup blend duration: over the first control ticks, linearly interpolate every joint
+# from the robot's initial measured pose into the policy's commanded target, so control
+# eases in without a snap on the first command.
+INIT_RAMP_S = 3.0
+
 
 def _extract_smpl_from_action(action: dict | None) -> np.ndarray | None:
     """Reassemble a (720,) SMPL window from ``smpl.{i}`` action keys, or None.
@@ -122,6 +132,27 @@ def _extract_vr3_from_action(action: dict | None) -> tuple[np.ndarray, np.ndarra
     return pos, orn
 
 
+def _extract_loco_from_action(action: dict | None) -> tuple[np.ndarray, np.ndarray] | None:
+    """Reassemble controller-stick locomotion from ``loco_axes.{i}`` / ``loco_btn.{i}``.
+
+    Returns ``(axes (4,) = [lx, ly, rx, ry], buttons (4,) = [A, B, X, Y])`` or None
+    when no locomotion state was sent this tick (sentinel: ``loco_axes.0``).
+    """
+    if not action or f"{LOCO_AXES_PREFIX}0" not in action:
+        return None
+    axes = np.fromiter(
+        (float(action.get(f"{LOCO_AXES_PREFIX}{i}", 0.0)) for i in range(LOCO_N_AXES)),
+        dtype=np.float32,
+        count=LOCO_N_AXES,
+    )
+    buttons = np.fromiter(
+        (float(action.get(f"{LOCO_BTN_PREFIX}{i}", 0.0)) for i in range(LOCO_N_BTN)),
+        dtype=np.float32,
+        count=LOCO_N_BTN,
+    )
+    return axes, buttons
+
+
 class SonicRuntime:
     """Shared SONIC control loop state (standalone demo + locomotion controller)."""
 
@@ -230,6 +261,14 @@ class SonicWholeBodyController:
         # motion set, replan, e-stop, WASD direction) fire once per physical press
         # instead of every 50 Hz tick while the key is held.
         self._prev_keys: set[str] = set()
+        # Edge state for the PICO A+B / X+Y locomotion-mode cycle (3-point teleop).
+        self._prev_loco_mode_pair: tuple[bool, bool] = (False, False)
+
+        # Startup blend: ease from the robot's initial pose into the first commanded
+        # policy targets over INIT_RAMP_S (captured on the first control tick).
+        self._init_ramp_steps = max(1, round(INIT_RAMP_S / CONTROL_DT))
+        self._init_step = 0
+        self._start_pose: dict[str, float] = {}
 
         # 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
@@ -370,6 +409,61 @@ class SonicWholeBodyController:
         if "-" in held:
             ms.height = max(0.1, (ms.height if ms.height >= 0 else DEFAULT_HEIGHT) - 0.005)
 
+    def _process_pico_loco(self, axes: np.ndarray, buttons: np.ndarray) -> None:
+        """Drive locomotion from the PICO controller sticks/buttons (encode_mode 1).
+
+        Mirrors gear_sonic's ``PlannerLoop`` VR-3PT tick: left/right sticks steer
+        movement/facing/height (via the shared :func:`apply_joystick_axes`, identical
+        to the G1 remote), and A+B / X+Y edge-cycle the locomotion mode within the
+        current motion set.
+        """
+        lx, ly, rx, ry = (float(v) for v in axes)
+        apply_joystick_axes(lx, ly, rx, ry, self.ms, self.controller)
+
+        # Mode cycling: step linearly through the LocomotionMode enum (A+B = next,
+        # X+Y = previous), exactly like gear_sonic's PlannerLoop — so the operator can
+        # reach squat/kneel/crawl, not just the modes in one UI motion set.
+        a, b, x, y = (v > 0.5 for v in buttons)
+        ab_now, xy_now = (a and b), (x and y)
+        ab_prev, xy_prev = self._prev_loco_mode_pair
+        mode = int(self.ms.mode)
+        if ab_now and not ab_prev:
+            mode = min(int(LM.INJURED_WALK), mode + 1)
+        elif xy_now and not xy_prev:
+            mode = max(int(LM.IDLE), mode - 1)
+        if mode != int(self.ms.mode):
+            self.ms.mode = LM(mode)
+            self.ms.needs_replan = True
+            self.controller.playing = True
+            logger.info("SONIC 3-point: locomotion mode -> %s", LM(self.ms.mode).name)
+        self._prev_loco_mode_pair = (ab_now, xy_now)
+
+    def _startup_blend(self, obs: dict, out: dict) -> dict:
+        """Ease into policy control at startup: for the first ``INIT_RAMP_S`` seconds,
+        interpolate between the robot's pose captured on the first tick and the policy's
+        live commanded target, so the handoff has no snap.
+
+        ``out`` is the policy's ``<joint>.q`` target dict for this tick; the blend ratio
+        climbs 0->1 over the ramp, after which the raw policy target passes through.
+        """
+        if self._init_step >= self._init_ramp_steps or not out:
+            return out
+        if self._init_step == 0:
+            # Capture the robot's actual pose as the interpolation start point.
+            self._start_pose = {
+                f"{m.name}.q": float(obs.get(f"{m.name}.q", DEFAULT_ANGLES[m.value]))
+                for m in G1_29_JointIndex
+            }
+        self._init_step += 1
+        ratio = min(1.0, self._init_step / self._init_ramp_steps)
+        blended = {
+            k: self._start_pose.get(k, float(tgt)) * (1.0 - ratio) + float(tgt) * ratio
+            for k, tgt in out.items()
+        }
+        if self._init_step >= self._init_ramp_steps:
+            logger.info("SONIC startup blend complete -> full policy control")
+        return blended
+
     def run_step(self, action: dict, lowstate) -> dict:
         if lowstate is None:
             return {}
@@ -387,13 +481,18 @@ class SonicWholeBodyController:
         smpl = _extract_smpl_from_action(action)
         root_quat = _extract_root_from_action(action)
         vr3 = _extract_vr3_from_action(action)
+        loco = _extract_loco_from_action(action)
         if smpl is None and vr3 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 self._smpl_stream.has_vr3:
-                    vr3 = (self._smpl_stream.vr3_pos, self._smpl_stream.vr3_orn)
+            # VR3 is independent of the SMPL window: the controller-state source
+            # (head + controllers only) sends 3-point targets with no SMPL frame.
+            elif self._smpl_stream.has_fresh_vr3:
+                vr3 = (self._smpl_stream.vr3_pos, self._smpl_stream.vr3_orn)
+                if self._smpl_stream.has_fresh_loco:
+                    loco = (self._smpl_stream.loco_axes, self._smpl_stream.loco_buttons)
 
         if smpl is not None:
             # Full-body whole-body tracking: SMPL drives the reference, not joystick.
@@ -405,9 +504,8 @@ class SonicWholeBodyController:
             # self-driven to avoid root-acceleration spikes from an unsmoothed
             # reference trajectory.
             self.controller.smpl_root_quat = root_quat if self.enable_smpl_root else None
-            return self._runtime.tick(obs, debug=False, use_joystick=False)
-
-        if vr3 is not None:
+            out = self._runtime.tick(obs, debug=False, use_joystick=False)
+        elif vr3 is not None:
             # 3-point VR teleop: upper body tracks the wrist/neck targets; the lower
             # body / locomotion keeps running off the planner, so the joystick (and
             # keyboard) still steer walking/turning underneath.
@@ -415,16 +513,29 @@ class SonicWholeBodyController:
                 self._enter_3point()
             self.controller.vr_3point_local_target = vr3[0]
             self.controller.vr_3point_local_orn_target = vr3[1]
-            return self._runtime.tick(obs, debug=False, use_joystick=True)
+            # Replicate the original encode_mode-1 handling: when the PICO controller
+            # sticks are forwarded, drive locomotion from them directly (and skip the
+            # wireless-remote joystick read). Otherwise leave the remote/keyboard path.
+            if loco is not None:
+                self._process_pico_loco(loco[0], loco[1])
+                out = self._runtime.tick(obs, debug=False, use_joystick=False)
+            else:
+                out = self._runtime.tick(obs, debug=False, use_joystick=True)
+        else:
+            # No (or stale) teleop reference: fall back to locomotion so the robot stays balanced.
+            if self.controller.encode_mode != 0:
+                self.controller.smpl_root_quat = None
+                self._exit_wholebody()
+            out = self._runtime.tick(obs, debug=False)
 
-        # No (or stale) teleop reference: fall back to locomotion so the robot stays balanced.
-        if self.controller.encode_mode != 0:
-            self.controller.smpl_root_quat = None
-            self._exit_wholebody()
-        return self._runtime.tick(obs, debug=False)
+        # Startup interpolation: blend from the robot's initial pose into the policy's
+        # commanded target over INIT_RAMP_S, regardless of mode.
+        return self._startup_blend(obs, out)
 
     def reset(self):
         self._runtime.reset()
+        self._init_step = 0  # re-run the startup blend after a reset
+        self._start_pose = {}
 
     def shutdown(self):
         if self._smpl_stream is not None:

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -375,13 +375,47 @@ class UnitreeG1(Robot):
         except Exception as e:
             logger.warning(f"Failed to send zero-torque on disconnect: {e}")
 
-    def disconnect(self):
-        # Put robot in passive mode before stopping threads
-        if not self.config.is_simulation:
-            self._send_zero_torque()
+    def _graceful_stop(self) -> None:
+        """Soft shutdown: hold the current pose and ramp joint stiffness (kp) to zero
+        over ``graceful_stop_s`` while keeping damping (kd), then go passive.
 
-        # Signal thread to stop and unblock any waits
+        Prevents the robot from collapsing the instant control ends (a bare
+        zero-torque command is kp=kd=0 ≈ free-fall). Must run after the controller
+        loop has stopped so the two aren't publishing at once.
+        """
+        if self.config.graceful_stop_s <= 0:
+            self._send_zero_torque()
+            return
+        with self._lowstate_lock:
+            lowstate = self._lowstate
+        if lowstate is None:
+            self._send_zero_torque()
+            return
+        q_hold = {f"{motor.name}.q": lowstate.motor_state[motor.value].q for motor in G1_29_JointIndex}
+        kp = np.array(self.kp, dtype=np.float32)
+        kd = np.array(self.kd, dtype=np.float32)
+        zeros = np.zeros(29, dtype=np.float32)
+        dt = self.controller.control_dt if self.controller is not None else self.config.control_dt
+        steps = max(1, int(self.config.graceful_stop_s / dt))
+        logger.info("Graceful stop: damping down over %.1fs", self.config.graceful_stop_s)
+        for i in range(steps):
+            ratio = (i + 1) / steps
+            self.publish_lowcmd(q_hold, kp=kp * (1.0 - ratio), kd=kd, tau=zeros)
+            time.sleep(dt)
+        self._send_zero_torque()
+
+    def disconnect(self):
+        # Stop the controller loop first so it isn't fighting the shutdown ramp.
         self._shutdown_event.set()
+        if self._controller_thread is not None:
+            self._controller_thread.join(timeout=2.0)
+            if self._controller_thread.is_alive():
+                logger.warning("Controller thread did not stop cleanly")
+
+        # Soft, damped settle instead of an instant limp (real robot only; the
+        # subscribe thread is still alive here to supply the current pose).
+        if not self.config.is_simulation:
+            self._graceful_stop()
 
         if self.controller is not None and hasattr(self.controller, "shutdown"):
             self.controller.shutdown()
@@ -392,12 +426,6 @@ class UnitreeG1(Robot):
             if self.subscribe_thread.is_alive():
                 logger.warning("Subscribe thread did not stop cleanly")
 
-        # Wait for controller thread to finish
-        if self._controller_thread is not None:
-            self._controller_thread.join(timeout=2.0)
-            if self._controller_thread.is_alive():
-                logger.warning("Controller thread did not stop cleanly")
-
         # Close simulation environment
         if self.config.is_simulation and self.sim_env is not None:
             try:

src/lerobot/teleoperators/pico_headset/pico_headset.py

@@ -24,6 +24,10 @@ from lerobot.types import RobotAction
 from ..teleoperator import Teleoperator
 from .config_pico_headset import PicoHeadsetConfig
 from .smpl_constants import (
+    LOCO_AXES_PREFIX,
+    LOCO_BTN_PREFIX,
+    LOCO_N_AXES,
+    LOCO_N_BTN,
     ROOT_ACTION_DIM,
     ROOT_ACTION_PREFIX,
     SMPL_ACTION_PREFIX,
@@ -60,6 +64,9 @@ class PicoHeadset(Teleoperator):
         if self.config.mode == "vr3":
             feats = {f"{VR3_POS_PREFIX}{i}": float for i in range(VR3_POS_DIM)}
             feats.update({f"{VR3_ORN_PREFIX}{i}": float for i in range(VR3_ORN_DIM)})
+            # Controller-stick locomotion travels alongside the VR targets.
+            feats.update({f"{LOCO_AXES_PREFIX}{i}": float for i in range(LOCO_N_AXES)})
+            feats.update({f"{LOCO_BTN_PREFIX}{i}": float for i in range(LOCO_N_BTN)})
             return feats
         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)})
@@ -107,16 +114,27 @@ class PicoHeadset(Teleoperator):
         # collapsed pose) and once the stream goes stale (so the controller falls
         # back to a safe standing/locomotion mode instead of freezing on the last
         # pose).
-        if not self._stream.has_data or self._stream.is_stale:
-            return {}
         if self.config.mode == "vr3":
-            # Sparse 3-point upper-body teleop (encode_mode 1). Needs the producer to
-            # be sending vr3_* fields; otherwise emit nothing and stay in locomotion.
-            if not self._stream.has_vr3:
+            # Sparse 3-point upper-body teleop (encode_mode 1). Gated on fresh vr3_*
+            # frames only (independent of the SMPL window), so the controller-state
+            # source (head + controllers, no body tracking) works. Emit nothing
+            # otherwise and stay in locomotion.
+            if not self._stream.has_fresh_vr3:
                 return {}
             action = {f"{VR3_POS_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.vr3_pos)}
             action.update({f"{VR3_ORN_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.vr3_orn)})
+            # Forward controller-stick locomotion when present, so the planner can
+            # steer walking/turning under the upper-body tracking (encode_mode 1).
+            if self._stream.has_fresh_loco:
+                action.update(
+                    {f"{LOCO_AXES_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.loco_axes)}
+                )
+                action.update(
+                    {f"{LOCO_BTN_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.loco_buttons)}
+                )
             return action
+        if not self._stream.has_data or self._stream.is_stale:
+            return {}
         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

src/lerobot/teleoperators/pico_headset/pico_publisher.py

@@ -49,8 +49,10 @@ import zmq
 
 from lerobot.teleoperators.pico_headset.smpl_fk import (
     SmplForwardKinematics,
+    ThreePointCalibrator,
     canonicalize_smpl_joints,
     compute_3point,
+    compute_3point_from_devices,
     root_quats_from_aa,
 )
 
@@ -66,6 +68,8 @@ def pack_message(
     root_transl: np.ndarray | None = None,
     vr3_pos: np.ndarray | None = None,
     vr3_orn: np.ndarray | None = None,
+    loco_axes: np.ndarray | None = None,
+    loco_buttons: np.ndarray | None = None,
 ) -> bytes:
     """Build the rt/smpl JSON message (single frame, topic embedded in payload).
 
@@ -86,6 +90,10 @@ def pack_message(
         data["vr3_pos"] = np.asarray(vr3_pos, np.float32).reshape(-1).tolist()
     if vr3_orn is not None:
         data["vr3_orn"] = np.asarray(vr3_orn, np.float32).reshape(-1).tolist()
+    if loco_axes is not None:
+        data["loco_axes"] = np.asarray(loco_axes, np.float32).reshape(-1).tolist()
+    if loco_buttons is not None:
+        data["loco_buttons"] = np.asarray(loco_buttons, np.float32).reshape(-1).tolist()
     return json.dumps({"topic": SMPL_TOPIC, "data": data}).encode("utf-8")
 
 
@@ -127,6 +135,18 @@ def main() -> None:
     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(
+        "--headset-source",
+        choices=["body", "devices"],
+        default="body",
+        help=(
+            "Live headset keypoint source: 'body' uses full-body tracking "
+            "(get_body_joints_pose, needs PICO Motion Trackers) and drives both SMPL "
+            "(encode_mode 2) and 3-point; 'devices' uses head + 2 controllers only "
+            "(get_headset_pose + get_*_controller_pose, no trackers) and emits 3-point "
+            "(encode_mode 1) exclusively."
+        ),
+    )
     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")
@@ -155,7 +175,11 @@ def main() -> None:
     ctx = zmq.Context.instance()
     sock = ctx.socket(zmq.PUB)
     sock.bind(f"tcp://*:{args.port}")
-    src_desc = f"motion-file {args.motion_file}" if clip else ("fake" if args.fake else "headset")
+    src_desc = (
+        f"motion-file {args.motion_file}"
+        if clip
+        else ("fake" if args.fake else f"headset:{args.headset_source}")
+    )
     print(
         f"[pico_publisher] '{SMPL_TOPIC}' bound to tcp://*:{args.port} @ {args.fps:.0f} Hz "
         f"[source: {src_desc}]"
@@ -165,11 +189,24 @@ def main() -> None:
 
     period = 1.0 / max(1.0, args.fps)
     frame_index = 0
+    last_tracked = -1
     t0 = time.time()
+
+    # 3-point operator calibration (device source only): map the operator's neutral
+    # rest pose onto the G1's neutral stance. Trigger a (re)capture with the A+B+X+Y
+    # controller combo, mirroring gear_sonic's ThreePointPose.calibrate_now.
+    calibrator = ThreePointCalibrator() if (xrt is not None and args.headset_source == "devices") else None
+    calib_combo_last = False
+    if calibrator is not None:
+        print(
+            "[pico_publisher] 3-point calibration: stand in a neutral rest pose and press "
+            "A+B+X+Y on the controllers to (re)calibrate."
+        )
     try:
         while True:
             loop_start = time.time()
             vr3_pos = vr3_orn = None
+            loco_axes = loco_buttons = None
             if clip is not None:
                 n = clip["joints"].shape[0]
                 if args.no_loop and frame_index >= n:
@@ -180,6 +217,53 @@ def main() -> None:
                 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()
+            elif not args.fake and args.headset_source == "devices":
+                # Controller-state 3-point path: head + 2 controllers only, no PICO
+                # Motion Trackers / body tracking. Emits encode_mode-1 targets only;
+                # the SMPL whole-body window is left as a zero placeholder.
+                head = np.asarray(xrt.get_headset_pose(), np.float32)
+                lc = np.asarray(xrt.get_left_controller_pose(), np.float32)
+                rc = np.asarray(xrt.get_right_controller_pose(), np.float32)
+                stamp_ns = int(xrt.get_time_stamp_ns())
+                if head.shape != (7,) or lc.shape != (7,) or rc.shape != (7,):
+                    time.sleep(0.005)
+                    continue
+                last_tracked = int(sum(np.linalg.norm(d[3:7]) > 1e-6 for d in (head, lc, rc)))
+                # Empty SMPL window: this source drives encode_mode 1 only, so the
+                # consumer must not mistake it for a (zero) whole-body reference.
+                joints = np.zeros((0, 3), np.float32)
+                root_quat = None
+                root_transl = None
+                vr3_pos, vr3_orn = compute_3point_from_devices(head, lc, rc)
+                # Edge-triggered (re)calibration on the A+B+X+Y combo.
+                combo_now = bool(
+                    xrt.get_A_button() and xrt.get_B_button() and xrt.get_X_button() and xrt.get_Y_button()
+                )
+                if combo_now and not calib_combo_last:
+                    calibrator.capture(vr3_pos, vr3_orn)
+                    print("\n[pico_publisher] 3-point calibration captured (neutral pose).")
+                calib_combo_last = combo_now
+                vr3_pos, vr3_orn = calibrator.apply(vr3_pos, vr3_orn)
+                # Controller-stick locomotion (encode_mode 1, replicated): left/right
+                # sticks + A/B/X/Y. The all-four combo is reserved for calibration, so
+                # suppress the button pairs on that frame to avoid a spurious mode cycle.
+                la = np.asarray(xrt.get_left_axis(), np.float32).reshape(-1)
+                ra = np.asarray(xrt.get_right_axis(), np.float32).reshape(-1)
+                loco_axes = np.array([la[0], la[1], ra[0], ra[1]], np.float32)
+                btn = (
+                    np.zeros(4, np.float32)
+                    if combo_now
+                    else np.array(
+                        [
+                            float(xrt.get_A_button()),
+                            float(xrt.get_B_button()),
+                            float(xrt.get_X_button()),
+                            float(xrt.get_Y_button()),
+                        ],
+                        np.float32,
+                    )
+                )
+                loco_buttons = btn
             else:
                 if args.fake:
                     body_poses = _fake_body_poses(loop_start - t0)
@@ -190,6 +274,13 @@ def main() -> None:
                     if body_poses.shape != (24, 7):
                         time.sleep(0.005)
                         continue
+                # How many joints are actually being tracked (non-zero-norm quat).
+                # If this stays near 0, the headset isn't streaming body data (e.g.
+                # "Full body"/"Send" not enabled, trackers uncalibrated, or a test
+                # device) and the reference will be static regardless of your motion.
+                # In that case, use --headset-source devices for head+controllers only.
+                quat_norms = np.linalg.norm(body_poses[:, 3:7], axis=1)
+                last_tracked = int(np.count_nonzero(quat_norms > 1e-6))
                 out = fk.compute(body_poses)
                 joints = out["smpl_joints_local"]
                 root_quat = out["root_quat"]
@@ -207,11 +298,18 @@ def main() -> None:
                     root_transl=root_transl,
                     vr3_pos=vr3_pos,
                     vr3_orn=vr3_orn,
+                    loco_axes=loco_axes,
+                    loco_buttons=loco_buttons,
                 )
             )
             frame_index += 1
             if frame_index % int(max(1, args.fps)) == 0:
-                print(f"[pico_publisher] sent {frame_index} frames", end="\r")
+                denom = 3 if (not args.fake and clip is None and args.headset_source == "devices") else 24
+                unit = "devices" if denom == 3 else "joints"
+                extra = f" | tracked {last_tracked}/{denom} {unit}" if last_tracked >= 0 else ""
+                if calibrator is not None:
+                    extra += " | calibrated" if calibrator.is_calibrated else " | UNCALIBRATED"
+                print(f"[pico_publisher] sent {frame_index} frames{extra}", end="\r")
 
             dt = time.time() - loop_start
             if dt < period:

src/lerobot/teleoperators/pico_headset/smpl_constants.py

@@ -49,3 +49,13 @@ VR3_ORN_DIM = VR3_N_POINTS * 4  # 12 (3 x wxyz)
 # Flat action-dict keys: ``vr3_pos.0 .. vr3_pos.8`` and ``vr3_orn.0 .. vr3_orn.11``.
 VR3_POS_PREFIX = "vr3_pos."
 VR3_ORN_PREFIX = "vr3_orn."
+
+# ── Controller-stick locomotion (SONIC encode_mode 1, replicated exactly) ────
+# In the original 3-point teleop the same tick that sends the VR targets also drives
+# locomotion from the PICO controller sticks/buttons (left stick -> move, right stick
+# -> yaw, A+B / X+Y -> cycle locomotion mode). We forward that raw controller state so
+# the consumer's planner can steer walking/turning underneath the upper-body tracking.
+LOCO_N_AXES = 4  # [left_x, left_y, right_x, right_y]
+LOCO_N_BTN = 4  # [A, B, X, Y]
+LOCO_AXES_PREFIX = "loco_axes."
+LOCO_BTN_PREFIX = "loco_btn."

src/lerobot/teleoperators/pico_headset/smpl_fk.py

@@ -30,6 +30,7 @@ plus the root orientation quaternion and pelvis translation.
 Quaternions are scalar-first (w, x, y, z) unless noted.
 """
 
+from dataclasses import dataclass, field
 from pathlib import Path
 
 import numpy as np
@@ -234,6 +235,19 @@ _VR3_OFFSETS = [
 _UNITY_TO_ROBOT = np.array([[-1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 1.0, 0.0]])
 
 
+def _safe_quat(quats: np.ndarray) -> np.ndarray:
+    """Replace zero-norm quaternions with the scalar-last identity.
+
+    The headset reports ``[0, 0, 0, 0]`` for joints it isn't currently tracking;
+    ``scipy.Rotation.from_quat`` rejects zero-norm quaternions, so we substitute the
+    identity ``[0, 0, 0, 1]`` (no rotation) for those rows to keep FK robust.
+    """
+    quats = np.asarray(quats, np.float64).copy()
+    bad = np.linalg.norm(quats, axis=-1) < 1e-8
+    quats[bad] = (0.0, 0.0, 0.0, 1.0)  # scalar-last identity
+    return quats
+
+
 def compute_3point(body_poses_np: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     """Extract the SONIC 3-point VR targets from headset body poses.
 
@@ -261,9 +275,10 @@ def compute_3point(body_poses_np: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     # and rotation-offset-corrected.
     positions = np.zeros((4, 3), np.float64)
     rotations: list[R] = []
+    quats = _safe_quat(body[:, 3:7])
     for out_i, j in enumerate((0, *_VR3_JOINTS)):
         positions[out_i] = q @ body[j, :3]
-        rot = R.from_quat(body[j, 3:7]).as_matrix()  # scalar-last input
+        rot = R.from_quat(quats[j]).as_matrix()  # scalar-last input
         rotations.append(R.from_matrix(q @ rot @ q.T) * _VR3_OFFSETS[out_i])
 
     root_inv = rotations[0].inv()
@@ -276,6 +291,159 @@ def compute_3point(body_poses_np: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     return pos, orn
 
 
+# ── 3-point VR teleop from raw device poses (no body trackers) ───────────────
+# PICO Y-up (X-right, Y-up, Z-back) -> robot Z-up world. Ported verbatim from
+# gear_sonic's controller path (``decoupled_wbc`` ``PicoStreamer.R_HEADSET_TO_WORLD``).
+_HEADSET_TO_WORLD = np.array([[0.0, 0.0, -1.0], [-1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
+
+
+def _device_pose_to_world(pose: np.ndarray) -> tuple[np.ndarray, R]:
+    """Convert a raw (7,) device pose ``[x, y, z, qx, qy, qz, qw]`` (PICO Y-up frame)
+    to a Z-up world ``(position, Rotation)`` pair.
+
+    Handles the all-zero quaternion the SDK emits when a device is momentarily
+    untracked by substituting the identity, matching ``PicoStreamer._process_xr_pose``.
+    """
+    pose = np.asarray(pose, np.float64)
+    xyz = _HEADSET_TO_WORLD @ pose[:3]
+    quat = pose[3:7]  # scalar-last
+    if np.linalg.norm(quat) < 1e-8:
+        quat = np.array([0.0, 0.0, 0.0, 1.0])
+    rot = _HEADSET_TO_WORLD @ R.from_quat(quat).as_matrix() @ _HEADSET_TO_WORLD.T
+    return xyz, R.from_matrix(rot)
+
+
+def compute_3point_from_devices(
+    head_pose: np.ndarray,
+    left_pose: np.ndarray,
+    right_pose: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Build the SONIC 3-point VR targets from raw head + controller poses.
+
+    This is the controller-state path (no PICO Motion Trackers / body tracking
+    required): the 3 keypoints are the left controller, right controller, and the
+    headset, each expressed relative to the **headset yaw frame** — mirroring
+    gear_sonic's ``decoupled_wbc`` ``PicoStreamer._process_xr_pose`` (Y-up -> Z-up,
+    then de-headed by the headset yaw). The headset stands in for the "neck" point,
+    so its root-relative position is ~0 and its orientation carries pitch/roll.
+
+    Args:
+        head_pose, left_pose, right_pose: (7,) ``[x, y, z, qx, qy, qz, qw]`` device
+            poses (scalar-last) from ``xrt.get_headset_pose()`` /
+            ``xrt.get_left_controller_pose()`` / ``xrt.get_right_controller_pose()``.
+
+    Returns:
+        (pos, orn):
+          - pos: (9,) float32, headset-yaw-relative ``[x, y, z]`` for [l-wrist, r-wrist, head]
+          - orn: (12,) float32, headset-yaw-relative ``[w, x, y, z]`` for the same order
+    """
+    head_pos, head_rot = _device_pose_to_world(head_pose)
+    left_pos, left_rot = _device_pose_to_world(left_pose)
+    right_pos, right_rot = _device_pose_to_world(right_pose)
+
+    # De-head: cancel the headset yaw so targets are expressed in a heading-local frame.
+    head_yaw = head_rot.as_euler("xyz")[2]
+    inv_yaw = R.from_euler("z", -head_yaw)
+
+    points = ((left_pos, left_rot), (right_pos, right_rot), (head_pos, head_rot))
+    pos = np.zeros(VR3_POS_DIM, np.float32)
+    orn = np.zeros(VR3_ORN_DIM, np.float32)
+    for k, (p_pos, p_rot) in enumerate(points):
+        pos[k * 3 : k * 3 + 3] = inv_yaw.apply(p_pos - head_pos)
+        orn[k * 4 : k * 4 + 4] = (inv_yaw * p_rot).as_quat(scalar_first=True)  # wxyz
+    return pos, orn
+
+
+# ── operator calibration for the 3-point targets ────────────────────────────
+# G1 neutral (zero-q) key-frame targets, pelvis-relative [x, y, z] in metres, from
+# MuJoCo FK on g1_29dof with gear_sonic's local offsets (wrists +0.18x ∓0.025y,
+# torso +0.35z). These are the poses the operator's rest pose is mapped onto so the
+# robot starts at its neutral stance. Orientations are identity at neutral.
+_G1_NEUTRAL_WRIST_POS = np.array([[0.3798, 0.1237, 0.0952], [0.3798, -0.1237, 0.0952]], np.float64)
+# Neck reconstruction chain (mirrors ThreePointPose._apply_calibration): torso link
+# +0.05 z, then +0.35 along the neck's local Z.
+_NECK_TORSO_OFFSET_Z = 0.05
+_NECK_LINK_LENGTH = 0.35
+
+
+@dataclass
+class ThreePointCalibrator:
+    """Aligns raw 3-point VR targets to the G1's neutral stance.
+
+    Ports gear_sonic ``ThreePointPose._capture_calibration`` / ``_apply_calibration``:
+    on :meth:`capture` (operator holding a neutral rest pose) it records (a) the
+    inverse of the head/neck orientation, used to de-tilt all points to upright, and
+    (b) per-wrist position + orientation offsets that map the corrected rest pose onto
+    the fixed G1 neutral wrist targets. :meth:`apply` then transforms every subsequent
+    frame by those offsets, and reconstructs the head/neck position from the calibrated
+    neck orientation via the torso->neck kinematic chain.
+
+    All quaternions are scalar-first (w, x, y, z), matching :func:`compute_3point`.
+    """
+
+    _neck_quat_inv: R | None = field(default=None, init=False)
+    _wrist_pos_offset: np.ndarray | None = field(default=None, init=False)
+    _wrist_rot_offset: list[R] = field(default_factory=list, init=False)
+
+    @property
+    def is_calibrated(self) -> bool:
+        return self._neck_quat_inv is not None
+
+    def reset(self) -> None:
+        self._neck_quat_inv = None
+        self._wrist_pos_offset = None
+        self._wrist_rot_offset = []
+
+    def capture(self, pos: np.ndarray, orn: np.ndarray) -> None:
+        """Capture calibration offsets from a neutral-pose frame.
+
+        Args:
+            pos: (9,) root-relative ``[x, y, z]`` for [l-wrist, r-wrist, head].
+            orn: (12,) root-relative ``[w, x, y, z]`` for the same order.
+        """
+        pos = np.asarray(pos, np.float64).reshape(3, 3)
+        orn = np.asarray(orn, np.float64).reshape(3, 4)
+        neck_rot = R.from_quat(orn[2], scalar_first=True)
+        neck_inv = neck_rot.inv()
+
+        self._wrist_pos_offset = np.zeros((2, 3), np.float64)
+        self._wrist_rot_offset = []
+        for k in range(2):
+            corrected_pos = neck_inv.apply(pos[k])
+            corrected_rot = neck_inv * R.from_quat(orn[k], scalar_first=True)
+            self._wrist_pos_offset[k] = corrected_pos - _G1_NEUTRAL_WRIST_POS[k]
+            self._wrist_rot_offset.append(corrected_rot.inv())  # g1 neutral rot = identity
+        self._neck_quat_inv = neck_inv
+
+    def apply(self, pos: np.ndarray, orn: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        """Apply the stored calibration; returns calibrated ``(pos (9,), orn (12,))``.
+
+        A no-op (returns the inputs unchanged) until :meth:`capture` has been called.
+        """
+        if self._neck_quat_inv is None:
+            return (
+                np.asarray(pos, np.float32).reshape(-1),
+                np.asarray(orn, np.float32).reshape(-1),
+            )
+        pos = np.asarray(pos, np.float64).reshape(3, 3)
+        orn = np.asarray(orn, np.float64).reshape(3, 4)
+        neck_inv = self._neck_quat_inv
+
+        out_pos = np.zeros((3, 3), np.float64)
+        out_orn = np.zeros((3, 4), np.float64)
+        for k in range(2):  # wrists
+            out_pos[k] = neck_inv.apply(pos[k]) - self._wrist_pos_offset[k]
+            corrected_rot = neck_inv * R.from_quat(orn[k], scalar_first=True)
+            out_orn[k] = (self._wrist_rot_offset[k] * corrected_rot).as_quat(scalar_first=True)
+
+        # Head/neck: orientation de-tilted, position from the torso->neck chain.
+        neck_rot = neck_inv * R.from_quat(orn[2], scalar_first=True)
+        out_orn[2] = neck_rot.as_quat(scalar_first=True)
+        neck_z = neck_rot.apply([0.0, 0.0, 1.0])
+        out_pos[2] = np.array([0.0, 0.0, _NECK_TORSO_OFFSET_Z]) + _NECK_LINK_LENGTH * neck_z
+        return out_pos.reshape(-1).astype(np.float32), out_orn.reshape(-1).astype(np.float32)
+
+
 class SmplForwardKinematics:
     """Rest-skeleton SMPL forward kinematics (no mesh, no torch)."""
 
@@ -323,7 +491,7 @@ class SmplForwardKinematics:
 
         # Global joint rotations from the headset (scalar-last), with the SMPL
         # +180 deg-about-Y frame fix, converted to per-joint local axis-angle.
-        global_rots = R.from_quat(body_poses_np[:, 3:7]) * R.from_euler("y", 180, degrees=True)
+        global_rots = R.from_quat(_safe_quat(body_poses_np[:, 3:7])) * R.from_euler("y", 180, degrees=True)
         gm = global_rots.as_matrix()  # (24, 3, 3)
 
         local_aa = np.zeros((24, 3), np.float64)

src/lerobot/teleoperators/pico_headset/smpl_stream.py

@@ -41,7 +41,16 @@ from collections import deque
 import numpy as np
 import zmq
 
-from .smpl_constants import JOINT_DIM, N_JOINTS, SMPL_OBS_DIM, VR3_ORN_DIM, VR3_POS_DIM, WINDOW
+from .smpl_constants import (
+    JOINT_DIM,
+    LOCO_N_AXES,
+    LOCO_N_BTN,
+    N_JOINTS,
+    SMPL_OBS_DIM,
+    VR3_ORN_DIM,
+    VR3_POS_DIM,
+    WINDOW,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -95,6 +104,12 @@ class SmplStream:
         self.vr3_pos = np.zeros(VR3_POS_DIM, np.float32)
         self.vr3_orn = np.tile([1.0, 0.0, 0.0, 0.0], VR3_ORN_DIM // 4).astype(np.float32)
         self._got_vr3 = False
+        self._last_vr3_t = 0.0
+        # Latest controller-stick locomotion (encode_mode 1): [lx, ly, rx, ry] + [A,B,X,Y].
+        self.loco_axes = np.zeros(LOCO_N_AXES, np.float32)
+        self.loco_buttons = np.zeros(LOCO_N_BTN, np.float32)
+        self._got_loco = False
+        self._last_loco_t = 0.0
         self._last_index = -1
         self._last_recv_t = 0.0
         self._warned_stale = False
@@ -121,6 +136,29 @@ class SmplStream:
         """True once the producer has sent at least one 3-point VR frame."""
         return self._got_vr3
 
+    @property
+    def has_fresh_vr3(self) -> bool:
+        """True when a 3-point VR frame arrived within ``stale_after_s``.
+
+        Unlike :attr:`has_data`, this is independent of the SMPL window, so the
+        controller-state source (head + controllers only, empty SMPL) still drives
+        ``encode_mode 1`` without a whole-body reference.
+        """
+        if not self._got_vr3:
+            return False
+        if not self.stale_after_s:
+            return True
+        return (time.time() - self._last_vr3_t) <= self.stale_after_s
+
+    @property
+    def has_fresh_loco(self) -> bool:
+        """True when controller-stick locomotion arrived within ``stale_after_s``."""
+        if not self._got_loco:
+            return False
+        if not self.stale_after_s:
+            return True
+        return (time.time() - self._last_loco_t) <= self.stale_after_s
+
     @property
     def seconds_since_last(self) -> float:
         """Wall-clock seconds since the last real frame (inf before the first)."""
@@ -143,6 +181,9 @@ class SmplStream:
         self._buf.clear()
         self._got_first = False
         self._got_vr3 = False
+        self._last_vr3_t = 0.0
+        self._got_loco = False
+        self._last_loco_t = 0.0
 
     # -- core ----------------------------------------------------------------
     def _drain_latest(self) -> np.ndarray | None:
@@ -155,6 +196,28 @@ class SmplStream:
         while dict(self._poller.poll(0)).get(self._sock) == zmq.POLLIN:
             payload = self._sock.recv()
             data = json.loads(payload.decode("utf-8")).get("data", {})
+
+            # Sparse 3-point VR targets (encode_mode 1). Parsed independently of the
+            # SMPL window so the controller-state source (head + controllers, empty
+            # SMPL) is still handled.
+            vp = data.get("vr3_pos")
+            vo = data.get("vr3_orn")
+            if vp is not None and vo is not None and len(vp) == VR3_POS_DIM and len(vo) == VR3_ORN_DIM:
+                self.vr3_pos = np.asarray(vp, np.float32)
+                self.vr3_orn = np.asarray(vo, np.float32)
+                self._got_vr3 = True
+                self._last_vr3_t = time.time()
+
+            # Controller-stick locomotion (encode_mode 1), also independent of SMPL.
+            la = data.get("loco_axes")
+            lb = data.get("loco_buttons")
+            if la is not None and lb is not None and len(la) == LOCO_N_AXES and len(lb) == LOCO_N_BTN:
+                self.loco_axes = np.asarray(la, np.float32)
+                self.loco_buttons = np.asarray(lb, np.float32)
+                self._got_loco = True
+                self._last_loco_t = time.time()
+
+            # SMPL whole-body window (encode_mode 2), optional on this stream.
             joints = np.asarray(data.get("smpl_joints_local", []), np.float32)
             if joints.size != N_JOINTS * JOINT_DIM:
                 continue
@@ -166,12 +229,6 @@ class SmplStream:
             rt = data.get("root_transl")
             if rt is not None and len(rt) == 3:
                 self.root_transl = np.asarray(rt, np.float32)
-            vp = data.get("vr3_pos")
-            vo = data.get("vr3_orn")
-            if vp is not None and vo is not None and len(vp) == VR3_POS_DIM and len(vo) == VR3_ORN_DIM:
-                self.vr3_pos = np.asarray(vp, np.float32)
-                self.vr3_orn = np.asarray(vo, np.float32)
-                self._got_vr3 = True
         return frame
 
     def step(self) -> np.ndarray:

src/lerobot/utils/import_utils.py

@@ -60,8 +60,18 @@ def is_package_available(
                     # If the package can't be imported, it's not available
                     package_exists = False
             else:
-                # For packages other than "torch", don't attempt the fallback and set as not available
-                package_exists = False
+                # The distribution may be published under a name that differs from the
+                # import name (e.g. ``onnxruntime`` imports from ``onnxruntime-gpu`` /
+                # ``onnxruntime-silicon``). Resolve the import name to its actual
+                # distribution(s) and read the version from there before giving up.
+                try:
+                    dists = importlib.metadata.packages_distributions().get(import_name, [])
+                    if dists:
+                        package_version = importlib.metadata.version(dists[0])
+                    else:
+                        package_exists = False
+                except importlib.metadata.PackageNotFoundError:
+                    package_exists = False
         logging.debug(f"Detected {pkg_name} version: {package_version}")
     if return_version:
         return package_exists, package_version