SinjaeKang/warmprior-dsrl

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},
}