WarpSAC: Towards the Pinnacle of Scalable Off-Policy Reinforcement Learning by Rethinking Exploration and Exploitation

This repository contains the official implementation of WarpSAC, together with the training scripts and environment integrations used in our experiments. WarpSAC extends FlashSAC with regime-aware replay and stabilization choices, improving sample efficiency and training wall-clock efficiency in the evaluated settings.

We thank Holiday Robotics for providing the FlashSAC code used in this project.

Demonstrations

<video controls muted loop playsinline width="720" src="https://github.com/user-attachments/assets/0dccd420-ed34-48cd-95bf-ebf9b65c6350">
Your browser does not support embedded video. Use the link above.
</video>

The framework supports JAX and PyTorch backends. The reported WarpSAC results use the JAX backend. The PPO implementation uses the PyTorch backend with mixed-precision execution and torch.compile; this is an optimized implementation built on top of the RSL-RL training pipeline, whose standard configuration does not provide these two optimizations.

Installation

Create the base environment:

conda env create -f environment.yml
conda activate warp_rl

The IsaacSim and IsaacLab stack is optional. Reserve at least 50 GB of free disk space for these packages:

conda activate warp_rl
python -m pip install --extra-index-url https://pypi.nvidia.com \
  "isaacsim[all,extscache]==5.1.0" \
  "isaaclab[isaacsim]==2.3.0"

Verify the installation with:

import jax
import torch

print("PyTorch:", torch.__version__)
print("CUDA:", torch.version.cuda, "available:", torch.cuda.is_available())
print("JAX:", jax.__version__, "devices:", jax.devices())

Training

All commands are run from the repository root. Use --seed to select an experiment seed.

WarpSAC on CPU-scale MuJoCo

python train_warpsac.py \
  --env-id Hopper-v4 \
  --env-type mujoco \
  --backend jax \
  --seed 1

WarpSAC on GPU-parallel IsaacLab

CUDA_VISIBLE_DEVICES=0 python train_warpsac.py \
  --env-id Isaac-Velocity-Flat-G1-v0 \
  --env-type isaaclab \
  --backend jax \
  --seed 1

Unitree G1 sim-to-real

The Unitree-G1-Flat profile uses the MJLab environment and the sim-to-real configuration defined by the training script.

CUDA_VISIBLE_DEVICES=0 python train_warpsac.py \
  --env-id Unitree-G1-Flat \
  --env-type mjlab \
  --backend jax \
  --seed 1

PPO on Unitree G1 sim-to-real

CUDA_VISIBLE_DEVICES=0 python train_ppo.py \
  --env-id Unitree-G1-Flat \
  --env-type mjlab \
  -- backend torch \
  --seed 1

Important WarpSAC options

The main controls that distinguish the WarpSAC sampling and normalization regimes are:

--decay-step 0
--actor-normalize-parameters
--critic-normalize-parameters

decay_step=0 disables the decay-based sampling bias and recovers uniform replay. Parameter normalization is enabled by default; for large GPU-parallel simulations, disabling actor and critic parameter normalization is recommended unless the experiment requires otherwise.

Citation

If WarpSAC is useful in your research, please cite:

@misc{wu2026warpsacpinnaclescalableoffpolicy,
      title={WarpSAC: Towards the Pinnacle of Scalable Off-policy RL by Rethinking Exploration and Exploitation}, 
      author={Zihao Wu and Hongyao Tang and Yi Ma and Huizhong Song and Pengyi Li and Yifu Yuan and Fei Ni and Jinyi Liu and Wei Wei and Jianrong Wang and Yan Zheng and Jianye Hao},
      year={2026},
      eprint={2608.24479},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2608.24479}, 
}

warprl

README