axial2d_batchacc_source.py

import random,json,time,numpy as np,torch
import torch.nn as nn,torch.nn.functional as F
S=20260927;random.seed(S);np.random.seed(S);torch.manual_seed(S);torch.cuda.manual_seed_all(S)
class Env:
 def __init__(self,s):self.r=np.random.default_rng(s);self.reset()
 def reset(self):self.x=float(self.r.uniform(-.15,.15));self.v=float(self.r.uniform(.3,.45));self.h=float(self.r.uniform(-.03,.03));self.t=0;self.os=[[float(self.r.uniform(-.8,.8)),float(self.r.uniform(.18,1)),float(self.r.uniform(-.025,.025))] for _ in range(5)];return self.obs()
 def obs(self):
  f=[self.x,self.v,self.h,self.t/120,1-self.x,1+self.x]
  for x,y,vx in sorted(self.os,key=lambda z:z[1]):f += [x,y,vx,y/max(self.v,.05)]
  return np.asarray((f+[0]*32)[:32],np.float32)
 def expert(self):
  th=[o for o in self.os if 0<o[1]<.34 and abs(o[0]-self.x)<.32]
  if th:
   o=min(th,key=lambda z:z[1]);return 0 if o[0]>=self.x and self.x>-.75 else (1 if self.x<.75 else 3)
  if self.x>.2:return 0
  if self.x<-.2:return 1
  return 2 if self.v<.72 else 3
 def step(self,a):
  self.h+=-.035 if a==0 else .035 if a==1 else 0;self.v+=.035 if a==2 else -.05 if a==3 else 0;self.h*=.82;self.v=float(np.clip(self.v,.05,.92));self.x+=self.h+float(self.r.normal(0,.005));self.t+=1
  for o in self.os:
   o[0]+=o[2];o[1]-=self.v*.055
   if o[1]<-.12:o[:]=[float(self.r.uniform(-.85,.85)),float(self.r.uniform(.65,1.2)),float(self.r.uniform(-.025,.025))]
  c=any(0<o[1]<.12 and abs(o[0]-self.x)<.18 for o in self.os);off=abs(self.x)>1;return self.obs(),c or off or self.t>=120,c,off
class Ax2(nn.Module):
 def __init__(self,D=96,H=6):
  super().__init__();self.n1=nn.LayerNorm(D);self.n2=nn.LayerNorm(D);self.n3=nn.LayerNorm(D);self.row=nn.MultiheadAttention(D,H,batch_first=True);self.col=nn.MultiheadAttention(D,H,batch_first=True);self.ff=nn.Sequential(nn.Linear(D,384),nn.GELU(),nn.Linear(384,D))
 def forward(self,x):
  B,R,C,D=x.shape;y=self.n1(x).reshape(B*R,C,D);y=self.row(y,y,y,need_weights=False)[0].reshape(B,R,C,D);x=x+y
  y=self.n2(x).permute(0,2,1,3);sh=y.shape;y=y.reshape(B*C,R,D);y=self.col(y,y,y,need_weights=False)[0].reshape(sh).permute(0,2,1,3);x=x+y
  return x+self.ff(self.n3(x))
class M(nn.Module):
 def __init__(self,D=96,K=8):
  super().__init__();self.e=nn.Linear(1,D);self.pos=nn.Parameter(torch.randn(1,4,8,D)*.02);self.pre=Ax2(D);self.chain=nn.ModuleList([Ax2(D) for _ in range(K)]);self.p=nn.Sequential(nn.Linear(D*(K+1),1536),nn.GELU(),nn.Linear(1536,768),nn.GELU(),nn.Linear(768,384),nn.GELU(),nn.Linear(384,4))
 def forward(self,o):
  x=self.e(o.reshape(-1,4,8,1))+self.pos;x=self.pre(x);zs=[x.mean((1,2))]
  for u in self.chain:x=u(x);zs.append(x.mean((1,2)))
  return self.p(torch.cat(zs,-1))
print('GENERATING_DATA',flush=True);O=[];Y=[]
for ep in range(3000):
 e=Env(100000+ep);o=e.reset();done=False
 while not done:a=e.expert();O.append(o);Y.append(a);o,done,_,_=e.step(a)
O=torch.tensor(np.asarray(O));Y=torch.tensor(Y);ix=torch.randperm(len(O));cut=int(.9*len(O));tr=ix[:cut];va=ix[cut:];m=M().cuda();opt=torch.optim.AdamW(m.parameters(),3e-4,weight_decay=1e-4);best=-1;be=0;sd=None
print('TRAIN_START','SAMPLES',len(O),'PARAMS',sum(p.numel() for p in m.parameters()),flush=True)
for ep in range(20):
 t=time.time();p=tr[torch.randperm(len(tr))];m.train()
 for j,i in enumerate(range(0,len(p),128)):
  b=p[i:i+128];loss=F.cross_entropy(m(O[b].cuda()),Y[b].cuda());opt.zero_grad(set_to_none=True);loss.backward();torch.nn.utils.clip_grad_norm_(m.parameters(),1);opt.step()
  if j%100==0:
   m.eval(); vv=va[:1024]; pp=[]
   with torch.no_grad():
    for q in range(0,len(vv),128):pp.append(m(O[vv[q:q+128]].cuda()).argmax(1).cpu())
   ba=(torch.cat(pp)==Y[vv]).float().mean().item();m.train();print('E',ep+1,'B',j,'LOSS',round(float(loss),6),'VAL_ACC',ba,flush=True)
 m.eval()
 with torch.no_grad():vb=va[:2048];pred=[]
 with torch.no_grad():
  for i in range(0,len(vb),128):pred.append(m(O[vb[i:i+128]].cuda()).argmax(1).cpu())
  v=(torch.cat(pred)==Y[vb]).float().mean().item()
 print('EPOCH',ep+1,'VAL_ACC',v,'SEC',round(time.time()-t,1),flush=True)
 if v>best:best=v;be=ep+1;sd={k:x.detach().cpu().clone() for k,x in m.state_dict().items()}
 torch.save({'state_dict':sd,'val_acc':best,'epoch':be,'seed':S,'architecture':'2D axial row+column + 8 chained axial blocks + policy'},'/out/axial2d_chain_policy.pt')
 open('/out/axial2d_chain_policy_acc.json','w').write(json.dumps({'best_val_acc':best,'best_epoch':be,'seed':S,'params':sum(p.numel() for p in m.parameters())},indent=2))
print('FINAL_BEST',best,'EPOCH',be,'PERSISTED',flush=True)