🤖 machine learning

FlowMPC: Improving Flow Matching policies with World Models

This paper introduces FlowMPC, a framework that enhances Flow Matching policies for robotic manipulation by integrating a learned world model to perform test-time MPPI planning, thereby improving task success rates without altering the original training objective.

Original authors: Chandon Hamel

Published 2026-06-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Chandon Hamel

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot arm to pick up a specific object, like a cube or a YCB tool, and place it exactly where you want it. This paper introduces a new way to make that robot smarter, called FlowMPC.

Here is the breakdown of how it works, using simple analogies:

1. The Problem: The "Talented but Naive" Apprentice

First, the researchers trained a robot using a method called Flow Matching (FM). Think of this FM policy as a very talented apprentice who has watched thousands of videos of a master craftsman doing the job.

What it's good at: Because it learned from so many examples, it is great at generating creative, varied ways to move its arm. It knows that there isn't just one way to pick up a cube; there are many successful paths.
The Weakness: The apprentice is strictly an "imitator." It only knows how to act based on what it saw in the videos. If the robot makes a tiny mistake (like a slight wobble) that wasn't in the training videos, the apprentice might panic or compound the error, causing the task to fail right at the finish line. It doesn't know how to "think ahead" to fix its own mistakes.

2. The Solution: Adding a "Crystal Ball" (The World Model)

To fix this, the researchers didn't retrain the apprentice. Instead, they gave the apprentice a Crystal Ball (a learned World Model) and a Planner.

The Crystal Ball (World Model): This is a simulator that predicts the future. If the robot moves its arm this way, the crystal ball says, "Okay, in one second, the object will be here, and you will likely succeed." If it moves that way, it says, "No, you'll drop it."
The Planner (MPPI): This is the decision-maker. It asks the apprentice: "Hey, give me 500 different ideas on how to finish this task."
- The apprentice (FM policy) quickly generates 24 of its best, most creative ideas based on what it learned from the videos.
- The planner fills the rest of the 500 ideas with random guesses.
- The planner then uses the Crystal Ball to simulate all 500 ideas into the future. It checks: "Which of these paths leads to a successful finish without dropping the object?"
- Finally, it picks the best path and executes the first move.

3. The Result: From "Good" to "Great"

The paper tested this on two tasks: picking up a cube and picking up a specific tool (PickSingleYCB).

The Outcome: The robot with the Crystal Ball didn't just reach the goal; it stayed there.
- PickCube: The success rate went from 93% to 97%.
- PickSingleYCB: The success rate went from 57% to 66%.
The Key Insight: The biggest improvement happened at the very end of the task. The FM apprentice alone could often get the object close to the target but would fumble it in the final seconds. The World Model acted like a safety net, allowing the robot to correct tiny errors and hold the object steady until the very last moment.

4. Why This Is Special

Usually, to make a robot better, you have to retrain it with a new, complex set of rules (Reinforcement Learning). This paper did something different:

It kept the apprentice's training exactly the same.
It just added a "thinking step" at the moment the robot actually performs the task.

The Metaphor:
Imagine a musician who has memorized a song perfectly (the FM policy). They play it well, but if they miss a note, they might lose their place.
FlowMPC is like giving that musician a conductor who listens to the next few seconds of the music in their head. If the musician is about to miss a beat, the conductor whispers, "Actually, try this slight variation instead," ensuring the song ends perfectly.

Summary

The paper shows that you can make a robot significantly more reliable by combining a learned imitator (who knows many ways to do a task) with a predictive simulator (who knows which of those ways will actually work). This allows the robot to fix its own small mistakes in real-time, leading to higher success rates, especially in the critical final moments of a task.

Technical Summary: FlowMPC

Problem Statement

Flow Matching (FM) has emerged as a powerful method for behavior cloning in multimodal action spaces, particularly for robotic manipulation where multiple distinct action trajectories can solve the same task. However, FM is fundamentally an imitation learning method trained to match expert behavior rather than to directly maximize expected return. Consequently, FM policies are limited by the quality and coverage of the demonstration data. They often struggle to recover from small errors, compounding mistakes, or states not well-represented in the expert dataset.

While recent work has attempted to address this by refining FM policies using reinforcement learning (RL), such approaches require modifying the FM training objective itself. This paper investigates a different question: Can a learned world model improve an FM policy without altering the FM training objective? The goal is to determine if model-based planning can complement flow-based imitation policies at test time.

Methodology: FlowMPC Framework

The proposed framework, FlowMPC, integrates an imitation-learned FM policy with a learned world model within a Model Predictive Path Integral (MPPI) planning loop. The core idea is to use the FM policy as a structured proposal mechanism for generating candidate action sequences, which are then evaluated and refined by the world model.

1. Training Paradigm

The training pipeline consists of five stages:

Expert Data Collection: A SAC agent with privileged state information generates expert demonstrations.
Trajectory Collection: Full trajectories (RGB, proprioception, actions, rewards) are collected.
FM Policy Training: An FM policy is trained on these expert trajectories to learn a multimodal action prior.
Visual Preprocessing: Expert trajectories are re-collected with RGB observations preprocessed by the FM policy's visual encoder to produce compact visual latents ( $z_{viz,t}$ ).
World Model Training: The world model is trained online using a replay buffer preloaded with expert data, utilizing the preprocessed visual latents combined with proprioceptive state ( $x_t = [s_t, z_{viz,t}]$ ).

2. Architectural Adaptations to TD-MPC2

The framework builds upon TD-MPC2 but replaces its learned actor with the imitation-learned FM policy. Several critical modifications were made to accommodate this hybrid:

Behavior-Cloning Action Predictor ( $\pi_{BC}$ ): A dedicated network maps the current model state to an action. This is used for terminal value bootstrapping during planning rollouts to ensure the value function evaluates in-distribution actions, stabilizing the planning process.
Unified Representation Interface: The world model consumes the FM visual latent directly concatenated with the proprioceptive state, rather than learning separate encoders. This allows the world model to leverage the FM policy's frozen visual features.
Simplified Objectives: Given the binary reward structure of the tasks, the complex discretized symlog-transformed rewards of standard TD-MPC2 were replaced with binary cross-entropy for rewards and standard mean-squared error for Q-value regression.
Dual World Model Variants:
- State-Space Variant: Predicts directly in the space of concatenated state and visual latent. Used for the PickCube task ( $H=1$ ).
- Latent-Space Variant: Follows standard TD-MPC2 with a learned latent representation and multi-step rollouts. Used for the more complex PickSingleYCB task ( $H=3$ ).

3. Inference and Planning

At test time, the planner operates as follows:

Proposal: Given the current observation, the planner samples $K$ candidate action trajectories. A subset ( $K_\pi$ ) is proposed by the FM policy, while the remainder are sampled randomly.
Rollout & Scoring: Each trajectory is rolled out in the learned world model. Trajectories are scored by summing predicted rewards and a terminal value estimate ( $\hat{Q}$ ) derived from the behavior-cloning predictor.
Refinement: MPPI reweights and refines the trajectories based on these scores. The first action of the optimized sequence is executed.
Efficiency: The FM visual encoding is computed once per observation and reused across all samples and integration steps.

Key Results

The approach was evaluated on two ManiSkill manipulation tasks: PickCube-v1 and PickSingleYCB-v1.

End-of-Episode Success: The most significant gains were observed in maintaining success through the end of the episode.
- PickCube: Improved from 93.14% (FM only) to 97.44% (FlowMPC).
- PickSingleYCB: Improved from 56.81% (FM only) to 66.41% (FlowMPC).
Anytime Success: Improvements were also seen in reaching a successful state at any point, though the gains were smaller.
- PickCube: 95.78% $\to$ 98.68%.
- PickSingleYCB: 68.77% $\to$ 69.78%.

The results indicate that the world model is particularly effective at correcting small action errors and maintaining stable control near task completion, rather than just helping the agent reach a transient success state.

Significance and Claims

The paper posits that world-model-based planning can effectively complement flow-based imitation policies without requiring a new FM training objective.

Hybrid Approach: The work demonstrates a viable hybrid between imitation learning and model-based RL. It uses imitation to learn a strong, multimodal action prior and planning to refine decisions at inference time.
No Retraining Required: The improvement comes from better utilization of the policy at test time, not from retraining the policy with a new RL objective.
Modest Scope: The authors explicitly state that the results are limited to two tasks and that the optimal world model design (state-space vs. latent-space, planning horizon) varied between tasks. They do not claim a universally optimal configuration but suggest that the combination of expressive generative policies and lookahead planning is a promising direction for robust policy improvement.
Performance Baseline: The method improved upon a strong FM baseline (which was already close to TD-MPC2 performance on one task), suggesting the planner "sharpens" a strong policy rather than merely rescuing a weak one.