code/stereo_core/role_observability_teacher.py

"""Train the action-only team-role teacher for the observability audit.

The teacher never sees images, task ids, agent ids, future observations, or
environment state.  It learns whether unordered expert action chunks come from
the same synchronized team instant.  Negatives replace exactly one arm with a
time-shifted chunk from the same episode, preserving the single-arm marginals.
"""
from __future__ import annotations

import argparse
import glob
import hashlib
import json
import math
import random
from collections import defaultdict
from pathlib import Path

import h5py
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset

from train_act import _stats, _trajectories, seed_everything


def episode_split(task: str, path: str, key: str, seed: int) -> str:
    digest = hashlib.sha256(f"{seed}|{task}|{path}|{key}".encode()).digest()
    bucket = int.from_bytes(digest[:8], "little") % 100
    return "train" if bucket < 70 else ("val" if bucket < 85 else "test")


class SynchronizedActionSets(Dataset):
    def __init__(self, trajectories, stats, split, seed, horizon=100, stride=8,
                 negative_mode="time_shift"):
        self.stats = stats
        self.horizon = horizon
        self.seed = seed
        self.negative_mode = negative_mode
        self.episodes = []
        self.by_task = defaultdict(list)
        for row in trajectories:
            path, key, n, present, task = row
            if episode_split(task, path, key, seed) != split or n < 2:
                continue
            item = (path, key, n, tuple(present), task)
            self.by_task[task].append(len(self.episodes))
            self.episodes.append(item)
        self.tasks = sorted(self.by_task)
        self.index = []
        for task in self.tasks:
            for episode_index in self.by_task[task]:
                n = self.episodes[episode_index][2]
                for t in range(0, n, stride):
                    self.index.append((episode_index, t))
        if not self.index:
            raise ValueError(f"no {split} synchronized episodes")

    def __len__(self):
        return len(self.index)

    def _chunk(self, action, t):
        chunk = np.asarray(action[t:t + self.horizon], np.float32)
        valid = len(chunk)
        padded = np.empty((self.horizon, action.shape[1]), np.float32)
        padded[:valid] = chunk
        padded[valid:] = chunk[-1]
        return (padded - self.stats["a_mean"]) / self.stats["a_std"]

    def __getitem__(self, index):
        episode_index, t = self.index[index]
        path, key, n, arms, task = self.episodes[episode_index]
        # Deterministic per-sample negative: same episode, one arm, nontrivial
        # temporal displacement.  This makes audit reruns bit-reproducible.
        rng = random.Random(self.seed * 1000003 + index)
        replaced = rng.randrange(len(arms))
        low = max(1, min(16, n // 8))
        high = max(low, min(96, max(1, n // 2)))
        delta = rng.randint(low, high)
        shifted_t = t + delta if t + delta < n else max(0, t - delta)
        positive, negative = [], []
        with h5py.File(path, "r") as h5:
            trajectory = h5[key]
            for arm_index, arm in enumerate(arms):
                action = trajectory["actions"][f"panda-{arm}"]
                local = self._chunk(action, t)
                positive.append(local)
                negative.append(self._chunk(action, shifted_t) if arm_index == replaced else local)
        action_dim = positive[0].shape[-1]
        pos = np.zeros((4, self.horizon, action_dim), np.float32)
        neg = np.zeros_like(pos)
        mask = np.zeros(4, np.bool_)
        pos[:len(arms)] = positive
        neg[:len(arms)] = negative
        use_duplicate = self.negative_mode == "role_duplicate" or (
            self.negative_mode == "mixed" and rng.random() < .5)
        if use_duplicate:
            neg[replaced] = pos[(replaced + 1) % len(arms)]
        mask[:len(arms)] = True
        return torch.from_numpy(pos), torch.from_numpy(neg), torch.from_numpy(mask)


class TeamRoleTeacher(nn.Module):
    def __init__(self, action_dim, width=192, role_dim=64):
        super().__init__()
        self.action_encoder = nn.Sequential(
            nn.Conv1d(action_dim, width, 5, padding=2), nn.GELU(),
            nn.Conv1d(width, width, 5, padding=2), nn.GELU(),
            nn.AdaptiveAvgPool1d(1),
        )
        self.role = nn.Sequential(
            nn.LayerNorm(width * 2), nn.Linear(width * 2, width), nn.GELU(),
            nn.Linear(width, role_dim),
        )
        self.q = nn.Linear(role_dim, role_dim, bias=False)
        self.k = nn.Linear(role_dim, role_dim, bias=False)

    def forward(self, actions, mask):
        batch, agents, horizon, action_dim = actions.shape
        encoded = self.action_encoder(
            actions.reshape(batch * agents, horizon, action_dim).transpose(1, 2)
        ).squeeze(-1).reshape(batch, agents, -1)
        weights = mask.to(encoded.dtype).unsqueeze(-1)
        total = (encoded * weights).sum(1, keepdim=True)
        peers = (total - encoded * weights) / (weights.sum(1, keepdim=True) - weights).clamp_min(1.0)
        role = F.normalize(self.role(torch.cat((encoded, peers), -1)), dim=-1)
        q, k = self.q(role), self.k(role)
        pair = torch.einsum("bid,bjd->bij", q, k) / math.sqrt(q.shape[-1])
        valid = mask.unsqueeze(1) & mask.unsqueeze(2)
        valid &= ~torch.eye(agents, dtype=torch.bool, device=mask.device).unsqueeze(0)
        score = (pair * valid).sum((1, 2)) / valid.sum((1, 2)).clamp_min(1)
        return score, role


def auc(labels, scores):
    order = torch.argsort(scores)
    ranks = torch.empty_like(order, dtype=torch.float32)
    ranks[order] = torch.arange(1, len(scores) + 1, dtype=torch.float32)
    positive = labels.bool()
    n_pos, n_neg = int(positive.sum()), int((~positive).sum())
    if not n_pos or not n_neg:
        return float("nan")
    return float((ranks[positive].sum() - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg))


@torch.no_grad()
def evaluate(model, loader, device):
    model.eval(); positive_scores, negative_scores = [], []
    for pos, neg, mask in loader:
        pos, neg, mask = pos.to(device), neg.to(device), mask.to(device)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            p, _ = model(pos, mask); n, _ = model(neg, mask)
        positive_scores.append(p.float().cpu())
        negative_scores.append(n.float().cpu())
    positive_scores = torch.cat(positive_scores)
    negative_scores = torch.cat(negative_scores)
    scores = torch.cat((positive_scores, negative_scores))
    labels = torch.cat((torch.ones_like(positive_scores), torch.zeros_like(negative_scores)))
    return {"auc": auc(labels, scores), "positive_mean": float(positive_scores.mean()),
            "negative_mean": float(negative_scores.mean()),
            "margin": float(positive_scores.mean() - negative_scores.mean())}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--normalization", required=True)
    parser.add_argument("--updates", type=int, default=5000)
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument("--seed", type=int, default=20260801)
    parser.add_argument("--workers", type=int, default=8)
    parser.add_argument("--negative-mode", choices=("time_shift", "role_duplicate", "mixed"),
                        default="time_shift")
    parser.add_argument("--checkpoint", default=None,
                        help="Evaluate an existing teacher checkpoint with corrected metrics")
    args = parser.parse_args()
    seed_everything(args.seed)
    paths = sorted({p for pattern in args.data.split(",") for p in glob.glob(pattern)})
    trajectories = _trajectories(paths, (0, 1, 2, 3))
    stats = torch.load(args.normalization, map_location="cpu", weights_only=False)["stats"]
    datasets = {split: SynchronizedActionSets(trajectories, stats, split, args.seed,
                                               negative_mode=args.negative_mode)
                for split in ("train", "val", "test")}
    loaders = {split: DataLoader(data, batch_size=args.batch_size, shuffle=(split == "train"),
                                 num_workers=args.workers, pin_memory=True, persistent_workers=args.workers > 0)
               for split, data in datasets.items()}
    sample = datasets["train"][0][0]
    device = torch.device("cuda")
    model = TeamRoleTeacher(sample.shape[-1]).to(device)
    out = Path(args.output); out.mkdir(parents=True, exist_ok=True)
    if args.checkpoint:
        checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
        model.load_state_dict(checkpoint["model"])
        metrics = {split: evaluate(model, loaders[split], device) for split in ("val", "test")}
        (out / "corrected_metrics.json").write_text(json.dumps(metrics, indent=2))
        print(json.dumps({"complete": True, "corrected_metrics": metrics}), flush=True)
        return
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
    split_manifest = {
        split: [{"path": p, "key": k, "task": task, "arms": list(arms), "length": n}
                for p, k, n, arms, task in data.episodes]
        for split, data in datasets.items()
    }
    (out / "episode_split.json").write_text(json.dumps(split_manifest, indent=2))
    iterator = iter(loaders["train"]); running = 0.0
    for update in range(1, args.updates + 1):
        try: pos, neg, mask = next(iterator)
        except StopIteration:
            iterator = iter(loaders["train"]); pos, neg, mask = next(iterator)
        pos, neg, mask = pos.to(device, non_blocking=True), neg.to(device, non_blocking=True), mask.to(device, non_blocking=True)
        model.train(); optimizer.zero_grad(set_to_none=True)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            p, _ = model(pos, mask); n, _ = model(neg, mask)
            loss = F.softplus(-p).mean() + F.softplus(n).mean()
        loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); optimizer.step()
        running += float(loss.detach())
        if update % 250 == 0 or update == args.updates:
            metrics = evaluate(model, loaders["val"], device)
            row = {"update": update, "train_loss": running / 250, "val": metrics}
            print(json.dumps(row), flush=True); running = 0.0
    metrics = {split: evaluate(model, loaders[split], device) for split in ("val", "test")}
    torch.save({"model": model.state_dict(), "stats": stats, "metrics": metrics,
                "config": vars(args), "split_manifest": split_manifest}, out / "teacher.pt")
    (out / "metrics.json").write_text(json.dumps(metrics, indent=2))
    print(json.dumps({"complete": True, "metrics": metrics}), flush=True)


if __name__ == "__main__":
    main()