run_e5_e6_e7.py

"""
E5: Analytic ProMP conditioning (Gaussian conditioning on target) — training-free baseline.
E6: Analytic IK + PID — sanity upper bound, no learning.
E7: Multi-threshold success rates (1/2/4 cm) for CL and PositionBC.
"""
import numpy as np
import torch
import time

from env.arm2d import Arm2DEnv
from experts.demo_generator import generate_dataset, ik
from models.promp import RBFBasis, ProMPPredictor, train_promp, decode_trajectory
from models.promp import evaluate_promp_closed_loop, compute_promp_weights

SEED = 42
N_BFS = 15
N_EVAL = 100
EPOCHS = 500
KP, KD = 80.0, 15.0


def run_e6(env, f):
    """Analytic IK + PD: solve target to joint angles, track with PD. No learning."""
    log = lambda *a: (print(*a, flush=True), f.write(" ".join(map(str, a)) + "\n"), f.flush())
    log("\n" + "=" * 70)
    log("E6: ANALYTIC IK + PD (no learning) sanity baseline")
    log("=" * 70)
    for n_episodes in [100]:
        successes = 0
        dists = []
        for _ in range(n_episodes):
            s = env.reset()
            q_goal = ik(env.target, env.l1, env.l2)
            done = False
            for t in range(200):
                q_des = q_goal  # constant goal
                u = KP * (q_des - env.q) - KD * env.qd
                s, _, done, info = env.step(u)
                if done:
                    break
            successes += 1 if info["success"] else 0
            dists.append(info["dist"])
        log(f"  IK+PD: success {successes/n_episodes:.3f}, dist {np.mean(dists):.3f}")


def run_e7(env, f):
    """Multi-threshold success for CL and PositionBC at 50 demos."""
    from run_e1_full import PositionBC, train_position_bc, evaluate_position_bc  # reuse
    log = lambda *a: (print(*a, flush=True), f.write(" ".join(map(str, a)) + "\n"), f.flush())
    log("\n" + "=" * 70)
    log("E7: MULTI-THRESHOLD SUCCESS (1/2/4 cm) at 50 demos")
    log("=" * 70)
    basis = RBFBasis(N_BFS)
    thresh = [0.01, 0.02, 0.04]
    for method in ["CL", "PosBC"]:
        log(f"\n  --- {method} ---")
        for seed in range(5):
            np.random.seed(SEED + seed)
            torch.manual_seed(SEED + seed)
            data = generate_dataset(env, 50, noise_levels=[0.03, 0.06, 0.09] * 17)
            if method == "CL":
                m = ProMPPredictor(10, N_BFS, 2, with_phase=True)
                train_promp(m, data, basis, epochs=EPOCHS, verbose=False, closed_loop=True)
            else:
                m = PositionBC(10, 2)
                train_position_bc(m, data, epochs=EPOCHS, lookahead=10)
            # custom eval collecting final distances
            m.eval()
            dists = []
            with torch.no_grad():
                for _ in range(N_EVAL):
                    s = env.reset()
                    for t in range(200):
                        s_t = torch.FloatTensor(s).unsqueeze(0)
                        if method == "CL":
                            w_flat = m(s_t, t / 200).squeeze(0).numpy()
                            w = w_flat.reshape(2, N_BFS)
                            traj = decode_trajectory(w, basis, 200)
                            q_des = traj[t]
                        else:
                            q_des = m(s_t, t / 200).squeeze(0).numpy()
                        u = KP * (q_des - env.q) - KD * env.qd
                        s, _, done, info = env.step(u)
                        if done:
                            break
                    dists.append(info["dist"])
            s1 = np.mean([d < 0.01 for d in dists])
            s2 = np.mean([d < 0.02 for d in dists])
            s4 = np.mean([d < 0.04 for d in dists])
            log(f"  seed{seed}: <1cm {s1:.3f} | <2cm {s2:.3f} | <4cm {s4:.3f} | mean_dist {np.mean(dists):.3f}")


def run_e5(env, f):
    """Analytic ProMP conditioning: fit Gaussian over demo weights, condition on target via regression."""
    log = lambda *a: (print(*a, flush=True), f.write(" ".join(map(str, a)) + "\n"), f.flush())
    log("\n" + "=" * 70)
    log("E5: ANALYTIC PROMP CONDITIONING (training-free) — fit weights, condition on target via linear regression")
    log("=" * 70)
    basis = RBFBasis(N_BFS)
    for b in [25, 50, 200]:
        log(f"\n  --- Budget {b} demos ---")
        for seed in range(5):
            np.random.seed(SEED + seed)
            torch.manual_seed(SEED + seed)
            data = generate_dataset(env, b, noise_levels=[0.03, 0.06, 0.09] * (b // 3 + 1))
            # Build feature->weight regression from demos
            # features: target (x,y), init q (from first state)
            Xr, Yw = [], []
            for d in data:
                states = d["states"]
                q0 = [np.arctan2(states[0][1], states[0][0]), np.arctan2(states[0][3], states[0][2])]
                Xr.append(np.concatenate([d["target"], q0]))
                cos_q1, sin_q1 = states[:, 0], states[:, 1]
                cos_q2, sin_q2 = states[:, 2], states[:, 3]
                q1 = np.arctan2(sin_q1, cos_q1)
                q2 = np.arctan2(sin_q2, cos_q2)
                w = compute_promp_weights(np.stack([q1, q2], 1), basis).ravel()
                Yw.append(w)
            Xr, Yw = np.array(Xr), np.array(Yw)
            # linear regression W = X A, A = (X^T X)^-1 X^T Y
            A = np.linalg.pinv(Xr.T @ Xr + 1e-4 * np.eye(Xr.shape[1])) @ Xr.T @ Yw
            # eval
            successes = 0
            dists = []
            for _ in range(N_EVAL):
                s = env.reset()
                q0 = [np.arctan2(s[1], s[0]), np.arctan2(s[3], s[2])]
                feat = np.concatenate([env.target, q0])
                w_flat = A.T @ feat
                w = w_flat.reshape(2, N_BFS)
                traj = decode_trajectory(w, basis, 200)
                done = False
                for t in range(200):
                    q_des = traj[t]
                    u = KP * (q_des - env.q) - KD * env.qd
                    s, _, done, info = env.step(u)
                    if done:
                        break
                successes += 1 if info["success"] else 0
                dists.append(info["dist"])
            log(f"  seed{seed}: analytic-ProMP success {successes/N_EVAL:.3f}, dist {np.mean(dists):.3f}")


def main():
    with open("results_e5e6e7.txt", "w") as f:
        f.write(f"E5/E6/E7 RESULTS — {time.strftime('%Y-%m-%d %H:%M')}\n")
        env = Arm2DEnv()
        run_e6(env, f)
        run_e7(env, f)
        run_e5(env, f)
    print("\nDone. results_e5e6e7.txt")


if __name__ == "__main__":
    main()