experiments/tasks/train_signal_forecast_priv.py

#!/usr/bin/env python3
"""Train + evaluate T8 v3 — privileged future-pressure conditioning (Option B).

Compared to train_signal_forecast.py:
  - Inputs: past 1.5s of `input_modalities` (e.g. just target modality)
            + future T_fut s of pressure (privileged side channel)
  - Output: future T_fut s of `target_modality`
  - Comparison baseline (A_priv): existing `_no_pressure` runs from T8 v2.
  - This run is the B_priv group; lift = skill(B_priv) - skill(A_priv).

If lift >> 0, future pressure trajectory carries information about future
kinematics that past kinematics alone do not encode. This directly tests
the Johansson 1984 hypothesis at the algorithmic level.
"""
from __future__ import annotations
import argparse
import json
import random
import sys
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

THIS = Path(__file__).resolve()
sys.path.insert(0, str(THIS.parent))
sys.path.insert(0, str(THIS.parents[1]))

from data.dataset_signal_forecast import (
    SignalForecastDataset, collate_signal_forecast,
    build_signal_train_test, EVENT_NAMES,
)
from nets.models_forecast_priv import DAFFuturePressure


def set_seed(seed: int):
    random.seed(seed); np.random.seed(seed)
    torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)


def train_epoch(model, loader, optimizer, device):
    model.train()
    total, n = 0.0, 0
    for x, y, y_last, fp, _et, _ in loader:
        x = {m: v.to(device) for m, v in x.items()}
        y = y.to(device)
        y_last = y_last.to(device).unsqueeze(1)
        fp = fp.to(device)
        residual_target = y - y_last
        optimizer.zero_grad()
        pred = model(x, fp)
        loss = ((pred - residual_target) ** 2).mean()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        total += loss.item() * y.numel()
        n += y.numel()
    return total / max(n, 1)


@torch.no_grad()
def evaluate(model, loader, device, t_fut, target_dim):
    model.eval()
    sse_m = np.zeros((5, t_fut), dtype=np.float64)
    sse_p = np.zeros((5, t_fut), dtype=np.float64)
    n_pairs = np.zeros((5, t_fut), dtype=np.int64)

    for x, y, y_last, fp, et, _ in loader:
        x = {m: v.to(device) for m, v in x.items()}
        y = y.to(device)
        y_last = y_last.to(device).unsqueeze(1)
        fp = fp.to(device)
        pred = model(x, fp)                                   # residual
        pred_full = pred + y_last
        persist = y_last.expand_as(y)
        m_err = ((pred_full - y) ** 2).mean(dim=-1)
        p_err = ((persist - y) ** 2).mean(dim=-1)
        et_np = et.numpy()
        m_np, p_np = m_err.cpu().numpy(), p_err.cpu().numpy()
        for k in range(m_np.shape[0]):
            e = int(et_np[k])
            sse_m[e] += m_np[k]; sse_p[e] += p_np[k]; n_pairs[e] += 1
            sse_m[4] += m_np[k]; sse_p[4] += p_np[k]; n_pairs[4] += 1

    out = {}
    for e in range(5):
        n = max(int(n_pairs[e].max()), 1)
        mse_m = (sse_m[e] / np.maximum(n_pairs[e], 1)).mean()
        mse_p = (sse_p[e] / np.maximum(n_pairs[e], 1)).mean()
        skill = 1.0 - (mse_m / mse_p) if mse_p > 1e-9 else 0.0
        per_h_skill = (1.0 - (sse_m[e] / np.maximum(n_pairs[e], 1)) /
                       np.maximum(sse_p[e] / np.maximum(n_pairs[e], 1), 1e-9)).tolist()
        name = EVENT_NAMES.get(e, "overall") if e < 4 else "overall"
        out[name] = {
            "n_anchors":  int(n),
            "mse_model":  float(mse_m),
            "mse_persist": float(mse_p),
            "skill_score": float(skill),
            "per_h_skill": [float(s) for s in per_h_skill],
        }
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input_modalities", required=True,
                    help="comma-separated; pressure NOT included unless you want past pressure too")
    ap.add_argument("--target_modality", required=True, choices=["imu", "emg", "mocap"])
    ap.add_argument("--t_obs", type=float, default=1.5)
    ap.add_argument("--t_fut", type=float, default=0.5)
    ap.add_argument("--anchor_stride", type=float, default=0.25)
    ap.add_argument("--per_event_max", type=int, default=8000)
    ap.add_argument("--epochs", type=int, default=25)
    ap.add_argument("--batch_size", type=int, default=64)
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--weight_decay", type=float, default=1e-4)
    ap.add_argument("--d_model", type=int, default=128)
    ap.add_argument("--dropout", type=float, default=0.1)
    ap.add_argument("--num_workers", type=int, default=2)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--patience", type=int, default=6)
    ap.add_argument("--output_dir", required=True)
    args = ap.parse_args()

    set_seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    inputs = args.input_modalities.split(",")
    print(f"device={device} seed={args.seed} model=DAF-priv "
          f"inputs={inputs} target={args.target_modality} "
          f"t_obs={args.t_obs} t_fut={args.t_fut}", flush=True)

    train_ds, test_ds = build_signal_train_test(
        input_modalities=inputs,
        target_modality=args.target_modality,
        t_obs_sec=args.t_obs, t_fut_sec=args.t_fut,
        anchor_stride_sec=args.anchor_stride,
        per_event_max=args.per_event_max,
        include_future_pressure=True,
        rng_seed=args.seed,
    )
    target_dim = train_ds.target_dim
    print(f"train={len(train_ds)} test={len(test_ds)} target_dim={target_dim}",
          flush=True)

    tr_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True,
                           num_workers=args.num_workers,
                           collate_fn=collate_signal_forecast, drop_last=False)
    te_loader = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False,
                           num_workers=args.num_workers,
                           collate_fn=collate_signal_forecast)

    model = DAFFuturePressure(
        train_ds.modality_dims, target_dim=target_dim,
        t_obs=train_ds.T_obs, t_fut=train_ds.T_fut,
        future_pressure_dim=50,
        d_model=args.d_model, dropout=args.dropout,
    ).to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"params={n_params:,}", flush=True)

    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr,
                                  weight_decay=args.weight_decay)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=args.epochs, eta_min=args.lr * 0.05
    )

    out_dir = Path(args.output_dir); out_dir.mkdir(parents=True, exist_ok=True)
    best_skill = -1e9
    best_epoch, best_eval = 0, None
    patience_counter = 0
    for ep in range(1, args.epochs + 1):
        t0 = time.time()
        tr_loss = train_epoch(model, tr_loader, optimizer, device)
        ev = evaluate(model, te_loader, device,
                      t_fut=train_ds.T_fut, target_dim=target_dim)
        sched.step()
        skill = ev["overall"]["skill_score"]
        print(f"  E{ep:2d} | tr_mse {tr_loss:.4f} | te_skill {skill:+.4f} "
              f"| pre {ev['pre-contact']['skill_score']:+.3f} "
              f"steady {ev['steady-grip']['skill_score']:+.3f} "
              f"release {ev['release']['skill_score']:+.3f} "
              f"non {ev['non-contact']['skill_score']:+.3f} "
              f"| {time.time()-t0:.1f}s", flush=True)
        if skill > best_skill:
            best_skill = skill
            best_epoch = ep
            best_eval = ev
            torch.save({k: v.cpu() for k, v in model.state_dict().items()},
                       out_dir / "model_best.pt")
            patience_counter = 0
        else:
            patience_counter += 1
        if patience_counter >= args.patience:
            print(f"  early stop at epoch {ep} (best {best_epoch})", flush=True)
            break

    out = {
        "method": "daf_priv",
        "input_modalities": inputs,
        "target_modality": args.target_modality,
        "future_pressure": True,
        "seed": args.seed, "n_params": n_params,
        "T_obs": train_ds.T_obs, "T_fut": train_ds.T_fut, "target_dim": target_dim,
        "best_epoch": int(best_epoch), "best_skill": float(best_skill),
        "eval": best_eval, "args": vars(args),
    }
    with open(out_dir / "results.json", "w") as f:
        json.dump(out, f, indent=2)
    print(f"\n[done] best skill={best_skill:+.4f} at epoch {best_epoch}", flush=True)


if __name__ == "__main__":
    main()