Imagination Augmented Agents
Description
TensorFlow implementation of Imagination Augmented Agents by DeepMind, published as a conference proceeding at NIPS 2017. This paper combines model free and model based RL in an agent that constructs implicit plans by learning to interpret predictions from a learned environment model.
The agent creates its own internal simulations using a learned environment model. This environment model is trained independently of the agent on future frame and reward prediction conditioned on action using a baseline actor (a2c) to play out a number of games. The environment model can then be used to simulate potential trajectories for the imagination augmented agent. To efficiently use these simulated trajectories, the agent learns an encoder that extracts information from these imaginations including both state and reward.
Above is the network architecture diagram from the paper. Below is the architecture of the environment model, also taken from the paper.
This implementation was done according to the specifications of the architecture in the appendix of the paper for Pacman. The game used in this was Pacman as used in the paper, however, any game could be substituted.
Code from OpenAI's Baseline is used for multiprocessing
(common/multiprocessing_env.py). DeepMind's implementation of MiniPacman
is used (common/deepmind.py). Implementation of common/minipacman.py is from @higgsfield.
Prerequisites
- TensorFlow 1.4.1
- NumPy 1.14.0
See requirements.txt for the complete list of requirements I used, however, only the two
listed above are important.
Usage
Train the actor critic model using the following. Training this is a necessary baseline and needs to be used to train the environment model.
python main.py a2c
Train the environment model using the following. Remember the a2c model must be trained already.
python env_model.py
Train the imagination augmented agent using the following. Remember the environment model must already be trained.
python main.py i2a
Evaluate any agent (in the terminal). Change the model checkpoint to whichever actor you would like to evaluate (either from i2a or a2c).
python evaluate.py
To see the agent play the game visually run the eval_actor_vis.ipynb Jupyter
Notebook.
To see the imagined states from the environment model visually run the
eval_env_model.ipynb Jupyter Notebook.
The hyperparameters of this model all work great for minipacman. The hyperparameters for the actor critic training can be found at the top of a2c.py. The hyperparamters can be found for the imagination part at the to top of i2a.py but they can be found at
the top of the i2a.py file, for the environment model at top of
env_model.py and for the overall training at the top of main.py. N_ENV is
the number of environments which are concurrently being simulated and trained
on. You can change this number based on the speed of your system.
NUM_ROLLOUTS is an interesting hyperparameter to play with as it corresponds to
the number of imagined states in the future.