import sys, json, random from pathlib import Path 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, train_model, evaluate, sweep_baseline, make_report SEED=2137 EPOCHS=10 NTR=1200 NTE=400 BATCH=128 BETA=.8 LAMBDA=.08 def H(z, delta, P, beta): # z,delta [B,S], P [B,S,S] m=np.min(z+delta,axis=0) return delta+beta*np.einsum('bxy,y->bx',P,m) def math_check(): rng=np.random.default_rng(SEED); B,S=4,7; beta=.73 P=rng.random((B,S,S)); P/=P.sum(2,keepdims=True); delta=rng.random((B,S)) z=rng.normal(size=(B,S)); w=rng.normal(size=(B,S)) ratio=np.max(abs(H(z,delta,P,beta)-H(w,delta,P,beta)))/np.max(abs(z-w)) star=np.zeros_like(delta) for _ in range(1000): new=H(star,delta,P,beta) if np.max(abs(new-star))<1e-13: break star=new cur=np.zeros_like(delta); errs=[] for _ in range(9): errs.append(float(np.max(abs(cur-star)))); cur=H(cur,delta,P,beta) ratios=[errs[i+1]/errs[i] for i in range(8) if errs[i]>1e-14] return {'operator_lipschitz_ratio':float(ratio),'beta':beta,'iteration_ratios':ratios, 'max_ratio_over_beta':float(max(ratios)/beta)} 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 branch_loss(net, x, y, beta=BETA): # Three nearby action branches: perturb only the most recent control u. xb=x.view(-1,8,3); offsets=torch.tensor([-0.35,0.,0.35],device=x.device) outs=[] for off in offsets: q=xb.clone(); q[:,-1,2]+=off; outs.append(net(q.reshape(len(x),-1))) pred=torch.cat(outs,1) # [n,3] # Local deficit marks relative to best candidate, with target supervision. delta=(pred-y[:,None]).pow(2); delta=delta-delta.min(1,keepdim=True).values # Learned transition consequence proxy: branch predictions are continuation outcomes. # Iterate finite Bellman map independently per sample (one-state branch pool). z=torch.zeros_like(delta) for _ in range(8): m=(z+delta).min(1,keepdim=True).values z=delta+beta*m scores=z+delta weights=torch.softmax(-scores/.12,dim=1) resolved=(weights*pred).sum(1) # Central action remains the task action; resolver consistency is auxiliary. return (pred[:,1]-y).pow(2).mean()+LAMBDA*(resolved-y).pow(2).mean(), weights.detach(), pred.detach() def train_idea(seed, lr, wd, collect=False): seed_all(seed); ds=get_dataset('dynamics',seed,n_train=NTR,n_test=NTE) net=make_model('rnn_small',tuple(ds['xtr'].shape[1:]),1) device='cuda' if torch.cuda.is_available() else 'cpu' try: net.to(device); x,y=ds['xtr'].to(device),ds['ytr'].to(device) opt=torch.optim.Adam(net.parameters(),lr=lr,weight_decay=wd) for _ in range(EPOCHS): net.train(); p=torch.randperm(len(x),device=device) for i in range(0,len(x),BATCH): idx=p[i:i+BATCH]; loss,_,_=branch_loss(net,x[idx],y[idx]) opt.zero_grad(); loss.backward(); opt.step() net.eval() with torch.no_grad(): xt,yt=ds['xte'].to(device),ds['yte'].to(device) loss,w,pred=branch_loss(net,xt,yt) metric=float((pred[:,1]-yt).pow(2).mean()) # behavior signature: smoothness of resolver weights under action perturbation smooth=float(torch.mean(torch.abs(w[:,2]-w[:,0]))) hard=float((pred.argmin(1)==0).float().mean()) if collect: return metric, {'resolver_weight_span':smooth,'hard_best_frequency':hard} return metric except RuntimeError: # CPU retry is explicit for shared-GPU failures. seed_all(seed); ds=get_dataset('dynamics',seed,n_train=NTR,n_test=NTE) net=make_model('rnn_small',tuple(ds['xtr'].shape[1:]),1).cpu(); x,y=ds['xtr'],ds['ytr'] opt=torch.optim.Adam(net.parameters(),lr=lr,weight_decay=wd) for _ in range(EPOCHS): p=torch.randperm(len(x)) for i in range(0,len(x),BATCH): idx=p[i:i+BATCH]; loss,_,_=branch_loss(net,x[idx],y[idx]); opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): pred=torch.cat([net(x.view(-1,8,3).clone().reshape(len(x),-1)) for _ in [0]],1) return float(((net(ds['xte'])[:,0]-ds['yte'])**2).mean()) def baseline_fn(cfg): def run(seed): seed_all(seed); ds=get_dataset('dynamics',seed,n_train=NTR,n_test=NTE) net=make_model('rnn_small',tuple(ds['xtr'].shape[1:]),1) _,metric,_=train_model(net,ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=cfg['wd'],log=lambda *_:None) return metric return run def main(): grid=[{'lr':lr,'wd':wd} for lr in [1e-3,3e-3,5e-3] for wd in [0.,1e-4]] base=sweep_baseline(baseline_fn,grid) # Same lr union is evaluated by baseline; idea has 3 nearby settings at best wd. wd=base['best_cfg']['wd']; idea_cfgs=[{'lr':lr,'wd':wd} for lr in [1e-3,3e-3,5e-3]] idea_cfg_results=[] for cfg in idea_cfgs: r=evaluate(lambda s,cfg=cfg:train_idea(s,cfg['lr'],cfg['wd']),seeds=range(8)) idea_cfg_results.append((cfg,r)) idea_cfg,best=min(idea_cfg_results,key=lambda q:q[1]['mean']) sig=[] for s in range(8): sig.append(train_idea(s,idea_cfg['lr'],idea_cfg['wd'],True)[1]) sigkeys=sig[0].keys(); signature={k:{'observed_mean':float(np.mean([a[k] for a in sig])),'per_seed': [float(a[k]) for a in sig]} for k in sigkeys} signature.update({'prediction':'resolver branch weights vary continuously with counterfactual action consequences; measured on trained models','confirmed':False}) rep=make_report('dynamics','rnn_small',base,best,{'math_check':math_check(),'trained_behavior':signature,'idea_sweep':[{'cfg':c,'mean':r['mean']} for c,r in idea_cfg_results]}) rep['protocol_notes']={'n_train':NTR,'n_test':NTE,'epochs':EPOCHS,'structural_match':'controlled pendulum dynamics; action branches are final-step control perturbations','baseline_knobs_swept':['lr','weight_decay']} Path('bench_report.json').write_text(json.dumps(rep,indent=2)) print(json.dumps(rep,indent=2)) if __name__=='__main__': main()