docs: Transformer 拟合避障轨迹规划训练方案（最终版）：无老师闭环部署、MeanFlow 动作头（WTA K=8 回退）、距离下界安全认证、APG 后训练与挖错循环、里程碑 M1–M8；附 HTML 版
This commit is contained in:
Hark
@@ -0,0 +1,123 @@
|
||||
# 2D 绕柱玩具实验：比较不同动作头在"左右(上下)双模态绕行"下的闭环表现
|
||||
import torch, torch.nn as nn, numpy as np, math, time, sys, json
|
||||
torch.manual_seed(0); np.random.seed(0); torch.set_num_threads(2)
|
||||
R=1.0; H=16; STEP=0.05
|
||||
def demo(s,g,mode):
|
||||
via=np.array([0.0, 1.35*mode])
|
||||
pts=[s, np.array([-0.9,0.95*mode]) if s[0]<-0.9 else s, via, np.array([0.9,0.95*mode]), g]
|
||||
P=np.array(pts); L=np.r_[0,np.cumsum(np.linalg.norm(np.diff(P,axis=0),axis=1))]
|
||||
n=int(L[-1]/STEP)+1; u=np.linspace(0,L[-1],n)
|
||||
x=np.interp(u,L,P[:,0]); y=np.interp(u,L,P[:,1]); Q=np.c_[x,y]
|
||||
# 平滑
|
||||
for _ in range(20): Q[1:-1]=0.5*Q[1:-1]+0.25*(Q[:-2]+Q[2:])
|
||||
return Q
|
||||
def sample_task(hard=False):
|
||||
sy=np.random.uniform(-0.3,0.3) if hard else np.random.uniform(-1.5,1.5)
|
||||
gy=np.random.uniform(-0.3,0.3) if hard else np.random.uniform(-1.5,1.5)
|
||||
return np.array([np.random.uniform(-2.5,-1.6),sy]), np.array([np.random.uniform(1.6,2.5),gy])
|
||||
def collide_path(Q):
|
||||
return np.any(np.linalg.norm(Q,axis=1)<R)
|
||||
def build(n=3000):
|
||||
O=[];A=[]
|
||||
for _ in range(n):
|
||||
s,g=sample_task(hard=np.random.rand()<0.4)
|
||||
d=s[1]+g[1]
|
||||
mode = np.sign(d) if abs(d)>0.3 else (1 if np.random.rand()<0.5 else -1)
|
||||
Q=demo(s,g,mode)
|
||||
if collide_path(Q): continue
|
||||
D=np.diff(Q,axis=0); D=np.r_[D,np.zeros((H,2))]
|
||||
for i in range(len(Q)-1):
|
||||
O.append(np.r_[Q[i],g]); A.append(D[i:i+H].reshape(-1)/STEP)
|
||||
return torch.tensor(np.array(O),dtype=torch.float32), torch.tensor(np.array(A),dtype=torch.float32)
|
||||
O,A=build(); print('samples',len(O),flush=True)
|
||||
AD=2*H
|
||||
def mlp(i,o,h=256): return nn.Sequential(nn.Linear(i,h),nn.SiLU(),nn.Linear(h,h),nn.SiLU(),nn.Linear(h,h),nn.SiLU(),nn.Linear(h,o))
|
||||
class MSE(nn.Module):
|
||||
def __init__(s): super().__init__(); s.f=mlp(4,AD)
|
||||
def loss(s,o,a): return ((s.f(o)-a)**2).mean()
|
||||
def sample(s,o,k=1): return s.f(o).unsqueeze(1).expand(-1,k,-1)
|
||||
class WTA(nn.Module):
|
||||
K=4
|
||||
def __init__(s): super().__init__(); s.f=mlp(4,AD*s.K+s.K)
|
||||
def out(s,o):
|
||||
y=s.f(o); return y[:,:AD*s.K].view(-1,s.K,AD), y[:,AD*s.K:]
|
||||
def loss(s,o,a):
|
||||
p,l=s.out(o); e=(p-a[:,None]).abs().mean(-1); k=e.argmin(1)
|
||||
return e.gather(1,k[:,None]).mean()+0.5*nn.functional.cross_entropy(l,k)
|
||||
def sample(s,o,k=1):
|
||||
p,l=s.out(o)
|
||||
if k==1: return p.gather(1,l.argmax(1)[:,None,None].expand(-1,1,AD))
|
||||
o_=l.argsort(1,descending=True); return p.gather(1,o_[:,:,None].expand(-1,-1,AD)) # 按概率排序
|
||||
class FM(nn.Module):
|
||||
def __init__(s,steps): super().__init__(); s.f=mlp(4+AD+1,AD); s.steps=steps
|
||||
def loss(s,o,a):
|
||||
e=torch.randn_like(a); t=torch.rand(len(a),1); z=(1-t)*a+t*e
|
||||
return ((s.f(torch.cat([o,z,t],1))-(e-a))**2).mean()
|
||||
def sample(s,o,k=1):
|
||||
B=len(o); oo=o.repeat_interleave(k,0); z=torch.randn(B*k,AD); dt=1/s.steps
|
||||
for i in range(s.steps):
|
||||
t=torch.full((B*k,1),1-i*dt); z=z-dt*s.f(torch.cat([oo,z,t],1))
|
||||
return z.view(B,k,AD)
|
||||
class MF(nn.Module): # MeanFlow (Geng et al. 2025)
|
||||
def __init__(s): super().__init__(); s.f=mlp(4+AD+2,AD)
|
||||
def u(s,o,z,r,t): return s.f(torch.cat([o,z,r,t],1))
|
||||
def loss(s,o,a):
|
||||
B=len(a); e=torch.randn_like(a)
|
||||
t=torch.rand(B,1); r=torch.rand(B,1); t,r=torch.maximum(t,r),torch.minimum(t,r)
|
||||
eq=(torch.rand(B,1)<0.75); r=torch.where(eq,t,r)
|
||||
z=(1-t)*a+t*e; v=e-a
|
||||
fn=lambda z_,r_,t_: s.u(o,z_,r_,t_)
|
||||
u,dudt=torch.func.jvp(fn,(z,r,t),(v,torch.zeros_like(r),torch.ones_like(t)))
|
||||
tgt=(v-(t-r)*dudt).detach()
|
||||
return ((u-tgt)**2).mean()
|
||||
def sample(s,o,k=1):
|
||||
B=len(o); oo=o.repeat_interleave(k,0); e=torch.randn(B*k,AD)
|
||||
x=e-s.u(oo,e,torch.zeros(B*k,1),torch.ones(B*k,1)); return x.view(B,k,AD)
|
||||
def train(m,iters=6000,bs=512):
|
||||
opt=torch.optim.AdamW(m.parameters(),1e-3,weight_decay=1e-4)
|
||||
sch=torch.optim.lr_scheduler.CosineAnnealingLR(opt,iters)
|
||||
for it in range(iters):
|
||||
idx=torch.randint(0,len(O),(bs,)); l=m.loss(O[idx],A[idx])
|
||||
opt.zero_grad(); l.backward(); opt.step(); sch.step()
|
||||
return m
|
||||
def seg_collide(p,q):
|
||||
for a in np.linspace(0,1,6):
|
||||
if np.linalg.norm(p+(q-p)*a)<R: return True
|
||||
return False
|
||||
@torch.no_grad()
|
||||
def rollout(m,s,g,k=1,safety=False,commit=3,T=500):
|
||||
p=s.copy(); hits=0; ups=[]; brakes=0
|
||||
for t in range(0,T,commit):
|
||||
o=torch.tensor(np.r_[p,g],dtype=torch.float32)[None]
|
||||
C=m.sample(o,k)[0].numpy().reshape(-1,H,2)*STEP
|
||||
chosen=None
|
||||
for c in C:
|
||||
path=p+np.cumsum(c[:commit],0)
|
||||
if not safety: chosen=c; break
|
||||
pts=np.r_[p[None],path]
|
||||
if not any(seg_collide(pts[i],pts[i+1]) for i in range(commit)): chosen=c; break
|
||||
if chosen is None: brakes+=1; continue
|
||||
for d in chosen[:commit]:
|
||||
q=p+d
|
||||
if seg_collide(p,q): hits+=1
|
||||
p=q
|
||||
if abs(p[0])<0.5: ups.append(np.sign(p[1]))
|
||||
if np.linalg.norm(p-g)<0.1: return dict(ok=1,hits=hits,brakes=brakes,steps=t,mode=np.mean(ups) if ups else 0)
|
||||
return dict(ok=0,hits=hits,brakes=brakes,steps=T,mode=np.mean(ups) if ups else 0)
|
||||
def evaluate(m,k,safety,n=150,hard=True):
|
||||
np.random.seed(123); res=[rollout(m,*sample_task(hard),k=k,safety=safety) for _ in range(n)]
|
||||
ok=np.mean([r['ok'] for r in res]); col=np.mean([r['hits']>0 for r in res])
|
||||
br=np.mean([r['brakes'] for r in res]); up=np.mean([r['mode']>0 for r in res if r['ok']]) if ok>0 else float('nan')
|
||||
return dict(reach=round(ok,3),collide_ep=round(col,3),brakes=round(br,1),frac_up=round(up,2))
|
||||
methods={'MSE':lambda:MSE(),'WTA4':lambda:WTA(),'FM10':lambda:FM(10),'FM1':lambda:FM(1),'MeanFlow1':lambda:MF()}
|
||||
which=sys.argv[1:] or list(methods)
|
||||
out={}
|
||||
for name in which:
|
||||
torch.manual_seed(0); t0=time.time(); m=train(methods[name]()); tt=time.time()-t0
|
||||
o=torch.tensor(np.r_[[-2,0],[2,0]],dtype=torch.float32)[None]
|
||||
t1=time.time(); [m.sample(o,8) for _ in range(50)]; lat=(time.time()-t1)/50*1000
|
||||
r={'train_s':round(tt),'infer_ms_k8':round(lat,2)}
|
||||
r['hard_raw_k1']=evaluate(m,1,False); r['hard_safe_k8']=evaluate(m,4 if name=='WTA4' else (1 if name=='MSE' else 8),True)
|
||||
r['easy_raw_k1']=evaluate(m,1,False,hard=False)
|
||||
out[name]=r; print(name,json.dumps(r,ensure_ascii=False),flush=True)
|
||||
json.dump(out,open(f'/workspace/work/toy/res_{"_".join(which)}.json','w'),indent=1)
|
||||
Reference in New Issue
Block a user