DIAMOND: Visual Details Matter in Atari and Diffusion for World Modeling

It was in 2018, when the thought of reinforcement learning within the context of a neural network world model was first introduced, and shortly, this fundamental principle was applied on world models. A number of the distinguished models that implement reinforcement learning were the Dreamer framework, which introduced reinforcement learning from the latent space of a recurrent state space model. The DreamerV2 demonstrated that the usage of discrete latents might lead to reduced compounding errors, and the DreamerV3 framework was capable of achieve human-like performance on a series of tasks across different domains with fixed hyperparameters. 

Moreover, parallels may be drawn between image generation models and world models indicating that the progress made in generative vision models may very well be replicated to profit the world models. Ever for the reason that use of transformers in natural language processing frameworks gained popularity, DALL-E and VQGAN frameworks emerged. The frameworks implemented discrete autoencoders to convert images into discrete tokens, and were capable of construct highly powerful and efficient text to image generative models by leveraging the sequence modeling abilities of the autoregressive transformers. At the identical time, diffusion models gained traction, and today, diffusion models have established themselves as a dominant paradigm for high-resolution image generation. Owing to the capabilities offered by diffusion models and reinforcement learning, attempts are being made to mix the 2 approaches, with the aim to reap the benefits of the flexibleness of diffusion models as trajectory models, reward models, planners, and as policy for data augmentation in offline reinforcement learning. 

World models offer a promising method for training reinforcement learning agents safely and efficiently. Traditionally, these models use sequences of discrete latent variables to simulate environment dynamics. Nonetheless, this compression can overlook visual details crucial for reinforcement learning. At the identical time, diffusion models have risen in popularity for image generation, difficult traditional methods that use discrete latents. Inspired by this shift, in this text, we’ll speak about DIAMOND (DIffusion As a Model Of eNvironment Dreams), a reinforcement learning agent trained inside a diffusion world model. We’ll explore the obligatory design decisions to make diffusion suitable for world modeling and show that enhanced visual details lead to higher agent performance. DIAMOND sets a brand new benchmark on the competitive Atari 100k test, achieving a mean human normalized rating of 1.46, the very best for agents trained entirely inside a world model. 

World models or Generative models of environments are emerging as one in every of the more essential components for generative agents to plan and reason about their environments. Although the usage of reinforcement learning has achieved considerable success in recent times, models implementing reinforcement learning are known for being sample inefficient, which significantly limits their real world applications. Then again, world models have demonstrated their ability to efficiently train reinforcement learning agents across diverse environments with a significantly improved sample efficiency, allowing the model to learn from real world experiences. Recent world modeling frameworks normally model environment dynamics as a sequence of discrete latent variables, with the model discretizing the latent space to avoid compounding errors over multi-step time horizons. Although the approach might deliver substantial results, it’s also related to a loss of data, resulting in lack of reconstruction quality and lack of generality. The loss of data might change into a major roadblock for real-world scenarios that require the data to be well-defined, like training autonomous vehicles. In such tasks, small changes or details within the visual input just like the color of the traffic light, or the turn indicator of the vehicle in front can change the policy of an agent. Although increasing the variety of discrete latents may also help avoid information loss, it shoots the computation costs significantly. 

Moreover, within the recent years, diffusion models have emerged because the dominant approach for high-quality image generation frameworks since frameworks built on diffusion models learn to reverse a noising process, and directly competes with a few of the more well-established approaches modeling discrete tokens, and due to this fact offers a promising alternative to eliminate the necessity for discretization in world modeling. Diffusion models are known for his or her ability to be easily conditioned and to flexibly model complex, multi-modal distributions without mode collapse. These attributes are crucial for world modeling, as conditioning enables a world model to accurately reflect an agent’s actions, resulting in more reliable credit task. Furthermore, modeling multimodal distributions offers a greater diversity of coaching scenarios for the agent, enhancing its overall performance. 

Constructing upon these characteristics, DIAMOND, (DIffusion As a Model Of eNvironment Dreams), a reinforcement learning agent trained inside a diffusion world model. The DIAMOND framework makes careful design decisions to make sure its diffusion world model stays efficient and stable over very long time horizons. The framework provides a qualitative evaluation to reveal the importance of those design decisions. DIAMOND sets a brand new state-of-the-art with a mean human normalized rating of 1.46 on the well-established Atari 100k benchmark, the very best for agents trained entirely inside a world model. Operating in image space allows DIAMOND’s diffusion world model to seamlessly substitute the environment, offering greater insights into world model and agent behaviors. Notably, the improved performance in certain games is attributed to higher modeling of critical visual details. The DIAMOND framework models the environment as a normal POMDP or Partially Observable Markov Decision Process with a set of states, a set of discrete actions, and a set of image observations. The transition functions describe the environment dynamics, and the reward function maps the transitions to scalar rewards. The commentary function describes the commentary probabilities, and emits image observations, which might be then utilized by the agents to see the environments, since they can not directly access the states. The first aim of the approach was to acquire a policy that maps observations to actions with the try and maximize the expected discount return with a reduction factor. World models are generative models of the environment, and world models may be used to create simulated environments to coach reinforcement learning agents in the true environment, and train reinforcement learning agents on this planet model environment. Figure 1 demonstrates the unrolling imagination of the DIAMOND framework over time. 

DIAMOND : Methodology and Architecture

At its core, diffusion models are a category of generative models that generate a sample by reversing the noising process, and draw heavy inspiration from non-equilibrium thermodynamics. The DIAMOND framework considers a diffusion process indexed by a continuous time variable with corresponding marginals and boundary conditions with a tractable unstructured prior distribution. Moreover, to acquire a generative model, which maps from noise to data, the DIAMOND framework must reverse the method, with the reversion process also being a diffusion process, running backwards in time. Moreover, at any given time limit, it is just not trivial to estimate the rating function for the reason that DIAMOND framework doesn’t access to the true rating function, and the model overcomes this hurdle by implementing rating matching objective, an approach that facilitates a framework to coach a rating model without knowing the underlying rating function. The score-based diffusion model provides an unconditional generative model. Nonetheless, a conditional generative model of environment dynamics is required to function a world model, and to serve this purpose, the DIAMOND framework looks at the final case of the POMDP approach, through which the framework could make use of past observations and actions to approximate the unknown Markovian state. As demonstrated in Figure 1., the DIAMOND framework makes use of this history to condition a diffusion model, to estimate and generate the following commentary directly. Although the DIAMOND framework can resort to any SDE or ODE solver in theory, there may be a trade-off between NFE or Variety of Function Evaluations, and sample quality that impacts the inference cost of diffusion models significantly. 

Constructing on the above learnings, allow us to now take a look at the sensible realization of the DIAMOND framework of a diffusion-based world model including the drift and diffusion coefficients corresponding to a specific selection of diffusion approach. As a substitute of choosing DDPM, a naturally suitable candidate for the duty, the DIAMOND framework builds on the EDM formulation, and considers a perturbation kernel with a real-valued function of diffusion time called the noise schedule. The framework selects the preconditioners to maintain the input and output variance for any voice level. The network training mixes signal and noise adaptively depending on the degradation level, and when the noise is low, and the goal becomes the difference between the clean and the perturbed signal, i.e. the added Gaussian noise. Intuitively, this prevents the training objective from becoming trivial within the low-noise regime. In practice, this objective is high variance on the extremes of the noise schedule, so the model samples the noise level from a log-normal distribution chosen empirically to be able to concatenate the training across the medium noise regions. The DIAMOND framework makes use of a normal U-Net 2D component for the vector field, and keeps a buffer of past observations and actions that the framework uses to condition itself. The DIAMOND framework then concatenates these past observations to the following noisy commentary, and input actions through adaptive group normalization layers within the residual blocks of the U-Net. 

DIAMOND: Experiments and Results

For comprehensive evaluation, the DIAMOND framework opts for the Atari 100k benchmark. The Atari 100k benchmark consists of 26 games designed to check a wide selection of agent capabilities. In each game, an agent is proscribed to 100k actions within the environment, which is roughly akin to 2 hours of human gameplay, to learn the sport before evaluation. For comparison, unconstrained Atari agents typically train for 50 million steps, representing a 500-fold increase in experience. We trained DIAMOND from scratch using 5 random seeds for every game. Each training run required around 12GB of VRAM and took roughly 2.9 days on a single Nvidia RTX 4090, amounting to 1.03 GPU years in total. The next table provides the rating for all games, the mean, and the IQM or interquartile mean of human-normalized scores. 

Following the constraints of point estimates, the DIAMOND framework provides stratified bootstrap confidence within the mean, and the IQM or interquartile mean of human-normalized scores together with performance profiles and extra metrics, as summed up in the next figure. 

The outcomes show that DIAMOND performs exceptionally well across the benchmark, surpassing human players in 11 games and achieving a superhuman mean HNS of 1.46, setting a brand new record for agents trained entirely inside a world model. Moreover, DIAMOND’s IQM is comparable to STORM and exceeds all other baselines. DIAMOND excels in environments where capturing small details is crucial, similar to Asterix, Breakout, and RoadRunner. Moreover, as discussed earlier, the DIAMOND framework has the flexibleness of implementing any diffusion model in its pipeline, even though it opts for the EDM approach, it might have been a natural selection to go for the DDPM model because it is already being implemented in quite a few image generative applications. To check the EDM approach against DDPM implementation, the DIAMOND framework trains each the variants with the identical network architecture on the identical shared static dataset with over 100k frames collected with an authority policy. The variety of denoising steps is directly related to the inference cost of the world model, and so fewer steps will reduce the price of coaching an agent on imagined trajectories. To make sure our world model stays computationally comparable with other baselines, similar to IRIS which requires 16 NFE per timestep, we aim to make use of not more than tens of denoising steps, preferably fewer. Nonetheless, setting the variety of denoising steps too low can degrade visual quality, resulting in compounding errors. To evaluate the steadiness of various diffusion variants, we display imagined trajectories generated autoregressively as much as t = 1000 timesteps in the next figure, using different numbers of denoising steps n ≤ 10. 

We observe that using DDPM (a), on this regime ends in severe compounding errors, causing the world model to quickly drift out of distribution. In contrast, the EDM-based diffusion world model (b) stays rather more stable over very long time horizons, even with a single denoising step. Imagined trajectories with diffusion world models based on DDPM (left) and EDM (right) are shown. The initial commentary at t = 0 is identical for each, and every row corresponds to a decreasing variety of denoising steps n. We observe that DDPM-based generation suffers from compounding errors, with smaller numbers of denoising steps resulting in faster error accumulation. In contrast, DIAMOND’s EDM-based world model stays rather more stable, even for n = 1. The optimal single-step prediction is the expectation over possible reconstructions for a given noisy input, which may be out of distribution if the posterior distribution is multimodal. While some games, like Breakout, have deterministic transitions that may be accurately modeled with a single denoising step, other games exhibit partial observability, leading to multimodal commentary distributions. In these cases, an iterative solver is obligatory to guide the sampling procedure towards a particular mode, as illustrated in the sport Boxing in the next figure. Consequently, The DIAMOND framework set n = 3 in all of our experiments.

The above figure compares single-step (top row) and multi-step (bottom row) sampling in Boxing. The movements of the black player are unpredictable, causing single-step denoising to interpolate between possible outcomes, leading to blurry predictions. In contrast, multi-step sampling produces a transparent image by guiding the generation towards a particular mode. Interestingly, for the reason that policy controls the white player, his actions are known to the world model, eliminating ambiguity. Thus, each single-step and multi-step sampling accurately predict the white player’s position.

Within the above figure, the trajectories imagined by DIAMOND generally exhibit higher visual quality and are more faithful to the true environment in comparison with those imagined by IRIS. The trajectories generated by IRIS contain visual inconsistencies between frames (highlighted by white boxes), similar to enemies being displayed as rewards and vice-versa. Although these inconsistencies may only affect just a few pixels, they’ll significantly impact reinforcement learning. As an illustration, an agent typically goals to focus on rewards and avoid enemies, so these small visual discrepancies could make it tougher to learn an optimal policy. The figure shows consecutive frames imagined with IRIS (left) and DIAMOND (right). The white boxes highlight inconsistencies between frames, which only appear in trajectories generated with IRIS. In Asterix (top row), an enemy (orange) becomes a reward (red) within the second frame, then reverts to an enemy within the third, and again to a reward within the fourth. In Breakout (middle row), the bricks and rating are inconsistent between frames. In Road Runner (bottom row), the rewards (small blue dots on the road) are inconsistently rendered between frames. These inconsistencies don’t occur with DIAMOND. In Breakout, the rating is reliably updated by +7 when a red brick is broken. 

Conclusion

In this text, we have now talked about DIAMOND, a reinforcement learning agent trained inside a diffusion world model. The DIAMOND framework makes careful design decisions to make sure its diffusion world model stays efficient and stable over very long time horizons. The framework provides a qualitative evaluation to reveal the importance of those design decisions. DIAMOND sets a brand new state-of-the-art with a mean human normalized rating of 1.46 on the well-established Atari 100k benchmark, the very best for agents trained entirely inside a world model. Operating in image space allows DIAMOND’s diffusion world model to seamlessly substitute the environment, offering greater insights into world model and agent behaviors. Notably, the improved performance in certain games is attributed to higher modeling of critical visual details. The DIAMOND framework models the environment as a normal POMDP or Partially Observable Markov Decision Process with a set of states, a set of discrete actions, and a set of image observations. The transition functions describe the environment dynamics, and the reward function maps the transitions to scalar rewards.