World Models Approach in RL

• World Models are unsupervised generative models that compress high-dimensional observations into low-dimensional latent spaces using VAEs and predict temporal dynamics with MDN-RNNs.
• The approach enables simulation-based policy training by using a simple controller to act on combined latent and temporal features, achieving robust transfer between simulated and real environments.
• Empirical results on tasks like Car Racing and VizDoom validate that training in 'dream' environments can yield efficient, sample-rich reinforcement learning with scalable computation.

A world model is an internal generative model that captures the compressed spatial and temporal structure of an environment, enabling an agent to simulate, plan, and act by leveraging learned representations of environment dynamics. The canonical "World Models Approach" is exemplified by an architecture in which high-dimensional environmental observations are encoded into a low-dimensional latent space, the temporal evolution of these latent features is predicted, and a lightweight controller utilizes these predictive representations for action selection. This approach supports policy training not only in the real environment but also within "hallucinated" environments generated by the learned world model, facilitating efficient simulation-based reinforcement learning and robust policy transfer.

1. Architectural Components: VAE and MDN-RNN

The approach is organized around two principal generative models:

• Variational Autoencoder (VAE, "V" model):
  • Function: Compresses each high-dimensional observation frame (e.g., $64 \times 64 \times 3$ RGB images) into a lower-dimensional latent vector $z$.
  • Operation: The encoder passes input through several convolutional layers to produce mean $\mu$ and variance $\sigma$ parameters, from which $z$ is sampled as $z \sim \mathcal{N}(\mu, \sigma^2 I)$. A deconvolutional decoder reconstructs the image from $z$.
  • Training Objective: Minimize the reconstruction loss plus a Kullback-Leibler (KL) divergence, enforcing a Gaussian latent prior.
• Mixture Density Network – Recurrent Neural Network (MDN-RNN, "M" model):
  • Function: Models the temporal dynamics in latent space, given by $p(z_{t+1} | z_t, a_t, h_t)$.
  • Formulation: At each step, the MDN-RNN (specifically using LSTM units) outputs parameters of a mixture of Gaussians, yielding:

  $$P(z_{t+1} \mid a_t, z_t, h_t) = \sum_{i} \pi_i \cdot \mathcal{N}(z_{t+1} \mid \mu_i, \sigma^2_i I)$$

  where the mixture weights $\pi_i$, means $\mu_i$, and variances $\sigma_i$ are learned at each time step. - Sampling: A temperature parameter $\tau$ modulates output uncertainty at sampling time.

• Controller (C model):

  • Structure: A deliberately simple model—often just a linear mapping—from concatenated VAE and MDN-RNN features. For input $[z_t ; h_t]$, the action is computed as $a_t = W_c [z_t ; h_t] + b_c$, where $W_c$ and $b_c$ are parameters optimized through evolution strategies such as CMA-ES.

2. Unsupervised Representation Learning

World model components are trained without requiring reward signals or explicit supervision:

• VAE Training:
  • Trained on large datasets of observation frames, typically collected through random agent rollouts.
  • Learns to reconstruct observations, with the KL term enforcing regularity in the latent space $z$.
  • Entire process is fully unsupervised.
• MDN-RNN Training:
  • Input sequences consist of $z$ vectors (from the VAE encoder) and corresponding action sequences.
  • Uses "teacher forcing" to train the RNN to predict the probability distribution of the next latent $z_{t+1}$.
  • Learning is unsupervised; the target sequence consists solely of the observed state evolution, independent of rewards.

This dual-stage unsupervised learning ensures the resulting world model captures both spatial regularities (through $z$) and temporal dependencies (through $h_t$).

3. Feature Extraction for Control

For policy learning, the agent's feature vector is formed by concatenating the latent encoding $z_t$ (from V) with the MDN-RNN's hidden state $h_t$ (from M):

• $z_t$ encodes the current observation's spatial features.
• $h_t$ contains information about recent temporal context and the RNN's predicted near-future.
• The controller operates directly on $[z_t ; h_t]$, providing an information-rich, low-dimensional basis for action selection. This design is critical for transferring policies between simulated and real environments, as $h_t$ abstracts history and $z_t$ grounds in the current state.

4. Policy Optimization in “Dream” Environments

• Controller Simplicity: The C-model is intentionally minimal (e.g., linear), shifting the burden of temporal and spatial modeling to the unsupervised world model. In the CarRacing-v0 benchmark, the controller has only 867 parameters.
• Training Inside the World Model (“the dream”):
  • The controller is trained via evolution strategies (CMA-ES) entirely within a hallucinated environment generated by the world model: actions are sampled, next states are simulated by the MDN-RNN, and cumulative rewards are assigned as in the real task.
  • Multiple controllers are evaluated in parallel in the dream, enabling large-scale, parallelizable policy search.
• Policy Transfer: Once a policy shows satisfactory performance in the simulated environment, it is deployed in the actual environment, keeping the same architecture and features for $[z_t ; h_t]$, demonstrating robust transfer (performance is preserved, sometimes improved, when shifting back to the “real” environment).

5. Empirical Results and Task Performance

• Car Racing:
  • Having the controller observe only $z_t$ yields moderate performance (score: $632 \pm 251$).
  • Adding a hidden layer increases performance ($788 \pm 141$).
  • Full model ($[z_t; h_t]$ supplied) achieves $906 \pm 21$ (over 100 trials), matching or exceeding state-of-the-art for that task.
• VizDoom: Take Cover:
  • The agent, trained solely within the hallucinated world, readily solves the “Take Cover” task (survival times $>1100$ steps, above the “solved” threshold of 750 and competitive with leaderboard agents when deployed back into VizDoom).
  • Additional modeling for the “done” signal enables the world model to simulate episodic terminations credibly.

These results demonstrate that features distilled from unsupervised models can encode sufficient task information for effective downstream control, and that hallucinated “dream” training can generalize to the real dynamics.

6. Mechanics and Deployment Implications

• Computation and Scaling: Both VAE and MDN-RNN operate on low-dimensional representations, greatly reducing the computation required for planning and simulation compared to full pixel-level models. Training is efficient and parallelizable.
• Limitations: Model capacity determines the fidelity of the simulated “dream” environment; errors in modeling can cause discrepancies between simulated and real-world transfer, but the approach’s success across tasks attests to broad robustness.
• Deployment Strategy: The architecture’s simplicity and reliance on unsupervised feature extraction make it suitable for rapid prototyping in novel environments, as data collection does not require extrinsic reward design or dense human intervention.

7. Interactive Exploration and Introspection

An interactive online version of the approach allows users to examine latent encodings, manipulate compressed representations, and observe the effect of temperature scaling on the uncertainty of simulated rollouts. This enables direct qualitative evaluation of spatial compressions, temporal modeling, and policy behavior in both synthetic and real environments.

---

In sum, the World Models Approach demonstrates that unsupervised generative modeling—via a combination of VAE-based spatial compression and RNN-based temporal prediction—can provide a succinct and predictive internal representation sufficient for optimizing control policies on challenging reinforcement learning benchmarks. Training policies within these hallucinated environments leads to robust results, and the approach offers a compelling paradigm for sample-efficient, modular, and interpretable reinforcement learning in complex sequential decision-making problems (Ha et al., 2018).