Overview

Research interest in world models has exploded in recent years, driven by their potential to enable scalable training of robotics. A world model that can accurately represent the environment and predict its future state given an action could allow robots to learn and refine policies within an imagined world, reducing the need for costly deployment and potentially unsafe real-world interaction. We explore how to train a DreamerV3 world model and policy through the lens of Flappy Birds.

Flappy Birds Environment

We chose Flappy Birds as our environment because of its minimal action set (flap, no-flap) and relatively simple objective (avoid the pipes for as long as possible). However, the game is deceptively unforgiving to play — as I have never reached a score above 20 😥. Every agent is trained on only visual input (pixels) and does not have access to the game's internal states. We downsample the number of pixels from 512x288 to 128x72 — a 4x compression rate. We reward the agent +1.0 for successfully passing a pipe, +0.1 every frame it stays alive, -1.0 for dying, and -0.5 each frame it touches the top of the screen.

Collecting a Behaviourally Diverse Dataset

Collecting a behaviourally diverse dataset is more challenging than naively collecting rollouts with random actions. A 50% flap probability seems reasonable, but you quickly realize that a flap action at approximately every other step quickly launches the bird into the top of the screen. Even if you considerably lower the flap probability, the large majority of rollouts will only include the bird dying on the first pipe. When the core of Flappy Birds gameplay is after the first pipe, the world model will need much deeper runs to accurately model the game.

We solve this by relying on a pre-trained PPO policy to collect longer rollouts. However, rollouts that are 100% guided by the PPO policy would only represent "correct" gameplay. In order for the world model to fully model the environment, it needs rollouts with a degree of randomness in actions and death reasons. We need rollouts that spam the flap action and rollouts that don't flap at all. We need rollouts that hit the first set of pipes at different locations, and we need deep runs. We need rollouts that die on the bottom pipe and rollouts that die on the top pipe. How can we represent all of this in a dataset?

To do this, we 1. induce various levels of randomness and 2. treat the internal game state representation as an embedding vector.

Formally, a rollout is a trajectory $\tau = (o_t, a_t, r_t)_{t=0}^{T-1}$ of observations, actions, and rewards, where $a_t \in \{0, 1\}$ (no-flap, flap) and $T$ is the step at which the bird dies. We collect a dataset $\mathcal{D}$ of $N = 10{,}000$ such rollouts, and we want $\mathcal{D}$ to cover as much of the space of possible playthroughs as we can.

Rollouts are generated by a behavior policy that mixes the pre-trained PPO policy $\pi_{\text{PPO}}$ with a Bernoulli flapping policy:

The flap probability of the random branch is drawn once per rollout as $p \sim \mathcal{U}(0, 0.06)$. We decided this range by collecting rollouts and found that there was approximately the same number of rollouts where the bird died on the top pipe vs. the bottom pipe. The mixing weight $q_t \in [0, 1]$ is the probability that we defer to PPO at step $t$. For the $n$-th rollout it is annealed linearly from 0% to 100% over the course of collection, and knocked back down to 0 with probability $\varepsilon = 0.01$ at each step:

The linear annealing enables a reasonable distribution of short rollouts vs. long rollouts, and the random dropout $z_t$ encourages more diverse deaths at the extremes of death heights (getting as high as possible between two pipes, hitting the ground between two pipes).

At each rollout collection step, we fire up a pool $\mathcal{P}$ of 64 playthroughs and keep the 16 most novel rollouts compared to all of the collected ones. We create a "behavior descriptor" of each rollout by concatenating the normalized internal state vectors $\hat{s}_t \in \mathbb{R}^d$ of the last $L = 50$ steps of the rollout:

We take the last 50 steps because they represent how the bird died, which is the most important part for the world model to represent. With the behavior descriptors, we can treat them as rollout embeddings and calculate a similarity score $\mathrm{sim}(\cdot, \cdot)$ between the rollout and all other already collected rollouts. The novelty of a candidate rollout is its similarity to its nearest neighbour in $\mathcal{D}$, so the rollout in the pool that is the most dissimilar is the most behaviourally diverse:

To train DreamerV3, we collect 10,000 rollouts. We show the histogram of rollout death height as well as rollout length. These histograms show we have a behaviourally diverse dataset of death heights and rollout lengths.

DreamerV3

We follow the implementation of DreamerV3. For more detailed information, please refer to their paper. We provide a figure below from the paper as a convenient overview of DreamerV3.

Hopefully, I can distil how DreamerV3 works into a simplified summary. Again, please refer to their paper for more information. DreamerV3 has two components: the world model and the behavior component. The world model can be simply understood as an RNN attached to an encoder and decoder of a latent representation (z). The latent (z), recurrent state (h), and action (a) are provided to the world model to predict the latent state of the next state. This predicted latent state can be decoded into an image that looks like a "dream". The behavior component uses an actor-critic model to predict the action that should be taken and the reward given for that action, as well as the next latent state.

Experiments

Baselines

We compare against the popular reinforcement learning baselines Q-Learning, DQN, and PPO.

Training & Results

We train a 12M parameter DreamerV3 world model for 20 epochs on our rollout dataset. For fair comparison, we train 1M parameter models for each method for 1 million real environment frames, with a learning rate of 1e-4 and a batch size of 64.

We generate 32 val seeds and 100 test seeds. We use the 32 val seeds during training to select the best checkpoints by highest mean score, and we use the test seeds to do final evaluation of those best checkpoints. We collect the mean, median, min, and max scores as well as the standard deviations, and provide the histograms of scores in the figures below. During test time, we stop the run after a score ceiling of 1,000.

DreamerV3 considerably outperforms all other baselines in both performance and consistency. Its performance distribution shows that only the bottom 25% of runs fail to achieve the maximum score of 1,000, resulting in a mean score of approximately 900 and a median of 1,000. Most notably, DreamerV3 achieves this performance after only 150,000 training steps, while all other agents are trained for 1 million steps.

PPO is also capable of reaching the maximum score of 1,000, but its performance is much less consistent. The majority of PPO runs achieve scores between 300 and 500. While this demonstrates that PPO can learn the task effectively in some runs, its overall performance is substantially less reliable than DreamerV3.

In contrast, DQN and Q-Learning struggle considerably with the task because they rely directly on the observation space to learn a value function. Since our agents receive RGB images as observations, the resulting observation space is extremely large. This makes it difficult for these methods to learn effective representations of the game state and, consequently, to learn Flappy Bird reliably.

DreamerV3's superior performance can be attributed in part to its ability to train using imagined experience generated by its learned world model. During training, DreamerV3 observes a real Flappy Bird frame and uses its latent world model to generate 1,024 imagined steps. These imagined experiences are then used for policy learning without requiring additional interaction with the real environment.

This makes DreamerV3 significantly more frame-efficient than the baseline methods. Although DreamerV3 is trained using only 150,000 real environment frames, each real observation can be used to generate up to 1,024 imagined steps. This corresponds to approximately 150,000 × 1,024 = 153.6 million imagined training frames. In comparison, the baseline agents require 1 million real environment steps. Thus, DreamerV3 achieves substantially better performance while requiring only 15% as many real environment interactions. This efficiency could be increased even more drastically if the policy was allowed to train for more than 1,024 imagined steps.

Conclusion

We show that training on latent representations is not only viable but may also be superior to training directly in the environment. I believe world models are well deserving of the excitement surrounding them and have tremendous potential for advancing robotics. I hope to explore world models further in my future research.

Acknowledgements

I'd like to thank the creators of flappy-bird-gymnasium and especially the authors of the DreamerV3 paper. This project would not have been possible without their incredible work.