Car racing World Model

This project is an implementation of the "World Models" architecture, designed to solve the CarRacing-v3 environment by learning a compressed spatial and temporal representation of the world. By dividing the agent into three distinct components—Vision (V), Memory (M), and Controller (C)—the model learns to navigate the track in a latent space, effectively allowing the agent to "dream" its own environment and train its policy within it [1].

Tool	Site
Weights & Biases	link

Installation

1. Clone the repository:

   git clone https://github.com/DeepLearningDS/car-racing-wm
   cd car-racing-wm

2. Create and activate a virtual environment:

   python3 -m venv venvCAR_RACING
   source venvCAR_RACING/bin/activate  # On Windows, use `venvCAR_RACING\Scripts\activate`

3. Install the dependencies:

   brew install swig # or apt-get install swig
   pip install -r requirements.txt

Test the Agent

You can easily test the trained agent in headless mode. This will quickly compute and print the final cumulative reward without rendering the environment:

python3 test.py model=controller

To watch the agent actually drive the car in real-time, enable the human render mode:

python3 test.py model=controller gymnasium.render_mode=human

Full Training Pipeline

If you want to reproduce the entire pipeline from scratch, follow these steps in order. The architecture must be trained sequentially: VAE $\rightarrow$ MDN-RNN $\rightarrow$ Controller.

1. Rollout Collection

To begin, you need to collect a dataset of environment interactions using a random policy. The following command launches a parallel generation process to collect 10,000 trajectories:

python3 scripts/collect_rollouts.py --rollouts 10000

Note

The rollouts dataset already exists on Kaggle.

2. VAE Training

Train the VAE. You can copy and paste the notebooks/02a_train_VAE.ipynb on Kaggle, or run:

python3 train.py model=vae

3. Encode Rollouts

Before training the RNN, you must compress the entire raw dataset using the newly trained VAE.

python3 scripts/encode_rollouts.py \
    --raw_dir data/raw/rollouts \
    --out_dir data/latent/rollouts \
    --vae_checkpoint checkpoints/vae/vae_final.pth

Note

The encoded rollouts dataset already exists on Kaggle.

4. MDN-RNN Training

Train the MDN-RNN. You can copy and paste the notebooks/03a_train_MDN_RNN.ipynb on Kaggle, or run:

python3 train.py model=mdn_rnn

5. Controller Training

Finally, train the Controller. You can copy and paste the notebooks/04b_train_controller.ipynb on Kaggle, or run:

python3 train.py model=controller

References

[1] Ha, D., & Schmidhuber, J. (2018). World Models. arXiv preprint arXiv:1803.10122.
Paper PDF (arXiv) | Interactive Article