Initial World Models checkpoint upload

README.md

@@ -0,0 +1,69 @@
+---
+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

config.json

@@ -0,0 +1,15 @@
+{
+  "architecture": "WorldModels",
+  "latent_size": 64,
+  "hidden_size": 256,
+  "action_size": 4,
+  "n_gaussians": 5,
+  "components": {
+    "vae": "Variational Autoencoder for visual compression",
+    "rnn": "MDRNN (Memory-Augmented Recurrent Neural Network)",
+    "controller": "CMA-ES optimized linear controller"
+  },
+  "trained_on": "Atari environment (Breakout)",
+  "upload_date": "2025-12-25T08:13:58.344484",
+  "framework": "PyTorch"
+}

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5e4139a4a780a9edcf8bb1b81f903d15570f33d8a4266584304898d9480c3f5f
+size 32753806