ml: 全流程代码与运行手册（数据打包/可微运动学与对拍/MeanFlow+WTA 模型/训练/闭环仿真/APG/挖错循环/ONNX 导出）；webapp 策略运行时 ik-policy.js + ?policy=1 接入；tests/policy_bench.cjs；ik-solver 增 pairList/clusterExport/criticalPairs/pairGap/leverTable
This commit is contained in:
Hark
@@ -0,0 +1,34 @@
|
||||
"""MeanFlow 损失（Geng et al. 2025；MP1 配方：25% 样本 r≠t，自适应损失加权）。
|
||||
约定：t=0 为数据、t=1 为噪声；z_t = (1−t)·x + t·ε，瞬时速度 v = ε − x；网络输出平均速度 u(z,r,t)。
|
||||
目标 u_tgt = v − (t − r)·du/dt，du/dt = v·∂_z u + ∂_t u 用 torch.func.jvp 一次前向得到（切向量 (v, 0, 1)）。"""
|
||||
import torch
|
||||
from torch.func import jvp
|
||||
|
||||
|
||||
def sample_rt(B, device, dtype, p_neq=0.25, mu=-0.4, sigma=1.0):
|
||||
t = torch.sigmoid(torch.randn(B, device=device, dtype=dtype) * sigma + mu)
|
||||
r = torch.sigmoid(torch.randn(B, device=device, dtype=dtype) * sigma + mu)
|
||||
t, r = torch.maximum(t, r), torch.minimum(t, r)
|
||||
eq = torch.rand(B, device=device) >= p_neq
|
||||
r = torch.where(eq, t, r)
|
||||
return r, t
|
||||
|
||||
|
||||
def meanflow_loss(model, mem, x, p_neq=0.25, ada_p=1.0, ada_c=1e-3):
|
||||
B = x.shape[0]
|
||||
r, t = sample_rt(B, x.device, x.dtype, p_neq)
|
||||
e = torch.randn_like(x)
|
||||
z = (1 - t)[:, None, None] * x + t[:, None, None] * e
|
||||
v = e - x
|
||||
fn = lambda z_, r_, t_: model.decode(mem, z_, r_, t_)
|
||||
u, dudt = jvp(fn, (z, r, t), (v, torch.zeros_like(r), torch.ones_like(t)))
|
||||
u_tgt = (v - (t - r)[:, None, None] * dudt).detach()
|
||||
err = (u - u_tgt).pow(2).mean(dim=(1, 2))
|
||||
w = 1.0 / (err.detach() + ada_c).pow(ada_p)
|
||||
return (w * err).mean(), err.mean().detach()
|
||||
|
||||
|
||||
def one_step(model, mem, e):
|
||||
B = e.shape[0]
|
||||
one = torch.ones(B, device=e.device, dtype=e.dtype)
|
||||
return e - model.decode(mem, e, torch.zeros_like(one), one)
|
||||
@@ -0,0 +1,159 @@
|
||||
"""策略网络（方案 §6.1）：条件特征 → 8 个条件 token → 4 层 encoder；16 个动作 token（噪声 z + 时间嵌入 (r,t)）
|
||||
→ 4 层 decoder（自注意力 + 交叉注意力，非自回归一次前向）→ 每 token 6 维（平均速度场 u）。
|
||||
价值头 / 子目标头只依赖观测（encoder 输出）。
|
||||
注意力用显式 matmul+softmax 实现：torch.func.jvp（MeanFlow 训练）需要前向模式自动微分，SDPA 融合内核不支持。"""
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from ml.common.features import F_DIM, H
|
||||
|
||||
|
||||
class Attn(nn.Module):
|
||||
def __init__(self, d, nh):
|
||||
super().__init__()
|
||||
self.nh, self.dh = nh, d // nh
|
||||
self.q, self.k, self.v, self.o = (nn.Linear(d, d) for _ in range(4))
|
||||
|
||||
def forward(self, x, mem=None):
|
||||
mem = x if mem is None else mem
|
||||
B, Lq, D = x.shape; Lk = mem.shape[1]
|
||||
q = self.q(x).view(B, Lq, self.nh, self.dh).transpose(1, 2)
|
||||
k = self.k(mem).view(B, Lk, self.nh, self.dh).transpose(1, 2)
|
||||
v = self.v(mem).view(B, Lk, self.nh, self.dh).transpose(1, 2)
|
||||
a = torch.softmax(q @ k.transpose(-1, -2) / math.sqrt(self.dh), dim=-1)
|
||||
return self.o((a @ v).transpose(1, 2).reshape(B, Lq, D))
|
||||
|
||||
|
||||
class FFN(nn.Module):
|
||||
def __init__(self, d, f):
|
||||
super().__init__(); self.a = nn.Linear(d, f); self.b = nn.Linear(f, d)
|
||||
|
||||
def forward(self, x):
|
||||
return self.b(F.gelu(self.a(x)))
|
||||
|
||||
|
||||
class EncBlock(nn.Module):
|
||||
def __init__(self, d, nh, f):
|
||||
super().__init__(); self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d); self.sa, self.ff = Attn(d, nh), FFN(d, f)
|
||||
|
||||
def forward(self, x):
|
||||
x = x + self.sa(self.n1(x))
|
||||
return x + self.ff(self.n2(x))
|
||||
|
||||
|
||||
class DecBlock(nn.Module):
|
||||
def __init__(self, d, nh, f):
|
||||
super().__init__()
|
||||
self.n1, self.n2, self.n3 = nn.LayerNorm(d), nn.LayerNorm(d), nn.LayerNorm(d)
|
||||
self.sa, self.ca, self.ff = Attn(d, nh), Attn(d, nh), FFN(d, f)
|
||||
|
||||
def forward(self, x, mem):
|
||||
x = x + self.sa(self.n1(x))
|
||||
x = x + self.ca(self.n2(x), mem)
|
||||
return x + self.ff(self.n3(x))
|
||||
|
||||
|
||||
def time_emb(t, dim):
|
||||
half = dim // 2
|
||||
fr = torch.exp(-math.log(1000.0) * torch.arange(half, device=t.device, dtype=t.dtype) / half)
|
||||
a = t[:, None] * 1000.0 * fr[None]
|
||||
return torch.cat([torch.sin(a), torch.cos(a)], -1)
|
||||
|
||||
|
||||
class Policy(nn.Module):
|
||||
def __init__(self, d=192, nh=8, ff=768, n_enc=4, n_dec=4, n_cond=8, f_dim=F_DIM, horizon=H, head='meanflow', K=8):
|
||||
super().__init__()
|
||||
self.d, self.n_cond, self.h, self.head, self.K = d, n_cond, horizon, head, K
|
||||
self.cond = nn.Sequential(nn.Linear(f_dim, 2 * d), nn.GELU(), nn.Linear(2 * d, n_cond * d))
|
||||
self.cpos = nn.Parameter(torch.zeros(1, n_cond, d))
|
||||
self.enc = nn.ModuleList(EncBlock(d, nh, ff) for _ in range(n_enc))
|
||||
self.enc_norm = nn.LayerNorm(d)
|
||||
self.ain = nn.Linear(6, d)
|
||||
self.apos = nn.Parameter(torch.zeros(1, horizon, d))
|
||||
self.temb = nn.Sequential(nn.Linear(2 * 64, d), nn.GELU(), nn.Linear(d, d))
|
||||
self.dec = nn.ModuleList(DecBlock(d, nh, ff) for _ in range(n_dec))
|
||||
self.dec_norm = nn.LayerNorm(d)
|
||||
self.aout = nn.Linear(d, 6)
|
||||
self.value = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, 1))
|
||||
self.sub = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, 13))
|
||||
if head == 'wta': # 预定回退：同骨干 WTA K 模式头（方案 §6.2）
|
||||
self.mode_emb = nn.Parameter(torch.randn(1, K, 1, d) * 0.02)
|
||||
self.mode_logit = nn.Linear(d, 1)
|
||||
nn.init.normal_(self.cpos, std=0.02); nn.init.normal_(self.apos, std=0.02)
|
||||
nn.init.zeros_(self.aout.weight); nn.init.zeros_(self.aout.bias)
|
||||
|
||||
def encode(self, feat):
|
||||
x = self.cond(feat).view(-1, self.n_cond, self.d) + self.cpos
|
||||
for b in self.enc: x = b(x)
|
||||
return self.enc_norm(x)
|
||||
|
||||
def heads(self, mem):
|
||||
p = mem.mean(1)
|
||||
return self.value(p).squeeze(-1), self.sub(p)
|
||||
|
||||
def decode(self, mem, z, r, t):
|
||||
"""平均速度场 u(z, r, t)；z [B,16,6]，r,t [B]"""
|
||||
te = self.temb(torch.cat([time_emb(t, 64), time_emb(t - r, 64)], -1))
|
||||
x = self.ain(z) + self.apos + te[:, None]
|
||||
for b in self.dec: x = b(x, mem)
|
||||
return self.aout(self.dec_norm(x))
|
||||
|
||||
def wta_modes(self, mem):
|
||||
"""WTA：K 个模式一次解码 → act [B,K,16,6], logits [B,K]"""
|
||||
B = mem.shape[0]
|
||||
x = (self.ain.bias + self.apos)[:, None] + self.mode_emb # [1,K,16,d]
|
||||
x = x.expand(B, -1, -1, -1).reshape(B * self.K, self.h, self.d)
|
||||
m = mem.repeat_interleave(self.K, 0)
|
||||
for b in self.dec: x = b(x, m)
|
||||
x = self.dec_norm(x)
|
||||
act = self.aout(x).view(B, self.K, self.h, 6)
|
||||
return act, self.mode_logit(x.mean(1)).view(B, self.K)
|
||||
|
||||
def sample(self, feat, z):
|
||||
"""部署接口（ONNX 同签名）：返回 (act [B,16,6] 归一化增量, value [B], sub [B,13])
|
||||
meanflow：一步生成 x = z − u(z, 0, 1)；
|
||||
wta：取第 0 行观测的 K 个模式，第 b 行输出模式 b mod K（部署时 8 行观测相同；前瞻只用各行的 value）。"""
|
||||
mem = self.encode(feat)
|
||||
v, s = self.heads(mem)
|
||||
B = feat.shape[0]
|
||||
if self.head == 'wta':
|
||||
act, _ = self.wta_modes(mem[:1])
|
||||
idx = torch.arange(B, device=feat.device) % self.K
|
||||
return act[0, idx] + 0.0 * z, v, s # z 不参与计算，保留为输入以统一 ONNX 签名
|
||||
one, zero = torch.ones(B, device=feat.device, dtype=feat.dtype), torch.zeros(B, device=feat.device, dtype=feat.dtype)
|
||||
x = z - self.decode(mem, z, zero, one)
|
||||
return x, v, s
|
||||
|
||||
def sample_k(self, feat, K, z=None):
|
||||
"""仿真 / APG 用：每个观测 K 条候选 → act [B,K,16,6], value [B], sub [B,13]"""
|
||||
mem = self.encode(feat)
|
||||
v, s = self.heads(mem)
|
||||
B = feat.shape[0]
|
||||
if self.head == 'wta':
|
||||
act, logit = self.wta_modes(mem)
|
||||
if K < self.K:
|
||||
top = logit.topk(K, -1).indices
|
||||
act = act.gather(1, top[:, :, None, None].expand(-1, -1, self.h, 6))
|
||||
return act, v, s
|
||||
z = torch.randn(B * K, self.h, 6, device=feat.device, dtype=feat.dtype) if z is None else z.reshape(B * K, self.h, 6)
|
||||
m = mem.repeat_interleave(K, 0)
|
||||
one = torch.ones(B * K, device=feat.device, dtype=feat.dtype)
|
||||
x = z - self.decode(m, z, torch.zeros_like(one), one)
|
||||
return x.view(B, K, self.h, 6), v, s
|
||||
|
||||
|
||||
def build(cfg):
|
||||
m = cfg.get('model', {})
|
||||
return Policy(d=m.get('d', 192), nh=m.get('heads', 8), ff=m.get('ff', 768), n_enc=m.get('enc', 4), n_dec=m.get('dec', 4),
|
||||
head=m.get('head', 'meanflow'), K=m.get('K', 8))
|
||||
|
||||
|
||||
def wta_loss(model, mem, y, w_all=0.05):
|
||||
"""WTA-L1 + 模式 CE；返回 (loss, 最优模式输出 [B,16,6])"""
|
||||
act, logit = model.wta_modes(mem)
|
||||
l1 = (act - y[:, None]).abs().mean(dim=(2, 3)) # [B,K]
|
||||
best = l1.argmin(-1)
|
||||
loss = l1.gather(1, best[:, None]).mean() + w_all * l1.mean() + 0.1 * F.cross_entropy(logit, best)
|
||||
return loss, act[torch.arange(y.shape[0]), best]
|
||||
Reference in New Issue
Block a user