Recurrent World Models Facilitate Policy Evolution

Problem

Framing

Model-based RL usually learns dynamics but still trains policies in the real environment, because learned simulators are exploitable. This paper closes that gap with a compact controller over a learned latent world model: a VAE for spatial compression, an MDN-RNN for stochastic dynamics, and CMA-ES for policy search. It solves CarRacing-v0 with $906 \pm 21$ and trains a VizDoom policy entirely inside the learned simulator, then transfers it to $1092 \pm 556$ in the real game.

Currently Used Methods

Foundational

• @kingmaVAE2013 — variational autoencoders learn compact latent codes from pixels.
  • Limitation in context: framewise compression alone misses temporal prediction needed for control.
• @hochreiterLSTM1997 — LSTM provides recurrent memory for sequential prediction.
  • Limitation in context: plain recurrence lacks stochastic output modeling for multimodal futures.
• @choGRU2014 — gated recurrent units popularize compact recurrent dynamics models.
  • Limitation in context: still not a probabilistic latent simulator for policy training.
• @mnihA3C_2016 — model-free deep RL learns control directly from observations.
  • Limitation in context: weaker sample efficiency and lower CarRacing scores from raw pixels.
• @goodfellowGAN2014 — adversarial training sharpens generative realism.
  • Limitation in context: no recurrent latent dynamics model for interactive control.

Proposed Method

Architecture

The agent factorizes into vision $V$, memory $M$, and controller $C$. $V$ is a ConvVAE that maps each frame to latent $\mathbf{z}_t$; $M$ is an LSTM-MDN that models the next-latent distribution; $C$ is a linear policy over $[\mathbf{z}_t; \mathbf{h}_t]$. CarRacing uses 256 LSTM units, Doom uses 512, and both use 5 Gaussian mixtures.

Loss / Objective

The method trains reconstruction, latent regularization, and next-state density modeling in separate stages.

Sampling Rule / Algorithm

Rollouts sample the next latent from the MDN-RNN, with temperature $\tau$ controlling simulator stochasticity.

Training Procedure

• Collect 10,000 random rollouts.
• Train ConvVAE on frames.
• Train MDN-RNN for 20 epochs with teacher forcing.
• CarRacing MDN-RNN: 256 LSTM units.
• Doom MDN-RNN: 512 LSTM units.
• Both tasks: 5 Gaussian mixtures.
• Controller optimization: CMA-ES.
• CMA-ES population size: 64.
• Fitness evaluation: 16 rollouts per agent.
• Doom transfer setting: $\tau = 1.15$.

Evaluation

Datasets

• CarRacing-v0, pixel-based continuous control.
• DoomTakeCover-v0, pixel-based survival task.
• 10,000 random-policy rollouts for model learning.

Metrics

• Average return over 100 trials.
• Survival time over 100 VizDoom rollouts.
• Transfer score from virtual to actual environment.

Headline results

• CarRacing-v0, V-only controller: $632 \pm 251$.
• CarRacing-v0, V-only with hidden layer: $788 \pm 141$.
• CarRacing-v0, full world model: $906 \pm 21$.
• DoomTakeCover-v0, virtual training at $\tau=1.15$: virtual $918 \pm 546$, actual $1092 \pm 556$.
• DoomTakeCover-v0, previous gym leader: actual $820 \pm 58$.

Ablations

• Remove memory $\mathbf{h}_t$: CarRacing drops from $906$ to $632$ mean return.
• Add one hidden layer to V-only controller: improves to $788$ but still fails the solve threshold.
• Sweep simulator temperature $\tau$: low $\tau$ induces exploitative policies that fail transfer.
• Raise $\tau$ to $1.15$: harder virtual training yields the best real-environment Doom score.

Method Strengths and Weaknesses

Strengths

• Solves CarRacing-v0 with $906 \pm 21$, exceeding the 900 solve threshold.
• Transfers a policy trained only in latent simulation to $1092 \pm 556$ on real VizDoom.
• Keeps control simple: only 867 controller parameters in CarRacing.
• Uses stochastic dynamics and temperature to reduce simulator exploitation.

Weaknesses

• Learned simulators remain exploitable; low $\tau$ policies collapse at transfer time.
• World model misses scene details, including exact monster counts in Doom.
• Policy receives simulator hidden state, richer than standard agent observations.
• Training is stage-wise, not end-to-end, which may limit task-specific representation learning.

Suggestions from the authors

• Train policies in generated environments with tuned uncertainty to improve transfer.
• Explore better defenses against controllers exploiting model imperfections.
• Study richer world models that capture game logic and rare events more faithfully.
• Extend latent-space training to more realistic, computationally expensive environments.

Links

Prior Papers

• @kingmaVAE2013 — provides the latent-variable encoder-decoder used as the vision module $V$.
• @hochreiterLSTM1997 — supplies the recurrent memory backbone used inside the MDN-RNN world model.
• @choGRU2014 — relevant recurrent-sequence precursor for learned latent dynamics over action-conditioned trajectories.
• @mnihA3C_2016 — serves as a model-free RL baseline family that underperforms the full world model on CarRacing.

Further Papers

• @hafnerDreamerV1_2019 — advances latent world models by optimizing behavior through imagined rollouts.
• @hafnerDreamerV2_2021 — extends the world-model paradigm to stronger visual control with improved latent planning.
• @hafnerDreamerV3_2023 — scales latent world models into a more general, robust control recipe.
• @schrittwieserMuZero2020 — builds a stronger model-based RL line around learned dynamics and policy/value optimization.