HiRE: Hindsight Reward Editing for Policy Finetuning

HiRE edits visual rewards using the outcomes of past interactions. Frozen foundation representations provide visual similarity, while successful and failed trajectories ground that similarity in task progress. HiRE attracts the policy toward successful states, suppresses visually plausible failure states, and incorporates the edited potential through reward shaping without training an additional reward model.

This repository contains the simulation training and evaluation code for HiRE. The experiments use DICE-RL to finetune diffusion and flow-matching policies on four manipulation tasks.

Use Linux with an NVIDIA GPU and EGL offscreen rendering. From the repository root:

conda create -n hire python=3.10 -y
conda activate hire

On NVIDIA Blackwell GPUs, install PyTorch's CUDA 12.8 wheels before the project packages:

python -m pip install torch==2.7.1 torchvision==0.22.1 --index-url https://download.pytorch.org/whl/cu128

RoboMimic 0.5.0 is available from its official Git tag. Install it alongside the project packages:

mkdir -p external
git clone --depth 1 --branch v0.5.0 https://github.com/ARISE-Initiative/robomimic.git external/robomimic
python -m pip install -e '.[vision]' -e './backends/dice_rl[sim]' -e external/robomimic
git clone https://github.com/NVlabs/mimicgen.git external/mimicgen
python -m pip install -e external/mimicgen
source scripts/set_path.sh

Data and outputs default to data_dir/ and log_dir/. Set HIRE_DATA_DIR and HIRE_LOG_DIR to use other locations. W&B logging is disabled by default. See setup for rendering and logging options.

The following commands run Stack Three. See the quick start for checkpoint locations and the data guide for the other tasks.

python scripts/dataset/prepare_mimicgen_datasets.py \
  --tasks stack_three --max-episodes 200 --jobs 1

This downloads the MimicGen demonstrations and prepares normalized pretraining and finetuning datasets. Tool Hang uses a separate data-conversion command with its camera settings, described in the data guide.

All four BC policies, including the Tool Hang flow-matching checkpoint at epoch 500, are available on Hugging Face. You can use them directly without BC pretraining:

python scripts/download_checkpoints.py --tasks stack_three

Omit --tasks to download all four policies (about 831 MB total). The script saves each checkpoint and its Hydra config under checkpoints/HiRE-release/. Set HIRE_CHECKPOINT_DIR to use another location.

python scripts/launch.py finetune stack_three --seed 42

The launcher selects the released checkpoint for the task. For online W&B logging, append --wandb online. The main-figure run table lists the recorded seeds and task-specific settings.

On first use, HiRE loads DINOv2 and builds a positive reference buffer from the demonstrations. During finetuning, successful and failed rollouts update the hindsight buffers used for reward computation.

For a sparse-reward comparison from the same base policy:

python scripts/launch.py finetune stack_three --reward sparse --seed 42

Use threading, three_piece_assembly, or tool_hang after preparing that task's data and downloading its checkpoint. For Tool Hang:

python scripts/download_checkpoints.py --tasks tool_hang
python scripts/launch.py finetune tool_hang --seed 42

python scripts/launch.py pretrain stack_three --seed 42
python scripts/launch.py finetune stack_three \
  --checkpoint /path/to/bc_run/checkpoint/state_100.pt --seed 42

Pretraining saves runs under log_dir/mimicgen-pretrain/. Keep .hydra/config.yaml with checkpoint/state_<epoch>.pt. An explicit --checkpoint overrides the released base policy.

Evaluate the finetuned policy and its base policy on the same seeded initial states:

python scripts/eval_rl_checkpoint.py \
  --ckpt_path /path/to/finetune_run/model_step_1000.pth \
  --num_eval_episodes 10 --eval_n_envs 10

This runs 100 episodes per policy (10 batches × 10 parallel environments) and reports sparse task success, return, and improvement over the base policy. Add --render to save rollout videos. Use --base_policy_path and --normalization_path if the checkpoint bundle or data has moved.

Task recipes are in configs/experiment/. Reward settings are in configs/reward/, and camera/environment settings are in configs/task/. For example, change the contrastive weight and negative-buffer capacity:

python scripts/launch.py finetune threading \
  reward.contrastive_weight=0.3 reward.negative_buffer.buffer_size=512

To use SigLIP, install '.[siglip]' and add encoder=siglip. Use --config-only to inspect a run's settings or --check to check its local inputs before starting it. See the training reference for defaults, ablations, and implementation details.

src/hire/       HiRE potential, reward shaping, visual encoders, and hindsight buffers
backends/       Policy training, replay, and environment integration
configs/        Task and experiment settings
scripts/        Data preparation, training, and evaluation commands
docs/           Setup, experiment instructions, and implementation reference
tests/          Reward, policy, and integration checks

For the method implementation, start with potential.py and shaping.py. See the implementation guide for how these components connect to the training loop.

If you find HiRE useful in your research, please cite:

@inproceedings{niu2026hire,
  title  = {HiRE: Hindsight Reward Editing for Policy Finetuning},
  author = {Niu, Haoyi and Han, Zhengtao and Ji, Yufeng and Li, Zhongyu and Sreenath, Koushil},
  booktitle = {Conference on Robot Learning (CoRL)},
  year   = {2026},
}

Our implementation builds on DICE-RL, DPPO, Diffusion Policy, RoboMimic, and MimicGen. We thank the authors for sharing their code. Please also cite DICE-RL when using the policy-learning backbone:

@article{sun2026prior,
  title={From Prior to Pro: Efficient Skill Mastery via Distribution Contractive RL Finetuning},
  author={Sun, Zhanyi and Song, Shuran},
  journal={arXiv preprint arXiv:2603.10263},
  year={2026}
}

Code is distributed under the MIT License. External datasets, pretrained encoders, and simulator packages retain their respective licenses.