JM2D checkpoints

Pre-trained checkpoints for Joint Model-based Model-free Diffusion for Planning with Constraints (CoRL 2025).

[Project page] [ArXiv] [Code]

Wonsuhk Jung*, Utkarsh A. Mishra*, Nadun Ranawaka Arachchige, Yongxin Chen, Danfei Xu, Shreyas Kousik — Georgia Institute of Technology

Each folder belongs to one branch of the code repository; that branch's scripts/download_checkpoints.py downloads its folder and checks the md5 sums below. The toy-domain checkpoints are small and ship with the public_toy branch. The hardware branch (public_hardware) is a guide for running JM2D on your own robot and has no checkpoint.

Files

Folder	Branch	Files	Size
pointmaze/	public_robomimic	maze2d_large_ep{1250,1500,1750,1900,2000}.pth: robomimic diffusion policies on D4RL maze2d-large-dense-v1 (the 5 checkpoints of paper Fig. 4)	5 × 524 MB
d3il/data/avoiding/	public_d3il	env_*.pkl (96 files) + MD5SUMS: the D3IL Avoiding demonstrations, needed for training and for evaluation (data normalizer)	2.6 MB
d3il/avoiding-d3il/H8_K20_Dmodels.GaussianDiffusion/<seed>/	public_d3il	state_best.pt + config pickles: Diffuser-style planners on D3IL-Avoiding, seeds 0–4 (paper Tables 2 and 3)	5 × 32 MB
pointmaze/dataset/	public_robomimic	maze2d_large_dense_v1.hdf5: the D4RL maze2d-large-dense-v1 dataset converted to robomimic format, used to train the PointMaze policies	303 MB

File	md5
pointmaze/maze2d_large_ep1250.pth	98614093985811a0dbcbb4a66b3ccfc4
pointmaze/maze2d_large_ep1500.pth	000b0f287307bd98289cc9b28479566c
pointmaze/maze2d_large_ep1750.pth	a7fd3fe32e1b35723dde1c566c61f1b9
pointmaze/maze2d_large_ep1900.pth	a756af86fe2070c2e77358ca380ccc44
pointmaze/maze2d_large_ep2000.pth	7f76546cb300ec7577c475d1f3b40f97
d3il/.../0/state_best.pt	8446a72feebe0e9b1839bc5cb48f448e
d3il/.../1/state_best.pt	c09dd9a85fcf435201d8e61835d053f0
d3il/.../2/state_best.pt	4c55388894fdf2ebeb22206678125ac3
d3il/.../3/state_best.pt	7d219c1b2df29557297bde3ab561b784
d3il/.../4/state_best.pt	7238d20b674872da4a8d88e8469a5a3b
pointmaze/dataset/maze2d_large_dense_v1.hdf5	aa4be41d7444c0bdc8fcf48707e11c84

Download

From a checkout of the matching branch:

python scripts/download_checkpoints.py

or directly:

hf download RoboWonsuhk/JM2D --include "pointmaze/*" --local-dir checkpoints_hf

The PointMaze dataset is a converted copy of the D4RL dataset (Fu et al., 2020, https://github.com/Farama-Foundation/D4RL), redistributed under its CC BY 4.0 license.

The D3IL Avoiding demonstrations are an unmodified copy from D3IL (Jia et al., 2024, https://github.com/ALRhub/d3il, MIT License).

The files are PyTorch pickles; load them only with the code of the matching branch.

Citation

@inproceedings{jung2025joint,
  title={Joint Model-based Model-free Diffusion for Planning with Constraints},
  author={Wonsuhk Jung and Utkarsh Aashu Mishra and Nadun Ranawaka Arachchige and Yongxin Chen and Danfei Xu and Shreyas Kousik},
  booktitle={9th Annual Conference on Robot Learning},
  year={2025},
  url={https://openreview.net/forum?id=E9t1ekt6W9}
}