🌐 Project Page  |  📄 Paper  |  🤗 Checkpoints  |  🤗 Data

Yuyang Liu*, Chuan Wen*, Yihang Hu, Dinesh Jayaraman, Yang Gao†

_{* Equal contribution    † Corresponding author}

If you find TimeRewarder useful, please consider giving it a ⭐ — it helps others find it!

This is the official codebase for TimeRewarder, reproducing the paper's MetaWorld results.

TimeRewarder learns a dense reward from passive (action-free) videos: it predicts the frame-wise temporal distance between two frames and uses that per-frame progress directly as the reward for downstream RL. The pipeline has two stages — (1) train TimeRewarder on expert videos, then (2) run RL with the learned reward.

TimeRewarder/
├── metaworld_generate_expert/   # collect MetaWorld scripted-policy expert videos
├── models/                      # TimeRewarder model definition (shared by training + RL)
├── training/                    # stage 1: train TimeRewarder on videos
└── RL/                          # stage 2: downstream RL with the TimeRewarder reward

conda env create -f environment.yml   # creates env "timerewarder" (Python 3.8 + PyTorch 2.4)
conda activate timerewarder
export MUJOCO_GL=egl

The conda env installs the Python packages (mujoco-py included). You still need the native MuJoCo binaries — set them up following the instructions in drqv2 — plus the system libraries libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf.

The default wheels target CUDA 12.1; for another CUDA, reinstall torch/torchvision from the matching index. MetaWorld is vendored under RL/.

• 🤗 CowAndSheep/timerewarder — the trained TimeRewarder checkpoints.
• 🤗 CowAndSheep/timerewarder-demos — the expert demo videos used in the paper.

huggingface-cli download CowAndSheep/timerewarder --local-dir models/ckpt
huggingface-cli download CowAndSheep/timerewarder-demos --repo-type dataset --local-dir demos

With the checkpoints in models/ckpt/, run RL directly — no expert data needed, since the reward comes from the pretrained model:

cd RL
bash run_example.sh             # drawer-open-v2
bash run_rl.sh window-close-v2  # any other task

Each run trains a DrQ-v2-style pixel agent for 200k env steps and logs success rate.

1. Get expert videos — download them (above), or generate them with the provided script:

cd metaworld_generate_expert
python generate_demo.py                       # built-in 10-task list
python generate_demo.py --task drawer-open-v2 # a single task (repeatable)

Per task this writes <task-id>/{videos/*.mp4, label.txt, label_val.txt, text.csv} (100 train + 100 held-out demos).

2. Train TimeRewarder (8 GPUs, 100 epochs) — from training/:

export DEMO_ROOT=/path/to/demos            # the directory holding <task-id>/ folders
cd training
bash train.sh drawer-open-v2               # exports models/ckpt/drawer_open_20bins.pth

3. Run downstream RL — from RL/:

cd ../RL
bash run_rl.sh drawer-open-v2              # auto-loads the checkpoint train.sh exported in step 2

run_rl.sh finds the checkpoint by the same name train.sh exports — models/ckpt/<task>_<BIN_NUM>bins[_text][_diffrep].pth — so just give both stages the matching task and flags (BIN_NUM, TEXT_INPUT, DIFF_REP) and it loads automatically. To point RL at a specific file instead, pass it explicitly: python train.py suite=metaworld suite/metaworld_task=<task> cost_encoder=timerewarder cost_encoder_ckpt=/path/to/ckpt.pth.

Text-conditioned variant. TimeRewarder is vision-only by default, but it also ships a language-conditioned model that runs the very same two-stage pipeline — just add TEXT_INPUT=1 to both stages (the demos already include a text.csv, so nothing else changes):

cd training && DEMO_ROOT=/path/to/demos TEXT_INPUT=1 bash train.sh drawer-open-v2   # stage 1 -> models/ckpt/drawer_open_text_20bins.pth
cd ../RL    && TEXT_INPUT=1 bash run_rl.sh drawer-open-v2                            # stage 2: downstream RL, logs success rate

drawer-open-v2 · door-open-v2 · window-open-v2 · window-close-v2 · button-press-topdown-v2 · lever-pull-v2 · plate-slide-v2 · stick-push-v2 · basketball-v3 · disassemble-v2

All scripts accept either the MetaWorld env id (e.g. drawer-open-v2) or the short name (e.g. drawer_open).

TimeRewarder isn't tied to MetaWorld — point stage 1 at your own videos to get a reward model for whatever downstream task you care about. Lay out one folder per task, set DEMO_ROOT to the directory holding them, and train:

cd training && DEMO_ROOT=/path/to/data bash train.sh <task-id>   # exports a reward model to models/ckpt/

Use clips where time tracks task progress — the reward is learned purely from the temporal ordering of frames, so each clip should run from start to goal. The dataset comes in two formats:

Vision-only (no language):

<task-id>/
├── videos/             # demo clips, one .mp4 each
├── label.txt           # train split:  one video filename per line
└── label_val.txt       # held-out split, same format

label.txt and label_val.txt list the clips in videos/ for the train and held-out splits — just one filename per line (the vision-only model needs no text id).

Text-conditioned (TEXT_INPUT=1) — add a text.csv and a real text id on each label.txt line:

<task-id>/
├── videos/
├── label.txt           # "<video_filename> <text_id>" — text_id indexes text.csv
├── label_val.txt
└── text.csv            # "id,name" rows, one description per id

The exported checkpoint is a standalone reward model — plug it into your own downstream pipeline however you like.

Each row toggles one setting; everything else stays at the paper default. Training knobs are appended to train.sh (as KEY Value pairs, or env vars where shown); RL knobs are Hydra overrides to run_rl.sh.

The checkpoint name encodes the structural settings (BIN_NUM, TEXT_INPUT, DIFF_REP) so RL loads a matching architecture. For any other ablation, set TAG=<name> on both stages to keep its checkpoint separate — e.g. TAG=uniform bash train.sh <task> DATA.WEIGHTED_SAMPLE False exports ..._20bins_uniform.pth, and TAG=uniform bash run_rl.sh <task> loads it.

DIFF_REP=1 predicts the temporal distance from the difference of each frame's separately-encoded features instead of the concatenated pair — the per-frame-encoding variant the paper uses to probe representation quality (Fig. 5). It is an architectural choice, so set it on both stages (training tags the checkpoint _diffrep, and RL must use the same flag to load it).

The reward-from-video baselines and an upper bound. These override the agent / reward source, so launch them with python train.py directly rather than run_rl.sh (which is hard-wired to the TimeRewarder reward). Run from RL/:

cd RL
# GT — environment-reward upper bound (no expert data needed)
python train.py suite=metaworld suite/metaworld_task=drawer_open cost_encoder=gt

# OT — ResNet features + Sinkhorn optimal transport
python train.py suite=metaworld suite/metaworld_task=drawer_open cost_encoder=resnet num_demos=100

# ADS — OT + Automatic Discount Scheduling
python train.py suite=metaworld suite/metaworld_task=drawer_open cost_encoder=resnet num_demos=100 adaptive_discount=true

# GAIfO — adversarial imitation from observations
python train.py suite=metaworld suite/metaworld_task=drawer_open agent=gaifo num_demos=100

# BC — behavior cloning
python train.py suite=metaworld suite/metaworld_task=drawer_open agent=bc num_demos_bc=100

The OT/ADS/GAIfO baselines score agent rollouts against an expert demo file (expert_dataset, default the matching task's expert_demos.pkl); BC clones a separate set (expert_dataset_bc). Generate the pickle alongside the videos with:

python generate_demo.py --task drawer-open-v2 --save-pkl   # writes <task-id>/expert_demos.pkl

Point expert_dataset= / expert_dataset_bc= at that file to use your own demos. GT needs no expert data — its reward is the environment reward itself. The main TimeRewarder line never uses expert_demos.pkl.

Built on IL_ADS and X-CLIP; environments from Meta-World.

@article{liu2025timerewarder,
  title={TimeRewarder: Learning Dense Reward from Passive Videos via Frame-wise Temporal Distance},
  author={Liu, Yuyang and Wen, Chuan and Hu, Yihang and Jayaraman, Dinesh and Gao, Yang},
  journal={arXiv preprint arXiv:2509.26627},
  year={2025}
}