run_decisive_experiments.py

"""
Decisive experiments for thesis revision — addressing peer review CRITICAL issues.

E1: Position-space BC baseline (separate trajectory representation from control interface)
E2: Raised actuator (acceleration) limit (test whether saturation is cause or symptom)
E3: Smoothness ablation with 20 seeds (rescue or downgrade the +7.6pp claim)
E7: Multi-threshold success rates (1/2/4 cm)

All experiments reuse the existing env/demo/models code. Same protocol as run_final.py
unless stated. Output is printed AND written to results_decisive.txt.
"""
import numpy as np
import torch
import time

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 (
    RBFBasis, ProMPPredictor, train_promp,
    evaluate_promp, evaluate_promp_closed_loop,
    evaluate_promp_cl_torque_aware, decode_trajectory,
)
from models.hybrid import BCResidual, train_hybrid, evaluate_hybrid_plan_once

import torch.nn as nn

SEED = 42
BUDGETS = [5, 10, 25, 50, 100, 200]
N_BFS = 15
N_EVAL = 100
EPOCHS = 500
KP, KD = 80.0, 15.0


def log(msg, f):
    print(msg)
    f.write(msg + "\n")
    f.flush()


def make_ds(env, n_demos, seed, noise_levels=(0.03, 0.06, 0.09)):
    np.random.seed(SEED + seed)
    torch.manual_seed(SEED + seed)
    return generate_dataset(env, n_demos, noise_levels=noise_levels * (n_demos // 3 + 1))


# ----------------------------------------------------------------------
# E1: Position-space BC
# ----------------------------------------------------------------------

class PositionBC(nn.Module):
    """MLP: state (+phase) -> absolute joint position q_des. Executed via same PD as ProMP."""

    def __init__(self, obs_dim, act_dim, hidden_dim=256, with_phase=True):
        super().__init__()
        self.with_phase = with_phase
        in_dim = obs_dim + (1 if with_phase else 0)
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, act_dim),
        )

    def forward(self, obs, phase=None):
        if self.with_phase:
            if phase is None:
                phase = 0.0
            if isinstance(phase, float):
                phase = torch.FloatTensor([[phase]])
            if phase.dim() < 2:
                phase = phase.unsqueeze(1) if phase.dim() == 1 else phase
            if obs.dim() == 1:
                obs = obs.unsqueeze(0)
            if phase.shape[0] != obs.shape[0]:
                phase = phase.expand(obs.shape[0], -1)
            obs = torch.cat([obs, phase], dim=-1)
        return self.net(obs)


def train_position_bc(model, dataset, epochs=EPOCHS, lr=1e-3, batch_size=64, lookahead=10):
    """Target = absolute joint position q_des at t+lookahead along the demonstrated
    trajectory (from cos/sin state). Predictive BC: model learns a smooth
    reference lookahead WITHOUT trajectory primitives, executed via the same PD.
    lookahead=0 reduces to naive next-step prediction (learns identity, fails);
    lookahead=10 predicts a smooth reference ahead, a fair no-primitive analogue."""
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()

    X, y, phases = [], [], []
    for d in dataset:
        states = d["states"]
        T = len(states)
        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)
        for t in range(T):
            t_tgt = min(t + lookahead, T - 1)
            X.append(states[t])
            phases.append(t / T)
            y.append([q1[t_tgt], q2[t_tgt]])

    X_t = torch.FloatTensor(np.array(X))
    y_t = torch.FloatTensor(np.array(y))
    ph_t = torch.FloatTensor(np.array(phases)).reshape(-1, 1)
    n = len(X_t)
    for epoch in range(epochs):
        idx = np.random.permutation(n)
        for start in range(0, n, batch_size):
            bi = idx[start:start + batch_size]
            pred = model(X_t[bi], ph_t[bi])
            loss = loss_fn(pred, y_t[bi])
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()


@torch.no_grad()
def evaluate_position_bc(model, env, n_episodes=N_EVAL, max_steps=200):
    model.eval()
    """Closed-loop execution: at each step predict q_des from state, track via PD."""
    successes = 0
    distances = []
    for _ in range(n_episodes):
        s = env.reset()
        done = False
        for t in range(max_steps):
            s_t = torch.FloatTensor(s).unsqueeze(0)
            phase = t / max_steps
            q_des = model(s_t, phase).squeeze(0).numpy()
            q_cur = env.q.copy()
            qd_cur = env.qd.copy()
            u = KP * (q_des - q_cur) - KD * qd_cur
            s, _, done, info = env.step(u)
            if done:
                break
        successes += 1 if info["success"] else 0
        distances.append(info["dist"])
    return successes / n_episodes, np.mean(distances), np.std(distances)


def run_e1(env, f):
    log("\n" + "=" * 70, f)
    log("E1: POSITION-SPACE BC vs raw-torque BC vs ProMP-CL", f)
    log("Purpose: separate trajectory representation from control interface", f)
    log("PositionBC = pointwise position prediction, SAME PD controller, NO trajectory primitives", f)
    log("=" * 70, f)

    env_bc = Arm2DEnv()
    basis = RBFBasis(N_BFS)
    for b in BUDGETS:
        log(f"\n--- Budget {b} ---", f)
        for m in ["bc", "posbc", "cl"]:
            log(f"  {m:6s} | succ(mean±std) | dist(mean±std)", f)
            break
        for seed in range(5):
            data = make_ds(env_bc, b, seed)
            if not data:
                continue
            # raw BC
            m = BCMLP(obs_dim=env_bc.obs_dim, act_dim=env_bc.act_dim)
            train_bc(m, data, epochs=EPOCHS, verbose=False)
            sr, dm, _ = evaluate_bc(m, env_bc, n_episodes=N_EVAL)
            # position BC
            mp = PositionBC(obs_dim=env_bc.obs_dim, act_dim=env_bc.act_dim)
            train_position_bc(mp, data, epochs=EPOCHS)
            srp, dmp, _ = evaluate_position_bc(mp, env_bc, n_episodes=N_EVAL)
            # ProMP-CL reference
            mcl = ProMPPredictor(obs_dim=env_bc.obs_dim, n_bfs=N_BFS, n_joints=env_bc.act_dim, with_phase=True)
            train_promp(mcl, data, basis, epochs=EPOCHS, verbose=False, closed_loop=True)
            src, dmc, _ = evaluate_promp_closed_loop(mcl, env_bc, basis, n_episodes=N_EVAL)
            log(f"  seed{seed}: BC {sr:.2f} ({dm:.3f}) | PosBC {srp:.2f} ({dmp:.3f}) | CL {src:.2f} ({dmc:.3f})", f)


def run_e2(env, f):
    log("\n" + "=" * 70, f)
    log("E2: RAISED ACCELERATION LIMIT — is saturation a cause or symptom?", f)
    log("Baseline: CL @50 demos, accel limit ±20. Raised: ±40, ±100. Also unconstrained (huge).", f)
    log("=" * 70, f)

    # Patch env to raise the clip limit
    basis = RBFBasis(N_BFS)
    for limit in [20.0, 40.0, 100.0, 1e6]:
        env_lim = Arm2DEnv()
        # override step to use a higher clip
        orig_step = env_lim.step

        def make_step(lim):
            def step(self, u):
                self.t += 1
                u = np.clip(u, -lim, lim)
                self.qd += u * self.dt
                self.qd = np.clip(self.qd, self.qd_lim[:, 0], self.qd_lim[:, 1])
                self.q += self.qd * self.dt
                self.q = np.clip(self.q, self.q_lim[:, 0], self.q_lim[:, 1])
                ee = self.fk()
                dist = np.linalg.norm(ee - self.target)
                success = dist < 0.04
                done = self.t >= self.max_steps or success
                reward = -dist - 0.01 * np.linalg.norm(u)
                return self._state(), reward, done, {"dist": dist, "success": success, "ee": ee}
            return step

        env_lim.step = make_step(limit)
        res = []
        for seed in range(5):
            data = make_ds(env_lim, 50, seed)
            m = ProMPPredictor(obs_dim=env_lim.obs_dim, n_bfs=N_BFS, n_joints=env_lim.act_dim, with_phase=True)
            train_promp(m, data, basis, epochs=EPOCHS, verbose=False, closed_loop=True)
            sr, dm, _ = evaluate_promp_closed_loop(m, env_lim, basis, n_episodes=N_EVAL)
            res.append(sr)
        log(f"  limit ±{limit:g}: CL@50 = {np.mean(res):.3f} ± {np.std(res):.3f}", f)


def run_e3(env, f):
    log("\n" + "=" * 70, f)
    log("E3: SMOOTHNESS ABLATION, 20 SEEDS, PAIRED", f)
    log("Purpose: is +7.6pp real? Paired per-seed baseline vs smoothness.", f)
    log("=" * 70, f)

    basis = RBFBasis(N_BFS)
    for b in [50, 200]:
        log(f"\n--- Budget {b} demos ---", f)
        for seed in range(20):
            data = make_ds(env, b, seed)
            # baseline CL
            m0 = ProMPPredictor(obs_dim=env.obs_dim, n_bfs=N_BFS, n_joints=env.act_dim, with_phase=True)
            train_promp(m0, data, basis, epochs=EPOCHS, verbose=False, closed_loop=True)
            sr0, _, _ = evaluate_promp_closed_loop(m0, env, basis, n_episodes=N_EVAL)
            # smoothness CL
            m1 = ProMPPredictor(obs_dim=env.obs_dim, n_bfs=N_BFS, n_joints=env.act_dim, with_phase=True)
            train_promp(m1, data, basis, epochs=EPOCHS, verbose=False, closed_loop=True,
                        torque_penalty_weight=0.001, penalty_type="smoothness")
            sr1, _, _ = evaluate_promp_closed_loop(m1, env, basis, n_episodes=N_EVAL)
            log(f"  seed{seed:2d}: base {sr0:.3f} | smooth {sr1:.3f} | delta {sr1-sr0:+.3f}", f)


def run_e7(env, f):
    log("\n" + "=" * 70, f)
    log("E7: MULTI-THRESHOLD SUCCESS (1/2/4 cm), CL @50 demos", f)
    log("=" * 70, f)

    basis = RBFBasis(N_BFS)
    thresholds = [0.01, 0.02, 0.04]
    for seed in range(5):
        data = make_ds(env, 50, seed)
        m = ProMPPredictor(obs_dim=env.obs_dim, n_bfs=N_BFS, n_joints=env.act_dim, with_phase=True)
        train_promp(m, data, basis, epochs=EPOCHS, verbose=False, closed_loop=True)
        # custom eval collecting final distances
        dists = []
        for _ in range(N_EVAL):
            s = env.reset()
            for t in range(200):
                s_t = torch.FloatTensor(s).unsqueeze(0)
                phase = t / 200
                w_flat = m(s_t, phase).squeeze(0).numpy()
                w = w_flat.reshape(env.act_dim, basis.n_bfs)
                traj = decode_trajectory(w, basis, 200)
                q_des = traj[t]
                u = KP * (q_des - env.q) - KD * env.qd
                s, _, done, info = env.step(u)
                if done:
                    break
            dists.append(info["dist"])
        s_1 = np.mean([d < 0.01 for d in dists])
        s_2 = np.mean([d < 0.02 for d in dists])
        s_4 = np.mean([d < 0.04 for d in dists])
        log(f"  seed{seed}: <1cm {s_1:.3f} | <2cm {s_2:.3f} | <4cm {s_4:.3f} | mean_dist {np.mean(dists):.3f}", f)


def main():
    with open("results_decisive.txt", "w") as f:
        log("DECISIVE EXPERIMENT RESULTS", f)
        log(f"Date: {time.strftime('%Y-%m-%d %H:%M')}", f)
        env = Arm2DEnv()
        run_e1(env, f)
        run_e2(env, f)
        run_e3(env, f)
        run_e7(env, f)
    print("\nDone. See results_decisive.txt")


if __name__ == "__main__":
    main()