1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6import torch.nn.functional as F
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import make_report
  9import obb_track
 10
 11SEEDS = tuple(range(8))
 12SWEEP_SEEDS = tuple(range(4))
 13LRS = (1e-3, 3e-3, 1e-2)
 14EPOCHS = 24
 15NTR, NTE = 400, 200
 16
 17def axes(y):
 18    u = torch.stack((torch.cos(y), torch.sin(y)), -1)
 19    return u, torch.stack((-torch.sin(y), torch.cos(y)), -1)
 20
 21def sat_margin(ego, obs, tau=None):
 22    # ego/obs: (...,2); observed boxes have fixed dimensions and headings.
 23    ce, co = ego, obs
 24    ye = torch.zeros_like(ce[..., 0]); yo = obs[..., 2]
 25    ue, ve = axes(ye); uo, vo = axes(yo)
 26    aa, ba = .5, .3
 27    ab, bb = .6, .4
 28    ax = torch.stack((ue, ve, uo, vo), -2)
 29    d = co[..., :2] - ce
 30    proj = (ax*d.unsqueeze(-2)).sum(-1).abs()
 31    re = aa*(ue.unsqueeze(-2)*ax).sum(-1).abs() + ba*(ve.unsqueeze(-2)*ax).sum(-1).abs()
 32    ro = ab*(uo.unsqueeze(-2)*ax).sum(-1).abs() + bb*(vo.unsqueeze(-2)*ax).sum(-1).abs()
 33    gaps = proj-re-ro
 34    if tau is None: return gaps.max(-1).values
 35    return tau*torch.logsumexp(gaps/tau, -1)
 36
 37def math_check():
 38    # Axis-aligned length-1 boxes cross at center separation 1.0.
 39    ds = torch.linspace(.5, 1.5, 201)
 40    e = torch.zeros(201,2); e[:,0] = 0
 41    o = torch.zeros(201,4); o[:,0] = ds
 42    exact = sat_margin(e,o)
 43    sm = sat_margin(e,o,.1)
 44    exact_zero = float(ds[torch.argmin(exact.abs())])
 45    smooth_zero = float(ds[torch.argmin(sm.abs())])
 46    gaps = torch.tensor([[1., .2, -.3, -.8]])
 47    bound_rows=[]
 48    for tau in (.2,.1,.05):
 49        err=float((tau*torch.logsumexp(gaps/tau,-1)-gaps.max()).item())
 50        bound_rows.append({'tau':tau,'error':err,'bound':tau*math.log(4)})
 51    x=torch.tensor([[0.,0.]],requires_grad=True); oo=torch.tensor([[1.5,0.,0.,0.]])
 52    sat_margin(x,oo,.1).backward()
 53    return {'predicted_exact_zero_m':1.1,'observed_exact_zero_m':exact_zero,
 54            'observed_smooth_zero_m':smooth_zero,'smoothing_bound':bound_rows,
 55            'active_axis_gradient_dx':float(x.grad[0,0]),
 56            'passed': abs(exact_zero-1)<.01 and all(abs(r['error'])<=r['bound']+1e-6 for r in bound_rows)}
 57
 58class Net(nn.Module):
 59    def __init__(self):
 60        super().__init__(); self.net=nn.Sequential(nn.Linear(4,32),nn.ReLU(),nn.Linear(32,32),nn.ReLU(),nn.Linear(32,16))
 61    def forward(self,x): return self.net(x)
 62
 63def train_one(seed, lr, barrier):
 64    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 65    d=obb_track.get_dataset(seed,NTR,NTE)
 66    dev=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 67    try:
 68        net=Net().to(dev); x=torch.tensor(d['xtr'],device=dev); y=torch.tensor(d['ytr'],device=dev)
 69        opt=torch.optim.Adam(net.parameters(),lr=lr)
 70        for _ in range(EPOCHS):
 71            pred=net(x).view(-1,8,2)
 72            loss=F.mse_loss(pred,y.view(-1,8,2))
 73            if barrier:
 74                obs=torch.tensor(d['xtr'][:,1:],device=dev)[:,None,:].expand(-1,8,-1)
 75                m=sat_margin(pred,obs,.1)
 76                loss=loss + .08*F.softplus((.05-m)/.05).mean()
 77            opt.zero_grad(); loss.backward(); opt.step()
 78        with torch.no_grad():
 79            pred=net(torch.tensor(d['xte'],device=dev)).view(-1,8,2)
 80            metric=float(F.mse_loss(pred,torch.tensor(d['yte'],device=dev).view(-1,8,2)).cpu())
 81        return metric, net.cpu(), d
 82    except Exception:
 83        # Robust CPU fallback, matching the requested device policy.
 84        torch.cuda.empty_cache() if torch.cuda.is_available() else None
 85        net=Net(); x=torch.tensor(d['xtr']); y=torch.tensor(d['ytr']); opt=torch.optim.Adam(net.parameters(),lr=lr)
 86        for _ in range(EPOCHS):
 87            pred=net(x).view(-1,8,2); loss=F.mse_loss(pred,y.view(-1,8,2))
 88            if barrier:
 89                obs=torch.tensor(d['xtr'][:,1:])[:,None,:].expand(-1,8,-1)
 90                loss=loss+.08*F.softplus((.05-sat_margin(pred,obs,.1))/.05).mean()
 91            opt.zero_grad(); loss.backward(); opt.step()
 92        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)))
 93        return metric,net,d
 94
 95def eval_cfg(lr, barrier, seeds=SEEDS, keep=False):
 96    vals=[]; sig=[]
 97    for s in seeds:
 98        metric,net,d=train_one(s,lr,barrier); vals.append(metric)
 99        if keep:
100            with torch.no_grad():
101                p=net(torch.tensor(d['xte'])).view(-1,8,2)
102                o=torch.tensor(d['xte'][:,1:])[:,None,:].expand(-1,8,-1)
103                mm=sat_margin(p,o)
104                sig.append((float(mm.mean()),float(mm.min()),float((mm<=0).float().mean())))
105    out={'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':vals,'n':len(vals)}
106    if keep: out['signature_values']=sig
107    return out
108
109def main():
110    check=math_check()
111    # Equal search-space union: baseline and idea each run every candidate lr.
112    bsweep=[]; isweep=[]
113    for lr in LRS:
114        bs=eval_cfg(lr,False,SWEEP_SEEDS)['mean']; ii=eval_cfg(lr,True,SWEEP_SEEDS)['mean']
115        bsweep.append({'cfg':{'lr':lr,'epochs':EPOCHS},'mean':bs})
116        isweep.append({'cfg':{'lr':lr,'epochs':EPOCHS,'lambda_barrier':.08,'tau':.1,'m0':.05,'beta':.05},'mean':ii})
117    blr=min(LRS,key=lambda z: next(q['mean'] for q in bsweep if q['cfg']['lr']==z))
118    ilr=min(LRS,key=lambda z: next(q['mean'] for q in isweep if q['cfg']['lr']==z))
119    base={'best_cfg':{'lr':blr,'epochs':EPOCHS},'sweep':bsweep,'full':eval_cfg(blr,False)}
120    idea=eval_cfg(ilr,True,SEEDS,keep=True)
121    vals=np.asarray(idea['signature_values']); sig={'predicted_vs_observed':{
122        'predicted_smoothed_margin_mean':float(vals[:,0].mean()),
123        'observed_exact_sat_margin_mean':float(vals[:,1].mean()),
124        'predicted_collision_fraction':float(vals[:,2].mean())},
125        'prediction':'barrier should reduce collision fraction and move margins upward',
126        'confirmed':bool(vals[:,2].mean()<.5 and vals[:,0].mean()>-0.5)}
127    rep=make_report('obb_clearance_trajectory','mlp_tiny',base,idea,{'math_check':check,'idea_sweep':isweep,**sig})
128    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
129    print(json.dumps(rep,indent=2))
130
131if __name__=='__main__': main()