cbb/Robot_Arm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 目标写法开关 mf_target(mf/lin/sq/mid)，默认 mf 且原地路径逐位不变；float32 接线用解析高斯场自证

Browse Source
This commit is contained in:
cbb
2026-10-20 21:48:28 +08:00
parent 78786881e7
commit 3b8682a345
3 changed files with 97 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ml/model/meanflow.py
+35 -6
View File
@@ -26,17 +26,46 @@ def sample_rt(B, device, dtype, p_neq=0.25, mu=-0.4, sigma=1.0, p_edge=0.0):
return r, t
def meanflow_loss(model, mem, x, p_neq=0.25, ada_p=1.0, ada_c=1e-3, return_aux=False, p_edge=0.0):
def avg_velocity_target(model, mem, x, e, r, t, target='mf', lam=0.5):
"""平均速度目标的三种写法（目标都 detach；u 带梯度）。target='mf' 与原实现逐位相同。
mf ：Eq.4 单段自举 v − (t−r)·du/dt，收敛解 = 平均速度（解析高斯场已验证偏差 <1e−12）
lin ：v − (t−r)·[(1−λ)·du(z_t,c,t)/dt + λ·du(z_c,r,c)/dc]，系数和=1 ⇒ 收敛解仍是平均速度
sq ：同上的 (1−λ)²/λ² 版（K-MF 论文摘要里的写法）⇒ 系数和 0.5，实测只线性收敛（结构性少修一半）
mid ：v − (t−r)·du(z_c,r,c)/dc，整条修正移到中点
"""
fn = lambda z_, r_, t_: model.decode(mem, z_, r_, t_)
z = (1 - t)[:, None, None] * x + t[:, None, None] * e
v = e - x
if target == 'mf':
u, dudt = jvp(fn, (z, r, t), (v, torch.zeros_like(r), torch.ones_like(t)))
return u, (v - (t - r)[:, None, None] * dudt).detach(), z
with torch.no_grad():
u = fn(z, r, t) # 这三种写法只要前向拿 u，导数全部进目标（sg）
c = (1 - lam) * t + lam * r
zc = (1 - c)[:, None, None] * x + c[:, None, None] * e
tv = (v, torch.zeros_like(r), torch.ones_like(t))
def jv(zz, rr, tt):
with torch.enable_grad():
return jvp(fn, (zz, rr, tt), tv)[1].detach()
if target == 'mid':
corr = jv(zc, r, c)
else:
d1 = jv(z, c, t)
d2 = jv(zc, r, c)
w = ((1 - lam) ** 2, lam ** 2) if target == 'sq' else ((1 - lam), lam)
corr = w[0] * d1 + w[1] * d2
return u, (v - (t - r)[:, None, None] * corr).detach(), z
def meanflow_loss(model, mem, x, p_neq=0.25, ada_p=1.0, ada_c=1e-3, return_aux=False, p_edge=0.0, target='mf', lam=0.5):
mem, x = mem.float(), x.float()
with torch.autocast('cuda', enabled=False):
B = x.shape[0]
r, t = sample_rt(B, x.device, x.dtype, p_neq, p_edge=p_edge)
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()
u, u_tgt, z = avg_velocity_target(model, mem, x, e, r, t, target, lam)
err = (u - u_tgt).pow(2).mean(dim=(1, 2))
w = 1.0 / (err.detach() + ada_c).pow(ada_p)
loss = (w * err).mean()
ml/train.py
+2 -1
View File
@@ -161,7 +161,8 @@ def main():
l_mf, xh = wta_loss(model, mem, Y); mf_raw = l_mf.detach(); xh = xh.float()
else:
p_neq = tc['p_neq'] if step >= warm else 0.0
l_mf, mf_raw, aux = meanflow_loss(model, mem, Y, p_neq=p_neq, return_aux=True, p_edge=tc.get('p_edge', 0.0))
l_mf, mf_raw, aux = meanflow_loss(model, mem, Y, p_neq=p_neq, return_aux=True, p_edge=tc.get('p_edge', 0.0),
target=tc.get('mf_target', 'mf'), lam=tc.get('mf_lam', 0.5))
sel = ((aux['r'] == aux['t']) & (aux['t'] <= tc.get('fk_tmax', 0.3))).nonzero(as_tuple=True)[0]
tt = aux['t'][sel][:, None, None]
xh = aux['z'][sel] - tt * aux['u'][sel] # 配对去噪估计（fp32）
tests/test_meanflow_targets.py
+60
View File
@@ -9,6 +9,7 @@
import unittest
import torch
from torch.func import jvp
LAM = 0.5
@@ -108,5 +109,64 @@ class MeanFlowTargetFixedPointTest(unittest.TestCase):
self.assertLess(half['mid'], full['mid'])
class AvgVelocityTargetWiringTest(unittest.TestCase):
"""把三种写法接进 ml/model/meanflow.py 之后，再用同一个解析场验一遍接线：
浮点从 float64 换到 float32，量级判据放宽，但"mf 无偏 / sq 只线性收敛"必须照样成立，
否则说明 z_c、切向量、r 槽换位这几处接错了。"""
SIG = 2.0
class Stub:
def __init__(self, sig2):
self.sig2 = sig2
def s2(self, t):
return (1 - t) ** 2 * self.sig2 + t ** 2
def decode(self, mem, z, r, t):
return z * (1 - torch.sqrt(self.s2(r) / self.s2(t)))[:, None, None] / (t - r)[:, None, None]
def bias(self, target, r, t, lam=0.5):
from ml.model.meanflow import avg_velocity_target
f = Field(self.SIG)
z = torch.tensor([[0.7] * 4, [0.7] * 4], dtype=torch.float32) # [2,4,6]? 形状自定，逐元素即可
z = torch.full((2, 4, 6), 0.7)
rt = torch.tensor([r, r]); tt = torch.tensor([t, t])
xh = (z * (1 - tt)[:, None, None] * self.SIG ** 2 / f.s2(tt)[:, None, None])
eh = (z * tt[:, None, None] / f.s2(tt)[:, None, None])
u, tgt, _ = avg_velocity_target(self.Stub(self.SIG ** 2), None, xh, eh, rt, tt, target, lam)
truth = f.u(torch.tensor(0.7, dtype=torch.float32), torch.tensor(r), torch.tensor(t))
return float(((tgt - truth).abs() / truth.abs()).mean()), float(truth)
def test_wiring_keeps_mf_unbiased_and_sq_first_order(self):
b_mf, _ = self.bias('mf', 0.8, 1.0)
self.assertLess(b_mf, 1e-4)
rel = []
for s in (0.2, 0.1, 0.05):
rel.append({k: self.bias(k, 1.0 - s, 1.0)[0] for k in ('sq', 'lin', 'mid', 'mf')})
print('float32 接线后的偏差：', [{k: round(v, 4) for k, v in d.items()} for d in rel])
r_sq = rel[1]['sq'] / rel[0]['sq']
r_mid = rel[1]['mid'] / rel[0]['mid']
self.assertGreater(r_sq, 0.3) # 只线性收敛 ⇒ 结构性偏差
self.assertLess(r_mid, 0.35) # 二次收敛 ⇒ 只剩积分误差
self.assertLess(rel[0]['lin'], rel[0]['sq'])
def test_default_path_is_bitwise_identical(self):
"""target='mf' 必须与重构前的原地写法逐位相同（默认训练路径不许被这次改动碰到）"""
from ml.model.meanflow import avg_velocity_target
torch.manual_seed(0)
m = self.Stub(2.0)
x = torch.randn(3, 5, 6); e = torch.randn(3, 5, 6)
r = torch.rand(3) * 0.4; t = 0.6 + torch.rand(3) * 0.4
u_new, tgt_new, z_new = avg_velocity_target(m, None, x, e, r, t, 'mf', 0.5)
fn = lambda a, b, c: m.decode(None, a, b, c)
z = (1 - t)[:, None, None] * x + t[:, None, None] * e
v = e - x
u0, dudt = jvp(fn, (z, r, t), (v, torch.zeros_like(r), torch.ones_like(t)))
tgt0 = (v - (t - r)[:, None, None] * dudt).detach()
for name, a_, b_ in (('u', u_new, u0), ('u_tgt', tgt_new, tgt0), ('z', z_new, z)):
with self.subTest(name=name):
torch.testing.assert_close(a_, b_, rtol=0, atol=0)
if __name__ == '__main__':
unittest.main()