WarmPrior | Prior-space RL Experiments
Frozen flow-matching policies and normalization statistics for the prior-space RL experiments of WarmPrior: Straightening Flow-Matching Policies with Temporal Priors (paper · code).
Each Robomimic task ships two backbones, which differ only in the source distribution they were trained under:
| Task | Baseline (N(0, I)) | WarmPrior | act_steps |
|---|---|---|---|
square | square_pre_fm_mlp_ta4 | square_pre_fm_warm_ac_ta8 | 4 |
transport | transport_pre_fm_mlp_ta8 | transport_pre_fm_warm_ac_ta16 | 8 |
The WarmPrior backbone predicts a chunk of length 2H: it executes the first half and reserves the
second half as the forecast that centers the source distribution at the next inference step. Its
horizon_steps is therefore twice the baseline's, while act_steps is the same.
Both are trained for 3000 epochs with the flow-matching objective on the Robomimic PH
demonstrations. normalization.npz holds the observation/action statistics and is required at
runtime by every config.
Layout
robomimic/{lift,can,square,transport}/normalization.npz
robomimic-pretrain/{task}_pre_fm_mlp_*/…/checkpoint/state_3000.pt # baseline
robomimic-pretrain/{task}_pre_fm_warm_ac_*/…/checkpoint/state_3000.pt # WarmPrior
The paths mirror dppo/log/ in the code release, so the download can be copied straight into place.
Usage
hf download SinjaeKang/warmprior-dsrl --local-dir /tmp/wp_ckpt
cd warmprior-dsrl
mkdir -p dppo/log/robomimic-pretrain dppo/log/robomimic
cp -r /tmp/wp_ckpt/robomimic-pretrain/* dppo/log/robomimic-pretrain/
cp -r /tmp/wp_ckpt/robomimic/* dppo/log/robomimic/
python train_dsrl.py --config-path=cfg/robomimic --config-name=dsrl_square_fm.yaml # DSRL baseline
python train_dsrl.py --config-path=cfg/robomimic --config-name=dsrl_square_fm_wp_ac.yaml # DSRL-WP
A _wp_ac config requires a warm_ac backbone and a baseline config requires an fm_mlp one; the
two families are not interchangeable.
Citation
@inproceedings{kang2026warmprior,
title = {WarmPrior: Straightening Flow-Matching Policies with Temporal Priors},
author = {Kang, Sinjae and Kim, Chanyoung and Wang, Kaixin and Zhao, Li and Lee, Kimin},
booktitle = {ICML Workshop on Structured Probabilistic Inference \& Generative Modeling},
year = {2026},
eprint = {2605.13959},
archivePrefix = {arXiv},
}