lerobot/examples/unitree_g1/sonic.py

#!/usr/bin/env python
"""
SONIC planner with full mode control.

Keyboard controls:
    N / P      - next / previous motion set
    1-8        - select mode within current set
    WASD       - movement direction
    Q / E      - rotate facing left / right
    9 / 0      - decrease / increase speed
    - / =      - decrease / increase height
    R          - force replan
    M          - toggle SMPL motion playback <-> locomotion (needs --motion-file)
    Space      - emergency stop -> IDLE
    Esc        - quit

Gamepad controls (Unitree wireless controller):
    Left stick Y  - speed (forward = fast, back = stop)
    Left stick X  - movement direction (offset from facing)
    Right stick X - facing direction (incremental rotation)
    Right stick Y - height (up = tall 0.8m, down = low 0.1m)
    Buttons       - unused (mode selection is keyboard-only)

For teleop integration use --robot.controller=SonicWholeBodyController instead.
"""

import argparse
import contextlib
import faulthandler
import gc
import os
import sys
import tempfile
import time

import numpy as np
from dataset_motion import DatasetJointMotion
from motion_loader import SmplMotion
from smpl_stream import DEFAULT_SMPL_HOST, DEFAULT_SMPL_PORT, SmplStream

from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
from lerobot.robots.unitree_g1.controllers.sonic_pipeline import (
    CONTROL_DT,
    DEFAULT_ANGLES,
    ENCODER_UPDATE_EVERY,
    LM,
    MOTION_SETS,
    MUJOCO_TO_ISAACLAB,
    RawKeyboard,
    compute_kp_kd,
    drain_keyboard,
)
from lerobot.robots.unitree_g1.controllers.sonic_whole_body import SonicRuntime
from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1


def _load_joint_trajectory(controller, joints: np.ndarray, velocities: np.ndarray) -> None:
    """Load a (T, 29) joint reference into the controller for encode_mode=0 tracking.

    The dataset provides joints in Unitree/G1_29_JointIndex order, but the SONIC
    encoder reference (motion_joint_positions) is in IsaacLab order. Reorder here.
    """
    joints = np.asarray(joints)[:, MUJOCO_TO_ISAACLAB]
    velocities = np.asarray(velocities)[:, MUJOCO_TO_ISAACLAB]
    t = joints.shape[0]
    with controller.motion_lock:
        cap = controller.motion_joint_positions.shape[0]
        if t > cap:
            controller.motion_joint_positions = np.zeros((t, 29), np.float64)
            controller.motion_joint_velocities = np.zeros((t, 29), np.float64)
            controller.motion_body_quats = np.zeros((t, 4), np.float64)
            controller.motion_body_quats[:, 0] = 1.0
            controller.motion_body_pos = np.zeros((t, 3), np.float64)
        controller.motion_joint_positions[:t] = joints
        controller.motion_joint_velocities[:t] = velocities
        controller.motion_body_quats[:t, 0] = 1.0
        controller.motion_body_quats[:t, 1:] = 0.0
        controller.motion_body_pos[:t] = 0.0
        controller.motion_timesteps = t
        controller.ref_cursor = 0
        controller.init_ref_quat = np.array([1, 0, 0, 0], np.float64)
    controller.encode_mode = 0
    controller.playing = True
    controller.first_motion = True  # triggers heading init on first obs
    controller.reinit_heading = True


def _tick_replay(runtime, obs: dict) -> dict:
    """One control tick for dataset replay: encode/decode + advance, no planner."""
    if not obs:
        runtime.step += 1
        return {}
    do_enc = runtime.step % ENCODER_UPDATE_EVERY == 0
    action = runtime.controller.step(obs, update_encoder=do_enc, debug=False)
    runtime.controller.advance_cursor()
    runtime.step += 1
    return action


def main():
    parser = argparse.ArgumentParser(description="SONIC planner with keyboard + gamepad control")
    parser.add_argument(
        "--ip",
        type=str,
        default=None,
        help="Robot IP for real hardware (e.g. 192.168.123.164). Omit for simulation.",
    )
    parser.add_argument(
        "--log-csv",
        action="store_true",
        help="Write /tmp/sonic_pose_log.csv (disabled by default for teleop perf)",
    )
    parser.add_argument(
        "--cpu",
        action="store_true",
        help="Force CPU ONNX Runtime (skip CUDA even if onnxruntime-gpu is installed)",
    )
    parser.add_argument(
        "--headless", action="store_true", help="Ignored for sim (stock UnitreeG1 uses hub MuJoCo defaults)"
    )
    parser.add_argument(
        "--gamepad",
        action="store_true",
        help="Read Unitree wireless gamepad in sim (default: keyboard-only in sim)",
    )
    parser.add_argument(
        "--keyboard-only", action="store_true", help="Ignore wireless gamepad (terminal keyboard only)"
    )
    parser.add_argument(
        "--motion-file",
        type=str,
        default=None,
        help="Play an SMPL motion clip (.npz) via SONIC whole-body mode "
        "(encode_mode=2) instead of locomotion planning.",
    )
    parser.add_argument(
        "--no-loop", action="store_true", help="With --motion-file, play once instead of looping"
    )
    parser.add_argument(
        "--smpl-stream",
        action="store_true",
        help="Use the live rt/smpl headset stream as the reference motion "
        "(SONIC whole-body, encode_mode=2), instead of a --motion-file clip.",
    )
    parser.add_argument(
        "--smpl-host",
        type=str,
        default=DEFAULT_SMPL_HOST,
        help=f"Host publishing rt/smpl (default: {DEFAULT_SMPL_HOST})",
    )
    parser.add_argument(
        "--smpl-port",
        type=int,
        default=DEFAULT_SMPL_PORT,
        help=f"Port for the rt/smpl stream (default: {DEFAULT_SMPL_PORT})",
    )
    parser.add_argument(
        "--replay-dataset",
        type=str,
        default=None,
        help="Replay a LeRobot dataset episode's 29-DoF observation.state as a SONIC "
        "encode_mode=0 joint reference (e.g. BitRobot/HIW-500-lerobot).",
    )
    parser.add_argument(
        "--episode", type=int, default=0, help="Episode index for --replay-dataset (default: 0)"
    )
    parser.add_argument(
        "--replay-frames",
        type=int,
        default=None,
        help="Cap the number of replayed frames (default: whole episode)",
    )
    args = parser.parse_args()

    exclusive = [bool(args.smpl_stream), bool(args.motion_file), bool(args.replay_dataset)]
    if sum(exclusive) > 1:
        parser.error("--smpl-stream, --motion-file and --replay-dataset are mutually exclusive")

    # Surface native crashes (onnxruntime / mujoco) with a real traceback, and
    # avoid losing buffered diagnostics if the process dies mid-loop.
    faulthandler.enable()
    with contextlib.suppress(Exception):
        sys.stdout.reconfigure(line_buffering=True)

    print("=" * 60)
    print("SONIC planner - full mode control")
    print("  N/P  cycle sets | 1-8 select mode | WASD move")
    print("  Q/E  rotate     | 9/0 speed       | -/= height")
    print("  R    replan     | Space IDLE       | Esc quit")
    if args.ip:
        print(f"  Robot IP: {args.ip}")
    else:
        print("  Mode: simulation")
    print("=" * 60 + "\n")

    cfg = UnitreeG1Config(controller=None)  # full-body SONIC; standalone loop owns publish
    if args.ip:
        cfg.is_simulation = False
        cfg.robot_ip = args.ip
    else:
        cfg.is_simulation = True
        if args.headless:
            print("[Note] --headless ignored: sim uses stock UnitreeG1 + hub env")
    robot = UnitreeG1(cfg)
    robot.connect()
    kp, kd = compute_kp_kd()
    robot.kp = kp.copy()
    robot.kd = kd.copy()

    runtime = SonicRuntime(force_cpu=args.cpu)
    controller = runtime.controller
    ms = runtime.ms

    # --replay-dataset drives SONIC mode 0: load a recorded 29-DoF joint trajectory
    # into the controller's reference buffers and let the policy try to track it,
    # bypassing the locomotion planner (which would otherwise overwrite the ref).
    replay = None
    if args.replay_dataset:
        replay = DatasetJointMotion(
            args.replay_dataset, episode=args.episode, max_frames=args.replay_frames
        )
        _load_joint_trajectory(controller, replay.joints, replay.velocities)
        dur = replay.num_frames / replay.fps
        print(f"\n[Replay] {args.replay_dataset} episode {args.episode} -> SONIC mode 0")
        print(
            f"  frames={replay.num_frames} fps={replay.fps:.0f} duration={dur:.1f}s "
            f"(src {replay.src_fps:.0f} fps, encode_mode=0, planner bypassed)"
        )

    motion = None
    if args.smpl_stream:
        motion = SmplStream(host=args.smpl_host, port=args.smpl_port)
        controller.smpl_motion = motion  # lets 'M' key toggle streaming
        controller.encode_mode = 2  # start in SONIC whole-body SMPL imitation
        print(f"\n[Motion] live SMPL stream (rt/smpl @ {args.smpl_host}:{args.smpl_port})")
        print("  Source: pico_manager_thread_server.py --manager on the publisher host.")
        print("  encode_mode=2. Press 'M' to toggle SMPL stream <-> locomotion at runtime.")
    elif args.motion_file:
        motion = SmplMotion(args.motion_file, loop=not args.no_loop)
        controller.smpl_motion = motion  # lets 'M' key toggle playback
        controller.encode_mode = 2  # start in SONIC whole-body SMPL imitation
        dur = motion.num_frames / motion.fps
        print(f"\n[Motion] SMPL whole-body playback: {args.motion_file}")
        print(
            f"  frames={motion.num_frames} fps={motion.fps:.1f} "
            f"duration={dur:.1f}s loop={not args.no_loop} encode_mode=2"
        )
        print("  Press 'M' to toggle SMPL playback <-> locomotion at runtime.")

    runtime.controller.print_input_diagnostics()

    print(f"\nStarting: {MOTION_SETS[0][0]} (default mode: {LM(ms.mode).name})")
    [print(f"  {i + 1}: {m.name}") for i, m in enumerate(MOTION_SETS[0][1])]
    print(
        "\n[Ready] Click THIS terminal, then W/A/S/D to move. 1-6 change mode, 9/0 speed, Esc quit.\n",
        flush=True,
    )

    # Sim hub publishes wireless_remote bytes that can fight terminal WASD.
    base_joystick = not args.keyboard_only and (args.gamepad or args.ip is not None)

    with RawKeyboard() as kb:
        try:
            gc.disable()
            gc_timer = 0.0
            robot.reset(CONTROL_DT, DEFAULT_ANGLES)
            time.sleep(1.0)

            last_status = time.time() - 2.1
            loop_t, enc_t, dec_t, obs_t, act_t = [], [], [], [], []
            slow_n = blend_n = 0
            stall_src = ""
            did_blend = False
            t_start = time.time()

            log_path = os.path.join(tempfile.gettempdir(), "sonic_pose_log.csv")
            jnames = [m.name for m in G1_29_JointIndex]
            log_ctx = open(log_path, "w") if args.log_csv else None  # noqa: SIM115
            if log_ctx:
                log_ctx.write(
                    "t,step,cursor,ts,blend,mode,"
                    + ",".join(f"q{i}" for i in range(29))
                    + ","
                    + ",".join(f"ref{i}" for i in range(29))
                    + ","
                    + ",".join(f"act{i}" for i in range(29))
                    + ",delta_max,action_norm,token_norm\n"
                )

            try:
                while not robot._shutdown_event.is_set():
                    t0 = time.time()
                    if drain_keyboard(kb, ms, controller):
                        break

                    obs = robot.get_observation()
                    t_obs = time.time()
                    obs_t.append(1000 * (t_obs - t0))
                    if not obs:
                        runtime.tick({}, use_joystick=False)
                        time.sleep(max(0.0, CONTROL_DT - (time.time() - t0)))
                        continue

                    # Dataset replay: SONIC tracks the recorded 29-DoF joint clip.
                    if replay is not None:
                        step_before = runtime.step
                        t_step = time.time()
                        action = _tick_replay(runtime, obs)
                        step_ms = 1000 * (time.time() - t_step)
                        (enc_t if step_before % 5 == 0 else dec_t).append(step_ms)
                        robot.send_action(action)
                        if controller.ref_cursor >= controller.motion_timesteps - 1:
                            print("\n[Replay] episode finished")
                            break
                        now = time.time()
                        loop_t.append(1000 * (now - t0))
                        time.sleep(max(0.0, CONTROL_DT - (now - t0)))
                        continue

                    # SMPL playback only while in whole-body mode; 'M' toggles it.
                    motion_active = motion is not None and controller.encode_mode == 2
                    if motion_active:
                        controller.smpl_joints_10frame_step1 = motion.step()
                        if motion.done:
                            print("\n[Motion] clip finished")
                            break

                    step_before = runtime.step
                    t_step = time.time()
                    action = runtime.tick(obs, use_joystick=base_joystick and not motion_active)
                    step_ms = 1000 * (time.time() - t_step)
                    do_enc = step_before % 5 == 0
                    (enc_t if do_enc else dec_t).append(step_ms)

                    t_act = time.time()
                    robot.send_action(action)
                    act_t.append(1000 * (time.time() - t_act))

                    if log_ctx and runtime.step % 5 == 0:
                        t_rel = time.time() - t_start
                        q_r = np.array([obs.get(f"{n}.q", 0) for n in jnames])
                        a_v = np.array([action.get(f"{n}.q", 0) for n in jnames])
                        cur, ts = controller.ref_cursor, controller.motion_timesteps
                        q_ref = (
                            controller.motion_joint_positions[min(cur, ts - 1)] if ts > 0 else np.zeros(29)
                        )
                        log_ctx.write(
                            f"{t_rel:.4f},{runtime.step},{cur},{ts},{int(did_blend)},{ms.mode},"
                            + ",".join(f"{v:.6f}" for v in q_r)
                            + ","
                            + ",".join(f"{v:.6f}" for v in q_ref)
                            + ","
                            + ",".join(f"{v:.6f}" for v in a_v)
                            + ","
                            + f"{np.max(np.abs(a_v - q_r)):.6f},"
                            f"{np.linalg.norm(a_v):.6f},"
                            f"{np.linalg.norm(controller.token):.6f}\n"
                        )
                        did_blend = False

                    now = time.time()
                    loop_ms = 1000 * (now - t0)
                    if loop_ms > 50:
                        stall_src = (
                            f"[STALL] {loop_ms:.0f}ms: "
                            f"obs={obs_t[-1]:.0f} step={step_ms:.0f} act={act_t[-1]:.0f}"
                        )
                    if loop_ms > CONTROL_DT * 1500:
                        slow_n += 1

                    if now - last_status > 2.0:

                        def _avg(lst):
                            return sum(lst) / len(lst) if lst else 0

                        hz = 1000 / _avg(loop_t) if _avg(loop_t) else 0
                        print(
                            f"\r  {ms.status_line()}  step={runtime.step} "
                            f"ref={controller.ref_cursor}/{controller.motion_timesteps} "
                            f"loop={_avg(loop_t):.1f}ms(max={max(loop_t, default=0):.1f}) hz={hz:.0f} "
                            f"enc={_avg(enc_t):.1f} dec={_avg(dec_t):.1f} obs={_avg(obs_t):.1f} "
                            f"slow={slow_n} blends={blend_n}",
                            end="",
                            flush=True,
                        )
                        if stall_src:
                            print(f"\n  {stall_src}")
                            stall_src = ""
                        last_status = now
                        loop_t, enc_t, dec_t, obs_t, act_t = [], [], [], [], []
                        slow_n = blend_n = 0

                    gc_timer += CONTROL_DT
                    if gc_timer >= 10.0:
                        gc.collect()
                        gc_timer = 0.0
                    loop_t.append(loop_ms)
                    time.sleep(max(0.0, CONTROL_DT - (time.time() - t0)))
            finally:
                if log_ctx:
                    log_ctx.close()

        except KeyboardInterrupt:
            pass
        finally:
            gc.enable()
            if args.log_csv:
                print(f"\n[Log] Saved to {log_path}")
            if motion is not None and hasattr(motion, "close"):
                motion.close()
            runtime.shutdown()
            print("\nStopping...")
            if robot.is_connected:
                robot.disconnect()
            print("Done.")


if __name__ == "__main__":
    main()