Robot_Arm/ml/model/meanflow.py

"""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)）。
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


def sample_rt(B, device, dtype, p_neq=0.25, mu=-0.4, sigma=1.0, p_edge=0.0):
    """2026-10-19（审查 review_model_train P0-3）：部署只在 (r=0, t=1) 推理一次，而 logit-normal 取 max/min 后
    P(t>0.95) 约 0.08%、P(r<0.05) 约 0.5%，这个角几乎从未被训练到（一步残差与噪声相关 −0.86 即外推失准）。
    p_neq > 0 时另以 p_edge 的概率把样本直接放在 (r=0, t=1)；预热期（p_neq=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)
    if p_neq > 0 and p_edge > 0:
        edge = torch.rand(B, device=device) < p_edge
        t = torch.where(edge, torch.ones_like(t), t); r = torch.where(edge, torch.zeros_like(r), r)
    return r, t


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)
        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()
    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):
    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)