Mastering Diverse Domains through World Models — part 3/4

Cleaned-text code points [30482, 46196). Bibliography omitted here; source.md includes it.

ℒ ⁢ ( θ ) ≐ − ∑ t = 1 T sg ⁡ ( ( R t λ − v ψ ⁢ ( s t ) ) / max ⁡ ( 1 , S ) ) ⁢ log ⁡ π θ ⁢ ( a t | s t ) + η ⁢ H ⁡ [ π θ ⁢ ( a t | s t ) ] approaches-limit ℒ 𝜃 superscript subscript 𝑡 1 𝑇 sg subscript superscript 𝑅 𝜆 𝑡 subscript 𝑣 𝜓 subscript 𝑠 𝑡 1 𝑆 subscript 𝜋 𝜃 conditional subscript 𝑎 𝑡 subscript 𝑠 𝑡 𝜂 H subscript 𝜋 𝜃 conditional subscript 𝑎 𝑡 subscript 𝑠 𝑡 \displaystyle\begin{aligned} \mathcal{L}(\theta)\doteq-\textstyle\sum_{t=1}^{T% }\operatorname{sg}\!\Big{(}\big{(}R^{\lambda}_{t}-v_{\psi}(s_{t})\big{)}/\max(% 1,\,S)\Big{)}\log\pi_{\theta}(a_{t}\;|\;s_{t})+\,\eta\operatorname{H}\big{[}% \pi_{\theta}(a_{t}\;\big{|}\;s_{t})\big{]}\end{aligned} start_ROW start_CELL caligraphic_L ( italic_θ ) ≐ - ∑ start_POSTSUBSCRIPT italic_t = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_T end_POSTSUPERSCRIPT roman_sg ( ( italic_R start_POSTSUPERSCRIPT italic_λ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT - italic_v start_POSTSUBSCRIPT italic_ψ end_POSTSUBSCRIPT ( italic_s start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) ) / roman_max ( 1 , italic_S ) ) roman_log italic_π start_POSTSUBSCRIPT italic_θ end_POSTSUBSCRIPT ( italic_a start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | italic_s start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) + italic_η roman_H [ italic_π start_POSTSUBSCRIPT italic_θ end_POSTSUBSCRIPT ( italic_a start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | italic_s start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) ] end_CELL end_ROW (6) The return distribution can be multi-modal and include outliers, especially for randomized environments where some episodes have higher achievable returns than others. Normalizing by the smallest and largest observed returns would then scale returns down too much and may cause suboptimal convergence. To be robust to these outliers, we compute the range from the 5 th to the 95 th return percentile over the return batch and smooth out the estimate using an exponential moving average:

S ≐ EMA ⁡ ( Per ⁡ ( R t λ , 95 ) − Per ⁡ ( R t λ , 5 ) , 0.99 ) approaches-limit 𝑆 EMA Per subscript superscript 𝑅 𝜆 𝑡 95 Per subscript superscript 𝑅 𝜆 𝑡 5 0.99 \displaystyle\begin{aligned} S\doteq\operatorname{EMA}\!\big{(}\operatorname{% Per}(R^{\lambda}_{t},95)-\operatorname{Per}(R^{\lambda}_{t},5),0.99\big{)}\end% {aligned} start_ROW start_CELL italic_S ≐ roman_EMA ( roman_Per ( italic_R start_POSTSUPERSCRIPT italic_λ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , 95 ) - roman_Per ( italic_R start_POSTSUPERSCRIPT italic_λ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , 5 ) , 0.99 ) end_CELL end_ROW (7) Previous work typically normalizes advantages 5 rather than returns, which puts a fixed amount of emphasis on maximizing returns over entropy regardless of whether rewards are within reach. Scaling up advantages when rewards are sparse can amplify noise that outweighs the entropy regularizer and stagnates exploration. Normalizing rewards or returns by standard deviation can fail under sparse rewards where their standard deviation is near zero, drastically amplifying rewards regardless of their size. Constrained optimization targets a fixed entropy on average across states 32 , 33 regardless of achievable returns, which is robust but explores slowly under sparse rewards and converges lower under dense rewards. We did not find stable hyperparameters across domains for these approaches. Return normalization with a denominator limit overcomes these challenges, exploring rapidly under sparse rewards and converging to high performance across diverse domains.

### Robust predictions

Reconstructing inputs and predicting rewards and returns can be challenging because the scale of these quantities can vary across domains. Predicting large targets using a squared loss can lead to divergence whereas absolute and Huber losses 7 stagnate learning. On the other hand, normalizing targets based on running statistics 5 introduces non-stationarity into the optimization. We suggest the symlog squared error as a simple solution to this dilemma. For this, a neural network f ⁢ ( x , θ ) 𝑓 𝑥 𝜃 f(x,\theta) italic_f ( italic_x , italic_θ ) with inputs x 𝑥 x italic_x and parameters θ 𝜃 \theta italic_θ learns to predict a transformed version of its targets y 𝑦 y italic_y . To read out predictions y ^ ^ 𝑦 \hat{y} over^ start_ARG italic_y end_ARG of the network, we apply the inverse transformation:

ℒ ⁢ ( θ ) ≐ 1 2 ⁢ ( f ⁢ ( x , θ ) − symlog ⁡ ( y ) ) 2 y ^ ≐ symexp ⁡ ( f ⁢ ( x , θ ) ) formulae-sequence approaches-limit ℒ 𝜃 1 2 superscript 𝑓 𝑥 𝜃 symlog 𝑦 2 approaches-limit ^ 𝑦 symexp 𝑓 𝑥 𝜃 \displaystyle\begin{aligned} \mathcal{L(\theta)}\doteq\textstyle\frac{1}{2}% \big{(}f(x,\theta)-\operatorname{symlog}(y)\big{)}^{2}\qquad\hat{y}\doteq% \operatorname{symexp}\!\big{(}f(x,\theta)\big{)}\end{aligned} start_ROW start_CELL caligraphic_L ( italic_θ ) ≐ divide start_ARG 1 end_ARG start_ARG 2 end_ARG ( italic_f ( italic_x , italic_θ ) - roman_symlog ( italic_y ) ) start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT over^ start_ARG italic_y end_ARG ≐ roman_symexp ( italic_f ( italic_x , italic_θ ) ) end_CELL end_ROW (8) Using the logarithm as transformation would not allow us to predict targets that take on negative values. Therefore, we choose a function from the bi-symmetric logarithmic family 34 that we name symlog as the transformation with the symexp function as its inverse:

symlog ⁡ ( x ) ≐ sign ⁡ ( x ) ⁢ ln ⁡ ( | x | + 1 ) symexp ⁡ ( x ) ≐ sign ⁡ ( x ) ⁢ ( exp ⁡ ( | x | ) − 1 ) formulae-sequence approaches-limit symlog 𝑥 sign 𝑥 𝑥 1 approaches-limit symexp 𝑥 sign 𝑥 𝑥 1 \displaystyle\begin{aligned} \operatorname{symlog}(x)\doteq\operatorname{sign}% (x)\ln\!\big{(}|x|+1\big{)}\qquad\operatorname{symexp}(x)\doteq\operatorname{% sign}(x)\big{(}\!\exp(|x|)-1\big{)}\end{aligned} start_ROW start_CELL roman_symlog ( italic_x ) ≐ roman_sign ( italic_x ) roman_ln ( | italic_x | + 1 ) roman_symexp ( italic_x ) ≐ roman_sign ( italic_x ) ( roman_exp ( | italic_x | ) - 1 ) end_CELL end_ROW (9) The symlog function compresses the magnitudes of both large positive and negative values. Unlike the logarithm, it is symmetric around the origin while preserving the input sign. This allows the optimization process to quickly move the network predictions to large values when needed. The symlog function approximates the identity around the origin so that it does not affect learning of targets that are already small enough.

For potentially stochastic targets, such as rewards or returns, we introduce the symexp twohot loss. Here, the network outputs the logits for a softmax distribution over exponentially spaced bins b i ∈ B subscript 𝑏 𝑖 𝐵 b_{i}\in B italic_b start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ∈ italic_B . Predictions are read out as the weighted average of the bin positions weighted by their predicted probabilities. Importantly, the network can output any continuous value in the interval because the weighted average can fall between the buckets:

y ^ ≐ softmax ( f ( x ) ) T B B ≐ symexp ( [ − 20 … + 20 ] ) \displaystyle\begin{aligned} \hat{y}\doteq\operatorname{softmax}(f(x))^{T}B% \qquad B\doteq\operatorname{symexp}(\begin{bmatrix}-20&...&+20\end{bmatrix})% \end{aligned} start_ROW start_CELL over^ start_ARG italic_y end_ARG ≐ roman_softmax ( italic_f ( italic_x ) ) start_POSTSUPERSCRIPT italic_T end_POSTSUPERSCRIPT italic_B italic_B ≐ roman_symexp ( [ start_ARG start_ROW start_CELL - 20 end_CELL start_CELL … end_CELL start_CELL + 20 end_CELL end_ROW end_ARG ] ) end_CELL end_ROW (10) The network is trained on twohot encoded targets 8 , 28 , a generalization of onehot encoding to continuous values. The twohot encoding of a scalar is a vector with | B | 𝐵 |B| | italic_B | entries that are all 0 0 except at the indices k 𝑘 k italic_k and k + 1 𝑘 1 k+1 italic_k + 1 of the two bins closest to the encoded scalar. The two entries sum up to 1 1 1 1 , with linearly higher weight given to the bin that is closer to the encoded continuous number. The network is then trained to minimize the categorical cross entropy loss for classification with soft targets. Note that the loss only depends on the probabilities assigned to the bins but not on the continuous values associated with the bin locations, decoupling the size of the gradients from the size of the targets:

ℒ ( θ ) ≐ − twohot ( y ) T log ⁢ softmax ( f ( x , θ ) ) \displaystyle\begin{aligned} \mathcal{L}(\theta)\doteq-\operatorname{twohot}(y% )^{T}\operatorname{log\,softmax}(f(x,\theta))\end{aligned} start_ROW start_CELL caligraphic_L ( italic_θ ) ≐ - roman_twohot ( italic_y ) start_POSTSUPERSCRIPT italic_T end_POSTSUPERSCRIPT start_OPFUNCTION roman_log roman_softmax end_OPFUNCTION ( italic_f ( italic_x , italic_θ ) ) end_CELL end_ROW (11) Applying these principles, Dreamer transforms vector observations using the symlog functions, both for the encoder inputs and the decoder targets and employs the synexp twohot loss for the reward predictor and critic. We find that these techniques enable robust and fast learning across many diverse domains. For critic learning, an alternative asymmetric transformation has previously been proposed 35 , which we found less effective on average across domains. Unlike alternatives, symlog transformations avoid truncating large targets 7 , introducing non-stationary from normalization 5 , or adjusting network weights when new extreme values are detected 36 .

## Results

We evaluate the generality of Dreamer across 8 domains—with over 150 tasks—under fixed hyperparameters. We designed the experiments to compare Dreamer to the best methods in the literature, which are often specifically designed and tuned for the benchmark at hand. We further compare to a high-quality implementation of PPO 5 , a standard reinforcement learning algorithm that is known for its robustness. We run PPO with fixed hyperparameters chosen to maximize performance across domains and that reproduce strong published results of PPO on ProcGen 37 . To push the boundaries of reinforcement learning, we apply Dreamer to the challenging video game Minecraft, comparing it to strong previous algorithms. Finally, we analyze the importance of individual components of Dreamer and its robustness to different model sizes and computational budgets. All Dreamer agents are trained on a single Nvidia A100 GPU each, making it reproducible for many research labs. A public implementation of Dreamer that reproduces all results is available on the project website.

#### Benchmarks

We perform an extensive empirical study across 8 domains that include continuous and discrete actions, visual and low-dimensional inputs, dense and sparse rewards, different reward scales, 2D and 3D worlds, and procedural generation. Figure 1 summarizes the benchmark results, showing that Dreamer outperforms a wide range of previous expert algorithms across diverse domains. Crucially, Dreamer substantially outperforms PPO across all domains.

• Atari This established benchmark contains 57 Atari 2600 games with a budget of 200M frames, posing a diverse range of challenges 38 . We use the sticky action simulator setting 39 . Dreamer outperforms the powerful MuZero algorithm 8 while using only a fraction of the computational resources. Dreamer also outperforms the widely-used expert algorithms Rainbow 40 and IQN 41 .

• ProcGen This benchmark of 16 games features randomized levels and visual distractions to test the robustness and generalization of agents 42 . Within the budget of 50M frames, Dreamer matches the tuned expert algorithm PPG 37 and outperforms Rainbow 42 , 40 . Our PPO agent with fixed hyperparameters matches the published score of the highly tuned official PPO implementation 37 .

• DMLab This suite of 30 tasks features 3D environments that test spatial and temporal reasoning 43 . In 100M frames, Dreamer exceeds the performance of the scalable IMPALA and R2D2+ agents 35 at 1B environment steps, amounting to a data-efficiency gain of over 1000%. We note that these baselines were not designed for data-efficiency but serve as a valuable comparison point for the performance previously achievable at scale.

• Atari100k This data-efficiency benchmark comntains 26 Atari games and a budget of only 400K frames, amounting to 2 hours of game time 17 . EfficientZero 44 holds the state-of-the-art by combining online tree search, prioritized replay, and hyperparameter scheduling, but also resets levels early to increase data diversity, making a comparison difficult. Without this complexity, Dreamer outperforms the best remaining methods, including the transformer-based IRIS and TWM agents, the model-free SPR, and SimPLe 45 .

• Proprio Control This benchmark contains 18 control tasks with continuous actions, proprioceptive vector inputs, and a budget of 500K environment steps 46 . The tasks range from classical control over locomotion to robot manipulation tasks, featuring dense and sparse rewards. Dreamer sets a new state-of-the-art on this benchmark, outperforming D4PG, DMPO, and MPO 33 .

• Visual Control This benchmark consists of 20 continuous control tasks where the agent receives only high-dimensional images as input and has a budget of 1M environment steps 46 . Dreamer establishes a new state-of-the-art on this benchmark, outperforming DrQ-v2 and CURL 47 , which are specialized to visual environments and leverage data augmentation.

• BSuite This benchmark includes 23 environments with a total of 468 configurations that are specifically designed to test credit assignment, robustness to reward scale and stochasticity, memory, generalization, and exploration 48 . Dreamer establishes a new state-of-the-art on this benchmark, outperforming Boot DQN and other methods 49 . Dreamer improves over previous algorithms especially in the scale robustness category.

#### Minecraft

Collecting diamonds in the popular game Minecraft has been a long-standing challenge in artificial intelligence 18 , 19 , 20 . Every episode in this game is set in a unique randomly generated and infinite 3D world. Episodes last until the player dies or up to 36000 steps equaling 30 minutes, during which the player needs to discover a sequence of 12 items from sparse rewards by foraging for resources and crafting tools. It takes about 20 minutes for experienced human players to obtain diamonds 20 . We follow the block breaking setting of prior work 19 because the provided action space would make it challenging for stochastic policies to keep a key pressed for a prolonged time.

Because of the training time in this complex domain, extensive tuning would be difficult for Minecraft. Instead, we apply Dreamer out of the box with its default hyperparameters. As shown in Figures 5 and 1 , Dreamer is the first algorithm to collect diamonds in Minecraft from scratch without using human data as was required by VPT 20 or adaptive curricula 19 . All the Dreamer agents we trained on Minecraft discover diamonds in 100M environment steps. While several strong baselines progress to advanced items such as the iron pickaxe, none of them discovers a diamond.

#### Ablations

In Figure 6 , we ablate the robustness techniques and learning signals on a diverse set of 14 tasks to understand their importance. The training curves of individual tasks are included in the supplementary material. We observe that all robustness techniques contribute to performance, most notably the KL objective of the world model, followed by return normalization and symexp twohot regression for reward and value prediction. In general, we find that each individual technique is critical on a subset of tasks but may not affect performance on other tasks.