DreamerV2 Pytorch

Pytorch implementation of DreamerV2 agent as in Mastering Atari with Discrete World Models, based heavily on the original danijar's Tensorflow 2 implementation. This implementation also follows closely the code structure as the original.

This repo intends to approximate as close as possible the results obtained by the original Tensorflow 2 implementation.

Install dependencies:

#install pytorch using conda or pip -> follow the instructions here https://pytorch.org

# install ffmpeg for saving agents gifs
# code for Ubuntu/debian
sudo apt install ffmpeg


#install other dependencies
pip3 install -r requirements.txt

Train the agent:

python3 dreamerv2/train.py --logdir ~/logdir/atari_pong/dreamerv2/1 --configs defaults atari --task atari_pong

Monitor results:

tensorboard --logdir ~/logdir

Testing environments (playing as a human)

python dreamerv2/play_test.py --env SpaceInvaders-v0 

Playing inside the agent dream!

python dreamerv2/play.py --configs defaults atari --task atari_space_invaders --logdir /logdir/space_invaders_logdir

---

Features:

• FP16 pytorch support
• Custom distributions and optimizers that closely match the tensorflow counterparts
• support logging to Wandb (you need to put imports on train.py though)

Not implemented (converted from original):

• Plan2Explore
• ram observation
• right now as implemented, doesn't support learning the actor by propagating the value network (torch.no_grad())
• the replay buffer dataloaders are synchronous for now (no workers)

Replication of results

Comparison of the Space Invaders Atari game learning performance over our implementation (Pytorch) vs the original implementation (Tensorflow), while using the same hyperparameters.

Results averaged over 5 randomly-seeded runs.

Pre-trained models

link: https://drive.google.com/drive/folders/1md-5Q5Ewh0a9EwCUb8LcQNSPE6IO8iPb?usp=sharing