import sys, json, math, random from pathlib import Path import numpy as np import torch import torch.nn as nn import torch.nn.functional as F sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import make_report import obb_track SEEDS = tuple(range(8)) SWEEP_SEEDS = tuple(range(4)) LRS = (1e-3, 3e-3, 1e-2) EPOCHS = 24 NTR, NTE = 400, 200 def axes(y): u = torch.stack((torch.cos(y), torch.sin(y)), -1) return u, torch.stack((-torch.sin(y), torch.cos(y)), -1) def sat_margin(ego, obs, tau=None): # ego/obs: (...,2); observed boxes have fixed dimensions and headings. ce, co = ego, obs ye = torch.zeros_like(ce[..., 0]); yo = obs[..., 2] ue, ve = axes(ye); uo, vo = axes(yo) aa, ba = .5, .3 ab, bb = .6, .4 ax = torch.stack((ue, ve, uo, vo), -2) d = co[..., :2] - ce proj = (ax*d.unsqueeze(-2)).sum(-1).abs() re = aa*(ue.unsqueeze(-2)*ax).sum(-1).abs() + ba*(ve.unsqueeze(-2)*ax).sum(-1).abs() ro = ab*(uo.unsqueeze(-2)*ax).sum(-1).abs() + bb*(vo.unsqueeze(-2)*ax).sum(-1).abs() gaps = proj-re-ro if tau is None: return gaps.max(-1).values return tau*torch.logsumexp(gaps/tau, -1) def math_check(): # Axis-aligned length-1 boxes cross at center separation 1.0. ds = torch.linspace(.5, 1.5, 201) e = torch.zeros(201,2); e[:,0] = 0 o = torch.zeros(201,4); o[:,0] = ds exact = sat_margin(e,o) sm = sat_margin(e,o,.1) exact_zero = float(ds[torch.argmin(exact.abs())]) smooth_zero = float(ds[torch.argmin(sm.abs())]) gaps = torch.tensor([[1., .2, -.3, -.8]]) bound_rows=[] for tau in (.2,.1,.05): err=float((tau*torch.logsumexp(gaps/tau,-1)-gaps.max()).item()) bound_rows.append({'tau':tau,'error':err,'bound':tau*math.log(4)}) x=torch.tensor([[0.,0.]],requires_grad=True); oo=torch.tensor([[1.5,0.,0.,0.]]) sat_margin(x,oo,.1).backward() return {'predicted_exact_zero_m':1.1,'observed_exact_zero_m':exact_zero, 'observed_smooth_zero_m':smooth_zero,'smoothing_bound':bound_rows, 'active_axis_gradient_dx':float(x.grad[0,0]), 'passed': abs(exact_zero-1)<.01 and all(abs(r['error'])<=r['bound']+1e-6 for r in bound_rows)} class Net(nn.Module): def __init__(self): super().__init__(); self.net=nn.Sequential(nn.Linear(4,32),nn.ReLU(),nn.Linear(32,32),nn.ReLU(),nn.Linear(32,16)) def forward(self,x): return self.net(x) def train_one(seed, lr, barrier): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) d=obb_track.get_dataset(seed,NTR,NTE) dev=torch.device('cuda' if torch.cuda.is_available() else 'cpu') try: net=Net().to(dev); x=torch.tensor(d['xtr'],device=dev); y=torch.tensor(d['ytr'],device=dev) opt=torch.optim.Adam(net.parameters(),lr=lr) for _ in range(EPOCHS): pred=net(x).view(-1,8,2) loss=F.mse_loss(pred,y.view(-1,8,2)) if barrier: obs=torch.tensor(d['xtr'][:,1:],device=dev)[:,None,:].expand(-1,8,-1) m=sat_margin(pred,obs,.1) loss=loss + .08*F.softplus((.05-m)/.05).mean() opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): pred=net(torch.tensor(d['xte'],device=dev)).view(-1,8,2) metric=float(F.mse_loss(pred,torch.tensor(d['yte'],device=dev).view(-1,8,2)).cpu()) return metric, net.cpu(), d except Exception: # Robust CPU fallback, matching the requested device policy. torch.cuda.empty_cache() if torch.cuda.is_available() else None net=Net(); x=torch.tensor(d['xtr']); y=torch.tensor(d['ytr']); opt=torch.optim.Adam(net.parameters(),lr=lr) for _ in range(EPOCHS): pred=net(x).view(-1,8,2); loss=F.mse_loss(pred,y.view(-1,8,2)) if barrier: obs=torch.tensor(d['xtr'][:,1:])[:,None,:].expand(-1,8,-1) loss=loss+.08*F.softplus((.05-sat_margin(pred,obs,.1))/.05).mean() opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): metric=float(F.mse_loss(net(torch.tensor(d['xte'])).view(-1,8,2),torch.tensor(d['yte']).view(-1,8,2))) return metric,net,d def eval_cfg(lr, barrier, seeds=SEEDS, keep=False): vals=[]; sig=[] for s in seeds: metric,net,d=train_one(s,lr,barrier); vals.append(metric) if keep: with torch.no_grad(): p=net(torch.tensor(d['xte'])).view(-1,8,2) o=torch.tensor(d['xte'][:,1:])[:,None,:].expand(-1,8,-1) mm=sat_margin(p,o) sig.append((float(mm.mean()),float(mm.min()),float((mm<=0).float().mean()))) out={'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':vals,'n':len(vals)} if keep: out['signature_values']=sig return out def main(): check=math_check() # Equal search-space union: baseline and idea each run every candidate lr. bsweep=[]; isweep=[] for lr in LRS: bs=eval_cfg(lr,False,SWEEP_SEEDS)['mean']; ii=eval_cfg(lr,True,SWEEP_SEEDS)['mean'] bsweep.append({'cfg':{'lr':lr,'epochs':EPOCHS},'mean':bs}) isweep.append({'cfg':{'lr':lr,'epochs':EPOCHS,'lambda_barrier':.08,'tau':.1,'m0':.05,'beta':.05},'mean':ii}) blr=min(LRS,key=lambda z: next(q['mean'] for q in bsweep if q['cfg']['lr']==z)) ilr=min(LRS,key=lambda z: next(q['mean'] for q in isweep if q['cfg']['lr']==z)) base={'best_cfg':{'lr':blr,'epochs':EPOCHS},'sweep':bsweep,'full':eval_cfg(blr,False)} idea=eval_cfg(ilr,True,SEEDS,keep=True) vals=np.asarray(idea['signature_values']); sig={'predicted_vs_observed':{ 'predicted_smoothed_margin_mean':float(vals[:,0].mean()), 'observed_exact_sat_margin_mean':float(vals[:,1].mean()), 'predicted_collision_fraction':float(vals[:,2].mean())}, 'prediction':'barrier should reduce collision fraction and move margins upward', 'confirmed':bool(vals[:,2].mean()<.5 and vals[:,0].mean()>-0.5)} rep=make_report('obb_clearance_trajectory','mlp_tiny',base,idea,{'math_check':check,'idea_sweep':isweep,**sig}) Path('bench_report.json').write_text(json.dumps(rep,indent=2)) print(json.dumps(rep,indent=2)) if __name__=='__main__': main()