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 meanflow_loss(model, mem, x, p_neq=0.25, ada_p=1.0, ada_c=1e-3, return_aux=False, p_edge=0.0):
    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()
        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)