https://github.com/brunoocs/world-model-2018

Modern PyTorch implementation of World Models with interactive notebooks for the Car Racing environment. Features VAE vision model, MDN-RNN memory system, and CMA-ES controller with visualization tools. Complete end-to-end reinforcement learning pipeline with clean, well-documented code.
https://github.com/brunoocs/world-model-2018

cma-es mdn pytorch rnn world-models

Last synced: 5 months ago
JSON representation

Modern PyTorch implementation of World Models with interactive notebooks for the Car Racing environment. Features VAE vision model, MDN-RNN memory system, and CMA-ES controller with visualization tools. Complete end-to-end reinforcement learning pipeline with clean, well-documented code.

• Host: GitHub
• URL: https://github.com/brunoocs/world-model-2018
• Owner: BrunooCS
• Created: 2025-05-05T17:26:17.000Z (5 months ago)
• Default Branch: main
• Last Pushed: 2025-05-05T17:28:16.000Z (5 months ago)
• Last Synced: 2025-05-05T18:45:07.892Z (5 months ago)
• Topics: cma-es, mdn, pytorch, rnn, world-models
• Language: Jupyter Notebook
• Homepage:
• Size: 59 MB
• Stars: 0
• Watchers: 1
• Forks: 0
• Open Issues: 0
• Metadata Files:
  • Readme: README.md

Awesome Lists containing this project

README

          


  World Models PyTorch Implementation 🏎️



---

[![Python](https://img.shields.io/badge/Python-3.12-blue?style=flat&logo=python)](https://www.python.org/)

[![PyTorch](https://img.shields.io/badge/PyTorch-2.0-orange?style=flat&logo=pytorch)](https://pytorch.org/)

[![Gymnasium](https://img.shields.io/badge/Gymnasium-0.29-green?style=flat&logo=openaigym)](https://gymnasium.farama.org/)

[![License](https://img.shields.io/badge/License-MIT-lightgrey?style=flat)](LICENSE)

A clean, interactive PyTorch implementation of ["World Models"](https://worldmodels.github.io/) by David Ha and Jürgen Schmidhuber.



  



## Overview

---

This project implements the complete World Models architecture for the CarRacing-v3 environment from Gymnasium. World Models consist of three components:

1. **Vision (V)**: A Variational Autoencoder (VAE) that compresses raw images into latent representations

2. **Memory (M)**: A Mixed Density Network with LSTM (MDN-RNN) that predicts future states

3. **Controller (C)**: A simple neural network policy trained with CMA-ES

## Interactive Notebooks

---

The implementation is organized into interactive notebooks that explain each component:

| Notebook | Description |

|----------|-------------|

| [1-Rollouts.ipynb](1-Rollouts.ipynb) | Generating dataset from environment interactions |

| [2-Vision (VAE).ipynb](2-Vision%20(VAE).ipynb) | Training the Variational Autoencoder |

| [3-Memory (rnn-mdn).ipynb](3-Memory%20(rnn-mdn).ipynb) | Building the MDN-RNN predictive model |

| [4-Controller (C).ipynb](4-Controller%20(C).ipynb) | Evolutionary training of the controller |

| [5-Videos.ipynb](5-Videos.ipynb) | Generating videos of model performance |

## Features

---

- **Pure PyTorch** implementation with clean, commented code

- **Interactive Visualization** of latent space and model predictions

- **End-to-End Pipeline** from data collection to agent training

- **Pre-trained Models** included in `checkpoints/` directory

- **Modular Design** allowing for experimentation with architectures

## Visualizations

### VAE Latent exploration tools



  



### Pygmae interactive game visualzation of trained models with keyboard controls



  



---

[@Ha2018WorldModels](https://doi.org/10.5281/zenodo.1207631)  

Ha, David and Schmidhuber, Jürgen. "World Models." Zenodo, 2018. [Link to paper](https://zenodo.org/record/1207631).  

Copyright: Creative Commons Attribution 4.0.

---

## License

This project is open-sourced under the [MIT License](LICENSE).

---