mujoco_linkerbot/tools/run_middle_step_experiment.py

#!/usr/bin/env python3
"""O6 中指 step 实验：真机 / 仿真共用同一套控制时序。

时序（默认，两边完全一致）：
  3s 张开 → 3s 握紧 → 3s 张开 → 2s 余量录制

用法：
  export ROS_DOMAIN_ID=31

  # 真机（需 can0 UP + SDK，或脚本尝试自动启 SDK）
  python3 tools/run_middle_step_experiment.py real

  # 仿真（自带 MuJoCo，无需 run_sim.sh）
  python3 tools/run_middle_step_experiment.py sim

  # 两次跑完后出对比图（自动找目录下最新 real/sim csv）
  python3 tools/run_middle_step_experiment.py compare

  # 指定 csv
  python3 tools/run_middle_step_experiment.py compare \\
    --real-csv reports/O6_middle_ros/ros_real_....csv \\
    --sim-csv  reports/O6_middle_ros/ros_sim_....csv
"""

from __future__ import annotations

import argparse
import glob
import os
import signal
import subprocess
import sys
import threading
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
OUT_DIR = ROOT / "reports/O6_middle_ros"

# 与 middle_step_cmd_pub.py 保持一致
PRE_WAIT = 3.0
HOLD = 3.0
RECORD_AFTER = 2.0
RATE_HZ = 10.0


def record_duration() -> float:
    return PRE_WAIT + HOLD + HOLD + RECORD_AFTER + 1.5


def setup_ros_env():
    os.environ.setdefault("ROS_DOMAIN_ID", "31")
    os.environ.setdefault("ROS_LOCALHOST_ONLY", "0")


def setup_mujoco_pythonpath():
    venv_site = ROOT / "venv/lib/python3.12/site-packages"
    if venv_site.is_dir():
        os.environ["PYTHONPATH"] = f"{venv_site}{os.pathsep}{os.environ.get('PYTHONPATH', '')}"


def run_real_experiment(args):
    setup_ros_env()
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    dur = record_duration()

    rec_cmd = [
        sys.executable,
        str(ROOT / "tools/ros_real_record_qvt.py"),
        "--hand",
        args.hand,
        "--finger",
        args.finger,
        "--label",
        "ros_real",
        "--out",
        str(OUT_DIR),
        "--duration",
        str(dur),
    ]
    print(f"[real] 启动录制 {dur:.1f}s ...")
    rec = subprocess.Popen(rec_cmd, cwd=str(ROOT))
    time.sleep(1.0)

    pub_cmd = [
        sys.executable,
        str(ROOT / "tools/middle_step_cmd_pub.py"),
        "--hand",
        args.hand,
        "--finger",
        args.finger,
        "--pre-wait",
        str(PRE_WAIT),
        "--hold",
        str(HOLD),
        "--rate",
        str(RATE_HZ),
    ]
    print("[real] 发送统一 step 指令 ...")
    pub = subprocess.run(pub_cmd, cwd=str(ROOT))
    if pub.returncode != 0:
        rec.send_signal(signal.SIGINT)
        rec.wait()
        return pub.returncode

    print("[real] 等待录制结束 ...")
    try:
        rec.wait(timeout=dur + 5)
    except subprocess.TimeoutExpired:
        rec.send_signal(signal.SIGINT)
        rec.wait()
    return rec.returncode


def run_sim_experiment(args):
    setup_ros_env()
    setup_mujoco_pythonpath()
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    dur = record_duration()

    rec_cmd = [
        sys.executable,
        str(ROOT / "tools/ros_mujoco_record_qvt.py"),
        "--hand",
        args.hand,
        "--finger",
        args.finger,
        "--label",
        "ros_sim",
        "--out",
        str(OUT_DIR),
        "--record-duration",
        str(dur),
    ]
    if args.viewer:
        rec_cmd.append("--viewer")

    print(f"[sim] 启动 MuJoCo 录制 {dur:.1f}s ...")
    rec = subprocess.Popen(rec_cmd, cwd=str(ROOT))
    time.sleep(1.0)

    pub_cmd = [
        sys.executable,
        str(ROOT / "tools/middle_step_cmd_pub.py"),
        "--hand",
        args.hand,
        "--finger",
        args.finger,
        "--pre-wait",
        str(PRE_WAIT),
        "--hold",
        str(HOLD),
        "--rate",
        str(RATE_HZ),
    ]
    print("[sim] 发送统一 step 指令 ...")
    pub = subprocess.run(pub_cmd, cwd=str(ROOT))
    if pub.returncode != 0:
        rec.send_signal(signal.SIGINT)
        rec.wait()
        return pub.returncode

    print("[sim] 等待录制结束 ...")
    try:
        rec.wait(timeout=dur + 10)
    except subprocess.TimeoutExpired:
        rec.send_signal(signal.SIGINT)
        rec.wait()
    return rec.returncode


def latest_csv(pattern: str) -> Path | None:
    files = sorted(glob.glob(str(OUT_DIR / pattern)), key=os.path.getmtime)
    return Path(files[-1]) if files else None


def run_compare(args):
    real_csv = Path(args.real_csv) if args.real_csv else latest_csv("ros_real_*middle*.csv")
    sim_csv = Path(args.sim_csv) if args.sim_csv else latest_csv("ros_sim_*middle*.csv")
    if not real_csv or not real_csv.is_file():
        print("[error] 找不到真机 csv，先跑: python3 tools/run_middle_step_experiment.py real")
        return 1
    if not sim_csv or not sim_csv.is_file():
        print("[error] 找不到仿真 csv，先跑: python3 tools/run_middle_step_experiment.py sim")
        return 1

    out = OUT_DIR / "ros_real_vs_sim_middle_left_compare_qvt.png"
    cmd = [
        sys.executable,
        str(ROOT / "tools/plot_real_sim_compare.py"),
        str(real_csv),
        str(sim_csv),
        "-o",
        str(out),
    ]
    print(f"[compare] real={real_csv.name}")
    print(f"[compare] sim ={sim_csv.name}")
    return subprocess.call(cmd, cwd=str(ROOT))


def main():
    parser = argparse.ArgumentParser(description="O6 middle step experiment (real/sim/compare)")
    parser.add_argument(
        "mode",
        choices=["real", "sim", "compare"],
        help="real=真机录制 | sim=MuJoCo录制 | compare=对比图",
    )
    parser.add_argument("--hand", default="left")
    parser.add_argument("--finger", default="middle")
    parser.add_argument("--viewer", action="store_true", help="sim 模式打开 MuJoCo 窗口")
    parser.add_argument("--real-csv", default="")
    parser.add_argument("--sim-csv", default="")
    args = parser.parse_args()

    print("=" * 60)
    print("O6 middle step 实验时序（真机 / 仿真相同）")
    print(f"  pre_wait={PRE_WAIT}s  hold={HOLD}s  record_after={RECORD_AFTER}s")
    print(f"  ROS_DOMAIN_ID={os.environ.get('ROS_DOMAIN_ID', '31')}")
    print("=" * 60)

    if args.mode == "real":
        print("\n[提示] 真机请先: sudo ip link set can0 up type can bitrate 1000000")
        print("       并运行 ./run_real_hand.sh（或本脚本前自行启动 SDK）\n")
        return run_real_experiment(args) or 0
    if args.mode == "sim":
        return run_sim_experiment(args) or 0
    return run_compare(args)


if __name__ == "__main__":
    raise SystemExit(main())