|
|
---
|
|
|
tags:
|
|
|
- reinforcement-learning
|
|
|
- world-models
|
|
|
- atari
|
|
|
- vae
|
|
|
- rnn
|
|
|
- cma-es
|
|
|
- pytorch
|
|
|
license: mit
|
|
|
---
|
|
|
|
|
|
# World Models - Atari Agent
|
|
|
|
|
|
This is a World Models implementation from Ha & Schmidhuber (2018) trained on Atari Breakout environment.
|
|
|
|
|
|
## Model Description
|
|
|
|
|
|
The World Models architecture consists of three main components:
|
|
|
|
|
|
1. **VAE (Variational Autoencoder)**: Compresses 64x64 RGB images into a 64-dimensional latent space
|
|
|
2. **RNN (Memory-Augmented Recurrent Neural Network - MDRNN)**: Predicts the next latent representation given current latent state and action
|
|
|
3. **Controller**: A linear controller optimized using CMA-ES to maximize cumulative reward
|
|
|
|
|
|
## Model Details
|
|
|
|
|
|
- **Latent Size**: 64
|
|
|
- **Hidden Size**: 256 (MDRNN)
|
|
|
- **Action Space**: 4 (Atari discrete actions)
|
|
|
- **Architecture**: Convolutional encoder/decoder for VAE, LSTM-based RNN
|
|
|
- **Optimization**: CMA-ES for controller training
|
|
|
|
|
|
## Usage
|
|
|
|
|
|
```python
|
|
|
import torch
|
|
|
from pathlib import Path
|
|
|
|
|
|
# Load checkpoint
|
|
|
checkpoint = torch.load('pytorch_model.bin')
|
|
|
|
|
|
# Access components
|
|
|
vae_state = checkpoint['vae_state_dict']
|
|
|
rnn_state = checkpoint['rnn_state_dict']
|
|
|
controller_state = checkpoint['controller_state_dict']
|
|
|
|
|
|
# Reconstruct models (see auto_train.py for architecture definitions)
|
|
|
# and load states into them
|
|
|
```
|
|
|
|
|
|
## Training Details
|
|
|
|
|
|
- **Environment**: Atari Breakout
|
|
|
- **Harvest Episodes**: 100
|
|
|
- **VAE Epochs**: 20
|
|
|
- **RNN Epochs**: 20
|
|
|
- **CMA-ES Generations**: 25
|
|
|
- **Population Size**: 32
|
|
|
- **Framework**: PyTorch
|
|
|
- **Training Script**: auto_train.py
|
|
|
|
|
|
## References
|
|
|
|
|
|
- Ha, D., & Schmidhuber, J. (2018). World Models. arXiv preprint arXiv:1803.10122.
|
|
|
- Original paper: https://worldmodels.github.io/
|
|
|
|
|
|
## License
|
|
|
|
|
|
MIT License
|
|
|
|
