R2-Dreamer: Redundancy-Reduced World Models without Decoders or Augmentation

This repository provides a PyTorch implementation of R2-Dreamer (ICLR 2026), a computationally efficient world model that achieves high performance on continuous control benchmarks. It also includes an efficient PyTorch DreamerV3 reproduction that trains ~5x faster than a widely used codebase, along with other baselines. Selecting R2-Dreamer via the config provides an additional ~1.6x speedup over this baseline.

Instructions

Install dependencies. This repository is tested with Ubuntu 24.04 and Python 3.11.

If you prefer Docker, follow docs/docker.md.

# Installing via a virtual env like uv is recommended.
pip install -r requirements.txt

Run training on default settings:

python3 train.py logdir=./logdir/test

Monitoring results:

tensorboard --logdir ./logdir

Switching algorithms:

# Choose an algorithm via model.rep_loss:
# r2dreamer|dreamer|infonce|dreamerpro
python3 train.py model.rep_loss=r2dreamer

For easier code reading, inline tensor shape annotations are provided. See docs/tensor_shapes.md.

Available Benchmarks

At the moment, the following benchmarks are available in this repository.

Environment	Observation	Action	Budget	Description
Meta-World	Image	Continuous	1M	Robotic manipulation with complex contact interactions.
DMC Proprio	State	Continuous	500K	DeepMind Control Suite with low-dimensional inputs.
DMC Vision	Image	Continuous	1M	DeepMind Control Suite with high-dimensional images inputs.
DMC Subtle	Image	Continuous	1M	DeepMind Control Suite with tiny task-relevant objects.
Atari 100k	Image	Discrete	400K	26 Atari games.
Crafter	Image	Discrete	1M	Survival environment to evaluates diverse agent abilities.
Memory Maze	Image	Discrete	100M	3D mazes to evaluate RL agents' long-term memory.

Use Hydra to select a benchmark and a specific task using env and env.task, respectively.

python3 train.py ... env=dmc_vision env.task=dmc_walker_walk

Headless rendering

If you run MuJoCo-based environments (DMC / MetaWorld) on headless machines, you may need to set MUJOCO_GL for offscreen rendering. Using EGL is recommended as it accelerates rendering, leading to faster simulation throughput.

# For example, when using EGL (GPU)
export MUJOCO_GL=egl
# (optional) Choose which GPU EGL uses
export MUJOCO_EGL_DEVICE_ID=0

More details: Working with MuJoCo-based environments

Code formatting

If you want automatic formatting/basic checks before commits, you can enable pre-commit:

pip install pre-commit
# This sets up a pre-commit hook so that checks are run every time you commit
pre-commit install
# Manual pre-commit run on all files
pre-commit run --all-files

Citation

If you find this code useful, please consider citing:

@inproceedings{
morihira2026rdreamer,
title={R2-Dreamer: Redundancy-Reduced World Models without Decoders or Augmentation},
author={Naoki Morihira and Amal Nahar and Kartik Bharadwaj and Yasuhiro Kato and Akinobu Hayashi and Tatsuya Harada},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://openreview.net/forum?id=Je2QqXrcQq}
}