MuDreamer: Learning Predictive World Models without Reconstruction

We would like to thank the reviewer for the feedback.

Please find below our answer to your comments.

1. About the reconstruction loss.

The reconstruction loss in equation 3 is optional and is only used to observe the latent dynamics imagined by the model. We explain in the above paragraph in section 4.1 that a "stop gradient" sg(.) operation is added to prevent the reconstruction loss gradients to be back-propagated to the rest of the network: "We optionally learn a decoder network to reconstruct the sequence of observations while using the stop gradient operator sg(.) to prevent the gradients from being back-propagated to other network parameters." The sg(.) operation is also shown in equation 3. The reconstruction loss is hence not required as the title says.

2. Paper novelty.

Inverse dynamics and continuation prediction loss are not contributions of our paper. We propose to learn hidden representations in a self-supervised manner using an action predictor network and a value predictor network. The action predictor predicts the sequence of action leading to the observed environment change while value predictor predicts the discounted sum of future rewards of the sampled trajectory. We also show that the use of batch normalization in the representation network is required to prevent a collapse of the model latent state to invariant features, where the standard deviation of the model categorical state logits converges to zero.

3. Training time.

The difference in training time is indeed relatively modest. The amount of parameters is also relatively small and MuDreamer requires less than 4GB of GPU memory during training. However, we anticipate MuDreamer to have a more significant impact on training time and memory requirement when scaling to larger encoder/decoder sizes.

We would like to thank the reviewer for the insightful feedback.

Please find below our answer to your questions.

1. Value predictors in this research are inspired by MuZero, and the inclusion of action prediction is a common practice in various model-based approaches. As a result, the novelty may be relatively constrained.

Our value predictor network is inspired by the MuZero value function. However, we didn't find similar works using an action predictor network to predict the sequence of preceding actions leading to the observed environment change. If this is the case, could you share with us literature papers using a similar method in model-based RL?

3. The experimental results are not satisfactory, as it appears that DreamerV3 performs better. I am aware that MuDreamer has fewer parameters, but it is important to analyze specifically where the differences lie. Please provide an analysis and remove the corresponding parts from DreamerV3 to assess the performance. Additionally, since the authors were inspired by MuZero in several aspects, it would be beneficial to compare this approach as well.

We evaluated MuDreamer on Atari100k in order to test our MuDreamer algorithm in another domain. While MuDreamer successfully solves DeepMind Control tasks, we observed that the agent was not able to achieve the same performance on Atari. This is especially the case for 'boxer' where MuDreamer fails to achieve good scores, significantly penalizing the final human normalized score. However, MuDreamer outperforms DreamerV3 in a subset of the games. This is why we categorized these early results as promising and plan to identify the reason of these failure in future works. One way would be to replace the value predictor lambda returns by n-step returns. Since the sampled batch includes trajectories with different policies, computing returns using a limited or dynamic number of TD-steps, depending on the training step the trajectory was generated, may provide better quality targets for representation learning.

4. Why is batch normalization used instead of other normalization techniques such as layer normalization? Can layer normalization achieve similar effects?

DreamerV3 uses layer normalization to stabilize learning and train larger models. In our case, Figure 7 and Table 3 show that using Layer Normalization create instabilities for some of the tasks. In some tasks like 'Walker Walk', layer normalization can be used instead of batch normalization. But this often create instabilities where the model categorical state logits become constant with an averaged std over all features near zero. This is a known issue in self-supervised learning for images and can be solved using batch normalization to regularize the output features.

5. Many explanations are not sufficiently in-depth. For example, in KL balancing, why does using a slight regularization of the representations toward the prior with $\beta_{rep}$ = 0.05 solve both of these issues?

The default regularization of $\beta_{rep}$ = 0.1 was too high to efficiently learn hidden representations. Using a $\beta_{rep}$ of 0.2, we observed that unnecessary information such as the environment floor and the agent shadow in DeepMind Control tasks were sometimes reconstructed as monochrome surfaces without the original details. Following BLAST, removing the representation loss $L_{rep}$ using a $\beta_{rep}$ of zero solved this problem. However, this resulted in instabilities with the prior latents sometimes diverging. We suppose this create difficulties for the prior to predict posterior representations since we do not use a slow moving teacher encoder network. The Representation network and Dynamics Predictor also do not share a common architecture while methods like SimSiam use a unique predictor network on top of the learned features. Using a smaller regularization of the representations toward the prior with $\beta_{rep}$=0.05 was sufficient to keep posterior and prior latents close to each other without impacting the learning of hidden representations.

We would like to thank the reviewer for the insightful feedback.

Please find below our answer to your questions.

1. Why is Figure 2 so good at reconstructing the observation?

Depending of the environment, some details of the observation like the background will not change or be very similar. These 'invariant' details do not need to be encoded inside the model latent state and can simply be learned by the decoder. In the other hand, you can see that details related to the agent body in Figure 2 are slightly more blurry than the original. These details progressively become less and less blurry along training, until becoming very close to the original image. While invariant background details are usually learned by the decoder in the early phase of training.

1.1 One would have expected to only capture what mattered for the task if the assumptions from the abstract/introduction were true?

Moving details that are not necessary for the agent are sometimes ignored and not reconstructed as the original. For instance, when setting ßrep to 0.1 and 0.2 or using less than 8 discrete latents for the model hidden state, we observed that unnecessary information such as the environment floor and the agent shadow in DeepMind Control were sometimes reconstructed as monochrome surfaces without the original details. A loss of performance was also observed. The model dynamics loss does not distinguish between relevant and irrelevant information. However, value and action prediction losses require encoding necessary information in the latent space, discarding the extra information if the model hidden state has not enough capacity.

1.2 Are there games where you have examples of “Dreamer failing to perceive crucial elements”, which MuDreamer does capture?

Dreamer sometimes fails to perceive small elements like the ball of the 'BallInCup' task in DeepMind Control. These details sometimes disappear and reappear at a different position during the imagination phase, which may harm behavior learning. In the case of BallInCup, this resulted in instabilities with training phases where the agent suddenly did not perceive the ball and was unable to correctly solve the task. We did not observe this kind of behavior for MuDreamer along training.

2. It would have been interesting to point to specific games where this effect should arise, and make a clear comparison between DreamerV3, MuDreamer and EfficientZero.

DreamerV3 performance is particularly impacted when visual distractions are present in the observation. In order to study the effect of visual distractions on DreamerV3 and MuDreamer performance, we experimented with the natural background setting, where the DeepMind Control tasks background is replaced with natural videos. You will find more details in section 'i' in the appendix of the revised paper. Figure 10 shows the comparison of MuDreamer with DreamerV3 under the natural background setting. MuDreamer successfully learn a policy while DreamerV3 fails on every task. Figure 11 shows the decoder reconstruction of observations by DreamerV3 and MuDreamer for the Walker Run and Finger Spin tasks. MuDreamer correctly reconstruct the agent body with a monochrome or blurry background while DreamerV3 focus on the background details, discarding the agent body and necessary information.

3. Did you explore using the Action predictor network directly for acting, instead of having another Actor network?

The actor predictor learns to predict the action that led to the observed environment change, so it cannot be used directly for planning. One way would be to generate the desired next model state to predict the action leading to this state. However, this seems to be a complex problem to solve compared to optimizing the actor and value networks.