run_e1_full.py

"""
E1 FULL: Position-space BC (lookahead reference prediction) vs raw-torque BC vs ProMP-CL.
Decisive test of the control-interface confound (review C1).
Budgets: 5,10,25,50,100,200 demos, 5 seeds, la10/la20.
"""
import numpy as np
import torch
import torch.nn as nn
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_closed_loop

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


class PositionBC(nn.Module):
    """MLP: state (+phase) -> absolute joint position q_des at lookahead horizon."""

    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=20):
    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()
    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 main():
    out = open("results_e1.txt", "w")
    def log(*a):
        msg = " ".join(str(x) for x in a)
        print(msg, flush=True)
        out.write(msg + "\n")
        out.flush()

    log("E1 FULL: Position-space BC (la10/la20) vs raw-torque BC vs ProMP-CL")
    log(f"Date: {time.strftime('%Y-%m-%d %H:%M')}")
    log("BUDGETS:", BUDGETS, "| SEEDS: 5 | EPOCHS:", EPOCHS)

    env = Arm2DEnv()
    basis = RBFBasis(N_BFS)

    for la in [10, 20]:
        log(f"\n===== LOOKAHEAD = {la} =====")
        log(f"{'Budget':>7} | {'BC':>8} | {'PosBC-la':>9} | {'CL':>8}")
        for b in BUDGETS:
            bc_s, pb_s, cl_s = [], [], []
            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))
                if not data:
                    continue
                m_bc = BCMLP(10, 2)
                train_bc(m_bc, data, epochs=EPOCHS, verbose=False)
                sr, _, _ = evaluate_bc(m_bc, env)
                bc_s.append(sr)
                mp = PositionBC(10, 2)
                train_position_bc(mp, data, epochs=EPOCHS, lookahead=la)
                srp, _, _ = evaluate_position_bc(mp, env)
                pb_s.append(srp)
                m_cl = ProMPPredictor(10, N_BFS, 2, with_phase=True)
                train_promp(m_cl, data, basis, epochs=EPOCHS, verbose=False, closed_loop=True)
                src, _, _ = evaluate_promp_closed_loop(m_cl, env, basis)
                cl_s.append(src)
            log(f"{b:>7} | {np.mean(bc_s):.3f}±{np.std(bc_s):.3f} | {np.mean(pb_s):.3f}±{np.std(pb_s):.3f} | {np.mean(cl_s):.3f}±{np.std(cl_s):.3f}")

    out.close()


if __name__ == "__main__":
    main()