ReBRAC-v2

Official JAX/Flax implementation and experiment configurations for ReBRAC-v2.

ReBRAC-v2 retains the algorithmic minimalism of ReBRAC while combining an end-to-end normalizing-flow actor, mixed behavior-cloning objectives, a distributional residual critic ensemble, staged training, and evaluation-time action refinement. The original ReBRAC implementation is available here.

Paper · Project page · Code

• algorithms/rebrac_v2_ogbench.py: OGBench training and evaluation.
• algorithms/rebrac_v2.py: D4RL training and evaluation.
• algorithms/nf_policy_action.py: affine-coupling normalizing-flow actor.
• algorithms/nf_policy_spline.py: neural-spline-flow actor used in ablations.
• algorithms/kron.py: Kron optimizer.
• configs/offline/rebrac-v2-ogbench/: final OGBench configurations.
• configs/offline/rebrac-v2/: final D4RL AntMaze and Adroit configurations.
• docs/: static GitHub Pages project website.
• results_data/: per-task scores used to build the paper tables and aggregate plots.

The supplied environment reproduces the Linux, Python 3.10, NVIDIA GPU, and CUDA 12 setup used for the experiments:

conda env create -f environment.yml
conda activate jax

Run all commands from the repository root. D4RL uses the legacy mujoco-py stack, so its system-level MuJoCo and OpenGL requirements must also be available. See the mujoco-py installation guide if environment creation succeeds but D4RL cannot initialize MuJoCo.

OGBench loads datasets through the installed ogbench package and caches them under ~/.ogbench/data by default. Override --ogbench_dataset_dir to use another location.

Run one task with a released configuration:

PYTHONPATH=. python3 algorithms/rebrac_v2_ogbench.py \
  --config_path configs/offline/rebrac-v2-ogbench/antmaze/giant_navigate_singletask_v0.yaml \
  --ogbench_task_id 1 \
  --train_seed 0

--ogbench_task_id accepts task IDs 1-5. Omit it to use the base environment name stored in the configuration. Any configuration field can be overridden from the command line.

To evaluate only after 800K, 900K, and 1M updates and log their running average:

PYTHONPATH=. python3 algorithms/rebrac_v2_ogbench.py \
  --config_path configs/offline/rebrac-v2-ogbench/antmaze/giant_navigate_singletask_v0.yaml \
  --ogbench_task_id 1 \
  --train_seed 0 \
  --track_checkpoint_average=true

Enabling track_checkpoint_average replaces periodic evaluation with evaluations at those three checkpoints. The final eval/checkpoint_average_normalized_score_mean metric is their arithmetic mean.

Run an AntMaze configuration:

PYTHONPATH=. python3 algorithms/rebrac_v2.py \
  --config_path configs/offline/rebrac-v2/antmaze/umaze_v2.yaml \
  --train_seed 0

Run an Adroit configuration:

PYTHONPATH=. python3 algorithms/rebrac_v2.py \
  --config_path configs/offline/rebrac-v2/pen/human_v1.yaml \
  --train_seed 0

The released YAML files define the training budget and evaluation protocol for each dataset.

The CSV files in results_data/ contain per-task means, standard deviations, seed-level values when available, and the benchmark labels used by the paper tables. See results_data/README.md for file-level provenance and protocol details.

Weights & Biases runs in offline mode by default. The scripts do not upload source code or collect Git metadata. To log to your own account, authenticate with W&B and set online mode:

WANDB_MODE=online PYTHONPATH=. python3 algorithms/rebrac_v2_ogbench.py \
  --config_path configs/offline/rebrac-v2-ogbench/scene/play_singletask_v0.yaml \
  --ogbench_task_id 1 \
  --train_seed 0

The W&B project, group, and name are regular configuration fields and can be overridden from the command line.

If you use ReBRAC-v2, please cite the paper:

@article{tarasov2026rebrac,
  title={ReBRAC-v2: The Return of the King},
  author={Tarasov, Denis and Katzschmann, Robert K},
  journal={arXiv preprint arXiv:2608.01205},
  year={2026}
}

This project is released under the Apache License 2.0.