124 lines
6.3 KiB
Python
124 lines
6.3 KiB
Python
# 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)
|