tests/scripts/dance_compare.py

"""Compare dance playback: Python's play_move vs Marionette-style stream.

Plays the same dance from pollen-robotics/reachy-mini-dances-library
via two different playback strategies, captures the actual robot
state during each run, dumps CSVs.

Strategies:
  - python_native : reachy.play_move(move, play_frequency=100). Internally
    that's three separate set_target_* commands per tick at 100 Hz.
  - marionette_style : one combined set_target() per tick at 50 Hz, with
    lead compensation (90 ms antennas, 205 ms head) applied to the
    sampled-from-motion time. This mirrors what the JS Marionette app
    does over WebRTC.

N runs per strategy; soft return to base between runs.
"""

import csv
import math
import sys
import threading
import time
from pathlib import Path

import numpy as np

from reachy_mini import ReachyMini
from reachy_mini.motion.recorded_move import RecordedMove, RecordedMoves

DATASET = "pollen-robotics/reachy-mini-dances-library"
DANCE = sys.argv[1] if len(sys.argv) > 1 else "head_tilt_roll"
N_RUNS = 5
CAPTURE_HZ = 50.0
CAPTURE_PERIOD = 1.0 / CAPTURE_HZ
HEAD_LEAD_S = 0.205
ANTENNA_LEAD_S = 0.090
MARIONETTE_STREAM_HZ = 50.0
MARIONETTE_STREAM_PERIOD = 1.0 / MARIONETTE_STREAM_HZ

OUT_DIR = Path("/Users/remi/Downloads/dance-compare")
OUT_DIR.mkdir(parents=True, exist_ok=True)


def sample_move_at(move: RecordedMove, t: float):
    """Evaluate the move at t (clipped) — returns (head 4x4, antennas, body_yaw)."""
    t_clipped = min(max(0.0, t), move.timestamps[-1] - 1e-3)
    return move.evaluate(t_clipped)


def capture_loop(mini: ReachyMini, samples: list, stop_event: threading.Event, t0: float):
    """Background thread: poll cached state at CAPTURE_HZ, append (t, head_flat16, antennas)."""
    tick = 0
    while not stop_event.is_set():
        now = time.perf_counter()
        t = now - t0
        try:
            head = mini.get_current_head_pose()      # 4x4 ndarray
            ant = mini.get_present_antenna_joint_positions()
            samples.append((t, head.flatten().tolist(), list(ant)))
        except Exception as e:
            print(f"  capture read error at t={t:.3f}: {e}")
        tick += 1
        next_tick = t0 + tick * CAPTURE_PERIOD
        sleep_for = next_tick - time.perf_counter()
        if sleep_for > 0:
            time.sleep(sleep_for)


def write_csv(path: Path, move: RecordedMove, samples: list):
    """CSV columns: t_s, cmd_head[16], cmd_ant_l, cmd_ant_r, act_head[16], act_ant_l, act_ant_r.

    Commanded values are evaluated analytically from the move at each
    measured t — same way the JS benchmark CSVs are structured.
    """
    headers = ["t_s"]
    headers += [f"cmd_h{i}" for i in range(16)]
    headers += ["cmd_ant_l_deg", "cmd_ant_r_deg"]
    headers += [f"act_h{i}" for i in range(16)]
    headers += ["act_ant_l_deg", "act_ant_r_deg"]

    with open(path, "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(headers)
        for t, head_actual_flat, ant_actual in samples:
            if t < 0 or t > move.duration + 0.2:
                continue
            head_cmd, ant_cmd, _ = sample_move_at(move, t)
            row = [f"{t:.4f}"]
            row += [f"{v:.6f}" for v in head_cmd.flatten()]
            row += [f"{math.degrees(ant_cmd[0]):.3f}", f"{math.degrees(ant_cmd[1]):.3f}"]
            row += [f"{v:.6f}" for v in head_actual_flat]
            row += [f"{math.degrees(ant_actual[0]):.3f}", f"{math.degrees(ant_actual[1]):.3f}"]
            w.writerow(row)


def reset_to_base(mini: ReachyMini):
    """Smooth goto base + small settle. Run between each measurement run."""
    base = np.eye(4)
    mini.goto_target(base, antennas=[-0.1745, 0.1745], duration=0.8)
    time.sleep(1.0)


def run_python_native(mini: ReachyMini, move: RecordedMove, run_idx: int):
    """play_move at 100 Hz; the SDK's stock playback path."""
    samples = []
    stop = threading.Event()
    t0 = time.perf_counter()
    capture_thread = threading.Thread(
        target=capture_loop, args=(mini, samples, stop, t0), daemon=True,
    )
    capture_thread.start()
    try:
        mini.play_move(move, play_frequency=100.0, initial_goto_duration=0.0)
        # Let the actuator settle for a beat so the tail of the motion
        # appears in capture.
        time.sleep(0.3)
    finally:
        stop.set()
        capture_thread.join(timeout=1.0)
    out = OUT_DIR / f"{DANCE}-python_native-run{run_idx:02d}.csv"
    write_csv(out, move, samples)
    return out


def run_marionette_style(mini: ReachyMini, move: RecordedMove, run_idx: int):
    """One combined set_target per tick at 50 Hz with lead compensation."""
    samples = []
    stop = threading.Event()
    t0 = time.perf_counter()
    capture_thread = threading.Thread(
        target=capture_loop, args=(mini, samples, stop, t0), daemon=True,
    )
    capture_thread.start()
    try:
        tick = 0
        while True:
            now = time.perf_counter()
            t = now - t0
            if t > move.duration:
                break
            t_head = min(t + HEAD_LEAD_S, move.duration - 1e-3)
            t_ant = min(t + ANTENNA_LEAD_S, move.duration - 1e-3)
            head_for_head, _, body_yaw_h = sample_move_at(move, t_head)
            _, ant_for_ant, _ = sample_move_at(move, t_ant)
            mini.set_target(
                head=head_for_head,
                antennas=list(ant_for_ant),
                body_yaw=float(body_yaw_h),
            )
            tick += 1
            next_tick = t0 + tick * MARIONETTE_STREAM_PERIOD
            sleep_for = next_tick - time.perf_counter()
            if sleep_for > 0:
                time.sleep(sleep_for)
        time.sleep(0.3)
    finally:
        stop.set()
        capture_thread.join(timeout=1.0)
    out = OUT_DIR / f"{DANCE}-marionette_style-run{run_idx:02d}.csv"
    write_csv(out, move, samples)
    return out


def main():
    print(f"Loading dance: {DANCE} from {DATASET}")
    library = RecordedMoves(DATASET)
    if DANCE not in library.list_moves():
        print(f"Dance '{DANCE}' not found. Available: {library.list_moves()}")
        sys.exit(2)
    move = library.get(DANCE)
    print(f"  duration: {move.duration:.2f} s, frames: {len(move.timestamps)}")

    print("Connecting to robot…")
    with ReachyMini() as mini:
        print("Connected. Initial reset to base…")
        reset_to_base(mini)

        all_outputs = []
        # Interleave the two methods so any drift in conditions (WiFi,
        # robot temp) affects both equally rather than biasing one.
        for run_idx in range(1, N_RUNS + 1):
            print(f"\n--- Run {run_idx}/{N_RUNS}: python_native ---")
            out = run_python_native(mini, move, run_idx)
            print(f"  → {out.name}")
            all_outputs.append(out)
            reset_to_base(mini)

            print(f"--- Run {run_idx}/{N_RUNS}: marionette_style ---")
            out = run_marionette_style(mini, move, run_idx)
            print(f"  → {out.name}")
            all_outputs.append(out)
            reset_to_base(mini)

    print(f"\nWrote {len(all_outputs)} CSVs to {OUT_DIR}")


if __name__ == "__main__":
    main()