path: root/diag/train_rec.py

"""Step-2: RRoG/TRM-on-GNN for ZINC ring-counting.

The graph is encoded once with a GIN encoder. A shared edge-free node-wise compute block then
refines hidden state over n_sup*T recurrent steps (TRM-style: carry latent detached between
deep-supervision steps). --grad_mode controls the LAST supervision step's recursion:
  full  : backprop through all T inner recursions (TRM)
  1step : backprop only the last inner recursion, first T-1 detached (HRM 1-step-gradient)

Run: PYTHONPATH=/home/yurenh2/rrog python3 diag/train_rec.py --grad_mode full --sigma 0 --K 1
"""
import argparse, json, os, time
import numpy as np
import torch
import torch.nn as nn
from torch_geometric.loader import DataLoader
from torch_geometric.data import Batch, Data
from torch_geometric.nn import (
    APPNP,
    ARMAConv,
    ChebConv,
    FiLMConv,
    GATv2Conv,
    GCNConv,
    GENConv,
    GINEConv,
    GINConv,
    GraphConv,
    MFConv,
    PNAConv,
    ResGatedGraphConv,
    SAGEConv,
    SGConv,
    TAGConv,
    TransformerConv,
    global_add_pool,
)
from torch_geometric.utils import degree
from diag.train_cycle import prepare

PROJECT_ROOT = os.environ.get(
    'RROG_ROOT',
    os.path.abspath(os.path.join(os.path.dirname(__file__), '..')),
)
OUT = os.environ.get('RROG_RUNS_DIR', os.path.join(PROJECT_ROOT, 'runs'))
SUPPORTED_VIEWS = [
    'gin', 'gine', 'gcn', 'graphsage', 'gatv2', 'graphconv', 'transformer', 'pna',
    'gen', 'film', 'resgated', 'tag', 'sgc', 'cheb', 'arma', 'mf', 'appnp',
]


def data_list(recs):
    return [Data(x=r['x'], edge_index=r['edge_index'], y=r['y'].view(1, 2),
                 num_nodes=r['x'].numel()) for r in recs]


def loader(recs, bs, shuffle, drop_last=False):
    data = recs if recs and isinstance(recs[0], Data) else data_list(recs)
    return DataLoader(data, batch_size=bs, shuffle=shuffle, drop_last=drop_last)


def degree_histogram(data):
    max_degree = 0
    degs = []
    for graph in data:
        deg = degree(graph.edge_index[1], num_nodes=graph.num_nodes, dtype=torch.long)
        degs.append(deg)
        if deg.numel():
            max_degree = max(max_degree, int(deg.max().item()))
    hist = torch.zeros(max_degree + 1, dtype=torch.long)
    for deg in degs:
        hist += torch.bincount(deg, minlength=hist.numel())
    return hist


def make_view_layer(view, hidden, deg):
    if view == 'gin':
        return GINConv(nn.Sequential(
            nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, hidden)), train_eps=True)
    if view == 'gine':
        return GINEConv(nn.Sequential(
            nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, hidden)),
            train_eps=True, edge_dim=hidden)
    if view == 'gcn':
        return GCNConv(hidden, hidden)
    if view == 'graphsage':
        return SAGEConv(hidden, hidden)
    if view == 'gatv2':
        return GATv2Conv(hidden, hidden, heads=4, concat=False)
    if view == 'graphconv':
        return GraphConv(hidden, hidden)
    if view == 'transformer':
        return TransformerConv(hidden, hidden, heads=4, concat=False)
    if view == 'pna':
        if deg is None:
            raise ValueError('PNA view requires a training-set degree histogram')
        return PNAConv(
            hidden, hidden,
            aggregators=['mean', 'min', 'max', 'std'],
            scalers=['identity', 'amplification', 'attenuation'],
            deg=deg,
        )
    if view == 'gen':
        return GENConv(hidden, hidden)
    if view == 'film':
        return FiLMConv(hidden, hidden)
    if view == 'resgated':
        return ResGatedGraphConv(hidden, hidden)
    if view == 'tag':
        return TAGConv(hidden, hidden, K=3)
    if view == 'sgc':
        return SGConv(hidden, hidden, K=2, cached=False)
    if view == 'cheb':
        return ChebConv(hidden, hidden, K=3)
    if view == 'arma':
        return ARMAConv(hidden, hidden, num_stacks=1, num_layers=2)
    if view == 'mf':
        return MFConv(hidden, hidden)
    if view == 'appnp':
        return APPNP(K=5, alpha=0.1)
    raise ValueError(f'unsupported view: {view}')


class RecGIN(nn.Module):
    def __init__(self, n_atom, hidden=128, T=3, n_sup=3, sigma=0.0, inner=2,
                 grad_mode='full', agg_layers=5, compute_layers=None, view='gin', deg=None):
        super().__init__()
        self.view = view
        self.agg_layers = agg_layers
        self.compute_layers = compute_layers or inner
        self.emb = nn.Embedding(n_atom, hidden)
        self.edge_emb = nn.Embedding(1, hidden) if view == 'gine' else None
        self.agg_convs = nn.ModuleList()
        for _ in range(agg_layers):
            self.agg_convs.append(make_view_layer(view, hidden, deg))
        self.agg_bns = nn.ModuleList([nn.BatchNorm1d(hidden) for _ in range(agg_layers)])
        core = []
        d = hidden
        for _ in range(self.compute_layers - 1):
            core += [nn.Linear(d, hidden), nn.GELU()]
            d = hidden
        core.append(nn.Linear(d, hidden))
        self.core_norm = nn.LayerNorm(hidden)
        self.core = nn.Sequential(*core)
        nn.init.zeros_(self.core[-1].weight)
        nn.init.zeros_(self.core[-1].bias)
        self.head = nn.Sequential(nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 2))
        self.qhead = nn.Sequential(nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 1))
        with torch.no_grad():
            self.qhead[-1].weight.zero_()
            self.qhead[-1].bias.fill_(-5.0)
        self.T, self.n_sup, self.sigma, self.grad_mode = T, n_sup, sigma, grad_mode

    def aggregate(self, x, ei):
        h = self.emb(x)
        for conv, bn in zip(self.agg_convs, self.agg_bns):
            if self.view == 'gine':
                edge_attr = self.edge_emb(torch.zeros(ei.size(1), dtype=torch.long, device=ei.device))
                h = bn(conv(h, ei, edge_attr)).relu()
            else:
                h = bn(conv(h, ei)).relu()
        return h

    def core_step(self, combined, state):
        """Shared TRM compute core. Deliberately edge-free."""
        return state + self.core(self.core_norm(combined))

    def _z_step(self, y, z, ctx, noise):
        z = self.core_step(ctx + y + z, z)
        if noise and self.sigma > 0:
            z = z + self.sigma * torch.randn_like(z)
        return z

    def _y_step(self, y, z, noise):
        y = self.core_step(y + z, y)
        if noise and self.sigma > 0:
            y = y + self.sigma * torch.randn_like(y)
        return y

    def recurse(self, y, z, ctx, noise, one_step=False):
        if self.T == 0:
            return y, z
        if one_step:                                  # HRM 1-step gradient
            with torch.no_grad():
                for _ in range(self.T - 1):
                    z = self._z_step(y, z, ctx, noise)
            z = z.detach()
            z = self._z_step(y, z, ctx, noise)          # only last inner carries grad
            y = self._y_step(y, z, noise)
            return y, z
        for _ in range(self.T):                        # TRM full recursion
            z = self._z_step(y, z, ctx, noise)
        y = self._y_step(y, z, noise)
        return y, z

    def predict(self, y, batch):
        pooled = global_add_pool(y, batch)
        return self.head(pooled), self.qhead(pooled).view(-1)

    def forward_trace(self, x, ei, batch, steps, noise=False):
        ctx = self.aggregate(x, ei)
        y = ctx
        z = torch.zeros_like(ctx)
        preds, q_logits = [], []
        for s in range(steps):
            y, z = self.recurse(y, z, ctx, noise, one_step=(self.grad_mode == '1step'))
            pred, q = self.predict(y, batch)
            preds.append(pred)
            q_logits.append(q)
            if s < steps - 1:
                y, z = y.detach(), z.detach()
        return preds, q_logits

    def forward(self, x, ei, batch, noise=False):
        ctx = self.aggregate(x, ei)
        y = ctx
        z = torch.zeros_like(ctx)
        preds = []
        for s in range(self.n_sup):
            if s < self.n_sup - 1:
                with torch.no_grad():
                    y, z = self.recurse(y, z, ctx, noise)
                y, z = y.detach(), z.detach()
            else:
                y, z = self.recurse(y, z, ctx, noise, one_step=(self.grad_mode == '1step'))
            pred, _ = self.predict(y, batch)
            preds.append(pred)
        _, q = self.predict(y, batch)
        return preds, q


@torch.no_grad()
def evaluate(model, ld, dev, ymu, ysd, K=1, select='none'):
    model.eval()
    ysd_d, ymu_d = ysd.to(dev), ymu.to(dev)
    ae = torch.zeros(2); ae_or = torch.zeros(2); n = 0
    for b in ld:
        b = b.to(dev)
        if K == 1:
            preds, _ = model(b.x, b.edge_index, b.batch, noise=model.sigma > 0)
            chosen = oracle = preds[-1]
        else:
            P, Q = [], []
            for _ in range(K):
                preds, q = model(b.x, b.edge_index, b.batch, noise=True)
                P.append(preds[-1]); Q.append(q)
            P = torch.stack(P); Q = torch.stack(Q)
            ar = torch.arange(P.size(1), device=dev)
            chosen = P[Q.argmax(0), ar] if select == 'bestq' else P.mean(0)
            oracle = P[(P - b.y.unsqueeze(0)).abs().sum(-1).argmin(0), ar]
        ae += ((chosen * ysd_d + ymu_d) - (b.y * ysd_d + ymu_d)).abs().sum(0).cpu()
        ae_or += ((oracle * ysd_d + ymu_d) - (b.y * ysd_d + ymu_d)).abs().sum(0).cpu()
        n += b.num_graphs
    return (ae / n).tolist(), (ae_or / n).tolist()


@torch.no_grad()
def evaluate_trace(model, ld, dev, ymu, ysd, steps, adaptive=False):
    model.eval()
    ysd_d, ymu_d = ysd.to(dev), ymu.to(dev)
    ae = torch.zeros(2)
    n = 0
    step_sum = 0.0
    for b in ld:
        b = b.to(dev)
        preds, q_logits = model.forward_trace(b.x, b.edge_index, b.batch, steps, noise=False)
        P = torch.stack(preds, dim=0)
        if adaptive:
            Q = torch.stack(q_logits, dim=0)
            halted = Q > 0
            any_halt = halted.any(dim=0)
            first_halt = halted.to(torch.int64).argmax(dim=0)
            fallback = torch.full_like(first_halt, steps - 1)
            idx = torch.where(any_halt, first_halt, fallback)
            chosen = P[idx, torch.arange(P.size(1), device=dev)]
            step_sum += (idx.to(torch.float32) + 1).sum().item()
        else:
            chosen = P[-1]
            step_sum += steps * b.num_graphs
        ae += ((chosen * ysd_d + ymu_d) - (b.y * ysd_d + ymu_d)).abs().sum(0).cpu()
        n += b.num_graphs
    return (ae / n).tolist(), step_sum / max(n, 1)


def _split_nodes(t, ptr):
    return [t[ptr[i].item():ptr[i + 1].item()].detach() for i in range(ptr.numel() - 1)]


def act_train_step(model, state, replacement_batch, opt, dev, args):
    replacement = replacement_batch.to_data_list()
    batch_size = len(replacement)
    if state is None:
        state = {
            'graphs': [None for _ in range(batch_size)],
            'y': [None for _ in range(batch_size)],
            'z': [None for _ in range(batch_size)],
            'steps': torch.zeros(batch_size, dtype=torch.long, device=dev),
            'halted': torch.ones(batch_size, dtype=torch.bool, device=dev),
        }

    halted_cpu = state['halted'].detach().cpu().tolist()
    for i, halted in enumerate(halted_cpu):
        if halted:
            state['graphs'][i] = replacement[i]

    b = Batch.from_data_list(state['graphs']).to(dev)
    ctx = model.aggregate(b.x, b.edge_index)
    ptr = b.ptr
    y_parts, z_parts = [], []
    for i in range(batch_size):
        start, end = ptr[i].item(), ptr[i + 1].item()
        if halted_cpu[i] or state['y'][i] is None:
            y_parts.append(ctx[start:end])
            z_parts.append(torch.zeros_like(ctx[start:end]))
        else:
            y_parts.append(state['y'][i].to(dev))
            z_parts.append(state['z'][i].to(dev))
    y = torch.cat(y_parts, dim=0)
    z = torch.cat(z_parts, dim=0)

    opt.zero_grad()
    y, z = model.recurse(y, z, ctx, noise=False, one_step=(model.grad_mode == '1step'))
    pred, q = model.predict(y, b.batch)
    per_graph_err = (pred - b.y).abs().mean(1)
    pred_loss = per_graph_err.mean()
    with torch.no_grad():
        if args.halt_target == 'binary':
            halt_target = (per_graph_err <= args.halt_norm_threshold).to(q.dtype)
        else:
            halt_target = torch.sigmoid((args.halt_norm_threshold - per_graph_err) / args.halt_temp)
    q_loss = nn.functional.binary_cross_entropy_with_logits(q, halt_target)
    loss = pred_loss + 0.5 * args.lam_q * q_loss
    y_det, z_det = y.detach(), z.detach()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()

    state['y'] = _split_nodes(y_det, ptr)
    state['z'] = _split_nodes(z_det, ptr)
    with torch.no_grad():
        was_halted = state['halted']
        steps = torch.where(was_halted, torch.zeros_like(state['steps']), state['steps']) + 1
        halted = (steps >= args.halt_max_steps) | (q.detach() > 0)
        if args.halt_exploration_prob > 0 and args.halt_max_steps > 1:
            explore = torch.rand_like(q) < args.halt_exploration_prob
            min_steps = torch.where(
                explore,
                torch.randint(2, args.halt_max_steps + 1, steps.shape, device=dev),
                torch.zeros_like(steps),
            )
            halted = halted & (steps >= min_steps)
        state['steps'] = steps
        state['halted'] = halted

    return state, {
        'loss': float(loss.detach().cpu()),
        'pred_loss': float(pred_loss.detach().cpu()),
        'q_loss': float(q_loss.detach().cpu()),
        'halted_frac': float(state['halted'].to(torch.float32).mean().detach().cpu()),
        'steps': float(state['steps'].to(torch.float32).mean().detach().cpu()),
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--grad_mode', choices=['full', '1step'], default='full')
    ap.add_argument('--sigma', type=float, default=0.0)
    ap.add_argument('--K', type=int, default=1)
    ap.add_argument('--select', choices=['none', 'bestq'], default='bestq')
    ap.add_argument('--T', type=int, default=3)
    ap.add_argument('--n_sup', type=int, default=3)
    ap.add_argument('--hidden', type=int, default=128)
    ap.add_argument('--agg_layers', type=int, default=5)
    ap.add_argument('--compute_layers', type=int, default=2)
    ap.add_argument('--view', choices=SUPPORTED_VIEWS, default='gin')
    ap.add_argument('--epochs', type=int, default=200)
    ap.add_argument('--lr', type=float, default=1e-3)
    ap.add_argument('--bs', type=int, default=128)
    ap.add_argument('--lam_q', type=float, default=1.0)
    ap.add_argument('--act', action='store_true',
                    help='train all recurrent depths up to halt_max_steps and train qhead as a halt head')
    ap.add_argument('--halt_max_steps', type=int, default=8)
    ap.add_argument('--halt_norm_threshold', type=float, default=0.30)
    ap.add_argument('--halt_temp', type=float, default=0.10)
    ap.add_argument('--halt_target', choices=['soft', 'binary'], default='soft')
    ap.add_argument('--halt_exploration_prob', type=float, default=0.1)
    ap.add_argument('--loss_mode', choices=['last', 'trace'], default='trace')
    ap.add_argument('--seed', type=int, default=0)
    ap.add_argument('--device', default='auto')
    args = ap.parse_args()
    torch.manual_seed(args.seed); np.random.seed(args.seed)
    dev = 'cuda' if args.device == 'auto' and torch.cuda.is_available() else (
        'cpu' if args.device == 'auto' else args.device)
    os.makedirs(OUT, exist_ok=True)

    tr, va, te = prepare('train'), prepare('val'), prepare('test')
    n_atom = int(max(r['x'].max() for r in tr + va + te)) + 1
    Ytr = torch.stack([r['y'] for r in tr]); ymu, ysd = Ytr.mean(0), Ytr.std(0) + 1e-8
    for recs in (tr, va, te):
        for r in recs:
            r['y'] = (r['y'] - ymu) / ysd
    train_data = data_list(tr)
    trl = loader(train_data, args.bs, True, drop_last=True)
    val, tel = loader(va, 256, False), loader(te, 256, False)

    deg = degree_histogram(train_data) if args.view == 'pna' else None
    model = RecGIN(n_atom, args.hidden, args.T, args.n_sup, args.sigma, grad_mode=args.grad_mode,
                   agg_layers=args.agg_layers, compute_layers=args.compute_layers,
                   view=args.view, deg=deg).to(dev)
    opt = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-5)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, args.epochs)
    l1 = nn.L1Loss()
    act_steps = max(1, args.halt_max_steps)

    t0 = time.time(); best_val = 9e9; best = {}; best_state = None; act_state = None
    for ep in range(args.epochs):
        model.train()
        act_metrics = []
        for b in trl:
            if args.act:
                act_state, metrics = act_train_step(model, act_state, b, opt, dev, args)
                act_metrics.append(metrics)
            else:
                b = b.to(dev); opt.zero_grad()
                if args.loss_mode == 'trace':
                    preds, q_logits = model.forward_trace(
                        b.x, b.edge_index, b.batch, args.n_sup, noise=model.sigma > 0)
                    q = q_logits[-1]
                else:
                    preds, q = model(b.x, b.edge_index, b.batch, noise=model.sigma > 0)
                loss = sum(l1(p, b.y) for p in preds) / len(preds)
                with torch.no_grad():
                    tq = -(preds[-1] - b.y).abs().mean(1)
                loss = loss + args.lam_q * nn.functional.mse_loss(q, tq)
                loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()
        sched.step()
        if (ep + 1) % 20 == 0 or ep == args.epochs - 1:
            if args.act:
                vm, _ = evaluate_trace(model, val, dev, ymu, ysd, act_steps, adaptive=False)
            else:
                vm, _ = evaluate(model, val, dev, ymu, ysd, args.K, args.select)
            if sum(vm) < best_val:
                best_val = sum(vm)
                if args.act:
                    tem, fixed_steps = evaluate_trace(model, tel, dev, ymu, ysd, act_steps, adaptive=False)
                    tea, adaptive_steps = evaluate_trace(model, tel, dev, ymu, ysd, act_steps, adaptive=True)
                    best = {'ep': ep + 1, 'val_mae': vm, 'test_mae': tem,
                            'test_mae_adaptive': tea, 'fixed_steps': fixed_steps,
                            'adaptive_steps': adaptive_steps}
                else:
                    tem, teo = evaluate(model, tel, dev, ymu, ysd, args.K, args.select)
                    best = {'ep': ep + 1, 'val_mae': vm, 'test_mae': tem, 'test_mae_oracle': teo}
                best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
            if args.act and act_metrics:
                hm = sum(m['halted_frac'] for m in act_metrics) / len(act_metrics)
                sm = sum(m['steps'] for m in act_metrics) / len(act_metrics)
                print(f"ep{ep+1} val_mae={[round(x,3) for x in vm]} halt={hm:.2f} train_steps={sm:.2f}", flush=True)
            else:
                print(f"ep{ep+1} val_mae={[round(x,3) for x in vm]}", flush=True)

    act_tag = f"_actfull{act_steps}_{args.halt_target}{args.halt_norm_threshold:g}_e{args.epochs}" if args.act else ""
    loss_tag = f"_{args.loss_mode}" if (not args.act and args.loss_mode != 'last') else ""
    view_tag = f"_{args.view}" if args.view != 'gin' else ""
    tag = f"rec_rrog{view_tag}_{args.grad_mode}_sig{args.sigma}_K{args.K}_{args.select}_T{args.T}_ns{args.n_sup}{loss_tag}{act_tag}_s{args.seed}"
    rep = {'dataset': 'ZINC-cycle56', 'tag': tag, **vars(args), 'sec': round(time.time() - t0, 1),
           'dev': dev, 'arch': 'rrog_once_agg_node_compute', 'y_std_raw': ysd.tolist(), **best}
    if args.act:
        print(f"[{tag}] test_mae={[round(x,3) for x in best.get('test_mae')]} "
              f"adaptive={[round(x,3) for x in best.get('test_mae_adaptive')]} "
              f"steps={best.get('adaptive_steps'):.2f}/{best.get('fixed_steps'):.2f} "
              f"@ep{best.get('ep')} ({rep['sec']}s)")
    else:
        print(f"[{tag}] test_mae={[round(x,3) for x in best.get('test_mae')]} "
              f"oracle@K={[round(x,3) for x in best.get('test_mae_oracle')]} @ep{best.get('ep')} ({rep['sec']}s)")
    with open(os.path.join(OUT, f"{tag}.json"), 'w') as f:
        json.dump(rep, f, indent=2)
    torch.save({'state': best_state or model.state_dict(),
                'cfg': {'n_atom': n_atom, 'hidden': args.hidden, 'T': args.T, 'n_sup': args.n_sup,
                        'sigma': args.sigma, 'grad_mode': args.grad_mode,
                        'agg_layers': args.agg_layers, 'compute_layers': args.compute_layers,
                        'view': args.view,
                        'loss_mode': args.loss_mode,
                        'act': args.act, 'act_impl': 'persistent_recycle' if args.act else 'none',
                        'halt_max_steps': act_steps,
                        'halt_exploration_prob': args.halt_exploration_prob,
                        'arch': 'rrog_once_agg_node_compute'},
                'ymu': ymu, 'ysd': ysd}, os.path.join(OUT, f"ckpt_{tag}.pt"))
    print("  wrote", os.path.join(OUT, f"ckpt_{tag}.pt"))


if __name__ == "__main__":
    main()