add 3 point calibration + waist coupling, remote controller and smoothed motion

src/lerobot/robots/unitree_g1/README.md

@@ -10,11 +10,17 @@ deploy stack (no `gear_sonic`/torch dependency).
 
 Selected with `--robot.controller=<ClassName>`:
 
-| Controller                     | Purpose                                                             |
-| ------------------------------ | ------------------------------------------------------------------- |
-| `SonicWholeBodyController`     | SONIC whole-body: locomotion, keyboard, and SMPL imitation (mode 2) |
-| `GrootLocomotionController`    | GR00T locomotion policy                                             |
-| `HolosomaLocomotionController` | Holosoma locomotion policy                                          |
+| Controller                     | Purpose                                                                               |
+| ------------------------------ | ------------------------------------------------------------------------------------- |
+| `SonicWholeBodyController`     | SONIC whole-body: locomotion (mode 0), 3-point VR teleop (mode 1), SMPL imitation (mode 2) |
+| `GrootLocomotionController`    | GR00T locomotion policy                                                               |
+| `HolosomaLocomotionController` | Holosoma locomotion policy                                                            |
+
+On startup the controller **interpolates** from the robot's measured pose into the
+policy's commanded target over ~3 s (no snap), and on disconnect (Ctrl-C) it performs a
+**graceful damped settle** — holding pose while ramping stiffness to zero over
+`--robot.graceful_stop_s` (default 1.5 s) instead of going instantly limp. Both apply in
+every mode.
 
 ## Requirements
 
@@ -50,15 +56,36 @@ lerobot-teleoperate \
 Controls: `WASD` move · `Q`/`E` turn · `1`–`8` mode · `9`/`0` speed · `-`/`=` height ·
 `R` replan · `Space` emergency-stop.
 
-**PICO headset teleop** (live SMPL whole-body):
+**PICO headset teleop — SMPL whole-body** (mode 2, needs PICO Motion Trackers):
 
 ```bash
+# 1) publisher (streams rt/smpl from full-body tracking)
+python -m lerobot.teleoperators.pico_headset.pico_publisher --fps 50
+# 2) controller
 lerobot-teleoperate \
   --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
   --teleop.type=pico_headset
 ```
 
-This requires the XRoboToolkit stack — see below.
+**PICO headset teleop — 3-point VR** (mode 1, head + 2 controllers only, **no trackers**):
+
+```bash
+# 1) publisher (head + controllers -> 3-point targets + stick locomotion)
+python -m lerobot.teleoperators.pico_headset.pico_publisher --fps 50 --headset-source devices
+# 2) controller
+lerobot-teleoperate \
+  --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
+  --teleop.type=pico_headset --teleop.mode=vr3
+```
+
+3-point controls: left stick move · right stick X turn · right stick Y height ·
+`A`+`B` / `X`+`Y` cycle locomotion mode (walk/run/squat/kneel/…) · hands+head track the
+upper body. **Calibration**: stand in a neutral rest pose and press `A`+`B`+`X`+`Y` — the
+publisher status line flips from `UNCALIBRATED` to `calibrated`. This maps your rest pose
+onto the G1's neutral stance and is required before the hands track well; the SMPL
+(mode 2) path is self-calibrating and needs no such step.
+
+Both require the XRoboToolkit stack — see below.
 
 ## PICO headset / XRoboToolkit install
 
@@ -75,7 +102,9 @@ Summary:
    - Jetson (aarch64): the arm64 `.deb`.
    - Windows (x64): the Windows PC Service build.
 2. **PICO app** — install `XRoboToolkit-PICO-*.apk` on the headset (see the guide),
-   enable Developer Mode, pair/calibrate the ankle motion trackers.
+   enable Developer Mode. For **SMPL whole-body** (mode 2) you also need the PICO Motion
+   Trackers paired/calibrated and "Full body" enabled; for **3-point** (mode 1,
+   `--headset-source devices`) only Head + Controller + Send are required — no trackers.
 3. **`xrobotoolkit_sdk`** — a pybind11/CMake build (not a pip package), from
    [`XRoboToolkit-PC-Service-Pybind`](https://github.com/XR-Robotics/XRoboToolkit-PC-Service-Pybind):
    - Linux x86_64: `pip install pybind11 cmake` then `bash setup_ubuntu.sh` (or the
@@ -85,7 +114,7 @@ Summary:
    - Windows x64: `pip install pybind11` then `setup_windows.bat` (needs git + an
      MSVC/CMake toolchain; uses the prebuilt `PXREARobotSDK.dll`/`.lib`).
 4. Connect PICO and workstation to the **same Wi-Fi**, open the XRoboToolkit app, enter
-   the PC IP, and enable Head/Controller/Full-body/Send.
+   the PC IP, and enable Head/Controller/Send (plus Full-body for SMPL mode 2).
 
 ### Platform support
 
@@ -112,7 +141,22 @@ python -m lerobot.teleoperators.pico_headset.pico_publisher --motion-file <clip>
 ## Notes
 
 - SMPL **root motion** into the mode-2 anchor is opt-in (`SonicWholeBodyController(enable_smpl_root=True)`);
-  it is off by default because an unsmoothed 30 Hz→50 Hz root trajectory can spike the
-  base acceleration.
+  it stays off by default (untested on hardware). When enabled, the per-frame root quat is
+  spherically smoothed (`root_smoothing_alpha`, default 0.15) before it reaches the anchor,
+  which removes the base-acceleration spikes the raw 30 Hz→50 Hz trajectory used to cause.
 - Direct `rt/smpl` subscription without the pico teleoperator is available via
   `SonicWholeBodyController(enable_smpl_stream=True, smpl_host=..., smpl_port=...)`.
+- 3-point (mode 1) uses the **headset-yaw frame** as its reference and the `A`+`B`+`X`+`Y`
+  calibration to align to the G1 neutral stance. Calibration maps the operator's rest pose
+  onto the G1's **standing** (`default_angles`) wrist/neck key-frame poses (position **and**
+  orientation) computed by FK — the `default_angles` stand-in for gear_sonic's live
+  measured-q recalibration, since the robot holds `default_angles` at calibration time.
+  Re-aligning the arms only (preserving the neck level) is available via the calibrator's
+  `recalibrate_wrists()`.
+- 3-point **locomotion** from the PICO sticks follows gear_sonic's `PlannerLoop` exactly:
+  a yaw accumulator on the right stick and **mode-dependent speed curves** on the left
+  (slow `0.1+0.5·mag`, run `1.5+3·mag`, walk = planner default). Stick signs replicate
+  gear_sonic's `get_controller_axes` usage (forward `+ly`, strafe `-lx`, turn `-rx`); since
+  the publisher forwards the same raw SDK axes, this is the correct convention by construction.
+- Startup interpolation and the graceful-stop settle are mode-agnostic; set
+  `--robot.graceful_stop_s=0` to restore the old instant zero-torque on disconnect.

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

@@ -1285,6 +1285,66 @@ def apply_joystick_axes(lx, ly, rx, ry, ms, controller=None):
         ms.height = max(0.1, min(1.0, (ms.height if ms.height >= 0 else DEFAULT_HEIGHT) + step))
 
 
+# gear_sonic PlannerLoop joystick constants (pico_manager_thread_server).
+PICO_YAW_GAIN = 1.5  # rad/s of facing rotation at full right-stick deflection
+PICO_JOY_DEADZONE = 0.15  # gear_sonic JOYSTICK_DEADZONE (distinct from the remote's)
+
+
+def apply_pico_loco_axes(lx, ly, rx, ry, ms, dt=CONTROL_DT):
+    """Map PICO controller sticks onto ``MovementState`` — the gear_sonic 3-point path.
+
+    Ports ``pico_manager_thread_server.PlannerLoop`` (encode_mode 1) rather than the
+    keyboard/remote-parity :func:`apply_joystick_axes`, so speed follows gear_sonic's
+    mode-dependent curves and facing comes from a yaw accumulator:
+
+    - Right stick X -> yaw accumulator (``+= PICO_YAW_GAIN * -rx * dt``), dead-zoned.
+    - Left stick -> magnitude ``mag`` (dead-zone-rescaled to 0..1) driving a
+      *mode-dependent* speed curve (slow ``0.1+0.5·mag``, run ``1.5+3·mag``, walk =
+      default), and a facing-rotated global movement vector -> ``movement_angle``.
+
+    Signs replicate gear_sonic's ``get_controller_axes`` usage exactly (forward = +ly,
+    strafe = -lx, turn = -rx), which resolves the PICO axis-sign question: because the
+    publisher forwards the same raw SDK axes gear_sonic reads, matching its sign usage
+    is by definition the correct mapping. ``ms.speed`` is later re-clamped into each
+    mode's valid range by :func:`clamp_mode_params`.
+    """
+    lx, ly, rx, ry = float(lx), float(ly), float(rx), float(ry)
+
+    # Facing: gear_sonic YawAccumulator — integrate only outside the dead zone.
+    if abs(rx) >= PICO_JOY_DEADZONE:
+        ms.facing_angle += PICO_YAW_GAIN * (-rx) * dt
+
+    raw_mag = min(1.0, math.hypot(lx, ly))
+    if raw_mag < PICO_JOY_DEADZONE:
+        ms.has_movement = False
+        ms.speed = -1.0
+        ms.joy_prev_active = False
+        return
+
+    mag = min(1.0, (raw_mag - PICO_JOY_DEADZONE) / (1.0 - PICO_JOY_DEADZONE))
+    m = LM(ms.mode)
+    if m == LM.SLOW_WALK:
+        ms.speed = 0.1 + 0.5 * mag  # 0.1 .. 0.6
+    elif m == LM.WALK:
+        ms.speed = -1.0  # planner default
+    elif m == LM.RUN:
+        ms.speed = 1.5 + 3.0 * mag  # 1.5 .. 4.5
+    else:
+        ms.speed = mag
+
+    # Facing-rotated movement vector: rotation_facing @ [-lx, ly] * scale.
+    scale = mag / raw_mag
+    mlx, mly = -lx * scale, ly * scale
+    fx, fy = math.cos(ms.facing_angle), math.sin(ms.facing_angle)
+    gx = -fy * mlx + fx * mly
+    gy = fx * mlx + fy * mly
+    if not ms.joy_prev_active:
+        ms.needs_replan = True
+    ms.has_movement = True
+    ms.joy_prev_active = True
+    ms.movement_angle = math.atan2(gy, gx)
+
+
 def process_joystick(obs, ms, controller=None):
     """Drive ``MovementState`` from the G1 wireless remote in ``obs``."""
     wr = obs.get("wireless_remote")

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

@@ -53,7 +53,7 @@ from .sonic_pipeline import (
     MovementState,
     PlannerController,
     SonicPlanner,
-    apply_joystick_axes,
+    apply_pico_loco_axes,
     clamp_mode_params,
     compute_kp_kd,
     make_ort_session_options,
@@ -240,6 +240,7 @@ class SonicWholeBodyController:
         force_cpu: bool = False,
         *,
         enable_smpl_root: bool = False,
+        root_smoothing_alpha: float = 0.15,
         enable_smpl_stream: bool = False,
         smpl_host: str | None = None,
         smpl_port: int | None = None,
@@ -252,10 +253,16 @@ class SonicWholeBodyController:
         self.ms = self._runtime.ms
 
         # When True, the per-frame SMPL root quaternion steers the mode-2 anchor.
-        # Off by default: feeding it currently produces root-acceleration spikes
-        # (NaN QACC at DOF 0) until the reference root trajectory is smoothed and
-        # rate-matched (30 Hz dataset -> 50 Hz control).
+        # Off by default: even with smoothing this changes the anchor/heading and is
+        # untested on hardware, so it stays opt-in. When enabled, the raw per-frame
+        # root quat (from a 30 Hz dataset resampled to a 50 Hz loop) is spherically
+        # smoothed by :meth:`_smooth_root_quat` before it reaches the anchor, which
+        # removes the root-acceleration spikes (NaN QACC at DOF 0) the unsmoothed
+        # trajectory caused. ``root_smoothing_alpha`` in (0, 1] is the per-tick blend
+        # toward the incoming quat (smaller = smoother/laggier, 1 = no smoothing).
         self.enable_smpl_root = enable_smpl_root
+        self._root_smoothing_alpha = float(np.clip(root_smoothing_alpha, 1e-3, 1.0))
+        self._smoothed_root_quat: np.ndarray | None = None
 
         # Tracks the previous keyboard held-key set so discrete controls (mode,
         # motion set, replan, e-stop, WASD direction) fire once per physical press
@@ -413,12 +420,12 @@ class SonicWholeBodyController:
         """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.
+        movement/facing/speed via :func:`apply_pico_loco_axes` (the faithful gear_sonic
+        yaw-accumulator + mode-dependent speed curves, not the keyboard-parity map), 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)
+        apply_pico_loco_axes(lx, ly, rx, ry, self.ms)
 
         # Mode cycling: step linearly through the LocomotionMode enum (A+B = next,
         # X+Y = previous), exactly like gear_sonic's PlannerLoop — so the operator can
@@ -438,6 +445,34 @@ class SonicWholeBodyController:
             logger.info("SONIC 3-point: locomotion mode -> %s", LM(self.ms.mode).name)
         self._prev_loco_mode_pair = (ab_now, xy_now)
 
+    def _smooth_root_quat(self, root_quat: np.ndarray | None) -> np.ndarray | None:
+        """Spherically smooth the per-frame SMPL root quaternion (mode-2 anchor).
+
+        The reference root trajectory is authored at ~30 Hz and consumed at 50 Hz, so
+        the raw per-tick quat steps unevenly and injects root-acceleration spikes into
+        the anchor. This keeps a persistent estimate and shortest-path nlerp-slerps it
+        toward each incoming (unit) quat by ``root_smoothing_alpha``, yielding a
+        continuous, rate-matched heading. Quaternions are scalar-first (w, x, y, z).
+        Returns ``None`` (leaving the anchor self-driven) for an invalid/zero input.
+        """
+        if root_quat is None:
+            self._smoothed_root_quat = None
+            return None
+        q = np.asarray(root_quat, np.float64)
+        n = np.linalg.norm(q)
+        if n < 1e-8:
+            return self._smoothed_root_quat
+        q = q / n
+        if self._smoothed_root_quat is None:
+            self._smoothed_root_quat = q
+        else:
+            prev = self._smoothed_root_quat
+            if np.dot(prev, q) < 0.0:  # shortest-path: quats double-cover SO(3)
+                q = -q
+            blended = prev + self._root_smoothing_alpha * (q - prev)
+            self._smoothed_root_quat = blended / (np.linalg.norm(blended) + 1e-12)
+        return self._smoothed_root_quat.astype(np.float32)
+
     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
@@ -500,10 +535,12 @@ class SonicWholeBodyController:
                 self._enter_wholebody()
             self.controller.smpl_joints_10frame_step1 = smpl
             # Root orientation steers the mode-2 anchor/heading, but only when
-            # explicitly enabled (see enable_smpl_root); otherwise the anchor stays
-            # 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
+            # explicitly enabled (see enable_smpl_root); the raw per-frame quat is
+            # spherically smoothed first so the 30->50 Hz resample doesn't spike the
+            # anchor. Disabled -> anchor stays self-driven.
+            self.controller.smpl_root_quat = (
+                self._smooth_root_quat(root_quat) if self.enable_smpl_root else 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
@@ -525,6 +562,7 @@ class SonicWholeBodyController:
             # 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._smoothed_root_quat = None
                 self._exit_wholebody()
             out = self._runtime.tick(obs, debug=False)
 
@@ -536,6 +574,7 @@ class SonicWholeBodyController:
         self._runtime.reset()
         self._init_step = 0  # re-run the startup blend after a reset
         self._start_pose = {}
+        self._smoothed_root_quat = None
 
     def shutdown(self):
         if self._smpl_stream is not None:

src/lerobot/teleoperators/pico_headset/smpl_fk.py

@@ -355,11 +355,28 @@ def compute_3point_from_devices(
 
 
 # ── 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)
+# G1 neutral key-frame targets, pelvis-relative [x, y, z] in metres, from MuJoCo FK on
+# g1_29dof at the robot's *standing* configuration (``default_angles`` — the pose the
+# robot actually holds at calibration time), with gear_sonic's local key-frame offsets
+# (``G1_KEY_FRAME_OFFSETS``: wrists +0.18x ∓0.025y, torso +0.35z). These are the poses
+# the operator's rest pose is mapped onto so the handoff starts at the robot's neutral
+# stance. This is the fixed-``default_angles`` stand-in for gear_sonic's
+# ``get_g1_key_frame_poses(q=measured_q)`` (see :meth:`ThreePointCalibrator.capture`):
+# since the robot stands at ``default_angles`` after the startup ramp, its measured q
+# equals this configuration, so these constants are the measured-q targets for the
+# nominal case. (Per-frame *live* measured q would need a reverse controller->publisher
+# feedback channel; not wired.)
+_G1_NEUTRAL_WRIST_POS = np.array(
+    [[0.2232, 0.2177, -0.1555], [0.2232, -0.2177, -0.1555]], np.float64
+)
+# Wrist neutral orientations (scalar-first w, x, y, z) at ``default_angles`` — NOT
+# identity: the wrists are rolled/pitched in the standing pose. Matches gear_sonic
+# using ``g1_lwrist_rot`` / ``g1_rwrist_rot`` from FK (not identity) as the rotation
+# calibration reference.
+_G1_NEUTRAL_WRIST_ROT = [
+    R.from_quat([0.9168, 0.0897, 0.3864, 0.0463], scalar_first=True),  # left
+    R.from_quat([0.9168, -0.0897, 0.3864, -0.0463], scalar_first=True),  # right
+]
 # 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
@@ -387,24 +404,41 @@ class ThreePointCalibrator:
 
     @property
     def is_calibrated(self) -> bool:
-        return self._neck_quat_inv is not None
+        return self._neck_quat_inv is not None and self._wrist_pos_offset is not None
 
     def reset(self) -> None:
         self._neck_quat_inv = None
         self._wrist_pos_offset = None
         self._wrist_rot_offset = []
 
+    def recalibrate_wrists(self) -> None:
+        """Clear only the wrist offsets, preserving the neck calibration.
+
+        Mirrors gear_sonic ``ThreePointPose.reset_with_measured_q``: the next
+        :meth:`capture` recomputes the wrist offsets (against the G1 neutral targets)
+        while keeping the already-captured neck orientation, so re-aligning the arms
+        doesn't force the operator to re-level their head.
+        """
+        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.
 
+        Neck calibration is captured once and then preserved across subsequent
+        captures (matching gear_sonic's ``if self._calibration_neck_quat_inv is
+        None``); call :meth:`reset` to clear it or :meth:`recalibrate_wrists` to
+        re-align only the arms.
+
         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()
+        if self._neck_quat_inv is None:
+            self._neck_quat_inv = R.from_quat(orn[2], scalar_first=True).inv()
+        neck_inv = self._neck_quat_inv
 
         self._wrist_pos_offset = np.zeros((2, 3), np.float64)
         self._wrist_rot_offset = []
@@ -412,15 +446,16 @@ class ThreePointCalibrator:
             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
+            # rot_offset maps the corrected rest orientation onto the G1 neutral wrist
+            # orientation: calibrated = rot_offset * (neck_inv * current).
+            self._wrist_rot_offset.append(_G1_NEUTRAL_WRIST_ROT[k] * corrected_rot.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:
+        if self._neck_quat_inv is None or self._wrist_pos_offset is None:
             return (
                 np.asarray(pos, np.float32).reshape(-1),
                 np.asarray(orn, np.float32).reshape(-1),