diff --git a/src/lerobot/robots/unitree_g1/config_unitree_g1.py b/src/lerobot/robots/unitree_g1/config_unitree_g1.py index fc4f76768..e768ed94b 100644 --- a/src/lerobot/robots/unitree_g1/config_unitree_g1.py +++ b/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 diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py index 99a4bfd92..f8605fd65 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py +++ b/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) 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 b99d8bfaf..6fb91441d 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/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 ``.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: diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 3d862d4f0..2146e293b 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/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: diff --git a/src/lerobot/teleoperators/pico_headset/pico_headset.py b/src/lerobot/teleoperators/pico_headset/pico_headset.py index 048180e74..3dbba0404 100644 --- a/src/lerobot/teleoperators/pico_headset/pico_headset.py +++ b/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 diff --git a/src/lerobot/teleoperators/pico_headset/pico_publisher.py b/src/lerobot/teleoperators/pico_headset/pico_publisher.py index 88462a047..d918ac552 100644 --- a/src/lerobot/teleoperators/pico_headset/pico_publisher.py +++ b/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: diff --git a/src/lerobot/teleoperators/pico_headset/smpl_constants.py b/src/lerobot/teleoperators/pico_headset/smpl_constants.py index 5d2169e95..c19b1295d 100644 --- a/src/lerobot/teleoperators/pico_headset/smpl_constants.py +++ b/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." diff --git a/src/lerobot/teleoperators/pico_headset/smpl_fk.py b/src/lerobot/teleoperators/pico_headset/smpl_fk.py index e60ea2e51..2d1db4bef 100644 --- a/src/lerobot/teleoperators/pico_headset/smpl_fk.py +++ b/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) diff --git a/src/lerobot/teleoperators/pico_headset/smpl_stream.py b/src/lerobot/teleoperators/pico_headset/smpl_stream.py index 36f0c321e..756e78081 100644 --- a/src/lerobot/teleoperators/pico_headset/smpl_stream.py +++ b/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: diff --git a/src/lerobot/utils/import_utils.py b/src/lerobot/utils/import_utils.py index 974dd5b51..be808787e 100644 --- a/src/lerobot/utils/import_utils.py +++ b/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