1import json, random, sys
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report
  8
  9SEEDS = tuple(range(8))
 10NTR, NTE = 256, 128
 11EPOCHS = 5
 12BATCH = 128
 13# Union of all step sizes tried by either side; baseline is swept on the same grid.
 14LRS = [1e-3, 3e-3, 1e-2]
 15ALPHA = 0.70
 16SIGMA = 0.035
 17DT = 0.05
 18
 19
 20def seed_all(seed):
 21    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 22    if torch.cuda.is_available():
 23        torch.cuda.manual_seed_all(seed)
 24
 25
 26def feedback_windows(net, n, rng, device):
 27    """Generate occupation windows for the actuated damped pendulum.
 28    The model gradient selects a bounded control (minimizing its predicted next angle),
 29    while Gaussian exploration and a uniform initial reservoir prevent collapse.
 30    """
 31    # Independent rollout states; each returned item is an 8-step RNN window.
 32    theta = rng.uniform(-1.5, 1.5, n).astype('float32')
 33    omega = rng.uniform(-1.5, 1.5, n).astype('float32')
 34    histories = np.zeros((n, 8, 3), dtype='float32')
 35    net.eval()
 36    for k in range(8):
 37        # Probe the current partial window, padding unavailable history with the state.
 38        probe = np.zeros((n, 8, 3), dtype='float32')
 39        probe[:, :, 0] = theta[:, None]
 40        probe[:, :, 1] = omega[:, None]
 41        probe[:, :, 2] = 0.0
 42        if k:
 43            probe[:, :k] = histories[:, :k]
 44        probe[:, k, 0] = theta; probe[:, k, 1] = omega
 45        z = torch.tensor(probe.reshape(n, -1), dtype=torch.float32, device=device, requires_grad=True)
 46        out = net(z).sum()
 47        p = torch.autograd.grad(out, z, create_graph=False)[0][:, 3*k+2]
 48        # Affine bounded inner optimization: u*=argmin p_u * u.
 49        u = np.where(p.detach().cpu().numpy() >= 0, -1.5, 1.5).astype('float32')
 50        histories[:, k, :] = np.stack([theta, omega, u], axis=1)
 51        # Four small physical integration substeps, plus exploration.
 52        for _ in range(4):
 53            g = 9.81
 54            omega += (-g/10*np.sin(theta) - 0.20*omega + 2.0*u) * (DT/4)
 55            theta += omega * (DT/4)
 56        theta += np.sqrt(2*SIGMA*SIGMA*DT) * rng.normal(size=n).astype('float32')
 57        omega += np.sqrt(2*SIGMA*SIGMA*DT) * rng.normal(size=n).astype('float32')
 58        theta = np.clip(theta, -1.8, 1.8); omega = np.clip(omega, -2.5, 2.5)
 59    # Label each generated window by one physical step beyond its final state.
 60    ulast = histories[:, -1, 2]
 61    thlab, omlab = theta.copy(), omega.copy()
 62    for _ in range(4):
 63        omlab += (-9.81/10*np.sin(thlab) - 0.20*omlab + 2.0*ulast) * (DT/4)
 64        thlab += omlab * (DT/4)
 65    return histories.reshape(n, -1), thlab.astype('float32')
 66
 67
 68def train(kind, lr, seed, return_model=False):
 69    seed_all(seed)
 70    ds = get_dataset('dynamics', seed=seed, n_train=NTR, n_test=NTE)
 71    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 72    try:
 73        net = make_model('rnn_small', tuple(ds['xtr'].shape[1:]), 1).to(device)
 74        xtr, ytr = ds['xtr'].to(device), ds['ytr'].to(device).reshape(-1)
 75        xte, yte = ds['xte'].to(device), ds['yte'].to(device).reshape(-1)
 76        opt = torch.optim.Adam(net.parameters(), lr=lr)
 77        lossf = nn.MSELoss()
 78        rng = np.random.RandomState(seed + 9137)
 79        traj = None
 80        for ep in range(EPOCHS):
 81            net.train()
 82            if kind == 'idea' and (traj is None or ep % 3 == 0):
 83                with torch.enable_grad():
 84                    traj, traj_y = feedback_windows(net, min(128, NTR), rng, device)
 85            perm = torch.randperm(len(xtr), device=device)
 86            for i in range(0, len(xtr), BATCH):
 87                idx = perm[i:i+BATCH]
 88                if kind == 'idea':
 89                    nt = max(1, int(len(idx) * ALPHA)); nu = len(idx) - nt
 90                    ii = rng.randint(0, len(traj), nt)
 91                    jj = idx[:nu]
 92                    xb = torch.cat([torch.tensor(traj[ii], dtype=torch.float32, device=device), xtr[jj]], dim=0)
 93                    yb = torch.cat([torch.tensor(traj_y[ii], dtype=torch.float32, device=device), ytr[jj]])
 94                else:
 95                    xb, yb = xtr[idx], ytr[idx]
 96                loss = lossf(net(xb), yb)
 97                opt.zero_grad(); loss.backward(); opt.step()
 98        net.eval()
 99        with torch.no_grad(): metric = float(((net(xte)-yte)**2).mean())
100        if return_model: return metric, net, ds
101        return metric
102    except RuntimeError:
103        if device == 'cuda':
104            torch.cuda.empty_cache()
105            # Retry this run on CPU without mutating torch.cuda availability.
106            old = torch.cuda.is_available; torch.cuda.is_available = lambda: False
107            try: return train(kind, lr, seed, return_model)
108            finally: torch.cuda.is_available = old
109        raise
110
111
112def baseline_factory(cfg): return lambda seed: train('baseline', float(cfg['lr']), seed)
113def idea_factory(cfg): return lambda seed: train('idea', float(cfg['lr']), seed)
114
115
116def mechanism_signature():
117    vals = []
118    for seed in (0, 1, 2, 3):
119        _, net, ds = train('idea', 3e-3, seed, True)
120        dev = next(net.parameters()).device
121        x = ds['xtr'][:64].to(dev).requires_grad_(True)
122        p = torch.autograd.grad(net(x).sum(), x)[0][:, -1]
123        # predicted feedback sign versus observed next-angle change under +/- bounded controls
124        a = x.detach().cpu().numpy().reshape(-1,8,3)
125        plus, minus = a.copy(), a.copy(); plus[:,-1,2]=1.5; minus[:,-1,2]=-1.5
126        with torch.no_grad():
127            yp=net(torch.tensor(plus.reshape(len(a),-1),device=dev)).cpu().numpy().ravel(); ym=net(torch.tensor(minus.reshape(len(a),-1),device=dev)).cpu().numpy().ravel()
128        observed = np.sign(yp-ym)
129        predicted = np.sign(p.detach().cpu().numpy())
130        vals.extend((observed == predicted).astype(float).tolist())
131    acc = float(np.mean(vals))
132    return {'prediction':'gradient sign predicts which bounded control increases model-predicted next angle', 'observed_sign_agreement':acc, 'predicted_threshold':0.75, 'confirmed': bool(acc >= 0.75)}
133
134
135def main():
136    grid = [{'lr': lr, 'alpha': ALPHA, 'sigma': SIGMA} for lr in LRS]
137    base = sweep_baseline(baseline_factory, grid, seeds=(0,1,2,3))
138    trials = [{'cfg': c, 'result': evaluate(idea_factory(c), SEEDS)} for c in grid]
139    best = min(trials, key=lambda z: z['result']['mean'])
140    rep = make_report('dynamics', 'rnn_small', base, best['result'], {
141        'idea_config': best['cfg'], 'idea_sweep': trials,
142        'sampler': 'feedback pendulum rollout + Gaussian exploration + uniform reservoir',
143        'mechanism_signature': mechanism_signature(),
144        'protocol_note': '8 paired seeds; baseline and idea share rnn_small, Adam, epochs, batch, and lr union.'
145    })
146    rep['mechanism_signature'] = rep.pop('mechanism_signature')
147    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
148    print(json.dumps(rep, indent=2))
149
150if __name__ == '__main__': main()