Towards Unraveling and Improving Generalization in World Models

Q5: — “An alternate method to improve robustness to perturbations is by training the model/value function on different augmentations of the states”

We agree that training the model/value function on different augmentations of the states is a valuable approach to improving robustness to perturbations. Following your suggestion, we added a new set of data augmentation experiments for comparison. We considered training with state images augmented with randomly-masked Gaussian noises. We evaluated it against 3 different perturbations: 1) modifying gravity constant g in the DMC walker environment 2) adding rotation to input image state 3) adding masked Gaussian noise to input image state.

The experiment results show that under varied gravity g and rotation perturbation, models trained with Jacobian regularization outperform state augmentation. Under masked Gaussian noise perturbation, Jacobian regularization outperforms state augmentation when the noise is more dense (mask with $\beta=70, 60$) and underperforms when $\beta=40, 50$. This may be due to the fact that state augmentation is trained with very similar Gaussian noise. State augmentation works well when the augmentation is consistent with the perturbation in inference, but does not generalize to unseen perturbations. Please see below table and Figure 3 in attached PDF for details.

The limitations of the rebuttal period prevent us from conducting a comprehensive empirical comparison at this time (e.g. more augmentation patterns), but we are excited to continue this line of research in our future works.

— “on the pros/cons of using the Jacobian regularization term over such data augmentation methods.”

Regarding the pros and cons of using the Jacobian regularization over data augmentation, we note the following:

Pros:

• Theoretical guarantees: Jacobian regularization is a principled approach grounded in theoretical results (Theorem 3.7 and Corollary 3.8), providing explicit control over model behavior in response to small error.

• Less reliant on data diversity: Unlike data augmentation, which heavily relies on diversity and relevance of augmented samples, Jacobian regularization targets the learning dynamics of WM.

• Less likelihood of overfitting: Jacobian regularization is less prone to overfitting as unlike data augmentation, which can lead to overfitting when overly reliant on certain perturbation patterns rather than general robustness.

• Bounding trajectory divergence: Jacobian regularization also mitigates error propagation during predictive rollouts (Theorem 4.1), a benefit not achieved by data augmentation.

Cons:

• Uncertainties with large error: The theoretical analysis for Jacobian regularization assumes the latent representation error is small. In cases when the encoder remains poorly learned, data augmentation may provide better model robustness.

• Computational overhead: computing Jacobian terms can introduce additional overhead

Q6: On the computational overhead of Jacobian regularization term

For every episode with 500 steps on an A100, training the model with Jacobian regularization took 28.702 seconds compared to 22.199 seconds for training without it (averaged over 20 episodes).

A7. Response to Limitation

We appreciate you raising this point. In the revision, with more page space, we will include a separate section addressing the limitations of our work. We will highlight that the main limitation of the SDE interpretation of world models (WM) is its restriction to a popular family of world models for theoretical soundness. As mentioned in the paper (lines 123-125), "we consider a popular class of world models, including Dreamer and PlaNet, where $\{z, \tilde{z}, \tilde{s}\}$ have distributions parameterized by neural networks’ outputs, and are Gaussian when the outputs are known."

Additionally, our results on the approximation error of latent representation focus on CNN (and similar) models, and future work is needed to generalize the study to models such as transformers. We are exploring this direction in our future work.

A5. On the clarification of experiment settings

We apologize for the confusion and will clarify further in the experiment section in the appendix. For the batch-size versus robustness experiment in Table 1, the considered perturbation methods are the same as those in the experiment on Jacobian regularization with perturbed states (Table 2 and Appendix D.2).

Regarding to your specific questions:

—”With regards to the perturbations - are they applied to every state in the trajectory?”

Yes, the perturbations are applied at every time step. This consistency is maintained, particularly in the robustness experiments with encoder noise, to align with our theoretical results, which studied continuously introduced encoder noise.

—”For masking, is the same mask used for every state, or is the mask also sampled randomly”

The masks are sampled randomly for each state, not the same mask for every state.

—”With regards to injecting encoder error - how to interpret the $\mu_t$ and $\sigma_t$ values?”

We consider a Gaussian noise process where at each time step, where $\mu_t$ amd $\sigma_t$ denotes mean and variance of the Gaussian process at time $t$ . An explanation is provided in the appendix (see pages 28, line 756, and pages 30, lines 764-765).

In the revised version, with more page space available, we will include further clarification of the experimental settings in Section 5 from the main text to avoid any confusion.

A6. Responses to Questions

Q1: On the gradient estimation errors in Table 1

Thank you for your insightful question. We used batch size experiments as a motivating example to introduce the relationship between error and generalization robustness. Batch-induced gradient estimation errors are relatively well-studied, whereas latent representation error is often overlooked. Our theoretical results on the effects of latent representation error interestingly align with the discovered empirical effects on gradient estimation error in the zero-drift case:

• Zero-mean error: Zero-mean errors, such as gradient estimation error and zero-drift latent regularization errors, have the potential for generalization improvement.
• Controlling the error: To harness these generalization gains, the batch size can influence gradient estimation error. When kept within a controlled range, this can enhance generalization—paralleling the control of latent representation error through Jacobian regularization.

In the revision, we will further clarify these points to avoid any confusion and provide a clearer explanation of the experimental settings of Table 1 in Section D in appendix.

Q3: Interpreting results for non-convex case

Thank you for the interesting question.

In the non-convex case, our most important finding from Corollary 3.8 remains relevant: the additional bias induced by the drift term of latent representation error can be controlled through the model’s Jacobian norm. While non-convex optimization presents more challenges due to the complicated gradient landscape, this control mechanism provides a valuable tool.

For further insights, non-convex optimization remains an open challenge, and we acknowledge the significant complexity of its gradient landscape. We are excited to consider future work that focuses on specific types of non-convex loss functions to extract more detailed insights and refine our understanding of the regularizing effects in these scenarios.

Q4: Extending the analysis on TD-MPC

For TD-MPC, which has a deterministic latent dynamics model, our SDE analysis can specialize to TD-MPC by setting all diffusion coefficient functions to zero in formulation equations (5) - (8), thus becoming systems of ODEs, provided that the learned MLP model in TD-MPC satisfies certain continuity assumptions to guarantee the unique existence of solutions.

Regarding the effects of latent representation error in TD-MPC, this would resemble a simplified version of Corollary 3.8 without diffusion terms. The result would similarly introduce a bias to the loss function, with its magnitude dependent on the error level. This effect can be modulated by the MLP’s input-output Jacobian norm. We are excited about future work to develop the theoretical details and experimental results in this context.

Thank you very much for your helpful feedback on our work! We respond to your questions as follow:

A1. On the analysis of compounding model error

Thank you for highlighting this important aspect. We agree that compounding model error is critical in world model analysis. During both training and predictive rollouts, even if the latent representation error is small at each time step, its accumulation through the dynamics model's predictions throughout the trajectory is significant. Our results show that in the training phase, the accumulated effects of zero-drift errors can act as implicit regularization (which is somewhat surprising but consistent with noise injection technique used in deep learning [1, 2]), while non-zero drift errors require additional regularization. During predictive rollouts, accumulated errors can lead to divergence if not properly regularized. Importantly, the degree of these effects depends on the dynamics model’s Jacobian. Hence, we propose Jacobian regularization.

Our theoretical results explicitly consider these compounding effects by interpreting world models as SDE dynamical systems to characterize error propagation. For instance, in Theorem 3.7 and Corollary 3.8, the derived regularization terms $\mathcal{R}$ and bias term $\tilde{\mathcal{R}}$ account for the accumulative effects of continuously introduced encoder errors on the sequence model $f$ and transition predictor $p$. This means that the effects of encoder errors from the initial time step to the entire trajectory of state predictions are captured in $\mathcal{R}$ and $\tilde{\mathcal{R}}$.

Regarding the compounding model error due to additional inaccuracies in the latent/state dynamics model predictions, our main results can be generalized to include errors from other model components, such as transition predictor error. For example, one could incorporate error terms $\sigma$ as $(\sigma_{\text{enc}}, 0, \sigma_{\text{pred}}, 0)$ and $\bar{\sigma}$ as $(\bar{\sigma}\_{\text{enc}}, 0, \bar{\sigma}\_{\text{pred}}, 0)$ in Theorem 3.7 and Corollary 4.2.

1. Alexander Camuto et al., "Explicit Regularisation in Gaussian Noise Injections," Advances in Neural Information Processing Systems 34 (NeurIPS 2020), 2020.

2. Soon Hoe Lim et al., "Noisy Recurrent Neural Networks," Proceedings of the 35th Conference on Neural Information Processing Systems (NeurIPS 2021), 2021.

A2. On the drift and diffusion components of error (and Q2)

Thank you for the question. In our case, the decomposition of latent representation error into drift and diffusion terms arises naturally as additive noise to the encoder, which is interpreted as an SDE process with drift and diffusion coefficient functions matching the stochastic design of the world model. In general, since the error process is a stochastic term, it is natural to consider its mean (drift) and variance (diffusion) decomposition, a common approach in stochastic control (see [3] for more details).

The case of zero-drift noise occurs when the learned encoder is unbiased but has variance. In contrast, the case of non-zero-drift noise corresponds to a more general situation where the learned encoder is biased and has variance. These are general settings that occur in world model learning.

3. Ramon van Handel, "Stochastic Calculus, Filtering, and Stochastic Control: Lecture Notes," Spring 2007.

A3. On the interpretation of latent representation error encouraging state exploration

Thank you for your insightful comment. We respectfully disagree with the reviewer’s assertion that erroneous latent states lead the agent to explore invalid states outside the state space of the MDP. We note that since the agent’s policy in the world model is conditioned on latent states, erroneous latent states $\tilde{z}$ lead the agent to take inaccurate actions $\tilde{a}$. The gap in the predicted value function due to latent representation error is captured in Corollary 4.2.

From the perspective of the task environment, the agent begins at a valid state $s$ and takes the inaccurate action $\tilde{a}$, but will still transition to a valid next state $\mathcal{T}(s, \tilde{a})$. This suboptimality can be interpreted as implicitly encouraging exploration by introducing a stochastic perturbation in the value function. Therefore, the agent remains within the valid state space while experiencing robustness similar to noise injection, as you suggested.

A4. On including low-level intuition for each result and overall readability

We appreciate your helpful suggestion. We will include a notation table in the appendix for easy reference. We will also provide more intuitive explanations for our theoretical results, expanding on the current ones, including

Theorem 3.7: “The presence of this Hessian-dependent term S, under latent representation error, implies a tendency towards wider minima in the loss landscape…” (page 6, 234-235)

Corollary 3.8: “The presence of $\tilde{\zeta}$ in ,$\tilde{Q}$ and $\tilde{S}$ induces a bias to the loss function with its magnitude dependent on the error level $\epsilon$, since $\tilde{zeta}$ is a non-zero term influenced on the drift term $\sigma$.” (page 7, 260-262)

Theorem 4.1: “... the expected divergence from error accumulation hinges on the expected error magnitude, the Jacobian norms within the latent dynamics model and the horizon length T.” (page 7, 300-302)

We believe these additions will greatly improve the readability and accessibility of our results.

Rebuttal will continue in the comments.

Thank you very much for your helpful comment!

• “The unseen images are produced via global/partial Gaussian noises and rotation, which seems more on the robustness side rather than the generalization of unseen images;”

Following your valuable feedback, we added a new set of experiments involving unseen dynamics to evaluate the robust generalization of our proposed Jacobian regularization. For Mujoco tasks, we varied the acceleration constants due to gravity (default being 9.8 m/s²). This tests the learned agent’s model and the latent dynamics model’s generalization capacity to unseen dynamics. Our results indeed validate that Jacobian regularization (both $\lambda = 0.01$ and $0.05$) outperforms the baseline. We report the mean and variants of eval returns across 5 eval runs in the following table. Please also see Figure 3 in the attached PDF for a plot of varied g vs eval returns.

We also acknowledge the challenges of setting up experiments for completely unseen/unrelated states and transitions in RL due to the complexity of defining such scenarios comprehensively. We considered local/global Gaussian noise, rotation, and varying acceleration constants due to gravity as perturbations. These were designed to best validate our theoretical findings, specifically Theorem 3.7 and Corollary 3.8, on implicit regularization, which link to favoring model solutions in regions of the loss landscape with improved generalization and robustness. Our approach is consistent with the literature's understanding of robust generalization [1, 2].

In the revised experiment section, we will clarify the distinction between robust generalization and more generic generalization to avoid any possible confusion.

1. Binghui Li et al., "Why Robust Generalization in Deep Learning is Difficult: Perspective of Expressive Power," Proceedings of the 36th Conference on Neural Information Processing Systems (NeurIPS 2022), 2022.

2. Soon Hoe Lim et al., "Noisy Recurrent Neural Networks," Proceedings of the 35th Conference on Neural Information Processing Systems (NeurIPS 2021), 2021.

Thank you for your positive feedback and suggestions! Based on your valuable feedback, we have improved the empirical analysis of our proposed regularization schemes with some new experiments and visualizations.

A1. Additional Experiments on Benchmark Environment

Thanks for the constructive suggestion. We intend to add 1-2 additional experiment environments to showcase robustness against perturbation brought by the Jacobian regularization in the final paper. Here we include a comparison of baseline method and Jacobian regularization in the challenging 2D environment -- Crafter [1] under noise perturbation. We add a Gaussian noise with 0 mean and 0.25 variance on the image state whose value ranges from -0.5 to 0.5 and apply masks with various percentages.

We would like to include a note that within the very limited rebuttal timeframe, the baseline method finished fewer steps than its Jacobian counterparts, thus earlier checkpoints are used for baseline in this comparison. We observed that the baseline performance has already plateaued according to the train/eval curves, although its performance can still improve with more training steps. We will keep the baseline experiment running and update this table should there be any change in the performance.

[1]: Danijar Hafner, 'Benchmarking the Spectrum of Agent Capabilities,' 2021

A2. Visualization of Reconstructed State Trajectories of Models Trained with/without Jacobian Regularization

Thanks for your great suggestions. We have added visualizations of reconstructed state trajectory samples in the revision to showcase the error propagation of exogenous zero-drift and non-zero drift error signals in latent states with and without Jacobian regularization. Please see Figure 1 & 2 in the attached PDF file.

As shown in the figure, the reconstructed states for the baseline model without Jacobian regularization appear fuzy, indicating the model has not correctly captured the dynamics of the environment; whereas the reconstructed states for model with Jacobian regularization are sharp and correctly reflect the dynamics of the environment. The visual comparison highlights the robustness brought by Jacobian regularization against latent noises.

A3. Responses to Questions

Q1. On the effects of $\lambda$ on model’s generalization

Thanks for raising the question. The choice of $\lambda$ can have subtle influence on the model's performance. The optimal value for $\lambda$ depends on the environment and the task. We suggest readers to try out $\lambda$ in the following ranges [0.1, 0.01, 0.001]. Please see table in A1, where we present crafter scores for 3 different values of $\lambda$ under a Gaussian noise with 0 mean and 0.25 variance on the input image states (where pixel values range from -0.5 to 0.5) with various masks.

We will include a section to discuss $\lambda$'s influence for different environments in the final version of the paper.

Q2. On the relation between $\lambda$ and task horizon

We did not observe a substantial relationship between $\lambda$ and task horizon in our current experiments. However, we hypothesize that as task horizon gets longer, larger regularization weights ($\lambda$ around 0.1) would be more beneficial due to their tighter control on error propagation. Due to time limitations during the rebuttal period, we were unable to conduct these additional experiments, but we plan to study this further in future works.

A4. Response to Limitation

We appreciate you raising this point. In the revision, with more page space, we will include a separate section addressing the limitations of our work. We will highlight that the main limitation of the SDE interpretation of world models (WM) is its restriction to a popular family of world models for theoretical soundness. As mentioned in the paper (lines 123-125), "we consider a popular class of world models, including Dreamer and PlaNet, where $\{z, \tilde{z}, \tilde{s}\}$ have distributions parameterized by neural networks’ outputs, and are Gaussian when the outputs are known."

In addition, our results on the approximation error of latent representation focus on CNN (and similar) models and future work is needed to generalize the study to models such as transformers. We are exploring this direction in our future work.