Hark
2a0a2e55ad
- M1 数据生成器：采样/老师/距离下界认证/时间律/统计/一致性门禁；确定性老师（碰撞次数计时 + resetSession）、pos 任务冻结 J6、--consensus 邻域共识过滤、Dr 类。 - 全流程代码与运行手册：数据打包/可微运动学与对拍/MeanFlow+WTA 模型/闭环仿真/APG/挖错循环/ONNX 导出；webapp 策略运行时 ik-policy.js、?policy=1 接入、tests/policy_bench.cjs；ik-solver 增 pairList/clusterExport/criticalPairs/pairGap/leverTable。
35 lines
1.5 KiB
Python
35 lines
1.5 KiB
Python
"""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)
|