world-modelz/dreamax

A scalable Dreamer implementation in JAX

Dreamax

This is a scalable implementation of the model-based RL algorithm Dreamer (Hafner et al., 2019) in JAX.

We used Yarden As' yardenas/jax-dreamer implementation as the starting point for our implementation.

For bug reports please use github issues.

Getting started

1. Build the docker container

To build the docker container run:

cd docker
./build.sh

2. Run the docker container

docker/start.sh

3. Start training with default configuration

Inside the docker container run:

python train.py

Tensorboard events will be written to the 'results' directory.

Papers

• Dreamer: Dream to Control: Learning Behaviors by Latent Imagination (explained by Yannic: Video)
• DreamerV2: Mastering Atari with Discrete World Models (explained by Yannic: Video)

Links

• Dreamer implementation in JAX: yardenas/jax-dreamer
• Original Dreamer v1 implementation (TF2): danijar/dreamer
• Original Dreamer v2 implementation (TF2): danijar/dreamerv2

#world-modelz

This is a community project of the #world-modelz channel of the Yannic Kilcher discord server. For more projects and the world models paper list visit our main repo: world-modelz/world-modelz.