WarmPrior: Straightening Flow-Matching Policies with Temporal Priors

This paper introduces WarmPrior, a temporally grounded prior derived from recent action history that replaces the standard Gaussian source in flow-matching policies to produce straighter probability paths, thereby consistently improving success rates, sample efficiency, and exploration in robotic manipulation tasks.

The Gist

Imagine you are teaching a robot arm to perform a delicate task, like stacking blocks or threading a needle. In the past, the most advanced way to teach these robots was to use a "generative" approach. Think of this like asking the robot to imagine a solution starting from total chaos.

The Old Way: Starting from Static

Standard methods tell the robot: "Start with a blank, random noise cloud (like static on an old TV). Now, slowly clean up that noise step-by-step until it looks like a perfect action."

The problem is that this "noise cloud" has no memory. It doesn't know what the robot was doing a second ago. If the robot is moving a cup, the noise doesn't know the cup is already halfway across the table. The robot has to rebuild the entire movement from scratch every single time, fighting against the randomness of the starting point. It's like trying to draw a perfect circle by starting with a pile of sand and hoping to sculpt it into a circle without ever looking at the previous stroke.

The New Idea: "WarmPrior" (The Warm Start)

The authors of this paper, WarmPrior, say: "Why start from cold, random static? Let's start from a 'warm' place."

Instead of starting with random noise, they start the robot's imagination with the last thing it actually did.

• The Analogy: Imagine you are walking down a path. The old method says, "Forget where you are; close your eyes, spin around, and try to guess the next step." The new method says, "Look at where your foot just landed. That's your starting point. Now, just take the next step from there."

They call this WarmPrior. It's a simple trick where the robot's "starting point" for its next move is anchored to its recent history.

Two Ways to Do It

The paper tests two simple versions of this idea:

1. The "Past" Version (WP-Past): The robot looks at the action it just finished and says, "Okay, I'm going to start my next guess right near where I just stopped." It's like a runner who knows their next stride will naturally follow the momentum of the last one.
2. The "Preview" Version (WP-Preview): This is a bit smarter. The robot tries to predict two steps ahead. It executes the first step, but it keeps the prediction for the second step in its head. When it's time to move again, it uses that "preview" of the future as its starting point. It's like a pianist who is already thinking about the next note while playing the current one.

Why It Works: Straightening the Path

The paper explains that this change makes the robot's "learning path" much straighter.

• The Curvy Road vs. The Highway: In the old method, because the robot starts from random noise, it has to take a winding, curvy path to get to the correct action. It's like driving from a random spot in the city to your house; you might take a detour.
• The Shortcut: With WarmPrior, the robot starts much closer to the destination. The path it learns is a straight line. This is like being dropped off right at the end of your driveway; you just walk straight to the door.

Because the path is straighter, the robot makes fewer mistakes, especially when it has to make decisions very quickly (using fewer "steps" to think).

The Results: Faster and Smarter

The researchers tested this on computer simulations and a real robot arm (a Franka Research 3).

• Better Success: The robot succeeded at tasks more often, especially on the hard ones.
• Faster Thinking: Even when the robot was allowed to think for a very short time (just one quick step), the "Warm" version worked much better than the "Cold" version.
• Reinforcement Learning: They also tried using this in a setting where the robot learns by trial and error (Reinforcement Learning). By starting with a "warm" guess, the robot learned new skills much faster because it didn't have to waste time searching random possibilities.

The Bottom Line

The paper argues that for a long time, robot designers ignored the "starting point" of their learning algorithms, treating it as a boring default setting. This paper shows that simply changing the starting point from "random noise" to "recent history" is a powerful, simple upgrade. It makes the robot's movements smoother, more consistent, and much more successful, without needing to change the complex brain (the neural network) underneath.

In short: Don't make the robot guess from scratch. Let it build on what it just did.

Technical Summary

Technical Summary: WarmPrior: Straightening Flow-Matching Policies with Temporal Priors

Problem Statement

Generative policies for robotic manipulation, particularly those based on diffusion and flow matching, have established a dominant paradigm for visuomotor control. In these frameworks, a neural field transports samples from a fixed source distribution (typically an isotropic Gaussian $N(0, I)$) to the data manifold of action chunks. While progress has been made in network architecture, interpolants, and integrators, the source distribution has remained largely static.

The authors identify a critical limitation in this standard approach: the isotropic Gaussian source is stateless and uninformative, ignoring the continuous, temporally correlated nature of robotic motion. Consequently, the policy is forced to rebuild every action chunk from scratch at each inference step. As denoising schedules shorten (reducing inference steps), the burden on the starting point increases, yet the "blind" source fails to provide the necessary temporal context, leading to curved probability paths and suboptimal performance, especially under low inference budgets or on complex tasks.

Methodology: WarmPrior

The paper introduces WarmPrior, a method that replaces the standard stateless source distribution with a temporally grounded prior constructed from readily available recent action history. This intervention is minimal, modifying only the source distribution $p_0$ while leaving the neural network, interpolant, and training objective untouched.

WarmPrior is instantiated in two variants:

1. WarmPrior-Past (WP-Past): Anchors the prior mean on the previously executed action chunk. The source is sampled as $a_0 = a_{\text{prev}} + \sigma \epsilon$, where $\epsilon \sim N(0, I)$.
2. WarmPrior-Preview (WP-Preview): Trains the policy to predict twice the chunk length ($2H$) at each step but executes only the first $H$. The second $H$ steps serve as a "preview" of the next chunk. At inference, the prior mean for the next step is set to the model's own previous forecast of the current chunk ($a_0 = \hat{a}_{\text{prev}}[H:2H] + \sigma \epsilon$).

In both variants, a residual Gaussian perturbation with scale $\sigma$ is added to ensure the source remains a proper stochastic distribution rather than a deterministic point mass. The cold region (unanchored positions) retains the standard Gaussian prior.

Key Contributions and Mechanisms

1. Geometric Straightening of Flow Paths

The primary mechanism for improvement is the straightening of probability paths. By initializing the transport close to the target manifold (the recent action history or forecast), the learned flow paths become shorter and straighter.

• Branching Cost Reduction: The authors formalize this via a "branching cost" analysis. Standard flow matching with independent couplings forces the velocity network to average over ambiguous endpoints, creating curved trajectories. WarmPrior acts as an implicit optimal-transport (OT) coupling, reducing the conditional variance of the endpoint given the intermediate point. This suppresses the irreducible residual error that causes path curvature.
• Empirical Evidence: Experiments show a significant reduction in pathwise curvature ($\kappa$) for WarmPrior variants compared to the baseline, correlating directly with success rate gains.

2. Tunable Temporal Consistency

WarmPrior introduces a continuous knob, $\sigma$, that balances temporal consistency and multimodal expressiveness.

• Implicit Commitment: A smaller $\sigma$ keeps the new action chunk within the basin of the previous mode, preventing "mode switching" (oscillating between valid action modes at chunk boundaries) without explicitly conditioning the network on full history.
• Robustness: This allows the policy to maintain consistency even when explicit action chunking is disabled ($H=1$), a regime where standard baselines typically fail catastrophically.

3. Enhanced Prior-Space Reinforcement Learning

The method extends to Reinforcement Learning (RL) in the prior space (e.g., DSRL). By using a WarmPrior-pretrained policy as a frozen base, the RL agent operates in a search space centered on a temporally grounded mean rather than the origin. This shrinks the exploration space, allowing the agent to learn a bounded residual correction around a competent anchor, leading to faster convergence and higher asymptotic performance.

Experimental Results

Simulation Benchmarks

Evaluated on Robomimic (state and image observations) and MimicGen (image observations) using Diffusion Policy backbones and the GR00T N1.5 VLA model:

• Success Rate: WarmPrior consistently outperforms the $N(0, I)$ baseline across all tasks.
• Inference Budget: Gains are most pronounced at low inference budgets (NFE = 1), where the curvature of the flow matters most. For example, on the difficult Transport-MH task with NFE=1, WP-Preview improved success rates from 23.3% to 34.5%.
• Real-Robot: Deployed on a Franka Research 3 with the GR00T N1.5 model, WarmPrior improved success rates on four tabletop manipulation tasks (Food Waste Disposal, Cup Stacking, Block Stacking, Cable Insertion), with the largest gains on precision-demanding tasks like Cable Insertion.

RL Efficiency

In prior-space RL experiments on Robomimic tasks, WarmPrior variants (WP-Past and WP-Preview) learned faster and reached higher asymptotic success rates (exceeding 99% on Square and ~97% on Transport) compared to vanilla DSRL baselines, which plateaued around 90%.

Comparison with Real-Time Chunking (RTC)

The paper compares WarmPrior with Real-Time Chunking (RTC), an inference-time procedure that clamps overlapping actions. The results indicate that while both methods improve performance, they address distinct failure modes: RTC fixes inter-chunk discontinuities at inference, while WarmPrior straightens the learned flow at training. Combining both yields additive improvements, confirming their complementary nature.

Significance and Claims

The paper claims that the source distribution is an important and underexplored design axis in generative robotic control. By shifting the paradigm from a "stateless default" to a "temporally grounded prior," WarmPrior achieves consistent performance gains without altering the core network architecture or training loss.

The authors position WarmPrior not as a replacement for complex architectural changes, but as a fundamental improvement to the generative process itself. It offers a simple, plug-and-play mechanism to:

1. Straighten flow trajectories, reducing the burden on the velocity network.
2. Provide a tunable mechanism for temporal consistency that works even without explicit action chunking.
3. Improve sample efficiency in downstream reinforcement learning by reshaping the exploration space.

The work suggests that future progress in generative robot control should look beyond the network and integrator to the construction of the prior space $p_0$ itself.