Project Page  |  arXiv:2606.28320

WARP-RM learns a dense, signed relative-progress signal from robot manipulation videos: per frame, how fast and in which direction the task is advancing. A frozen DINOv3 backbone feeds a bidirectional Transformer with a categorical progress-velocity head, trained fully self-supervised via the WARP time-warp augmentation (variable playback speed + reversals, sampled from a smooth AR(1) process). The output — v̂ₜ ≈ 1 at expert pace, ≈ 0 while stalling, < 0 while regressing — is a per-frame reward used downstream to filter and reweight behavior-cloning action chunks (WARP-BC).

This repository is the reward-model: training, scoring, dense inference, and annotation injection. The downstream WARP-BC chunk reweighting and simulation evaluator live in the public companion repositories linked below.

git clone https://github.com/uynitsuj/WARP-RM.git && cd WARP-RM
uv sync                       # core (training + scoring)
uv sync --extra wandb         # optional: Weights & Biases logging

Requires Python ≥ 3.10 and a CUDA GPU. Torch is pinned per CUDA runtime (cu128 default; run scripts/setup_torch.sh to auto-pick cu128/cu130).

WARP-RM trains and scores on LeRobot v2.1 / v3.0 datasets of successful demonstrations (see docs/dataset_schema.md for the expected layout). The defaults are the recipe:

# Real T-shirt folding recipe: bidirectional attention, fs3/sss45, window 32, 15k steps.
python scripts/train.py --lerobot-repo /path/to/your/lerobot/dataset

# Paper simulation results repro RM checkpoint configuration (WARP-RM sss15):
# selected checkpoint was step 14,400 of this 20k-step no-abs run.
python scripts/train.py --ablation no_abs --feature-stride 1 \
    --source-standard-stride 15 --max-steps 20000 \
    --lerobot-repo /path/to/your/lerobot/dataset

# List ablation configs:
python scripts/train.py --list-configs

Batch size defaults to 1024 (needs a ≥40 GB GPU); pass --batch-size 256 on smaller cards. Vision backbone features are cached on first use under ~/.cache/warp_rm/features/ (override with WARP_RM_FEATURE_CACHE); standalone caching via scripts/data/precompute_features.py.

Then score the dataset with the trained checkpoint and inject the per-frame reward columns back into it:

python scripts/data/write_warp_rm_annotations.py \
    --checkpoint checkpoints/best_model_<tag>.pt \
    --lerobot-repo /path/to/your/lerobot/dataset
# writes per-frame `warp_rm_signed_magnitude` (+ `warp_rm_progress`) columns

See docs/recipe.md for the full data → supervision walkthrough and docs/metrics_glossary.md for every metric in the training logs.

The paper's public sim reproduction uses these public artifacts beyond this repo (the RM itself trains and scores any compatible LeRobot dataset):

• MuJoCo-Warp data + full-state supplement: uynitsuj/sim-bottles-mjwarp-v1. Derived dataset (re-rendered at 30 Hz) from ABC's public bottles simulation data. It contains the LeRobot dataset, portable scene XML, and full qpos for the 2,436 re-renderable episodes.
• Paper RM head: uynitsuj/warp-rm-sim-bottles-sss15. The frozen DINOv3 backbone is obtained separately under Meta's terms.
• ABC evaluator: uynitsuj/abc-rabc@release-candidate provides the MuJoCo-Warp evaluator and deterministic score_bottles.py.
• OpenPI repository: uynitsuj/openpi provides optional pi0 + RABC policy support. Its upstream Pi0 terms apply to the corresponding published policy parameters.
• Policy checkpoints and canonical traces: uynitsuj/paper-sim-policy-checkpoints and uynitsuj/paper-sim-n128-traces.

For the deterministic n=128 audit, download the public trace artifact and run the public ABC scorer:

hf download uynitsuj/paper-sim-n128-traces --repo-type dataset \
  --local-dir ../traces/paper-sim-n128
git clone --branch release-candidate https://github.com/uynitsuj/abc-rabc.git
cd abc-rabc
python score_bottles.py --trace-dir ../traces/paper-sim-n128 --self-test

The self-test verifies every published table value from the canonical traces. It uses only the published data, checkpoints, traces, and source repositories; fresh rollouts are evaluated against tolerances rather than expected to be trajectory-identical.

WARP-RM/
├── scripts/
│   ├── train.py                      # training entrypoint (defaults = the recipe)
│   ├── config.py                     # shared hyperparameters
│   ├── data/{precompute_features,write_warp_rm_annotations,inject_rm_column}.py
│   ├── eval/{score_episodes,render}.py
│   └── visualize_samplers.py         # visualize the WARP time-warp sampler
├── warp_rm/                          # the package
│   ├── data/      {dataset, samplers, labelers, lerobot_dataset, video_reader, ...}.py
│   ├── models/    aggregators/transformer.py, backbones/dinov3.py
│   ├── core/      {trainer, loss, metrics, ablation, ...}.py
│   ├── eval/      scoring.py
│   ├── visualization/ {inference, renderer, plotting}.py
│   └── utils/     {caching, environment, schema}.py
└── docs/{recipe,dataset_schema,metrics_glossary}.md

@article{yu2026warp,
  title={WARP-RM: A Warp-Augmented Relative Progress Reward Model for Data Curation},
  author={Yu, Justin and Goldberg, Andrew and Kondap, Kavish and El-Refai, Karim and Ransing, Ethan and Chen, Qianzhong and Schwager, Mac and Shentu, Fred and Wu, Philipp and Goldberg, Ken},
  journal={arXiv preprint arXiv:2606.28320},
  year={2026}
}

This project is available under the MIT License.