train_weights.py

import random,json,numpy as np,torch,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)
def mlp(s):
 q=[]
 for i in range(len(s)-1):q+=[nn.Linear(s[i],s[i+1])]+([nn.GELU()] if i<len(s)-2 else [])
 return nn.Sequential(*q)
class E:
 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.o()
 def o(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.array((f+[0]*32)[:32],np.float32)
 def ex(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))]
  return self.o(),self.t>=120 or abs(self.x)>1 or any(0<o[1]<.12 and abs(o[0]-self.x)<.18 for o in self.os)
class U(nn.Module):
 def __init__(self):super().__init__();self.f=mlp([256,768,768,256]);self.n=nn.LayerNorm(256)
 def forward(self,x):return self.n(x+self.f(x))
class M(nn.Module):
 def __init__(self):super().__init__();self.enc=mlp([32,512,256]);self.chain=nn.ModuleList([U() for _ in range(8)]);self.policy=mlp([2304,2048,1024,512,4])
 def forward(self,o):
  z=self.enc(o);zs=[z]
  for u in self.chain:z=u(z);zs.append(z)
  return self.policy(torch.cat(zs,-1))
O=[];Y=[]
for ep in range(3000):
 e=E(100000+ep);o=e.reset();done=False
 while not done:
  a=e.ex();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));c=int(.9*len(O));tr=ix[:c];va=ix[c:];m=M().cuda();opt=torch.optim.AdamW(m.parameters(),3e-4,weight_decay=1e-4);best=-1;sd=None;be=0
print('START GPU',torch.cuda.get_device_name(0),'PARAMS',sum(p.numel() for p in m.parameters()),flush=True)
for ep in range(20):
 p=tr[torch.randperm(len(tr))]
 for i in range(0,len(p),512):
  b=p[i:i+512];loss=F.cross_entropy(m(O[b].cuda()),Y[b].cuda());opt.zero_grad();loss.backward();torch.nn.utils.clip_grad_norm_(m.parameters(),1);opt.step()
 with torch.no_grad():v=(m(O[va[:10000]].cuda()).argmax(1).cpu()==Y[va[:10000]]).float().mean().item()
 print('EPOCH',ep+1,'VAL_ACC',v,flush=True)
 if v>best:best=v;be=ep+1;sd={k:x.cpu().clone() for k,x in m.state_dict().items()}
torch.save({'state_dict':sd,'val_acc':best,'epoch':be,'seed':S},'/out/model.pt');open('/out/acc.json','w').write(json.dumps({'best_val_acc':best,'best_epoch':be,'seed':S},indent=2));print('PERSISTED',best,be,flush=True)