File size: 1,794 Bytes
5216407 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
---
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
|