fix(train): 首次正式预训练(21850 步闭环到达 0%)排查后修正——FK/平滑改加在配对去噪估计上(原加在随机噪声一步采样上=回归条件均值,多峰塌成单峰,实测候选标准差 0.030/幅度 0.151);时间嵌入最高角频率 1000→50、MF 分支 fp32、p_neq 前 5000 步为 0(实测 du/dt 均方根 6–12 vs v 1.1);增广 40%/3°/12 步收回;新增 dagger 阶段(策略闭环状态请老师重规划 ×2 轮微调);开训前老师轨迹回放自检;评估加开环误差与候选多样性,best/早停按到达率+开环误差

This commit is contained in:
Hark
2026-10-18 15:56:39 +00:00
parent c2732fb3cd
commit 979a5814c1
19 changed files with 386 additions and 61 deletions
+24 -13
View File
@@ -1,6 +1,12 @@
"""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))。"""
目标 u_tgt = v − (t − r)·du/dt,du/dt = v·∂_z u + ∂_t u 用 torch.func.jvp 一次前向得到(切向量 (v, 0, 1))。
2026-10-18 首次正式预训练后改:
- 整个 MF 分支(解码器前向 + JVP)在 fp32 下算:目标里含网络自己的导数,bf16 下 JVP 误差直接变成目标噪声;
- return_aux=True 时返回 (z, t, r, u),训练用 r==t 且 t 小的样本做"配对去噪估计" x̂ = z − t·u 来加 FK / 平滑损失。
原来把 FK 加在 one_step(随机噪声) 上、对齐这一条样本的老师 Y:对任意噪声都回归同一个 Y,最优解是条件均值 E[Y|X],
多峰头被训成单峰(实测 last.pt 8 个候选间标准差 0.030,动作幅度 0.151),左绕/右绕取平均正好穿柱。
x̂ 由 z_t 里带着的那一条 x 决定(t 小时 z_t 几乎就是 x),不会把不同模式拉到一起。"""
import torch
from torch.func import jvp
@@ -14,18 +20,23 @@ def sample_rt(B, device, dtype, p_neq=0.25, mu=-0.4, sigma=1.0):
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 meanflow_loss(model, mem, x, p_neq=0.25, ada_p=1.0, ada_c=1e-3, return_aux=False):
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)
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)
loss = (w * err).mean()
if return_aux:
return loss, err.mean().detach(), {'z': z, 't': t, 'r': r, 'u': u}
return loss, err.mean().detach()
def one_step(model, mem, e):
+9 -2
View File
@@ -55,10 +55,17 @@ class DecBlock(nn.Module):
return x + self.ff(self.n3(x))
TEMB_WMAX = 50.0 # 时间嵌入最高角频率(rad / 单位 t)
def time_emb(t, dim):
"""2026-10-18 首次正式预训练后改(step 21850 闭环到达 0%):原实现照扩散模型写成 sin(1000·t·f),f∈[1e-3, 1],
最高角频率 1000 —— t∈[0,1] 上网络对 t 极其敏感,实测 last.pt 的 ∂u/∂t 均方根 6–12,而 v 只有约 1.1,
MeanFlow 目标 u_tgt = v − (t−r)·du/dt 被自激的 JVP 项主导:MF 残差 2 万步一直停在 7–8、17k–19k 尖峰到 96、梯度 100% 被裁剪。
改为角频率在 [0.5, 50] 几何分布(最快一维在 [0,1] 上约转 8 圈,足够分辨 t,导数量级降 20 倍)。"""
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]
w = torch.exp(torch.linspace(math.log(0.5), math.log(TEMB_WMAX), half, device=t.device, dtype=t.dtype))
a = t[:, None] * w[None]
return torch.cat([torch.sin(a), torch.cos(a)], -1)