run_final.py

"""
Final comparison: BC vs ProMP-open vs ProMP-FB vs ProMP-CL (closed-loop) vs Hybrid.
All budgets, multiple seeds, standardized setup.
"""
import numpy as np
import torch
from env.arm2d import Arm2DEnv
from experts.demo_generator import generate_dataset
from models.bc import BCMLP, train_bc, evaluate_bc
from models.promp import ProMPPredictor, train_promp, evaluate_promp, RBFBasis
from models.promp import evaluate_promp_feedback_corrected, evaluate_promp_closed_loop
from models.hybrid import BCResidual, train_hybrid, evaluate_hybrid_plan_once

SEED = 42
BUDGETS = [5, 10, 25, 50, 100, 200]
N_BFS = 15
N_EVAL = 100
N_SEEDS = 5
EPOCHS = 500

np.random.seed(SEED)
torch.manual_seed(SEED)

env = Arm2DEnv()
basis = RBFBasis(n_bfs=N_BFS)

labels = {"bc": "BC", "po": "ProMP-open", "pfb": "ProMP-FB",
          "cl": "ProMP-CL", "hyb": "Hybrid"}
methods = list(labels)

results = {k: {b: {"s": [], "d": []} for b in BUDGETS} for k in methods}

for n_demos in BUDGETS:
    print(f"\n{'='*60}")
    print(f"Budget: {n_demos} demos  ({N_SEEDS} seeds)")
    print(f"{'='*60}")

    for seed in range(N_SEEDS):
        np.random.seed(SEED + seed)
        torch.manual_seed(SEED + seed)
        noise = [0.03, 0.06, 0.09]
        data = generate_dataset(env, n_demos, noise_levels=noise * (n_demos // 3 + 1))
        if len(data) == 0:
            continue

        # BC
        m = BCMLP(obs_dim=env.obs_dim, act_dim=env.act_dim)
        train_bc(m, data, epochs=EPOCHS, verbose=False)
        sr, dm, _ = evaluate_bc(m, env, n_episodes=N_EVAL)
        results["bc"][n_demos]["s"].append(sr); results["bc"][n_demos]["d"].append(dm)

        # ProMP open-loop
        m_po = ProMPPredictor(obs_dim=env.obs_dim, n_bfs=N_BFS, n_joints=env.act_dim)
        train_promp(m_po, data, basis, epochs=EPOCHS, verbose=False, closed_loop=False)
        sr, dm, _ = evaluate_promp(m_po, env, basis, n_episodes=N_EVAL)
        results["po"][n_demos]["s"].append(sr); results["po"][n_demos]["d"].append(dm)

        # ProMP feedback-corrected (plan once + PID)
        m_pfb = ProMPPredictor(obs_dim=env.obs_dim, n_bfs=N_BFS, n_joints=env.act_dim)
        train_promp(m_pfb, data, basis, epochs=EPOCHS, verbose=False, closed_loop=False)
        sr, dm, _ = evaluate_promp_feedback_corrected(m_pfb, env, basis, n_episodes=N_EVAL)
        results["pfb"][n_demos]["s"].append(sr); results["pfb"][n_demos]["d"].append(dm)

        # ProMP closed-loop (remaining-traj + phase, per-step replan)
        m_cl = ProMPPredictor(obs_dim=env.obs_dim, n_bfs=N_BFS, n_joints=env.act_dim, with_phase=True)
        train_promp(m_cl, data, basis, epochs=EPOCHS, verbose=False, closed_loop=True)
        sr, dm, _ = evaluate_promp_closed_loop(m_cl, env, basis, n_episodes=N_EVAL)
        results["cl"][n_demos]["s"].append(sr); results["cl"][n_demos]["d"].append(dm)

        # Hybrid: plan-once + BC residual
        m_ph = ProMPPredictor(obs_dim=env.obs_dim, n_bfs=N_BFS, n_joints=env.act_dim)
        m_res = BCResidual(obs_dim=env.obs_dim, act_dim=env.act_dim, with_phase=True)
        train_hybrid(m_ph, m_res, data, basis, epochs=EPOCHS)
        sr, dm, _ = evaluate_hybrid_plan_once(m_ph, m_res, env, basis, n_episodes=N_EVAL)
        results["hyb"][n_demos]["s"].append(sr); results["hyb"][n_demos]["d"].append(dm)

    # Per-budget print
    for k in methods:
        s = np.mean(results[k][n_demos]["s"])
        s_std = np.std(results[k][n_demos]["s"])
        d = np.mean(results[k][n_demos]["d"])
        d_std = np.std(results[k][n_demos]["d"])
        print(f"  {labels[k]:15s} succ={s:.2f}±{s_std:.2f}  dist={d:.3f}±{d_std:.3f}")

# Final summary
print("\n" + "=" * 70)
print("SUCCESS RATES")
header = f"{'Method':15s}"
for b in BUDGETS: header += f" | {b:>2d}dem"
print(header + " | Scaling")
print("-" * 70)
for k in methods:
    line = f"{labels[k]:15s}"
    vals = [np.mean(results[k][b]["s"]) for b in BUDGETS]
    for v in vals:
        line += f" | {v:.2f}"
    scaling = (vals[-1] - vals[0]) / vals[0] if vals[0] > 0 else float('inf')
    line += f" | {scaling:+.1f}x"
    print(line)

print("\n" + "=" * 70)
print("DISTANCES")
header = f"{'Method':15s}"
for b in BUDGETS: header += f" | {b:>2d}dem"
print(header)
print("-" * 70)
for k in methods:
    line = f"{labels[k]:15s}"
    for b in BUDGETS:
        d = np.mean(results[k][b]["d"])
        line += f" | {d:.3f}"
    print(line)