This repo reproduces the original implementation of World Models. This implementation uses TensorFlow 2.2.
The easiest way to handle dependencies is with Nvidia-Docker. Follow the instructions below to generate and attach to the container.
docker image build -t wm:1.0 -f docker/Dockerfile.wm .
docker container run -p 8888:8888 -v /home/ubuntu/world-models:/app --rm -it wm:1.0
To visualize the environment from the agents perspective or generate synthetic observations use the visualizations jupyter notebook. It can be launched from your container with the following:
jupyter notebook --no-browser --port=8888 --ip=0.0.0.0 --allow-root
| Real Frame Sample | Reconstructed Real Frame |
|---|---|
 |
 |
| Real Frame Sample | Reconstructed Real Frame |
|---|---|
 |
 |
To visualize the VAE/GAN loss and distribution plots as well as the random/reconstucted images, while in the docker container run:
tensorboard --logdir='./WorldModels/vae_gan/logs' --port=6006
These instructions assume a machine with a 4 core cpu.
To reproduce results for CarRacing-v0 run the following bash script
bash launch_scripts/carracing.bash
We have not run this for long enough(~45 days wall clock time) to verify that we produce the same results on CarRacing-v0 as the original implementation. Our average run time for the VAE and VAE/GAN training was approximately 2/3 days.
Average return curves comparing the original implementation and ours. The shaded area represents a standard deviation above and below the mean.
| Generation | Returns Actual Environment | ||
|---|---|---|---|
| D. Ha Original | 2000 | 868 +/- 511 | |
| Eager | 100 | 501 +/- 94 |