Offline Reinforcement Learning
with Universal Horizon Models

Paper | Project Page

Official Github repository for "Offline Reinforcement Learning with Universal Horizon Models".
ICML 2026

Overview

The Universal Horizon Model (UHM) is a general framework for future-predictive models, enabling efficient value learning in offline RL.

Installation

Our codebase requires Python 3.9+ and is based on JAX. To install the necessary dependencies, run the following commands:

pip install -r requirements.txt

Usage

The main implementation of UHM is in agents/uhm.py, and our implementations of additional baselines (DTD($\lambda$), MBTD($\lambda$), GHM) can also be found in the same directory:

# for UHM
python main.py --agent=agents/uhm.py --env_name=antmaze-large-navigate-singletask-task1-v0

# for DTD($\lambda$)
python main.py --agent=agents/dtd.py --env_name=antmaze-large-navigate-singletask-task1-v0

# for MBTD($\lambda$)
python main.py --agent=agents/mbtd.py --env_name=antmaze-large-navigate-singletask-task1-v0

# for GHM
python main.py --agent=agents/ghm.py --env_name=antmaze-large-navigate-singletask-task1-v0

Reproducing the main results

run.sh provides the complete list of the exact command-line flags used to produce the main results in the paper.

Acknowledgments

This codebase is implemented on the top of Flow Q Learning.