OMP: One-step Meanflow Policy with Directional Alignment¶

会议: ICML2026
arXiv: 2512.19347
代码: 待确认
领域: 机器人 / 具身智能 / 生成式策略
关键词: MeanFlow, 单步策略, 方向对齐, JVP 有限差分, 机器人操作

一句话总结¶

本文针对将 MeanFlow 范式直接搬到机器人操作时暴露出的三个理论病灶（频谱偏差、低速区梯度饥饿、嵌套 JVP 内存爆炸），提出 OMP：用一项 cosine-style 方向对齐损失把预测平均速度与真实平均速度方向"锁死"，再用有限差分 DDE 近似 Jacobian-Vector Product 解耦前后向，让单步（NFE=1）生成策略在 Adroit/Meta-World 上以 6.8ms 级延迟做到比 MP1 平均高 3.4%、在 Meta-World Very Hard 任务高 10.6% 的成功率。

研究背景与动机¶

领域现状：生成式机器人策略当前主流是把动作生成建模为概率去噪过程，Diffusion Policy / DP3 这类扩散策略靠 10 步左右的迭代去噪拿到了高成功率，但 NFE=10 带来的推理延迟挤死了高频闭环控制；为提速，FlowPolicy、ManiFlow 这类基于 flow matching 或一致性蒸馏的方法把推理压到单步，但训练侧要靠分段直线流或显式一致性约束，过强的架构约束又损失了泛化。

现有痛点：MeanFlow（2025）从理论上给出了一条更干净的"单步化"路径——直接学习区间平均速度 \(u(z_t, r, t)\)，绕开 ODE 求解器；其在机器人侧的落地 MP1 把延迟压到 6.8 ms。但作者发现把 MeanFlow 照搬到机器人，会暴露三个 image generation 场景看不见的病灶。

核心矛盾：图像生成里像素级动态范围大、梯度信号充足，掩盖了 MeanFlow 目标自身的频谱与几何缺陷；而机器人动作空间维度低、精细任务里真实平均速度 \(\|v_0\|\) 接近 0，三个理论病灶集中爆发：(1) 频谱偏差——时间积分相当于除以 \(i\omega\)，目标 PSD 按 \(1/\omega^2\) 衰减，等价于低通滤波器，把精细操作里的高频方向调整全压掉；(2) 梯度饥饿——MSE 损失对角度误差的梯度为 \(2\rho\rho^*\sin\alpha\)，与目标幅度 \(\rho^*\) 乘性耦合，\(\rho^*\to 0\) 时模型干脆把自己输出收缩到 0 而不去对齐方向；(3) 内存复杂度——MeanFlow Identity 里的总导数展开后含 JVP \(\nabla_z u\cdot dz/dt\)，对其求 \(\nabla_\theta\) 等价于嵌套 Forward-AD + Reverse-AD，必须同时保存 primal/tangent/adjoint 三套激活，大点云骨干训不动。

本文目标：(a) 拆掉 MSE 把方向和幅度强耦合的设定，让方向监督在低速区不消失；(b) 用一个不需要符号微分的近似替换 JVP，把训练内存压回标准 backprop 水平；(c) 保持 NFE=1 的推理速度。

切入角度：既然根本病因是 MSE 把方向和幅度绑死、又是 JVP 的解析展开把内存炸开，那就直接绕过去——用 cosine 把方向单拎出来当一项独立损失，用中心差分近似时间导数。

核心 idea：用方向对齐损失把预测平均速度的"指向"显式锁到真实平均速度 \(v_0\) 上，叠加一个 \(O(\epsilon)\) 中心差分代替 JVP 解耦前后向。

方法详解¶

整体框架¶

OMP 建在 MeanFlow 框架上：输入 3D 点云观测（FPS 下采到 512 或 1024 点）+ 2 步观测历史，输出长度为 4 的动作序列、每次执行其中 3 步。训练时模型 \(u_\theta(z_t, r, t \mid c)\) 学习时刻 \(r,t\) 之间的平均速度，遵循 MeanFlow Identity

\[u(z_t,r,t|c)=v(z_t,t|c)-(t-r)\dfrac{d}{dt}u(z_t,r,t|c)\]

把右边当 target \(u_{tgt}\)。推理时一次正向直接从噪声 \(z_T\sim\mathcal{N}(0,I)\) 走到动作 \(z_0\)，沿用 \(v_0 \triangleq z_T - z_0\) 作为真实平均速度。OMP 在 MP1 的 \(\mathcal{L}_{mse} + \lambda_{Disp}\mathcal{L}_{Disp}\) 基础上叠了第三项 \(\mathcal{L}_{DA}\) 解决方向对齐，并把 Identity 里的 \(\frac{d}{dt}u\) 实现替成 DDE 中心差分，得到 OMP-JVP（保留解析 JVP）和 OMP-DDE（用差分近似）两个版本。

关键设计¶

方向对齐损失 \(\mathcal{L}_{DA}\):
- 功能：把预测平均速度 \(u(z_t,r,t|c)\) 的方向显式拉到真实平均速度 \(v_0=z_T-z_0\) 上，解掉 MSE 在低 \(\rho^*\) 区域的梯度饥饿和频谱偏差。
- 核心思路：先算 cosine 相似度 \(\cos\alpha = \dfrac{v_0\cdot u}{\|v_0\|\cdot\|u\|}\)，再用对数形式 \(\mathcal{L}_{DA}=-\log\!\big(\frac{\cos\alpha+1}{2}\big)\) 当损失。这个形式有几个好处：对 \(\cos\alpha=-1\) 处发散给最强惩罚、对 \(\cos\alpha=1\) 处趋零、整个损失只取决于方向不取决于幅度，所以即使 \(\|v_0\|\to 0\) 方向梯度也不会塌缩；分母处加 \(\epsilon_{dir}\approx 10^{-6}\) 防数值除零。
- 设计动机：作者在 §4.2.2 用法-余弦定理拆 MSE 得到 \(\partial\mathcal{L}_{MSE}/\partial\alpha = 2\rho\rho^*\sin\alpha\)，证明角度梯度被 \(\rho^*\) 乘性压制；机器人精细操作 \(\rho^*\approx 0\)，MSE 干脆鼓励 \(\rho\to 0\) 拿到"静止策略"。\(\mathcal{L}_{DA}\) 把方向和幅度解耦，弹道阶段 \(\mathcal{L}_{mse}\) 主导走幅度，接触阶段 \(\mathcal{L}_{DA}\) 主导走方向。同时显式对齐 \(v_0\) 还顺手跳出了 \(1/\omega^2\) 低通滤波——目标不再经过时间积分。
Differential Derivation Equation (DDE):
- 功能：用中心差分替掉 MeanFlow Identity 里的解析 \(du_\theta/dt\)，避免对 JVP 再做反向传播带来的嵌套 AD 内存爆炸。
- 核心思路：把时间导数近似为 \(\dfrac{du_\theta(z_t,t,r|c)}{dt}\approx\dfrac{u_\theta(z_{t+\epsilon},t+\epsilon,r|c)-u_\theta(z_{t-\epsilon},t-\epsilon,r|c)}{2\epsilon}\)，\(\epsilon\) 是小扰动常数（敏感度分析见 §E.2）。这样训练图里只剩两次普通 forward + 一次普通 backward，不再需要存 tangent 激活，回到标准 backprop 的内存量级。
- 设计动机：§4.2.3 算了一下，对 JVP \(\nabla_z u_\theta\cdot v\) 再求 \(\nabla_\theta\) 等价于二阶混合偏导 \(\partial^2 u/\partial\theta\partial z\)，PyTorch/JAX 里要嵌套 Forward-AD 在 Reverse-AD 外面，必须同时存原始激活 \(X\)、tangent \(\delta X\)、tangent 的 adjoint 三套图，PointNet++/Transformer 这种点云骨干根本喂不进 4090。差分近似精度只是 \(O(\epsilon^2)\) 的截断误差，对成功率影响可控但内存收益巨大。
组合损失与 JVP/DDE 双版本:
- 功能：把方向、幅度、特征判别三种信号融成一个训练目标，并把内存优化做成独立开关。
- 核心思路：最终训练目标是 \(\mathcal{L}=\mathcal{L}_{mse}+\lambda_{Disp}\mathcal{L}_{Disp}+\lambda_{DA}\mathcal{L}_{DA}\)，其中 \(\mathcal{L}_{Disp}\) 沿用 MP1 的 dispersive loss 让特征空间更可分。把 \(\frac{d}{dt}u\) 的实现拆成两套——OMP-JVP 保留解析 JVP 拿最佳精度，OMP-DDE 用 DDE 近似换显存——这样可以根据实际任务规模（点云尺寸、动作 horizon）按需切换，而不用为了内存被迫永久牺牲精度。
- 设计动机：弹道段（大平移）和精细接触段（角度对齐）对损失的需求其实是分时段的，靠加权和能让模型在两阶段都拿到非零梯度；JVP/DDE 双版本是承认"内存-精度"是个真实 trade-off，OMP-JVP 给学术对照、OMP-DDE 给实际部署。

损失函数 / 训练策略¶

损失：\(\mathcal{L}=\mathcal{L}_{mse}+\lambda_{Disp}\mathcal{L}_{Disp}+\lambda_{DA}\mathcal{L}_{DA}\)；\(\mathcal{L}_{DA}=-\log\!\big(\frac{\cos\alpha+1}{2}\big)\)；DDE 时间步长 \(\epsilon\) 在 §E.2 做了敏感度扫描。
数据：每个仿真任务 10 条专家演示；点云 FPS 到 512 或 1024 点；图像 84×84；观测 history=2、prediction horizon=4、execution horizon=3。
训练：AdamW，lr=1e-4，batch=128；Adroit 训 3000 epoch，Meta-World 训 1000 epoch，每 200 epoch 评估一次，最终成功率取前 5 高的平均、再跨种子 (0/10/20) 平均；硬件单卡 RTX 4090。

实验关键数据¶

主实验：Adroit + Meta-World 37 任务平均¶

方法	NFE	Adroit Pen	MW Medium	MW Hard	MW Very Hard	总平均
DP (RSS'23)	10	13±2	11.0±2.5	5.25±2.5	22.0±5.0	35.2±5.3
DP3 (RSS'24)	10	46±10	44.5±8.7	32.7±7.7	39.4±9.0	68.7±4.7
FlowPolicy (AAAI'25)	1	54±4	58.2±7.9	40.2±4.5	52.2±5.0	71.6±3.5
MP1 (AAAI'26)	1	58±5	68.0±3.1	58.1±5.0	67.2±2.7	78.9±2.1
OMP-JVP	1	60±4	77.4±2.2	62.5±3.1	77.8±3.0	82.3±1.6
OMP-DDE	1	64±3	76.4±2.7	61.0±3.0	70.6±4.9	80.8±2.2

OMP-JVP 在总平均上比 MP1 高 3.4%、比 FlowPolicy 高 10.7%；越难的任务 OMP 增益越大——Meta-World Medium +9.4%、Very Hard +10.6%。MP1 在 Easy 子集（21/37 任务）已逼近上限（88%+），把绝对增量主要拉低到了 1.5%。

真机实验（3 任务，成功率 %）¶

方法	Place	Clean	Slip Ring
DP3	65	60	50
FlowPolicy	60	50	40
MP1	70	65	55
OMP	80	75	70

最难的 Slip Ring 上 OMP 比 MP1 高 15%，验证了方向对齐在"真实低速精细操作"上的核心收益。

消融 + 内存¶

配置	总平均成功率	说明
OMP-JVP (Full)	82.3	完整模型
− \(\mathcal{L}_{Disp}\)	81.2	去 dispersive，掉 1.1%（小）
− \(\mathcal{L}_{DA}\)	78.9	去方向对齐，掉 3.4%，回到 MP1 水平
− \(\mathcal{L}_{Disp}\) − \(\mathcal{L}_{DA}\)	78.3	全去，Adroit Pen 60→48
OMP-DDE (Full)	80.8	差分近似版
− \(\mathcal{L}_{DA}\) (DDE)	77.2	同样验证方向对齐是核心

任务 / Horizon	OMP-JVP 显存	OMP-DDE 显存
Adroit Hammer / H=4	6.60 GB	5.35 GB
Place Bottle / H=4	23.49 GB	18.33 GB
Adroit Hammer / H=16	7.69 GB	6.12 GB
Place Bottle / H=16	26.71 GB	19.19 GB

关键发现¶

方向对齐是主功臣：去 \(\mathcal{L}_{DA}\) 直接掉 3.4–3.6%，去 \(\mathcal{L}_{Disp}\) 只掉 0.7–1.1%，证明病灶根源在 MSE 的几何耦合而不是特征判别。
OMP 增益与任务难度正相关：Easy 任务上 MP1 已经把 ceiling 顶到 88%+，OMP 拉不开差距；Very Hard 上 +10.6% 说明方向对齐主要救的是低速精细任务，和理论预期完全吻合。
JVP→DDE 是精度-内存交易：DDE 平均掉 1.5%（Very Hard 掉 7.2% 偏多），换来 Place Bottle/H=16 上 26.71 GB→19.19 GB 的 28% 显存下降；点云越大、horizon 越长，DDE 性价比越高。
训练曲线更稳：Figure 5 显示 OMP 的成功率曲线方差远小于 FlowPolicy/MP1 的剧烈震荡，方向对齐顺带提升了训练稳定性。

亮点与洞察¶

把"频谱偏差 + 梯度饥饿 + 内存爆炸"打包成一个理论叙事：作者没有只丢一个 cosine 损失，而是先用 PSD 频域分析、再用法-余弦定理给 MSE 角度梯度的闭式表达、再用 AD 图分析三套激活，把三个看似分离的问题串成同一个"MeanFlow 不适合机器人"的故事，给了 \(\mathcal{L}_{DA}\) + DDE 强动机，论文的说服力主要来自这三段分析而不是数字本身。
cosine 损失的对数形式：写成 \(-\log((\cos\alpha+1)/2)\) 而不是直接 \(1-\cos\alpha\)，是个可复用的 trick——对数让 \(\cos\alpha=-1\) 附近梯度发散，模型在"完全走反"时被强惩罚；而 \(1-\cos\alpha\) 在反向时梯度反而最小，会陷在反向局部最优。
DDE 的真正价值不是"近似精度"，是"解耦计算图"：跨任务可复用——任何需要 \(du_\theta/dt\) 又要对 \(\theta\) 反传的场景（不止 MeanFlow，还包括 score matching 的二阶变体、一些 NeuralODE 训练）都可以借这一招把内存压回标准 backprop 量级。
Easy/Hard 任务分桶汇报：作者明确指出"在 Easy 子集 MP1 已经接近上限"，把总平均增益拆到 Easy/Medium/Hard/Very Hard 分别看，给出了任务难度 → 增益曲线，让方法的"适用边界"很清晰，值得任何"打榜接近饱和"的工作借鉴。

局限与展望¶

真机实验规模偏小：只有 3 个任务、每任务 20 次试验（10% 的颗粒度可见），统计强度有限，没给标准差；Slip Ring +15% 这种数字需要更大样本支撑。
DDE 的 \(\epsilon\) 是手调超参：作者把敏感度分析放进附录但没给自适应方案，不同任务可能需要重调；理想的下一步是用 trajectory curvature 自适应或 second-order trust region 选 \(\epsilon\)。
没有和近期 OneDP/ManiFlow 等做基于蒸馏的单步方法详细对比：表里 baseline 主要是 MP1/FlowPolicy 同代，缺一个 well-distilled DP3 的强基线（理论上多步教师 + 一致性蒸馏可能拿到接近 OMP 的精度）。
方向对齐对 multimodal 动作分布的影响未讨论：精细操作里可能有多个等价方向（左手/右手）能完成任务，cosine 强制单一方向是否会丢失模式多样性、值得在 dexterous 任务上做 mode coverage 度量。
缺乏对模仿数据规模的扫描：只用 10 条 demo，10→100 demo 增长曲线、\(\mathcal{L}_{DA}\) 收益是否随 demo 增加而衰减没有给。

评分¶

新颖性: ⭐⭐⭐⭐ 单一 cosine 损失 + 中心差分都不算新数学，但用频谱/几何/AD 三段分析把它们组装成"MeanFlow 在机器人侧的修复套件"是真正的洞察。
实验充分度: ⭐⭐⭐⭐ 37 个仿真任务 × 3 种子 + 3 个真机任务 + 完整消融 + 显存对照表，但真机统计强度有限、缺与蒸馏类单步方法的强对照。
写作质量: ⭐⭐⭐⭐ 理论分析（§4.2 三个小节）结构清晰、动机推得很顺，主表难度分桶呈现是亮点；公式用了 v3 自动转 LaTeX 的格式偶尔有冗余。
价值: ⭐⭐⭐⭐ 给社区一个"MeanFlow 在低维动作空间的修复模板"，方向对齐损失和 DDE 都可独立复用到其他单步生成式策略框架。