matthewh6/cspi0_bridge

CSPi0 — WidowX / Bridge Context-Smoothed π₀ Prior

Context-smoothed π₀ (CSPi0) policy prior for the real-world WidowX / BridgeData V2 setup. This is the pretrained prior used for the real-robot 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

openpi-format JAX/Flax checkpoint:

• params/ — orbax checkpoint (π₀ backbone, ~3.24B params)
• 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 against a WidowX robot:

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

# timestep-modulated RL (TMRL)
python3 sac_robot.py dataset=widowx method=tmrl ckpt=cspi0_bridge

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

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