bundle_v1/scripts/cifar_repro.py

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "torch",
#   "torchvision",
#   "numpy",
#   "kron-torch",
#   "huggingface_hub",
# ]
# ///
"""Faithful reproduction of the non-stationary CIFAR-10 experiment (Table 1 /
Figs. 3 & 8) of "Stable Deep Reinforcement Learning via Isotropic Gaussian
Representations" (arXiv:2602.19373, code https://github.com/asahebpa/IsoGaussian-DRL).

Model, SIGReg regularizer, non-stationarity protocol, and hyperparameters are
copied from the authors' cifar.py (CNN -> MLP trunk (default/medium/medium,
LayerNorm) -> linear head; 10k CIFAR-10 images; labels re-permuted at epochs
20/40/60/80; batch 256; lr 2.5e-4; 100 epochs; SIGReg with 16 slices, 8
frequencies, t_max 5).

Additions relative to the authors' script (which logs only accuracy + SIGReg
loss): per-epoch feature-rank (RankMe) and dormant-neuron measurement on a
fixed 1024-image probe batch in eval mode, using the metric definitions from
the authors' utils/representation_dynamics.py and utils/utils.py; covariance
eigenspectra and random-projection moment snapshots (for the Claim-4 mechanism
audit); CSV/NPZ outputs; seeding; no wandb.

Deviations documented in the logbook: probe-based metric evaluation (paper
does not specify its probe), no test-set pass (Table 1 uses train metrics),
mlp_type assumed "default".
"""
import argparse
import csv
import math
import os
import random
import time

import numpy as np
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset


# ---------------------- SIGReg (verbatim from authors' cifar.py) ----------------------
class SIGReg(nn.Module):
    def __init__(self, embedding_dim, num_slices=16, num_t=8, t_max=5.0):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.num_slices = num_slices
        self.num_t = num_t
        self.t_max = t_max
        self.register_buffer("t_grid", torch.linspace(-t_max, t_max, steps=num_t))

    def forward(self, embeddings):
        B, D = embeddings.shape
        a = torch.randn(self.num_slices, D, device=embeddings.device)
        a = a / (a.norm(dim=1, keepdim=True) + 1e-12)
        s = torch.matmul(embeddings, a.t()).t()
        t = self.t_grid.to(embeddings.device)
        loss = 0.0
        for ti in range(self.num_t):
            tt = t[ti]
            cos_ts = torch.cos(tt * s)
            sin_ts = torch.sin(tt * s)
            re = cos_ts.mean(dim=1)
            im = sin_ts.mean(dim=1)
            target = math.exp(-0.5 * (tt.item() ** 2))
            loss += ((re - target) ** 2 + (im ** 2)).mean()
        return loss / float(self.num_t)


# ---------------------- Non-stationary dataset (verbatim) ----------------------
class NonStationaryDataset(Dataset):
    def __init__(self, dataset):
        self.dataset = dataset
        self.permutation = list(range(len(dataset.targets)))

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

    def __getitem__(self, idx):
        image, _ = self.dataset[idx]
        label = self.dataset.targets[self.permutation[idx]]
        return image, label

    def reshuffle_labels(self):
        n = len(self.dataset.targets)
        self.permutation = np.random.permutation(n).tolist()


# ---------------------- MLP trunk ("default" type, from authors' models/mlp.py) ----------------------
def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
    torch.nn.init.orthogonal_(layer.weight, std)
    torch.nn.init.constant_(layer.bias, bias_const)
    return layer


class MLP(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers,
                 use_ln=False, last_act=True):
        super().__init__()
        mlp = []
        if num_layers == 1:
            hidden_size = output_size
        for i in range(num_layers):
            if i == 0:
                mlp.append(layer_init(nn.Linear(input_size, hidden_size)))
            elif i == num_layers - 1:
                mlp.append(layer_init(nn.Linear(hidden_size, output_size)))
            else:
                mlp.append(layer_init(nn.Linear(hidden_size, hidden_size)))
            if i < num_layers - 1:
                if use_ln:
                    mlp.append(nn.LayerNorm(hidden_size))
                mlp.append(nn.ReLU())
            elif i == num_layers - 1:
                if use_ln:
                    mlp.append(nn.LayerNorm(output_size))
                if last_act:
                    mlp.append(nn.ReLU())
        self.net = nn.Sequential(*mlp)

    def forward(self, x):
        return self.net(x)


# ---------------------- CNN+MLP model (verbatim structure from authors' cifar.py) ----------------------
class CIFARClassifier(nn.Module):
    def __init__(self, mlp, num_classes, use_ln=False):
        super().__init__()
        layers = [
            nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2, 2), nn.Dropout2d(0.2),
            nn.Conv2d(64, 64, 3, padding=1), nn.ReLU(),
            nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2, 2), nn.Dropout2d(0.2),
            nn.Conv2d(128, 128, 3, padding=1), nn.ReLU(),
            nn.Conv2d(128, 128, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2, 2), nn.Dropout2d(0.2),
        ]
        if use_ln:
            layers.insert(5, nn.GroupNorm(8, 64))
            layers.insert(12, nn.GroupNorm(16, 128))
            layers.insert(19, nn.GroupNorm(16, 128))
        self.cnn = nn.Sequential(*layers)
        self.network = nn.Sequential(self.cnn, nn.Flatten())
        self.trunk = mlp
        self.classifier = nn.Linear(512, num_classes)

    def forward(self, x):
        features = self.network(x)
        trunk_features = self.trunk(features)
        outputs = self.classifier(trunk_features)
        return outputs, trunk_features


# ---------------------- Metrics (from authors' utils) ----------------------
@torch.no_grad()
def rankme(features):
    """exp(entropy of normalized eigenvalues of E[phi phi^T]) — authors'
    'feature rank' (representation_dynamics.py, compute_ppo_metrics)."""
    h = features.detach().cpu().double()  # cpu: fp64 unsupported on MPS
    cov = h.T @ h / h.shape[0]
    eig = torch.linalg.eigvalsh(cov).clamp_min(0)
    p = eig / eig.sum() + 1e-6
    entropy = -(p * p.log()).sum()
    return float(entropy.exp()), eig.cpu().numpy()


@torch.no_grad()
def dormant_neurons(model, images):
    """Fraction of ReLU units with max activation <= 0 over the probe batch
    (authors' utils.get_dormant_neurons)."""
    acts = {}
    hooks = []

    def mk(name):
        def hook(_m, _i, out):
            acts[name] = out.detach()
        return hook

    for name, module in model.named_modules():
        if isinstance(module, nn.ReLU):
            hooks.append(module.register_forward_hook(mk(name)))
    model(images)
    for h in hooks:
        h.remove()
    fracs = {}
    for name, a in acts.items():
        if a.dim() == 4:
            r = a.permute(1, 0, 2, 3).reshape(a.shape[1], -1)
        else:
            r = a.t()
        fracs[name] = float((r.max(dim=1).values <= 0).sum().item() / a.shape[1])
    mlp_f = [v for k, v in fracs.items() if "trunk" in k]
    cnn_f = [v for k, v in fracs.items() if "cnn" in k]
    return (float(np.mean(mlp_f)), float(np.mean(cnn_f)),
            float(np.mean(list(fracs.values()))))


@torch.no_grad()
def projection_moments(features, gen, n_dirs=64):
    """Skewness / excess kurtosis of random 1-d projections (Claim 4 probe)."""
    d = features.shape[1]
    v = torch.randn(n_dirs, d, generator=gen).to(features.device)
    v = v / v.norm(dim=1, keepdim=True)
    z = features @ v.T  # [B, n_dirs]
    z = z - z.mean(0, keepdim=True)
    std = z.std(0, keepdim=True).clamp_min(1e-12)
    zn = z / std
    skew = (zn ** 3).mean(0)
    kurt = (zn ** 4).mean(0) - 3.0
    return float(skew.abs().mean()), float(kurt.mean())


def get_optimizer(name, params, lr):
    if name == "adam":
        return torch.optim.Adam(params, lr=lr)
    if name == "radam":
        return torch.optim.RAdam(params, lr=lr)
    if name == "kron":
        from kron_torch import Kron
        return Kron(params, lr=lr)
    raise ValueError(name)


class TensorCIFAR:
    """In-memory replacement for DataLoader(NonStationaryDataset(...)):
    identical math (ToTensor + Normalize(0.5, 0.5) transform precomputed once;
    label = targets[permutation[idx]]; fresh shuffle each epoch), but pure
    tensor indexing — no worker processes (macOS DataLoader workers hang at
    shutdown, and a hung exit would burn the GPU-job timeout)."""

    def __init__(self, data_uint8, targets):
        x = torch.from_numpy(np.ascontiguousarray(data_uint8)).float() / 255.0
        x = (x - 0.5) / 0.5  # Normalize((0.5,)*3, (0.5,)*3)
        self.x = x.permute(0, 3, 1, 2).contiguous()
        self.targets = torch.as_tensor(list(targets), dtype=torch.long)
        self.permutation = torch.arange(len(self.targets))

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

    def reshuffle_labels(self):
        self.permutation = torch.from_numpy(np.random.permutation(len(self.targets)))

    def epoch_batches(self, batch_size):
        order = torch.randperm(len(self.targets))
        labels = self.targets[self.permutation]
        for i in range(0, len(order), batch_size):
            idx = order[i:i + batch_size]
            yield self.x[idx], labels[idx]


def build_data(data_root, npz_path=""):
    n = 10000  # authors' num_datapoints
    if npz_path:
        arr = np.load(npz_path)
        data, targets = arr["data"][:n], arr["labels"][:n].tolist()
    else:
        base = torchvision.datasets.CIFAR10(root=data_root, train=True,
                                            download=True)
        data, targets = base.data[:n], base.targets[:n]
    return TensorCIFAR(data, targets)


def train_one(optimizer_name, lambda_sig, seed, args, device):
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)

    trainset = build_data(args.data_root, npz_path=args.npz_path)
    probe_x = trainset.x[:args.probe_size].to(device)

    mlp = MLP(input_size=2048, hidden_size=args.mlp_width,
              output_size=512, num_layers=args.mlp_layers, use_ln=args.use_ln)
    model = CIFARClassifier(mlp, num_classes=10, use_ln=args.use_ln).to(device)
    sigreg = SIGReg(embedding_dim=512).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = get_optimizer(optimizer_name, model.parameters(), args.lr)
    proj_gen = torch.Generator().manual_seed(12345)  # fixed probe directions across runs

    tag = f"{optimizer_name}_lam{lambda_sig}_seed{seed}"
    os.makedirs(args.out_dir, exist_ok=True)
    csv_path = os.path.join(args.out_dir, f"cifar_{tag}.csv")
    spectra = {}
    rows = []
    t_run = time.time()
    for epoch in range(args.epochs):
        if epoch in (20, 40, 60, 80):
            trainset.reshuffle_labels()
            print(f"[{tag}] reshuffling labels at epoch {epoch}", flush=True)
        model.train()
        t0 = time.time()
        total_sig, correct, total, n_batches = 0.0, 0, 0, 0
        for images, labels in trainset.epoch_batches(args.batch_size):
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs, trunk = model(images)
            loss_ce = criterion(outputs, labels)
            loss_sig = sigreg(trunk)
            loss = loss_ce + lambda_sig * loss_sig
            loss.backward()
            optimizer.step()
            total_sig += loss_sig.item()
            correct += (outputs.argmax(1) == labels).sum().item()
            total += labels.size(0)
            n_batches += 1
        train_acc = correct / total
        train_sig = total_sig / n_batches

        model.eval()
        with torch.no_grad():
            _, probe_feats = model(probe_x)
            rank, eig = rankme(probe_feats)
            d_mlp, d_cnn, d_all = dormant_neurons(model, probe_x)
            sig_probe = float(sigreg(probe_feats).item())
            skew, kurt = projection_moments(probe_feats, proj_gen)
            top_share = float(eig[-1] / max(eig.sum(), 1e-12))
        if epoch % 5 == 0 or epoch in (19, 20, 39, 40, 59, 60, 79, 80, 99):
            spectra[f"epoch_{epoch}"] = eig
        dt = time.time() - t0
        rows.append(dict(epoch=epoch, train_acc=train_acc, train_sigreg=train_sig,
                         probe_sigreg=sig_probe, rankme=rank,
                         dormant_mlp=d_mlp, dormant_cnn=d_cnn, dormant_all=d_all,
                         proj_abs_skew=skew, proj_excess_kurt=kurt,
                         top_eig_share=top_share, epoch_time_s=dt))
        print(f"[{tag}] epoch {epoch}: acc={train_acc:.4f} sig={train_sig:.4f} "
              f"rank={rank:.1f} dorm_mlp={d_mlp*100:.1f}% ({dt:.1f}s)", flush=True)

    with open(csv_path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0]))
        w.writeheader()
        w.writerows(rows)
    np.savez_compressed(os.path.join(args.out_dir, f"spectra_{tag}.npz"), **spectra)

    auc = {k: float(np.mean([r[k] for r in rows]))
           for k in ("train_acc", "train_sigreg", "rankme", "dormant_mlp", "dormant_all")}
    print(f"[{tag}] DONE in {time.time()-t_run:.0f}s | AUC(mean/epoch): "
          f"acc={auc['train_acc']*100:.1f} sig={auc['train_sigreg']*100:.1f} "
          f"rank={auc['rankme']:.1f} dorm_mlp={auc['dormant_mlp']*100:.1f} "
          f"dorm_all={auc['dormant_all']*100:.1f}", flush=True)
    return auc


def maybe_upload(out_dir, dataset_repo, prefix):
    """Push results to a dataset repo (the Python API supports only
    model/dataset/space repo types — buckets are CLI-only, validated 2026-07-16)."""
    if not dataset_repo:
        return
    try:
        from huggingface_hub import HfApi
        HfApi().upload_folder(folder_path=out_dir, path_in_repo=prefix,
                              repo_id=dataset_repo, repo_type="dataset")
        print(f"uploaded {out_dir} -> dataset {dataset_repo}/{prefix}", flush=True)
    except Exception as e:  # keep training results even if upload hiccups
        print(f"WARNING: results upload failed: {e}", flush=True)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--optimizer", default="adam", choices=["adam", "radam", "kron"])
    p.add_argument("--lambda_sig", type=float, default=1.0)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--epochs", type=int, default=100)
    p.add_argument("--batch_size", type=int, default=256)
    p.add_argument("--lr", type=float, default=0.00025)
    p.add_argument("--mlp_width", type=int, default=1536)   # authors' "medium"
    p.add_argument("--mlp_layers", type=int, default=8)     # authors' "medium" default-type depth
    p.add_argument("--use_ln", action="store_true", default=True)
    p.add_argument("--probe_size", type=int, default=1024)
    p.add_argument("--data_root", default="./cifar10_data")
    p.add_argument("--npz_path", default="",
                   help="load the 10k-image subset from this npz (data/labels) "
                        "instead of downloading via torchvision")
    p.add_argument("--npz_hf", default="",
                   help="HF dataset repo holding cifar10k.npz; downloaded to "
                        "--npz_path location at startup (for HF Jobs)")
    p.add_argument("--out_dir", default="./outputs_cifar")
    p.add_argument("--grid", action="store_true",
                   help="run the full 3-optimizer x {0,1} x seeds grid")
    p.add_argument("--grid_seeds", type=int, default=2)
    p.add_argument("--time_budget_s", type=int, default=9000,
                   help="grid mode: stop launching new runs when the estimated "
                        "next-run time would exceed this budget")
    p.add_argument("--results_repo", default="",
                   help="HF dataset repo to upload outputs to (grid mode)")
    p.add_argument("--bucket_prefix", default="claim5")
    p.add_argument("--device", default="auto",
                   help="auto|cuda|mps|cpu (kron_torch needs cuda or cpu: its "
                        "torch.compile backend does not support mps)")
    args = p.parse_args()

    if args.device != "auto":
        device = torch.device(args.device)
    else:
        device = torch.device("cuda" if torch.cuda.is_available()
                              else "mps" if torch.backends.mps.is_available() else "cpu")

    if args.npz_hf:
        from huggingface_hub import hf_hub_download
        args.npz_path = hf_hub_download(repo_id=args.npz_hf, repo_type="dataset",
                                        filename="cifar10k.npz")
        print(f"downloaded npz from {args.npz_hf} -> {args.npz_path}", flush=True)
    print(f"device: {device}", flush=True)

    if not args.grid:
        train_one(args.optimizer, args.lambda_sig, args.seed, args, device)
        return

    t0 = time.time()
    run_times = []
    results = {}
    skipped = []
    # seed-major order: a full 6-config grid at seed s completes before seed s+1
    for seed in range(args.grid_seeds):
        for opt in ("adam", "radam", "kron"):
            for lam in (0.0, 1.0):
                elapsed = time.time() - t0
                est_next = max(run_times) if run_times else 0.0
                if run_times and elapsed + est_next * 1.3 > args.time_budget_s:
                    skipped.append((opt, lam, seed))
                    continue
                t1 = time.time()
                auc = train_one(opt, lam, seed, args, device)
                run_times.append(time.time() - t1)
                results[f"{opt}_lam{lam}_seed{seed}"] = auc
                maybe_upload(args.out_dir, args.results_repo, args.bucket_prefix)
    if skipped:
        print(f"SKIPPED (time budget): {skipped}", flush=True)
    print("GRID SUMMARY (AUC = mean per epoch; acc/sig/dorm x100):", flush=True)
    for k, v in results.items():
        print(f"  {k}: acc={v['train_acc']*100:.1f} sig={v['train_sigreg']*100:.1f} "
              f"rank={v['rankme']:.1f} dorm_mlp={v['dormant_mlp']*100:.1f} "
              f"dorm_all={v['dormant_all']*100:.1f}", flush=True)
    with open(os.path.join(args.out_dir, "grid_summary.csv"), "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["run", "acc_auc", "sigreg_auc_x100", "rankme_auc",
                    "dormant_mlp_auc", "dormant_all_auc"])
        for k, v in results.items():
            w.writerow([k, v["train_acc"] * 100, v["train_sigreg"] * 100,
                        v["rankme"], v["dormant_mlp"] * 100, v["dormant_all"] * 100])
    maybe_upload(args.out_dir, args.results_repo, args.bucket_prefix)


if __name__ == "__main__":
    main()