run_generalization.py

"""
Generalization test: train on targets in right half-plane, test on left half-plane.
Measures whether trajectory-level (ProMP) generalizes better than BC.
"""
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
from models.hybrid import BCResidual, train_hybrid, evaluate_hybrid_plan_once

SEED = 42
N_BFS = 15
N_EVAL = 100
N_DEMOS = 50  # fixed budget
N_SEEDS = 5
EPOCHS = 200

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", "hyb": "Hybrid"}
methods = list(labels)

# Modified target sampling for generalization test
def sample_right_target(env):
    """Target in right half-plane: theta in [-pi/2, pi/2]."""
    r = np.random.uniform(0.1, env.l1 + env.l2 - 0.05)
    theta = np.random.uniform(-np.pi/2, np.pi/2)
    return np.array([r * np.cos(theta), r * np.sin(theta)])

def sample_left_target(env):
    """Target in left half-plane: theta in [pi/2, 3pi/2]."""
    r = np.random.uniform(0.1, env.l1 + env.l2 - 0.05)
    theta = np.random.uniform(np.pi/2, 3*np.pi/2)
    return np.array([r * np.cos(theta), r * np.sin(theta)])

# Results: {method: {"ind": {"s": [...], "d": [...]}, "ood": {"s": [...], "d": [...]}}}
results = {k: {"ind": {"s": [], "d": []}, "ood": {"s": [], "d": []}} for k in methods}

print("="*70)
print("Generalization test: train on RIGHT targets, test on LEFT targets")
print(f"Budget: {N_DEMOS} demos, {N_SEEDS} seeds")
print("="*70)

for seed in range(N_SEEDS):
    np.random.seed(SEED + seed)
    torch.manual_seed(SEED + seed)
    noise = [0.03, 0.06, 0.09]

    # Override target sampler for training data
    orig_sampler = env._sample_reachable_target
    env._sample_reachable_target = lambda: sample_right_target(env)
    data = generate_dataset(env, N_DEMOS, noise_levels=noise * (N_DEMOS // 3 + 1))
    env._sample_reachable_target = orig_sampler  # restore

    if len(data) == 0:
        continue

    print(f"\nSeed {seed}: {len(data)} right-half demos")

    # --- In-distribution eval (right targets) ---
    # BC
    m_bc = BCMLP(obs_dim=env.obs_dim, act_dim=env.act_dim)
    train_bc(m_bc, data, epochs=EPOCHS, verbose=False)
    sr, dm, _ = evaluate_bc(m_bc, env, n_episodes=N_EVAL)
    results["bc"]["ind"]["s"].append(sr); results["bc"]["ind"]["d"].append(dm)

    # ProMP open
    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"]["ind"]["s"].append(sr); results["po"]["ind"]["d"].append(dm)

    # ProMP-FB
    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"]["ind"]["s"].append(sr); results["pfb"]["ind"]["d"].append(dm)

    # Hybrid
    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"]["ind"]["s"].append(sr); results["hyb"]["ind"]["d"].append(dm)

    # OOD eval (left targets)
    env._sample_reachable_target = lambda: sample_left_target(env)

    sr, dm, _ = evaluate_bc(m_bc, env, n_episodes=N_EVAL)
    results["bc"]["ood"]["s"].append(sr); results["bc"]["ood"]["d"].append(dm)

    sr, dm, _ = evaluate_promp(m_po, env, basis, n_episodes=N_EVAL)
    results["po"]["ood"]["s"].append(sr); results["po"]["ood"]["d"].append(dm)

    sr, dm, _ = evaluate_promp_feedback_corrected(m_pfb, env, basis, n_episodes=N_EVAL)
    results["pfb"]["ood"]["s"].append(sr); results["pfb"]["ood"]["d"].append(dm)

    sr, dm, _ = evaluate_hybrid_plan_once(m_ph, m_res, env, basis, n_episodes=N_EVAL)
    results["hyb"]["ood"]["s"].append(sr); results["hyb"]["ood"]["d"].append(dm)

    env._sample_reachable_target = orig_sampler

    # Print seed results
    for k in methods:
        ind_s = results[k]["ind"]["s"][-1]
        ood_s = results[k]["ood"]["s"][-1]
        gap = ind_s - ood_s
        print(f"  {labels[k]:15s} in-dist={ind_s:.2f}  ood={ood_s:.2f}  gap={gap:.2f}")

# Final summary
print("\n" + "="*70)
print("GENERALIZATION SUMMARY")
print(f"{'Method':15s} | {'In-dist':>12} | {'OOD':>12} | {'Gap':>8}")
print("-"*70)
for k in methods:
    ind_s = np.mean(results[k]["ind"]["s"])
    ood_s = np.mean(results[k]["ood"]["s"])
    ind_d = np.mean(results[k]["ind"]["d"])
    ood_d = np.mean(results[k]["ood"]["d"])
    ind_std = np.std(results[k]["ind"]["s"])
    ood_std = np.std(results[k]["ood"]["s"])
    print(f"{labels[k]:15s} | {ind_s:.2f}±{ind_std:.2f} | {ood_s:.2f}±{ood_std:.2f} | {ind_s-ood_s:.2f}")
print()
print(f"{'Method':15s} | {'Dist-in':>12} | {'Dist-ood':>12}")
print("-"*70)
for k in methods:
    ind_d = np.mean(results[k]["ind"]["d"])
    ood_d = np.mean(results[k]["ood"]["d"])
    print(f"{labels[k]:15s} | {ind_d:.3f} | {ood_d:.3f}")