TS-DP: Drafters and Schedulers

Released weights for TS-DP: Reinforcement Speculative Decoding for Temporal Adaptive Diffusion Policy Acceleration (paper, code).

TS-DP accelerates the official Diffusion Policy transformer checkpoints. This repository only contains the TS-DP components; the base models are downloaded from the Diffusion Policy website.

Each folder <name>/ holds one setting:

• drafter.ckpt: the distilled one-block Transformer drafter (drafter weights only);
• scheduler.pth: the PPO scheduler that adapts the speculative parameters at run time.

Download

pip install -U huggingface_hub
hf download childtang/TS-DP --local-dir data/ts_dp                    # all settings
hf download childtang/TS-DP --include "can/*" --local-dir data/ts_dp  # a single setting

With older versions of huggingface_hub, use huggingface-cli download instead of hf download.

Settings

A drafter is distilled from one specific base model and only works with it. Base model paths are relative to the Diffusion Policy experiment logs.

<name>	--task	Base model
lift	lift	image/lift_ph/diffusion_policy_transformer/train_0/checkpoints/latest.ckpt
can	can	image/can_ph/diffusion_policy_transformer/train_0/checkpoints/latest.ckpt
square	square	image/square_ph/diffusion_policy_transformer/train_0/checkpoints/latest.ckpt
transport	transport	image/transport_ph/diffusion_policy_transformer/train_0/checkpoints/latest.ckpt
tool_hang	tool_hang	image/tool_hang_ph/diffusion_policy_transformer/train_0/checkpoints/epoch=0100-test_mean_score=0.773.ckpt
pusht	pusht	image/pusht/diffusion_policy_transformer/train_1/checkpoints/latest.ckpt
lift_mh	lift	image/lift_mh/diffusion_policy_transformer/train_0/checkpoints/latest.ckpt
can_mh	can	image/can_mh/diffusion_policy_transformer/train_0/checkpoints/latest.ckpt
square_mh	square	image/square_mh/diffusion_policy_transformer/train_0/checkpoints/latest.ckpt
transport_mh	transport	image/transport_mh/diffusion_policy_transformer/train_0/checkpoints/latest.ckpt
kitchen	kitchen	low_dim/kitchen/diffusion_policy_transformer/train_1/checkpoints/epoch=2700-test_mean_score=0.574.ckpt
block_pushing	block_pushing	low_dim/block_pushing/diffusion_policy_transformer/train_0/checkpoints/latest.ckpt

Usage

python scripts/eval_scheduler.py --task can \
    --checkpoint data/checkpoints/image/can_ph/diffusion_policy_transformer/train_0/checkpoints/latest.ckpt \
    --drafter-checkpoint data/ts_dp/can/drafter.ckpt \
    --ppo-checkpoint data/ts_dp/can/scheduler.pth \
    --output-dir data/eval/can

See the code repository for installation and the full evaluation protocol. Due to the stochasticity of diffusion sampling and the simulators, results on a different machine may deviate slightly from those reported in the paper.

Citation

@article{li2025ts,
  title         = {TS-DP: Reinforcement Speculative Decoding For Temporal Adaptive Diffusion Policy Acceleration},
  author        = {Ye Li and Jiahe Feng and Yuan Meng and Kangye Ji and Chen Tang and Jiajun Fan and Xinwan Wen and Jiawei Li and Wenwu Zhu and Shu-Tao Xia and Xin Wang and Zhi Wang},
  journal       = {arXiv preprint arXiv:2512.15773},
  eprint        = {2512.15773},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  doi           = {10.48550/arXiv.2512.15773},
  year          = {2025}
}