import json, random, sys import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) NTR, NTE = 256, 128 EPOCHS = 5 BATCH = 128 # Union of all step sizes tried by either side; baseline is swept on the same grid. LRS = [1e-3, 3e-3, 1e-2] ALPHA = 0.70 SIGMA = 0.035 DT = 0.05 def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def feedback_windows(net, n, rng, device): """Generate occupation windows for the actuated damped pendulum. The model gradient selects a bounded control (minimizing its predicted next angle), while Gaussian exploration and a uniform initial reservoir prevent collapse. """ # Independent rollout states; each returned item is an 8-step RNN window. theta = rng.uniform(-1.5, 1.5, n).astype('float32') omega = rng.uniform(-1.5, 1.5, n).astype('float32') histories = np.zeros((n, 8, 3), dtype='float32') net.eval() for k in range(8): # Probe the current partial window, padding unavailable history with the state. probe = np.zeros((n, 8, 3), dtype='float32') probe[:, :, 0] = theta[:, None] probe[:, :, 1] = omega[:, None] probe[:, :, 2] = 0.0 if k: probe[:, :k] = histories[:, :k] probe[:, k, 0] = theta; probe[:, k, 1] = omega z = torch.tensor(probe.reshape(n, -1), dtype=torch.float32, device=device, requires_grad=True) out = net(z).sum() p = torch.autograd.grad(out, z, create_graph=False)[0][:, 3*k+2] # Affine bounded inner optimization: u*=argmin p_u * u. u = np.where(p.detach().cpu().numpy() >= 0, -1.5, 1.5).astype('float32') histories[:, k, :] = np.stack([theta, omega, u], axis=1) # Four small physical integration substeps, plus exploration. for _ in range(4): g = 9.81 omega += (-g/10*np.sin(theta) - 0.20*omega + 2.0*u) * (DT/4) theta += omega * (DT/4) theta += np.sqrt(2*SIGMA*SIGMA*DT) * rng.normal(size=n).astype('float32') omega += np.sqrt(2*SIGMA*SIGMA*DT) * rng.normal(size=n).astype('float32') theta = np.clip(theta, -1.8, 1.8); omega = np.clip(omega, -2.5, 2.5) # Label each generated window by one physical step beyond its final state. ulast = histories[:, -1, 2] thlab, omlab = theta.copy(), omega.copy() for _ in range(4): omlab += (-9.81/10*np.sin(thlab) - 0.20*omlab + 2.0*ulast) * (DT/4) thlab += omlab * (DT/4) return histories.reshape(n, -1), thlab.astype('float32') def train(kind, lr, seed, return_model=False): seed_all(seed) ds = get_dataset('dynamics', seed=seed, n_train=NTR, n_test=NTE) device = 'cuda' if torch.cuda.is_available() else 'cpu' try: net = make_model('rnn_small', tuple(ds['xtr'].shape[1:]), 1).to(device) xtr, ytr = ds['xtr'].to(device), ds['ytr'].to(device).reshape(-1) xte, yte = ds['xte'].to(device), ds['yte'].to(device).reshape(-1) opt = torch.optim.Adam(net.parameters(), lr=lr) lossf = nn.MSELoss() rng = np.random.RandomState(seed + 9137) traj = None for ep in range(EPOCHS): net.train() if kind == 'idea' and (traj is None or ep % 3 == 0): with torch.enable_grad(): traj, traj_y = feedback_windows(net, min(128, NTR), rng, device) perm = torch.randperm(len(xtr), device=device) for i in range(0, len(xtr), BATCH): idx = perm[i:i+BATCH] if kind == 'idea': nt = max(1, int(len(idx) * ALPHA)); nu = len(idx) - nt ii = rng.randint(0, len(traj), nt) jj = idx[:nu] xb = torch.cat([torch.tensor(traj[ii], dtype=torch.float32, device=device), xtr[jj]], dim=0) yb = torch.cat([torch.tensor(traj_y[ii], dtype=torch.float32, device=device), ytr[jj]]) else: xb, yb = xtr[idx], ytr[idx] loss = lossf(net(xb), yb) opt.zero_grad(); loss.backward(); opt.step() net.eval() with torch.no_grad(): metric = float(((net(xte)-yte)**2).mean()) if return_model: return metric, net, ds return metric except RuntimeError: if device == 'cuda': torch.cuda.empty_cache() # Retry this run on CPU without mutating torch.cuda availability. old = torch.cuda.is_available; torch.cuda.is_available = lambda: False try: return train(kind, lr, seed, return_model) finally: torch.cuda.is_available = old raise def baseline_factory(cfg): return lambda seed: train('baseline', float(cfg['lr']), seed) def idea_factory(cfg): return lambda seed: train('idea', float(cfg['lr']), seed) def mechanism_signature(): vals = [] for seed in (0, 1, 2, 3): _, net, ds = train('idea', 3e-3, seed, True) dev = next(net.parameters()).device x = ds['xtr'][:64].to(dev).requires_grad_(True) p = torch.autograd.grad(net(x).sum(), x)[0][:, -1] # predicted feedback sign versus observed next-angle change under +/- bounded controls a = x.detach().cpu().numpy().reshape(-1,8,3) plus, minus = a.copy(), a.copy(); plus[:,-1,2]=1.5; minus[:,-1,2]=-1.5 with torch.no_grad(): 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() observed = np.sign(yp-ym) predicted = np.sign(p.detach().cpu().numpy()) vals.extend((observed == predicted).astype(float).tolist()) acc = float(np.mean(vals)) 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)} def main(): grid = [{'lr': lr, 'alpha': ALPHA, 'sigma': SIGMA} for lr in LRS] base = sweep_baseline(baseline_factory, grid, seeds=(0,1,2,3)) trials = [{'cfg': c, 'result': evaluate(idea_factory(c), SEEDS)} for c in grid] best = min(trials, key=lambda z: z['result']['mean']) rep = make_report('dynamics', 'rnn_small', base, best['result'], { 'idea_config': best['cfg'], 'idea_sweep': trials, 'sampler': 'feedback pendulum rollout + Gaussian exploration + uniform reservoir', 'mechanism_signature': mechanism_signature(), 'protocol_note': '8 paired seeds; baseline and idea share rnn_small, Adam, epochs, batch, and lr union.' }) rep['mechanism_signature'] = rep.pop('mechanism_signature') with open('bench_report.json','w') as f: json.dump(rep,f,indent=2) print(json.dumps(rep, indent=2)) if __name__ == '__main__': main()