---
license: mit
tags:
- robotics
- diffusion-policy
- speculative-decoding
- reinforcement-learning
---

# TS-DP: Drafters and Schedulers

Released weights for **TS-DP: Reinforcement Speculative Decoding for Temporal Adaptive Diffusion Policy Acceleration**
([paper](https://arxiv.org/abs/2512.15773), [code](https://github.com/ChildTang/TS-DP)).

TS-DP accelerates the official [Diffusion Policy](https://github.com/real-stanford/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

```bash
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](https://diffusion-policy.cs.columbia.edu/data/experiments/).

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

```bash
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](https://github.com/ChildTang/TS-DP) 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

```bibtex
@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}
}
```