1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6import torch.nn.functional as F
  7
  8SEED = 2746
  9np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
 10
 11def get_device():
 12    return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 13
 14def boxes_from_tensors(cx, cy, yaw, length, width):
 15    c = torch.stack((cx, cy), -1)
 16    u = torch.stack((torch.cos(yaw), torch.sin(yaw)), -1)
 17    v = torch.stack((-torch.sin(yaw), torch.cos(yaw)), -1)
 18    return c, u, v, length / 2, width / 2
 19
 20def sat_barrier(ca, ya, la, wa, cb, yb, lb, wb, tau=0.1):
 21    ca, ua, va, aa, ba = boxes_from_tensors(ca[...,0], ca[...,1], ya, la, wa)
 22    cb, ub, vb, ab, bb = boxes_from_tensors(cb[...,0], cb[...,1], yb, lb, wb)
 23    axes = torch.stack((ua, va, ub, vb), -2)
 24    d = cb - ca
 25    proj = (axes * d.unsqueeze(-2)).sum(-1).abs()
 26    def radius(u, v, a, b):
 27        pu = (u.unsqueeze(-2) * axes).sum(-1).abs()
 28        pv = (v.unsqueeze(-2) * axes).sum(-1).abs()
 29        return a.unsqueeze(-1)*pu + b.unsqueeze(-1)*pv
 30    gaps = proj - radius(ua, va, aa, ba) - radius(ub, vb, ab, bb)
 31    margin = tau * torch.logsumexp(gaps / tau, dim=-1)
 32    return margin, gaps
 33
 34def hard_margin(ca, ya, la, wa, cb, yb, lb, wb):
 35    _, gaps = sat_barrier(ca, ya, la, wa, cb, yb, lb, wb, tau=1e-4)
 36    return gaps.max(-1).values
 37
 38def soft_barrier(margin, m0=0.05, beta=0.05):
 39    return F.softplus((m0-margin)/beta)
 40
 41def exact_signed_clearance(ca, ya, la, wa, cb, yb, lb, wb):
 42    # Polygon signed clearance for separated rectangles, with negative overlap
 43    # represented by the SAT maximum margin (the exact positive clearance).
 44    return hard_margin(ca, ya, la, wa, cb, yb, lb, wb)
 45
 46def math_sweeps():
 47    out = {}
 48    dev = torch.device('cpu')
 49    z = torch.tensor([[0., 0.]], device=dev)
 50    one = torch.tensor([0.], device=dev)
 51    dims = [torch.tensor([2.]), torch.tensor([1.])]
 52    ds = np.linspace(0.5, 5.0, 91)
 53    exact=[]; smooth=[]
 54    for d in ds:
 55        cb=torch.tensor([[float(d),0.]])
 56        ex=hard_margin(z,one,dims[0],dims[1],cb,one,dims[0],dims[1]).item()
 57        sm=sat_barrier(z,one,dims[0],dims[1],cb,one,dims[0],dims[1],tau=.1)[0].item()
 58        exact.append(ex); smooth.append(sm)
 59    # Prediction 1: axis-aligned SAT crossing is d=length sum / 2 = 2m.
 60    i=np.argmin(np.abs(np.asarray(exact)))
 61    j=np.argmin(np.abs(np.asarray(smooth)))
 62    out['zero_crossing']={'predicted_exact_m':2.0,'observed_exact_m':float(ds[i]),'observed_smooth_zero_m':float(ds[j]),'smooth_bias_m':float(smooth[j]-exact[j])}
 63    # Prediction 2: logsumexp smoothing is bounded above by tau log(4), and approaches max as tau shrinks.
 64    gaps=torch.tensor([[1.0, .2, -.3, -.8]])
 65    maxg=gaps.max().item(); bounds=[]
 66    for tau in [.2,.1,.05,.01]:
 67        val=(tau*torch.logsumexp(gaps/tau,-1)).item()
 68        bounds.append({'tau':tau,'value':val,'error':val-maxg,'bound':tau*math.log(4)})
 69    out['smoothing_bound']={'predicted_error_leq_tau_log4':bounds}
 70    # Prediction 3: at a separated axis-aligned configuration the active-axis gradient is ~1.
 71    grad=[]
 72    for tau in [.5,.2,.1,.05,.01]:
 73        x=torch.tensor([[3.0,0.]],requires_grad=True)
 74        m,_=sat_barrier(z,one,dims[0],dims[1],x,one,dims[0],dims[1],tau=tau)
 75        m.backward(); grad.append({'tau':tau,'dx':float(x.grad[0,0])})
 76    out['gradient']={'predicted_dx_near_1':grad}
 77    return out
 78
 79class Planner(nn.Module):
 80    def __init__(self, horizon=8):
 81        super().__init__(); self.net=nn.Sequential(nn.Linear(3,32),nn.Tanh(),nn.Linear(32, horizon*2)); self.h=horizon
 82    def forward(self,x): return self.net(x).view(-1,self.h,2)
 83
 84def train_compare():
 85    torch.manual_seed(SEED)
 86    dev=get_device(); n=96; h=8
 87    x=torch.linspace(-1,1,n,device=dev).unsqueeze(1)
 88    inp=torch.cat((x, torch.sin(2*x), torch.cos(2*x)),1)
 89    target=torch.stack((torch.linspace(0,4,h,device=dev).repeat(n,1), .7*x.repeat(1,h)), -1)
 90    # Obstacle box centered at x=2, y=0; desired paths cross it for x near zero.
 91    obs_c=torch.tensor([2.,0.],device=dev); obs_y=torch.tensor(0.,device=dev)
 92    results={}
 93    for use_barrier in [False,True]:
 94        torch.manual_seed(SEED); model=Planner(h).to(dev); opt=torch.optim.Adam(model.parameters(),lr=.02)
 95        for _ in range(350):
 96            pred=model(inp); imitation=((pred-target)**2).mean()
 97            if use_barrier:
 98                ca=torch.zeros((n,h,2),device=dev); ca[...,0]=pred[...,0]; ca[...,1]=pred[...,1]
 99                cb=obs_c.view(1,1,2).expand(n,h,2)
100                m,_=sat_barrier(ca, torch.zeros((n,h),device=dev), torch.full((n,h),2.,device=dev), torch.full((n,h),1.,device=dev), cb, obs_y.expand(n,h), torch.full((n,h),2.,device=dev), torch.full((n,h),1.,device=dev),tau=.1)
101                loss=imitation + .25*soft_barrier(m,m0=.05,beta=.05).mean()
102            else: loss=imitation
103            opt.zero_grad(); loss.backward(); opt.step()
104        with torch.no_grad():
105            pred=model(inp); ca=pred; cb=obs_c.view(1,1,2).expand(n,h,2)
106            margins=hard_margin(ca,torch.zeros((n,h),device=dev),torch.full((n,h),2.,device=dev),torch.full((n,h),1.,device=dev),cb,obs_y.expand(n,h),torch.full((n,h),2.,device=dev),torch.full((n,h),1.,device=dev))
107            results['barrier' if use_barrier else 'baseline']={'mse':float(((pred-target)**2).mean()),'min_margin':float(margins.min()),'collision_fraction':float((margins<=0).float().mean())}
108    return results
109
110def main():
111    report={'seed':SEED,'device':str(get_device()),'math':math_sweeps(),'mini_experiment':train_compare()}
112    Path('results.json').write_text(json.dumps(report,indent=2))
113    print(json.dumps(report,indent=2))
114
115if __name__=='__main__': main()