Mastering Diverse Domains through World Models — part 2/4

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

The prediction loss trains the decoder and reward predictor via the symlog squared loss described later, and the continue predictor via logistic regression. The dynamics loss trains the sequence model to predict the next representation by minimizing the KL divergence between the predictor p ϕ ⁢ ( z t | h t ) subscript 𝑝 italic-ϕ conditional subscript 𝑧 𝑡 subscript ℎ 𝑡 p_{\phi}(z_{t}\;|\;h_{t}) italic_p start_POSTSUBSCRIPT italic_ϕ end_POSTSUBSCRIPT ( italic_z start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | italic_h start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) and the next stochastic representation q ϕ ⁢ ( z t | h t , x t ) subscript 𝑞 italic-ϕ conditional subscript 𝑧 𝑡 subscript ℎ 𝑡 subscript 𝑥 𝑡 q_{\phi}(z_{t}\;|\;h_{t},x_{t}) italic_q start_POSTSUBSCRIPT italic_ϕ end_POSTSUBSCRIPT ( italic_z start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | italic_h start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) . The representation loss, in turn, trains the representations to become more predictable allowing us to use a factorized dynamics predictor for fast sampling during imagination training. The two losses differ in the stop-gradient operator sg ⁡ ( ⋅ ) sg ⋅ \operatorname{sg}(\cdot) roman_sg ( ⋅ ) and their loss scale. To avoid a degenerate solution where the dynamics are trivial to predict but fail to contain enough information about the inputs, we employ free bits 26 by clipping the dynamics and representation losses below the value of 1 nat ≈ \approx ≈ 1.44 bits. This disables them while they are already minimized well to focus learning on the prediction loss:

ℒ pred ⁢ ( ϕ ) ≐ − ln ⁡ p ϕ ⁢ ( x t | z t , h t ) − ln ⁡ p ϕ ⁢ ( r t | z t , h t ) − ln ⁡ p ϕ ⁢ ( c t | z t , h t ) ℒ dyn ⁢ ( ϕ ) ≐ max ( 1 , KL [ sg ( q ϕ ( z t | h t , x t ) ) ∥ p ϕ ( z t | h t ) ] ) ℒ rep ⁢ ( ϕ ) ≐ max ( 1 , KL [ q ϕ ( z t | h t , x t ) ∥ sg ( p ϕ ( z t | h t ) ) ] ) \displaystyle\begin{aligned} \mathcal{L_{\mathrm{pred}}}(\phi)&\doteq-\ln p_{% \phi}(x_{t}\;|\;z_{t},h_{t})-\ln p_{\phi}(r_{t}\;|\;z_{t},h_{t})-\ln p_{\phi}(% c_{t}\;|\;z_{t},h_{t})\\ \mathcal{L_{\mathrm{dyn}}}(\phi)&\doteq\max\bigl{(}1,\operatorname{KL}\!\big{[% }\operatorname{sg}(q_{\phi}(z_{t}\;|\;h_{t},x_{t}))\;\big{\|}\;\hskip 13.77771% ptp_{\phi}(z_{t}\;|\;h_{t})\hphantom{)}\big{]}\bigr{)}\\ \mathcal{L_{\mathrm{rep}}}(\phi)&\doteq\max\bigl{(}1,\operatorname{KL}\!\big{[% }\hskip 13.77771ptq_{\phi}(z_{t}\;|\;h_{t},x_{t})\hphantom{)}\;\big{\|}\;% \operatorname{sg}(p_{\phi}(z_{t}\;|\;h_{t}))\big{]}\bigr{)}\end{aligned} start_ROW start_CELL caligraphic_L start_POSTSUBSCRIPT roman_pred end_POSTSUBSCRIPT ( italic_ϕ ) end_CELL start_CELL ≐ - roman_ln italic_p start_POSTSUBSCRIPT italic_ϕ end_POSTSUBSCRIPT ( italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | italic_z start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , italic_h start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) - roman_ln italic_p start_POSTSUBSCRIPT italic_ϕ end_POSTSUBSCRIPT ( italic_r start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | italic_z start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , italic_h start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) - roman_ln italic_p start_POSTSUBSCRIPT italic_ϕ end_POSTSUBSCRIPT ( italic_c start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | italic_z start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , italic_h start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) end_CELL end_ROW start_ROW start_CELL caligraphic_L start_POSTSUBSCRIPT roman_dyn end_POSTSUBSCRIPT ( italic_ϕ ) end_CELL start_CELL ≐ roman_max ( 1 , roman_KL [ roman_sg ( italic_q start_POSTSUBSCRIPT italic_ϕ end_POSTSUBSCRIPT ( italic_z start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | italic_h start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) ) ∥ italic_p start_POSTSUBSCRIPT italic_ϕ end_POSTSUBSCRIPT ( italic_z start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | italic_h start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) ] ) end_CELL end_ROW start_ROW start_CELL caligraphic_L start_POSTSUBSCRIPT roman_rep end_POSTSUBSCRIPT ( italic_ϕ ) end_CELL start_CELL ≐ roman_max ( 1 , roman_KL [ italic_q start_POSTSUBSCRIPT italic_ϕ end_POSTSUBSCRIPT ( italic_z start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | italic_h start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) ∥ roman_sg ( italic_p start_POSTSUBSCRIPT italic_ϕ end_POSTSUBSCRIPT ( italic_z start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | italic_h start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) ) ] ) end_CELL end_ROW (3) Previous world models require scaling the representation loss differently based on the visual complexity of the environment 21 . Complex 3D environments contain details unnecessary for control and thus prompt a stronger regularizer to simplify the representations and make them more predictable. In games with static backgrounds and where individual pixels may matter for the task, a weak regularizer is required to extract fine details. We find that combining free bits with a small representation loss resolves this dilemma, allowing for fixed hyperparameters across domains. Moreover, we transform vector observations using the symlog function described later, to prevent large inputs and large reconstruction gradients, further stabilizing the trade-off with the representation loss.

We occasionally observed spikes the in KL losses in earlier experiments, consistent with reports for deep variational autoencoders 27 . To prevent this, we parameterize the categorical distributions of the encoder and dynamics predictor as mixtures of 1% uniform and 99% neural network output, making it impossible for them to become deterministic and thus ensuring well-behaved KL losses. Further model details and hyperparameters are included in the supplementary material.

### Critic learning

The actor and critic neural networks learn behaviors purely from abstract trajectories of representations predicted by the world model 14 . For environment interaction, we select actions by sampling from the actor network without lookahead planning. The actor and critic operate on model states s t ≐ { h t , z t } approaches-limit subscript 𝑠 𝑡 subscript ℎ 𝑡 subscript 𝑧 𝑡 s_{t}\doteq\{h_{t},z_{t}\} italic_s start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ≐ { italic_h start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , italic_z start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT } and thus benefit from the Markovian representations learned by the recurrent world model. The actor aims to maximize the return R t ≐ ∑ τ = 0 ∞ γ τ ⁢ r t + τ approaches-limit subscript 𝑅 𝑡 superscript subscript 𝜏 0 superscript 𝛾 𝜏 subscript 𝑟 𝑡 𝜏 R_{t}\doteq\textstyle\sum_{\tau=0}^{\infty}\gamma^{\tau}r_{t+\tau} italic_R start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ≐ ∑ start_POSTSUBSCRIPT italic_τ = 0 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ∞ end_POSTSUPERSCRIPT italic_γ start_POSTSUPERSCRIPT italic_τ end_POSTSUPERSCRIPT italic_r start_POSTSUBSCRIPT italic_t + italic_τ end_POSTSUBSCRIPT with a discount factor γ = 0.997 𝛾 0.997 \gamma=0.997 italic_γ = 0.997 for each model state. To consider rewards beyond the prediction horizon T = 16 𝑇 16 T=16 italic_T = 16 , the critic learns to approximate the distribution of returns 28 for each state under the current actor behavior:

Actor: a t ∼ π θ ⁢ ( a t | s t ) Critic: v ψ ⁢ ( R t | s t ) missing-subexpression Actor: missing-subexpression similar-to subscript 𝑎 𝑡 subscript 𝜋 𝜃 conditional subscript 𝑎 𝑡 subscript 𝑠 𝑡 Critic: missing-subexpression subscript 𝑣 𝜓 conditional subscript 𝑅 𝑡 subscript 𝑠 𝑡 \displaystyle\begin{aligned} &\text{Actor:}\quad&&a_{t}\sim\pi_{\theta}(a_{t}% \;|\;s_{t})\qquad&\text{Critic:}\quad&&v_{\psi}(R_{t}\;|\;s_{t})\end{aligned} start_ROW start_CELL end_CELL start_CELL Actor: end_CELL start_CELL end_CELL start_CELL italic_a start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ∼ italic_π start_POSTSUBSCRIPT italic_θ end_POSTSUBSCRIPT ( italic_a start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | italic_s start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) end_CELL start_CELL Critic: end_CELL start_CELL end_CELL start_CELL italic_v start_POSTSUBSCRIPT italic_ψ end_POSTSUBSCRIPT ( italic_R start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | italic_s start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) end_CELL end_ROW (4) Starting from representations of replayed inputs, the world model and actor generate a trajectory of imagined model states s 1 : T subscript 𝑠 : 1 𝑇 s_{1:T} italic_s start_POSTSUBSCRIPT 1 : italic_T end_POSTSUBSCRIPT , actions a 1 : T subscript 𝑎 : 1 𝑇 a_{1:T} italic_a start_POSTSUBSCRIPT 1 : italic_T end_POSTSUBSCRIPT , rewards r 1 : T subscript 𝑟 : 1 𝑇 r_{1:T} italic_r start_POSTSUBSCRIPT 1 : italic_T end_POSTSUBSCRIPT , and continuation flags c 1 : T subscript 𝑐 : 1 𝑇 c_{1:T} italic_c start_POSTSUBSCRIPT 1 : italic_T end_POSTSUBSCRIPT . Because the critic predicts a distribution, we read out its predicted values v t ≐ E [ v ψ ( ⋅ | s t ) ] v_{t}\doteq\operatorname{E}[v_{\psi}(\,\cdot\;|\;s_{t})] italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ≐ roman_E [ italic_v start_POSTSUBSCRIPT italic_ψ end_POSTSUBSCRIPT ( ⋅ | italic_s start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) ] as the expectation of the distribution. To estimate returns that consider rewards beyond the prediction horizon, we compute bootstrapped λ 𝜆 \lambda italic_λ -returns 29 that integrate the predicted rewards and the values. The critic learns to predict the distribution of the return estimates R t λ subscript superscript 𝑅 𝜆 𝑡 R^{\lambda}_{t} italic_R start_POSTSUPERSCRIPT italic_λ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT using the maximum likelihood loss:

ℒ ⁢ ( ψ ) ≐ − ∑ t = 1 T ln ⁡ p ψ ⁢ ( R t λ | s t ) R t λ ≐ r t + γ ⁢ c t ⁢ ( ( 1 − λ ) ⁢ v t + λ ⁢ R t + 1 λ ) R T λ ≐ v T formulae-sequence approaches-limit ℒ 𝜓 superscript subscript 𝑡 1 𝑇 subscript 𝑝 𝜓 conditional subscript superscript 𝑅 𝜆 𝑡 subscript 𝑠 𝑡 formulae-sequence approaches-limit subscript superscript 𝑅 𝜆 𝑡 subscript 𝑟 𝑡 𝛾 subscript 𝑐 𝑡 1 𝜆 subscript 𝑣 𝑡 𝜆 subscript superscript 𝑅 𝜆 𝑡 1 approaches-limit subscript superscript 𝑅 𝜆 𝑇 subscript 𝑣 𝑇 \displaystyle\begin{aligned} \mathcal{L}(\psi)\doteq-\textstyle\sum_{t=1}^{T}% \ln p_{\psi}(R^{\lambda}_{t}\;|\;s_{t})\qquad R^{\lambda}_{t}\doteq r_{t}+% \gamma c_{t}\Big{(}(1-\lambda)v_{t}+\lambda R^{\lambda}_{t+1}\Big{)}\qquad R^{% \lambda}_{T}\doteq v_{T}\end{aligned} start_ROW start_CELL caligraphic_L ( italic_ψ ) ≐ - ∑ start_POSTSUBSCRIPT italic_t = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_T end_POSTSUPERSCRIPT roman_ln italic_p start_POSTSUBSCRIPT italic_ψ end_POSTSUBSCRIPT ( italic_R start_POSTSUPERSCRIPT italic_λ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | italic_s start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) italic_R start_POSTSUPERSCRIPT italic_λ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ≐ italic_r start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT + italic_γ italic_c start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ( ( 1 - italic_λ ) italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT + italic_λ italic_R start_POSTSUPERSCRIPT italic_λ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t + 1 end_POSTSUBSCRIPT ) italic_R start_POSTSUPERSCRIPT italic_λ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT ≐ italic_v start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT end_CELL end_ROW (5) While a simple choice would be to parameterize the critic as a Normal distribution, the return distribution can have multiple modes and vary by orders of magnitude across environments. To stabilize and accelerate learning under these conditions, we parameterize the critic as categorical distribution with exponentially spaced bins, decoupling the scale of gradients from the prediction targets as described later. To improve value prediction in environments where rewards are challenging to predict, we apply the critic loss both to imagined trajectories with loss scale β val = 1 subscript 𝛽 val 1 \beta_{\mathrm{val}}=1 italic_β start_POSTSUBSCRIPT roman_val end_POSTSUBSCRIPT = 1 and to trajectories sampled from the replay buffer with loss scale β repval = 0.3 subscript 𝛽 repval 0.3 \beta_{\mathrm{repval}}=0.3 italic_β start_POSTSUBSCRIPT roman_repval end_POSTSUBSCRIPT = 0.3 . The critic replay loss uses the imagination returns R t λ subscript superscript 𝑅 𝜆 𝑡 R^{\lambda}_{t} italic_R start_POSTSUPERSCRIPT italic_λ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT at the start states of the imagination rollouts as on-policy value annotations for the replay trajectory to then compute λ 𝜆 \lambda italic_λ -returns over the replay rewards.

Because the critic regresses targets that depend on its own predictions, we stabilize learning by regularizing the critic towards predicting the outputs of an exponentially moving average of its own parameters. This is similar to target networks used previously in reinforcement learning 7 but allows us to compute returns using the current critic network. We further noticed that the randomly initialized reward predictor and critic networks at the start of training can result in large predicted rewards that can delay the onset of learning. We thus initialize the output weight matrix of the reward predictor and critic to zeros, which alleviates the problem and accelerates early learning.

### Actor learning

The actor learns to choose actions that maximize return while exploring through an entropy regularizer 30 . However, the correct scale for this regularizer depends both on the scale and frequency of rewards in the environment. Ideally, we would like the agent to explore more if rewards are sparse and exploit more if rewards are dense or nearby. At the same time, the exploration amount should not be influenced by arbitrary scaling of rewards in the environment. This requires normalizing the return scale while preserving information about reward frequency.

To use a fixed entropy scale of η = 3 × 10 − 4 𝜂 3 superscript 10 4 \eta=3\times 10^{-4} italic_η = 3 × 10 start_POSTSUPERSCRIPT - 4 end_POSTSUPERSCRIPT across domains, we normalize returns to be approximately contained in the interval [ 0 , 1 ] 0 1 [0,1] [ 0 , 1 ] . In practice, substracting an offset from the returns does not change the actor gradient and thus dividing by the range S 𝑆 S italic_S is sufficient. Moreover, to avoid amplifying noise from function approximation under sparse rewards, we only scale down large return magnitudes but leave small returns below the threshold of L = 1 𝐿 1 L=1 italic_L = 1 untouched. We use the Reinforce estimator 31 for both discrete and continuous actions, resulting in the surrogate loss function: