Extending the applicability of world modeling with DreamerX

DreamerX is an implementation of model-based DreamerV3 with minor optimizations and novel training adjustments. DreamerX learns a compact latent representation of environment dynamics and estimates transitions based on sequential actions. This is known as the world model. Then, the actor-critic trains exclusively on imagined latent states from the world model, allowing learning using inferred dynamics through backpropagation using the world model weights.

The library is designed to be flexible and user-friendly, allowing researchers and practitioners to easily swap in custom environments, architectural variants, or benchmark new implementations against a validated baseline. DreamerX has been tested using Gymnasium and MuJoCo environments, and is constantly expanding to include more environments and techniques. Please check out the documentation to get started.

Note

This repository is in active development. Please check back later for more information, features, and examples.

Features

• Explainable Code: The flow, sections, and logic of the training process and individual components are well-documented with references to relevant source material.
• Modular Design: Model components, distributions, and losses are easily interchangable with easy-to-understand and documented interfaces for custom implementations.
• Novel Optimizations: Minor optimizations and novel discretization strategies, including two hot discretization, are available, with more coming soon.
• Detailed Logging: Logging, checkpointing, and evaluation are available out of the box with pre-assembled configuration files readily available.
• Gymnasium API Support: DreamerX is broadly applicable to environments complying to the Gymnasium API, with optional rendering support.

Trained examples

MuJoCo Walker2d-v5
100K Steps	200K Steps	500K Steps	1M Steps
MuJoCo Ant-v5
10K Steps	50K Steps	100K Steps	200K Steps
MuJoCo Hopper-v5
10K Steps	20K Steps	30K Steps	40K Steps
BipedalWalker-v3
10K Steps	50K Steps	100K Steps	200K Steps
LunarLander-v3
1K Steps	10K Steps	25K Steps	50K Steps

Installation

To install the library, first clone the repository

git clone https://github.com/Oafish1/DreamerX
cd DreamerX

Then, install DreamerX and dependencies

make install  # Only install necessary libraries, equivalent to `pip install -e .`
make install-dev  # Also install optional libraries, equivalent to `pip install -e .[dev,gym,extras]`

Usage

Please see the Dreamer notebook in the examples folder for usage examples. A quick-start tutorial will be available shortly. In the meantime, please refer to the documentation and don't hesitate to post an issue if you have any questions or concerns.