Pytorch implementation of the "WorldModels"
Paper: Ha and Schmidhuber, "World Models", 2018. https://doi.org/10.5281/zenodo.1207631. For a quick summary of the paper and some additional experiments, visit the github page.
Prerequisites
The implementation is based on Python3 and PyTorch, check their website here for installation instructions. The rest of the requirements is included in the requirements file, to install them:
pip3 install -r requirements.txt
Running the worldmodels
The model is composed of three parts:
- A Variational Auto-Encoder (VAE), whose task is to compress the input images into a compact latent representation.
- A Mixture-Density Recurrent Network (MDN-RNN), trained to predict the latent encoding of the next frame given past latent encodings and actions.
- A linear Controller (C), which takes both the latent encoding of the current frame, and the hidden state of the MDN-RNN given past latents and actions as input and outputs an action. It is trained to maximize the cumulated reward using the Covariance-Matrix Adaptation Evolution-Strategy (CMA-ES) from the
cma
python package.
In the given code, all three sections are trained separately, using the scripts trainvae.py
, trainmdrnn.py
and traincontroller.py
.
Training scripts take as argument:
- --logdir : The directory in which the models will be stored. If the logdir specified already exists, it loads the old model and continues the training.
- --noreload : If you want to override a model in logdir instead of reloading it, add this option.
1. Data generation
Before launching the VAE and MDN-RNN training scripts, you need to generate a dataset of random rollouts and place it in the datasets/carracing
folder.
Data generation is handled through the data/generation_script.py
script, e.g.
python data/generation_script.py --rollouts 1000 --rootdir datasets/carracing --threads 8
Rollouts are generated using a brownian random policy, instead of the white noise random action_space.sample()
policy from gym, providing more consistent rollouts.
2. Training the VAE
The VAE is trained using the trainvae.py
file, e.g.
python trainvae.py --logdir exp_dir
3. Training the MDN-RNN
The MDN-RNN is trained using the trainmdrnn.py
file, e.g.
python trainmdrnn.py --logdir exp_dir
A VAE must have been trained in the same exp_dir
for this script to work.
4. Training and testing the Controller
Finally, the controller is trained using CMA-ES, e.g.
python traincontroller.py --logdir exp_dir --n-samples 4 --pop-size 4 --target-return 950 --display
You can test the obtained policy with test_controller.py
e.g.
python test_controller.py --logdir exp_dir
Notes
When running on a headless server, you will need to use xvfb-run
to launch the controller training script. For instance,
xvfb-run -s "-screen 0 1400x900x24" python traincontroller.py --logdir exp_dir --n-samples 4 --pop-size 4 --target-return 950 --display
If you do not have a display available and you launch traincontroller
without
xvfb-run
, the script will fail silently (but logs are available in
logdir/tmp
).
Be aware that traincontroller
requires heavy gpu memory usage when launched
on gpus. To reduce the memory load, you can directly modify the maximum number
of workers by specifying the --max-workers
argument.
If you have several GPUs available, traincontroller
will take advantage of
all gpus specified by CUDA_VISIBLE_DEVICES
.
Authors
- Corentin Tallec - ctallec
- Léonard Blier - leonardblier
- Diviyan Kalainathan - diviyan-kalainathan
License
This project is licensed under the MIT License - see the LICENSE.md file for details