FlowMPC: Improving Flow Matching policies with World Models

Abstract: Flow Matching (FM) is a powerful approach for behavior cloning in multimodal action spaces [Jiang et al., 2025], but because it is not trained to directly maximize expected return, there is still room to improve how FM policies act at test time. This work investigates whether a learned world model can improve FM policies by enabling Model Predictive Path Integral (MPPI) planning over candidate action sequences proposed by the policy. Building on TD-MPC2 [Hansen et al., 2024], I introduce FlowMPC, a framework that combines an imitation-learned FM policy with a learned world model for test-time planning in ManiSkill manipulation tasks [Tao et al., 2025]. Across PickCube and PickSingleYCB, adding the world model improved performance over the FM policy alone, with especially clear gains in end-of-episode success. These results suggest that world-model-based planning can effectively complement flow-based imitation policies without modifying the FM training objective.

• The paper introduces FlowMPC, which uses a frozen flow-matching imitation policy to propose multimodal action trajectories while a TD-MPC2-style world model and MPPI planner select higher-return actions at inference time.
• The method improves end-of-episode success from 93.14% to 97.44% on PickCube and from 56.81% to 66.41% on PickSingleYCB, while producing smaller gains in anytime success.
• The results show that planning can stabilize competent imitation policies without changing their training objective, although performance depends on task-specific world-model representations and planning horizons.

Overview

FlowMPC investigates whether a learned world model can improve a Flow Matching (FM) imitation policy at test time without modifying the FM training objective (2606.16286). The framework builds on TD-MPC2's model-based planning backbone, replacing its learned actor with an FM policy that serves as a proposal distribution for Model Predictive Path Integral (MPPI) planning. The central claim is that the FM policy supplies a rich, multimodal prior over action trajectories, while the world model evaluates and refines those proposals according to predicted return—yielding gains over pure behavior cloning on two ManiSkill manipulation tasks.

Motivation and positioning

FM has proven effective for behavior cloning in multimodal action spaces, which is common in robotic manipulation where many distinct trajectories solve the same task. However, as an imitation method, FM is bounded by demonstration quality and coverage; it is not trained to maximize expected return and can compound errors under distribution shift. Prior remedies refine flow policies with reinforcement learning objectives—for example, Flow Matching Policy Gradients [mcallister2026flow] and RL for flow-matching policies (Pfrommer et al., 20 Jul 2025)—but these alter the learning objective itself. FlowMPC takes a complementary route: it keeps the policy frozen after imitation training and improves decision making purely through inference-time planning, treating the FM policy as an action prior rather than an optimization target.

Relative to TD-MPC2 (Hansen et al., 2023), which trains both a world model and an actor jointly, FlowMPC decouples these roles: imitation provides the proposal mechanism, and model-based planning provides selection. The author notes this specific combination had been relatively underexplored.

Method

The training pipeline proceeds in five stages: (1) train a SAC expert with privileged state access; (2) collect expert trajectories containing RGB observations, proprioceptive state, actions, and rewards; (3) train a conditional 1D temporal U-Net FM policy on those trajectories; (4) re-encode the expert data through the FM policy's visual encoder to produce compact visual latents $z_{\mathrm{viz},t}$; and (5) train the world model online from a replay buffer preloaded with expert data. The world model operates on $x_t = [s_t, z_{\mathrm{viz},t}]$, reusing the frozen FM visual representation rather than learning from raw pixels—a design choice the paper identifies as important for tractability, particularly for the state-space variant.

Several modifications to TD-MPC2 were required:

• Behavior-cloning action head: a predictor $\pi_{\mathrm{BC}}(x_t)$ maps model states to actions, used during training and for terminal value bootstrapping at the end of planning rollouts. Without it, evaluating terminal states against arbitrary sampled actions destabilized planning.
• Simplified reward/value targets: since the tasks use binary rewards, the discretized symlog machinery of TD-MPC2 was replaced with binary cross-entropy for rewards and MSE regression for Q-values.
• Two world model variants: a latent-space version following standard TD-MPC2 (learned representation, multi-step rollouts, consistency loss), and a state-space version that predicts directly in $[s_t, z_{\mathrm{viz},t}]$ space with one-step SAC-style transitions, decoupling the heads from any shared encoder.

At inference, the planner samples $K = 512$ candidate trajectories per step, of which $K_\pi = 24$ come from the FM policy and the remainder are random. Each candidate is scored by summed predicted rewards plus a terminal value estimate $$\hat{Q}(\hat{x}_{t+H+1}, \pi_{\mathrm{BC}}(\hat{x}_{t+H+1}))$$, refined via MPPI, and the first action executed. The FM visual encoding is computed once per observation and reused across all samples and Euler integration steps for efficiency.

Notably, the best configuration differed by task: PickCube favored the state-space model with horizon $H=1$; PickSingleYCB favored the latent model with $H=3$. The cross-variant swap performed worse in preliminary experiments, indicating the method does not yet have a single universally preferred configuration.

Experimental results

Evaluation uses 50-step episodes in ManiSkill3 [taomaniskill3] on PickCube-v1 and PickSingleYCB-v1, with binary rewards and two metrics: anytime success (task succeeds at any point) and end success (still successful at the final timestep). Results are averaged over 4096 episodes across 3 seeds.

Metric	Task	FM policy	FlowMPC
End success (%)	PickCube	93.14 ± 0.77	97.44 ± 0.48
End success (%)	PickSingleYCB	56.81 ± 1.52	66.41 ± 1.45
Anytime success (%)	PickCube	95.78 ± 0.62	98.68 ± 0.35
Anytime success (%)	PickSingleYCB	68.77 ± 1.42	69.78 ± 1.41

The strongest gains are in end-of-episode success: +4.30 points on PickCube and +9.60 points on PickSingleYCB, with non-overlapping confidence intervals. Anytime success improves more modestly, and on PickSingleYCB the confidence intervals overlap—the only metric where the improvement is not statistically clear. The pattern suggests planning primarily helps correct small errors and maintain stable control through episode completion rather than enabling task completion that would otherwise fail entirely.

A further point in favor of the result is baseline strength: the FM policy alone already approaches the roughly 60% performance reported for TD-MPC2 on PickSingleYCB, and the world-model-augmented variant exceeds it. The planner therefore sharpens a competent policy rather than rescuing a weak one.

Limitations and open questions

The paper is explicit about scope constraints. Evaluation covers only two ManiSkill tasks, and the optimal world model architecture and planning horizon differ between them, so the reported gains depend meaningfully on representation choice and horizon selection; robustness of these design choices across a broader task distribution remains untested. The expert demonstrations derive from a privileged-state SAC agent, tying data quality to that particular source. The near-null effect on PickSingleYCB anytime success leaves open whether planning confers benefits beyond end-of-episode stabilization on harder tasks. The paper also does not compare against RL-refined flow policies directly, so the relative merits of test-time planning versus objective modification remain an open empirical question.

Conclusion

FlowMPC demonstrates that coupling a frozen FM imitation policy with a TD-MPC2-style world model and MPPI planning yields consistent improvements on manipulation tasks, most clearly in sustained end-of-episode success, without altering the FM training objective. The contribution is a hybrid framework connecting generative imitation policies with model-based planning, supported by controlled comparisons over large evaluation sets, though its generality beyond the two evaluated tasks and its sensitivity to world-model design choices remain to be established.