1from __future__ import annotations
  2import sys, json, math, itertools
  3from pathlib import Path
  4import numpy as np
  5import torch
  6import torch.nn as nn
  7
  8sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  9from bench import make_model, train_model, sweep_baseline, make_report
 10from bench.protocol import evaluate
 11
 12META = {"name":"tangent_evasion_dynamics", "domain":"dynamics", "description":"2D evader policy with pursuer exclusion disk and tangent-branch safety structure."}
 13RADIUS = .8
 14CENTER = np.array([0., 0.], dtype=np.float32)
 15GOAL = np.array([3., 0.], dtype=np.float32)
 16DT = .08
 17VMAX = 1.0
 18
 19
 20def rot(theta):
 21    c,s=np.cos(theta),np.sin(theta)
 22    return np.array([[c,-s],[s,c]], dtype=np.float32)
 23
 24def tangent(center, target, radius, side):
 25    z=center-target; d=np.linalg.norm(z)
 26    if d <= radius: return None
 27    a=np.arcsin(np.clip(radius/d,-1,1))
 28    return target + np.sqrt(d*d-radius*radius)/d * (rot(-side*a) @ z)
 29
 30def label_velocity(x, y, heading):
 31    """Analytic branch reference, used only to generate task observations/labels."""
 32    p=np.array([x,y], dtype=np.float32); h=np.array([np.cos(heading),np.sin(heading)], dtype=np.float32)
 33    candidates=[]
 34    for side in (-1,1):
 35        s=tangent(CENTER, GOAL, RADIUS, side)
 36        if s is None: continue
 37        vec=s-p; ang=abs(np.arctan2(np.cross(h,vec),np.dot(h,vec)))
 38        candidates.append((ang+np.linalg.norm(GOAL-s),s))
 39    if np.linalg.norm(p-CENTER) < 1.35 and candidates:
 40        s=min(candidates,key=lambda z:z[0])[1]
 41        v=s-p
 42    else: v=GOAL-p
 43    n=np.linalg.norm(v)
 44    return (v/max(n,1e-6)*VMAX).astype(np.float32)
 45
 46def get_dataset(seed, n_train, n_test):
 47    rng=np.random.RandomState(seed)
 48    def make(n):
 49        X=np.empty((n,24),np.float32); Y=np.empty((n,2),np.float32)
 50        for i in range(n):
 51            # Include safe and near-boundary states, with heading as the third GRU channel.
 52            a=rng.uniform(-np.pi,np.pi); rad=rng.uniform(.82,2.0)
 53            x,y=rad*np.cos(a),rad*np.sin(a)
 54            heading=rng.uniform(-np.pi,np.pi)
 55            state=np.tile(np.array([x,y,heading],np.float32),(8,1))
 56            state[:,2]+=rng.normal(0,.025,8).astype(np.float32)
 57            X[i]=state.reshape(-1)
 58            Y[i]=label_velocity(x,y,heading)
 59        return X,Y
 60    xtr,ytr=make(n_train); xte,yte=make(n_test)
 61    return {"xtr":xtr,"ytr":ytr,"xte":xte,"yte":yte,"task":"regression","metric":"mse","input_shape":(24,),"out_dim":2}
 62
 63
 64def project_np(pos, vel):
 65    p=np.asarray(pos,float); v=np.asarray(vel,float).copy(); d=np.linalg.norm(p-CENTER)
 66    if d>1e-9:
 67        n=(p-CENTER)/d
 68        if d <= RADIUS + DT*np.linalg.norm(v) and np.dot(v,n)<0: v-=np.dot(v,n)*n
 69    z=np.linalg.norm(v)
 70    return v if z<=VMAX else v/z*VMAX
 71
 72def tangent_project_torch(x, out):
 73    # Network input is (x,y,heading) repeated; compute active-disk half-space projection.
 74    p=x[:, -1, :2]; v=out
 75    d=torch.linalg.vector_norm(p,dim=1,keepdim=True).clamp_min(1e-6)
 76    n=p/d
 77    inward=(v*n).sum(1,keepdim=True)
 78    active=(d <= RADIUS + DT*torch.linalg.vector_norm(v,dim=1,keepdim=True)).float()
 79    v=v-active*torch.minimum(inward,torch.zeros_like(inward))*n
 80    speed=torch.linalg.vector_norm(v,dim=1,keepdim=True).clamp_min(1e-6)
 81    return v*torch.minimum(torch.ones_like(speed), VMAX/speed)
 82
 83def run(seed, lr, idea, epochs=15):
 84    torch.manual_seed(10000+int(seed)); np.random.seed(10000+int(seed))
 85    ds=get_dataset(int(seed),400,200)
 86    net=make_model("rnn_small", ds["input_shape"], ds["out_dim"])
 87    device="cuda" if torch.cuda.is_available() else "cpu"
 88    try:
 89        net=net.to(device); opt=torch.optim.Adam(net.parameters(),lr=lr)
 90        xtr=torch.as_tensor(ds["xtr"],device=device); ytr=torch.as_tensor(ds["ytr"],device=device)
 91        lossf=nn.MSELoss()
 92        for _ in range(epochs):
 93            net.train(); perm=torch.randperm(len(xtr),device=device)
 94            for j in range(0,len(xtr),128):
 95                ix=perm[j:j+128]; raw=net(xtr[ix])
 96                pred=tangent_project_torch(xtr[ix].view(-1,8,3),raw) if idea else raw
 97                loss=lossf(pred,ytr[ix]); opt.zero_grad(); loss.backward(); opt.step()
 98        net.eval(); xt=torch.as_tensor(ds["xte"],device=device)
 99        with torch.no_grad():
100            raw=net(xt); pred=tangent_project_torch(xt.view(-1,8,3),raw) if idea else raw
101            metric=float(lossf(pred,torch.as_tensor(ds["yte"],device=device)))
102        return metric
103    except Exception:
104        # Explicit CPU fallback, matching the bench's robust-device requirement.
105        torch.cuda.empty_cache() if torch.cuda.is_available() else None
106        torch.manual_seed(10000+int(seed)); net=make_model("rnn_small",ds["input_shape"],2)
107        opt=torch.optim.Adam(net.parameters(),lr=lr); xtr=torch.as_tensor(ds["xtr"]); ytr=torch.as_tensor(ds["ytr"])
108        for _ in range(epochs):
109            for j in range(0,400,128):
110                raw=net(xtr[j:j+128]); pred=tangent_project_torch(xtr[j:j+128].view(-1,8,3),raw) if idea else raw
111                loss=nn.functional.mse_loss(pred,ytr[j:j+128]); opt.zero_grad(); loss.backward(); opt.step()
112        with torch.no_grad():
113            raw=net(torch.as_tensor(ds["xte"])); pred=tangent_project_torch(torch.as_tensor(ds["xte"]).view(-1,8,3),raw) if idea else raw
114            return float(nn.functional.mse_loss(pred,torch.as_tensor(ds["yte"])))
115
116def main():
117    # Equal union search space: every idea LR is also evaluated by baseline.
118    grid=[{"lr":1e-3},{"lr":3e-3},{"lr":1e-2}]
119    base=sweep_baseline(lambda c: lambda s: run(s,c["lr"],False),grid)
120    # Idea is run at best baseline LR and two nearby grid settings; all are already baseline-swept.
121    idea_cfgs=[{"lr":c["lr"]} for c in grid]
122    vals=[]
123    for c in idea_cfgs:
124        r=evaluate(lambda s,c=c: run(s,c["lr"],True),tuple(range(8)))
125        vals.append((r["mean"],c,r))
126    _,best_cfg,idea=min(vals,key=lambda z:z[0])
127    # Trained-model behavior signature: one-step disk penetration on held-out states.
128    def penetration(seed, isidea):
129        ds=get_dataset(seed,400,200); x=torch.as_tensor(ds["xte"]); torch.manual_seed(10000+seed); net=make_model("rnn_small",(24,),2); opt=torch.optim.Adam(net.parameters(),lr=best_cfg["lr"])
130        for _ in range(15):
131            for j in range(0,400,128):
132                raw=net(torch.as_tensor(ds["xtr"])[j:j+128]); p=tangent_project_torch(torch.as_tensor(ds["xtr"])[j:j+128].view(-1,8,3),raw) if isidea else raw
133                l=nn.functional.mse_loss(p,torch.as_tensor(ds["ytr"])[j:j+128]); opt.zero_grad();l.backward();opt.step()
134        with torch.no_grad(): out=net(x); out=tangent_project_torch(x.view(-1,8,3),out) if isidea else out
135        pos=x[:,-3:-1].numpy(); nxt=pos+DT*out.numpy(); return float(np.maximum(0,RADIUS-np.linalg.norm(nxt-CENTER,axis=1)).max())
136    bp=float(np.mean([penetration(s,False) for s in range(8)])); ip=float(np.mean([penetration(s,True) for s in range(8)]))
137    sig={"predicted":"projection keeps one-step disk penetration at numerical zero","observed_baseline_max_penetration":bp,"observed_idea_max_penetration":ip,"confirmed":ip < 1e-6 and ip <= bp+1e-8}
138    rep=make_report("tangent_evasion_dynamics","rnn_small",base,idea,{**sig,"custom_track":{"name":"tangent_evasion_dynamics","file":"bench_tangent_stage2.py","domain":"dynamics"},"idea_sweep":[{"cfg":c,"mean":m} for m,c,_ in vals]})
139    Path("bench_report.json").write_text(json.dumps(rep,indent=2))
140    print(json.dumps(rep,indent=2))
141if __name__=="__main__": main()