matthewh6/cspi0_libero

CSPi0 — LIBERO Context-Smoothed π₀ Prior

Context-smoothed π₀ (CSPi0) policy prior for the LIBERO simulation benchmark. This is the pretrained prior used for the LIBERO experiments in TMRL: Diffusion Timestep-Modulated Pretraining Enables Exploration for Efficient Policy Finetuning (paper).

"Context smoothing" finetunes π₀ with a timestep-modulated (CSPi0) architecture, so the diffusion timestep can be steered at inference time. This is what enables sample-efficient RL finetuning (TMRL) on top of the prior — the high-level RL policy modulates the prior's denoising timestep as its action.

Format

PyTorch checkpoint:

• model.safetensors — π₀ weights
• metadata.pt — checkpoint metadata
• assets/ — normalization statistics

The optimizer state is not included (inference/RL-finetuning only).

Usage

Use with the TMRL codebase. Download into checkpoints/, then run high-level RL on LIBERO:

hf download matthewh6/cspi0_libero --repo-type model --local-dir checkpoints/cspi0_libero

# libero-90 (timestep-modulated RL / TMRL)
python3 sac_libero.py dataset=libero_90 method=tmrl ckpt=cspi0_libero

# libero-goal
python3 sac_libero.py dataset=libero_goal perturbation=task method=tmrl ckpt=cspi0_libero

ckpt=cspi0_libero is mapped to the CSPi0 architecture by tmrl/utils/common.py:load_pi0_model. See the TMRL repo for full environment setup.

Citation

@inproceedings{hong2026tmrl,
    title  = {TMRL: Diffusion Timestep-Modulated Pretraining Enables Exploration for Efficient Policy Finetuning},
    author = {Hong, Matthew M. and Zhang, Jesse and Nagabandi, Anusha and Gupta, Abhishek},
    booktitle = {Robotics: Science and Systems (RSS)},
    year   = {2026}
}