WarmPrior:
Straightening Flow-Matching Policies with Temporal Priors

Abstract

Generative policies based on diffusion and flow matching have become a dominant paradigm for visuomotor robotic control, capturing complex, multimodal action distributions by iteratively transforming samples drawn from a Gaussian prior. In this work, we introduce WarmPrior, a new source distribution for generative policies. Our key idea is simple: rather than a zero-mean prior, we center the Gaussian on the action sequence predicted at the previous inference step. This yields straighter flow trajectories, mirroring the effect of OT couplings in Rectified Flow. Across simulation benchmarks (Robomimic and MimicGen with a Diffusion Policy backbone) and on a real Franka arm with the GR00T N1.5 vision-language-action model, WarmPrior consistently improves success rates. Our results identify the source distribution as an important yet underexplored design axis in generative robot policies.

Method

WarmPrior modifies only the source distribution. The network, interpolant, and training objective are untouched. We instantiate it in two minimal variants that differ in how the prior mean is anchored to the agent's own action history.

Main Results

Success rate (%) with the Diffusion Policy (ChiTransformer) backbone. Parentheses: gain over \(\mathcal{N}(\mathbf{0}, \mathbf{I})\).
(Green: non-overlapping 1σ seed intervals. Bold: best per (task, NFE).)

Why It Works: Straighter Flows

The baseline paths curve noticeably as the network pulls samples from a random origin onto the action manifold; WarmPrior paths, already starting close to the manifold, are visibly straighter and more parallel. Because fewer flows cross one another, the network spends less capacity realigning samples from the random base distribution and can devote more to refining actions, exactly where it matters for downstream success rate.

Pathwise curvature (lower is straighter)

Task	N(0, I)	WP-Past	WP-Preview
Square-PH	1.000	0.823	0.803
Square-MH	1.000	0.705	0.559
Transport-PH	1.000	0.720	0.692
Transport-MH	1.000	0.695	0.637
Tool-Hang-PH	1.000	0.806	0.807

Tasks with the largest curvature reduction (Square-MH, Transport-MH) also show the largest success-rate gain.

Flow paths on Square-MH. Normalized action coordinate vs. denoising time \(t \in [0, 1]\). WarmPrior paths are straighter and more parallel than the baseline.

We make this quantitative through a branching cost bound:

A better prior mean \(\mu\) tightens the first term, while the residual \(\sigma^2\) governs the second, and the two interact through a non-trivial trade-off: too small a \(\sigma\) concentrates the source onto \(\mu\) and is fragile when \(\mu\) is imperfect, while too large a \(\sigma\) leaves the field less tightly straightened. See Appendix B of the paper for the full derivation.

A Tunable Knob for Temporal Consistency

Setup. A 1D navigation toy where the observation \(o\) is the agent's horizontal position and the action is its vertical height. Demonstrations split evenly between passing above and below each obstacle, so the conditional \(p(a \mid o)\) is multimodal at every position. A regression policy collapses to the mean and drives straight through the obstacle (panel b).

Mode switching. A flow-matching policy recovers both modes per inference (panel c), but the standard objective places no constraint linking consecutive chunks: at every chunk boundary the policy can flip between modes. We call this pathology mode switching. Action chunking commits within a chunk but does not stop the oscillation. Naive history conditioning (panel d) eliminates the flips only by pinning each rollout to a single mode, sacrificing multimodality, and is further fragile to inference-time distributional shift over the action history.

σ as a knob. Because the WarmPrior mean is correlated with the previous chunk, a small \(\sigma\) keeps the new chunk in the nearby mode's basin and prevents crossings between distant modes (panel e); a large \(\sigma\) broadens the source and recovers the multimodal distribution at every step (panel f). \(\sigma\) thus becomes a continuous regulator between temporal consistency and multimodal expressiveness, an implicit alternative to history conditioning that does not pin the policy to a single mode.

Real-Robot Experiments

Four tabletop tasks on a Franka Research 3 using the DROID setup, with the GR00T N1.5 VLA backbone.
30 demos per task, 50 trials per seed, 3 seeds, NFE = 4.

Four tabletop tasks: Food Waste Disposal, Cup Stacking, Block Stacking, Cable Insertion.

WarmPrior consistently improves success rate. Gains are largest on the harder tasks (Cable Insertion, Block Stacking) where endpoint ambiguity is highest, mirroring the simulation pattern.

WarmPrior Improves Prior-Space RL Efficiency

DSRL-WP: residual parameterization. We adapt DSRL to the WarmPrior policy by reparameterizing the prior sample as a residual around the WarmPrior mean \(\mu\) and conditioning the RL agent on that mean. The augmented observation is \(\tilde{o} = [o, \mu]\), and the RL policy outputs a bounded offset \(\Delta\) that is added to \(\mu\):

We use \(\delta = 1.5\) on Robomimic Square and \(\delta = 1.0\) on Transport, matching vanilla DSRL. Conditioning \(\Delta\) on \(\mu\) lets it adapt to the current mean anchor and keeps the MDP Markovian. We instantiate DSRL-NA with this residual parameterization and call the resulting setup DSRL-WP.

Findings. We fine-tune on Square and Transport, the two Robomimic tasks not saturated under BC, with a frozen WarmPrior backbone pretrained by behavior cloning for 3000 epochs. DSRL-WP learns faster, converges more stably, and reaches a higher asymptote than both DSRL baselines, consistently exceeding 99.0 on Square and attaining 98.1 ± 0.54 on Transport, where DSRL-NA and DSRL-SAC plateau around 90.0. Two further findings on the Transport task accompany this gain. DSRL-WP attains \(Q^W\) values at the actor's deterministic output that are consistently higher than the DSRL-NA baseline. DSRL-WP also achieves a critic MSE TD-error roughly an order of magnitude lower throughout training, with tighter per-step fluctuations. We read both as signs that WarmPrior makes the RL policy search a localized region of the prior space where samples already decode to good actions, rather than the random noise space, consistent with the faster convergence reported above. This action-aligned prior space also yields the stronger asymptotic performance.

Prior-space RL. DSRL baselines vs. DSRL-WP on Square and Transport, averaged over 3 seeds (±1σ shading).

WarmPrior improves critic learning. DSRL-WP achieves higher \(Q^W(s, z_{\text{actor}})\) across 4k demonstration states (left) and ~10× lower critic loss with smaller fluctuations (right) on Transport.