Towards Understanding Robustness and Generalization in World Models

Thank you for your feedback and questions!

A1: “What are the experimental settings in Table 1? … “

We clarify that Table 1 is extracted from experimental studies presented in the experiment section in Section 5 with details presented 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).

A2: “There are various types of errors in latent dynamics models. In my opinion, the most crucial is the error in the transition models. However, this error is not addressed in equations (5-6). Could the authors provide insights into why they focus solely on representation errors?”

Thank you for highlighting this perspective. In this work, we focus on latent representation errors because they are unique to world models that learn a low-dimensional latent state representation, unlike conventional model-based RL approaches. We also want to emphasize that the most significant impact—whether from latent representation errors or transition model errors—often arises from error accumulation over time rather than from the errors themselves at each individual step. During both training and predictive rollouts, even small errors can compound significantly across the trajectory of the dynamics model's predictions.

Our main results can, in fact, be extended to address inaccuracies from the transition model's predictions. Specifically, to handle additional errors from the transition model, one could generalize our framework by incorporating these errors into Theorem 3.7 and Corollary 4.2. 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)$. We are happy to elaborate further if needed.

A3: “What is the ground truth probability measure P introduced in Line 232?”

We assume that P is a probability measure supported on Z with its Radon-Nikodym derivative p w.r.t $\mu_\mathcal{Z}$. In practice, P is typically an easy-to-sample distribution, such as a uniform distribution, which is determined by the specific choice of encoder-decoder architecture (Lines 232-235).

“In Dreamer, latent representations are learned using a per-sample reconstruction loss of states and a prediction loss. How do these implementation losses relate to the two minimizations of distribution distance mentioned in Line 238?”

As discussed in Lines 242-246, we use Wasserstein-1 distance to measure the discrepancy between probability measures. Specifically, for the Dreamer-type loss function that uses KL-divergence, we note that the squared W1 distance between two probability measures can be upper bounded by their KL-divergence up to a constant (Gibbs and Su, 2002). This implies that minimizing KL-divergence in the model could reasonably be expected to also reduce the W1 distance.

“Additionally, how does the distance relate to the scale of errors?”

The Wasserstein-1 distance measures the discrepancy between the target distribution and the represented distribution, which indicates the scale of approximation error in latent representation.

A4: “How is Corollary 4.2 regarding Q functions in Section 4 defined and derived? There do not appear to be any reward variables in the SDE formulation.”

Corollary 4.2 pertains to the currently learned Q function in policy learning phase in world models, which does not need to be explicitly defined in terms of reward variables. The corollary quantifies the impact of latent representation errors on this Q function during the inference phase of the latent dynamics model. Specifically, it shows that latent representation errors can implicitly encourage exploration of unseen states by introducing stochastic perturbations into the value function.

A5: “How is the Jacobian regularization term in equation (20) implemented in the experiments? I could not find implementation details in the Appendix.”

As mentioned in Lines 1574-1575, we use a computationally efficient approximation of the Jacobian norm for the sequence model. Specifically, we follow the method described in Hoffman et al. (2019), employing a single projection to approximate the Jacobian regularization term.

Thank you for your feedback and suggestions!

A1: Readability of the paper

We respectfully disagree with the statement. In fact, we made deliberate efforts to introduce and contextualize these concepts early in the text. Specifically, we describe LREs as errors incurred by latent encoders (in Line 71), and we explain their influence on robustness and generalization (Lines 48-49: "... investigate how latent representation errors can enhance robustness against perturbations, which in turn often improves generalization"). We also outline how LREs differ from conventional noise injection schemes (Lines 71-87), emphasizing their inherent nature and the challenges of error propagation during rollouts. Furthermore, the introduction establishes the relevance of zero/non-zero drift cases and the role of Jacobian regularization in mitigating the adverse effects of biased latent errors (Lines 99-103: For the case where latent representation errors exhibit non-zero drifts, … propose to add Jacobian regularization to mitigate these effects.").

Nonetheless, we are open to further clarifying these points to enhance readability.

A2: “Secondly, while the results look promising, they are not very extensive. Does the regularization give performance gains in other mujoco environments too? How does the performance compare to other ways of regularizing the dynamics model, like Repo [1] or RPC [2], or just L2 regularization on dynamics weights?”

We emphasize that our work focuses on deepening the theoretical understanding of the generalization capabilities of world models—an area previously observed in empirical studies but lacking systematic theoretical analysis. Our experiments are specifically designed to support and validate the theoretical insights presented. For example, our "Robustness against Encoder Errors" experiment directly corroborates Corollary 3.8, which demonstrates how the model’s Jacobian norm can control the impact of errors on the loss $\mathcal{L}$. We acknowledge that comparing our method with other regularization techniques, such as L2 regularization, Repo, or RPC, would be valuable for a future work with a broader empirical scope. We want to respectfully emphasize that the primary goal of this work is to advance the theoretical understanding of latent representation error within the specific context of world models, rather than to conduct a broad empirical comparison.

Thank you for your feedback!

A1: “The presentation is quite dense and difficult to follow, with repetitive statements and inadequate motivation for why this analysis approach is needed.”

We respectfully disagree with this statement. (As a matter of fact, all three NeurIPS reviewers think this presentation is easy to follow.) We intentionally structured the introduction to provide a comprehensive and layered explanation of our approach, aiming to make our contributions clear and accessible to a broad audience. Specifically, we motivate our analysis approach by emphasizing the practical importance of understanding the robustness and generalization capabilities of world models (Lines 31-45) and explaining how our work addresses open questions in this area (Lines 47-70). Additionally, we succinctly summarize our main contributions in bullet points (Lines 94-111) to ensure clarity and to highlight the novelty and relevance of our work.

That said, we greatly value the reviewer's perspective and try best to improve the clarity of our presentation. We would sincerely appreciate more specific suggestions or examples of sections that feel overly dense or repetitive, or guidance on how the motivation could be more effectively articulated.

A2: “The assumption of requiring C3 functions with bounded Lipschitz derivatives for theoretical analysis can be very strong and potentially unrealistic in practice.”

We respectfully note that assuming $\mathcal{C}^3$ with bounded Lipschitz derivatives is a common practice in theoretical analyses of general neural networks (see, for example, Assumption A in [1]). This assumption encompasses widely used activation functions, such as tanh, which is frequently employed in RNNs for world models. While the assumption does exclude certain unbounded activation functions, those functions are prone to issues like exploding gradients and are therefore less commonly used in practical settings.

A3: “Empirical evaluation is limited to only Walker and Quadruped environments, leading to imcomprehensive evaluations.”

We emphasize that our work focuses on deepening the theoretical understanding of the generalization capabilities of world models—an area previously observed in empirical studies but lacking systematic theoretical analysis. Our experiments are specifically designed to support and validate the theoretical insights presented. For example, our "Robustness against Encoder Errors" experiment directly corroborates Corollary 3.8, which demonstrates how the model’s Jacobian norm can control the impact of errors on the loss $\mathcal{L}$.

A4: “The hyperparameter sensitivity for the regularization weight, $\lambda$, is expected to significantly affect the resulting model performance, but its sensitivity is not discussed (either theoretically or empirically) in the manuscript.”

We conducted experiments with \lambda values of 0.01, 0.05, and 0.1, and the results (in Appendix 6) indicate that the model's performance remained relatively stable across this range. Specifically, we did not observe significant variations in performance metrics, suggesting that our method is robust to moderate changes in \lambda. We are open to expanding our discussion of $\lambda$ sensitivity, including both theoretical considerations and additional empirical studies, in the future work.

A5: “There is a lack of comparison with existing approaches on addressing robustness and generalisability in world models.”

To the best of our knowledge, our work is the first to provide a theoretical understanding of the robustness and generalizability of world models. We discuss how our approach differs from existing research on robustness and generalizability in deep RL (Lines 124-133), specifically focusing on the unique challenges posed by latent representation errors.in world models.

That said, if the reviewer could provide specific related theoretical work in world models on robustness and generalizability, we would be more than happy to include additional comparisons.

[1] Soon Hoe Lim, N Benjamin Erichson, Liam Hodgkinson, and Michael W Mahoney. Noisy recurrent neural networks. Advances in Neural Information Processing Systems, 34:5124–5137, 2021.

Thank you for your feedback and questions!

A1. Interpretation and Presentation of Theoretical Results

• “The theorems are weakly interpreted in the discussion sections. Specifically, the primary variables and parameters within the theorems are not sufficiently explained, making it challenging for readers to fully understand the implications of the results.”

We respectfully clarify that we have already provided detailed interpretations of the theoretical results throughout the paper. For instance:

Theorem 3.6: “Theorem 3.6 indicates that with an appropriate CNN configuration, the W1 approximation error can be made to reside in a small region.. this allows us to apply the perturbation analysis of the dynamical system defined in Equations (5 - 8) in the following sections.” (Line 266 - 269)

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… ng… this insight resonates with the empirical results in Table 1 that model’s robustness gain is not significant when the error induced by large batch sizes is too small.” (Line 320 - 327)

Corollary 3.8: “This observation also aligns with the theoretical insights in Lim et al. (2021) that the introduction of Brownian motion.. that could modulate the new bias term R˜ for stabilized implicit regularization.” (Line 355 - 367)

Theorem 4.1: “Theorem 4.1 correlates with the empirical findings … the Jacobian norms within the latent dynamics model and the horizon length.” (406 - 409)

Corollary 4.2: “ This corollary reveals that latent representation errors implicitly encourage exploration of unseen states … Intuitively, the stochasticity in the LDM also encourages greater exploration compared to its deterministic counterparts.” (Line 424 - 427)

• “In Theorem 3.6, the authors suggest that it can be assumed to be small, but they do not provide details on the relevant parameters or the conditions under which this assumption is valid.”

We respectfully note that Theorem 3.6 proves (instead of assume) that the W1 approximation error can indeed be confined to a small region, under an appropriate CNN configuration of the encoder-decoder architecture, rather than an arbitrary assumption of small error.

We are happy to include further details to improve accessibility and are open to specific suggestions from the reviewer.

A2. Experimental comparison to other robustness techniques

We note that our comparison with data augmentation methods is detailed in Appendix D.6, due to space limitations in the main text. In a nutshell, our experiments show that models trained with Jacobian regularization outperform those using augmentation techniques, especially when confronted with perturbations not seen during augmentation. While state augmentation is effective for inference perturbations similar to those used in training, it would struggle to generalize to novel perturbations. In contrast, Jacobian regularization offers robustness that is less dependent on the diversity of augmented samples..

A3. On expanding experiments to more experiments

We emphasize that our work focuses on deepening the theoretical understanding of the generalization capabilities of world models—an area previously observed in empirical studies but lacking systematic theoretical analysis. Our experiments are specifically designed to support and validate the theoretical insights presented so in this way experimental studies are complementary. For example, our "Robustness against Encoder Errors" experiment directly corroborates Corollary 3.8, which demonstrates how the model’s Jacobian norm can control the impact of errors on the loss $\mathcal{L}$.

A4. Responses to Questions

Q1: Connection of results in table (1) with main theoretical results

Our somewhat surprising theoretical results indicate that moderate latent representation errors can enhance the world model’s ability to handle perturbations, thereby improving both exploration and generalization. This behavior is analogous to the effect of gradient estimation errors observed during batch training, where a certain level of noise can act as implicit regularization.

Specifically, we note that when the error $\bar{\sigma}_t(\cdot)$ is too small, the impact of the term S as a form of implicit regularization becomes less pronounced. Intuitively, this aligns with the empirical observations in Table 1, where the robustness gains are limited when errors induced by large batch sizes are insufficient to drive significant regularization effects. We further discuss this connection in Lines 324-327.

Q2: “In (4), is the decoder intentionally conditioned on $\tilde{z}_t$?”

The decoder is actually conditioned on $z_t$.Thank you for pointing out this typo.

A1. Interpretation and Presentation of Theoretical Results

We respectfully clarify that we have already provided detailed interpretations of the theoretical results throughout the paper. For instance: ...

Thank you for sharing examples of your interpretations. However, the reviewer has carefully reviewed these sections of the paper, which led to the use of the term "weak interpretation." The reviewer has also provided a specific example to clarify what is meant by "weak interpretation".

We respectfully note that Theorem 3.6 proves (instead of assume) that the W1 approximation error can indeed be confined to a small region, under an appropriate CNN configuration of the encoder-decoder architecture, rather than an arbitrary assumption of small error.

It is recommended that the authors carefully revisit both their paper and the review comments to ensure clarity in their responses. The reviewer's wording aligns directly with the authors' explanation of Theorem 3.6, particularly the statement on line 268: "In particular, this result indicates that the error magnitude ε in SDE (5) can be assumed to be small."

A2. Experimental comparison to other robustness techniques

A3. On expanding experiments to more experiments

We note that our comparison with data augmentation methods is detailed in Appendix D.6, due to space limitations in the main text. In a nutshell, our experiments show that models trained with Jacobian regularization outperform those using augmentation techniques, especially when confronted with perturbations not seen during augmentation. While state augmentation is effective for inference perturbations similar to those used in training, it would struggle to generalize to novel perturbations. In contrast, Jacobian regularization offers robustness that is less dependent on the diversity of augmented samples..

We emphasize that our work focuses on deepening the theoretical understanding of the generalization capabilities of world models—an area previously observed in empirical studies but lacking systematic theoretical analysis. Our experiments are specifically designed to support and validate the theoretical insights presented so in this way experimental studies are complementary. For example, our "Robustness against Encoder Errors" experiment directly corroborates Corollary 3.8, which demonstrates how the model’s Jacobian norm can control the impact of errors on the loss .

The paper claims to propose a Jacobian regularization technique, supported by extensive experimental studies to demonstrate improvements in stability and robustness. However, without additional comparisons to similar baselines and experiments across more environments, it becomes challenging to substantiate this claim. This is particularly critical given that two of the main contributions emphasize the practical benefits of Jacobian regularization. Furthermore, while Jacobian regularization is presented as a main contribution, the idea was previously introduced in another cited work.

Thank you for further clarification on the questions!

We greatly appreciate your prompt response to our rebuttal. We appreciate this further opportunity that allows us to improve the clarity of our work.

A1. Interpretation and Presentation of Theoretical Results

Thank you for pointing out the statements where our interpretation of theoretical results could be improved. To clarify, Theorem 3.6 proves (rather than assumes) that W1 approximation error can indeed be confined to a small region, given an appropriate CNN configuration for the encoder-decoder architecture. This result provides a justification for treating the perturbation as small as needed. In the revised version, we incorporate this clarification, particularly on Line 268 (now on Line 270-272, highlighted in red), to better connect the theorem to its implications for the subsequent perturbation analysis.

A3. On expanding experiments to more experiments

Thank you for your detailed observations regarding the positioning of Jacobian regularization in our work. We would like to clarify that the primary contribution emphasized in our paper is the theoretical understanding of the impacts of non-zero drift error on instability in world models, which shed light on how the well-known Jacobian regularization technique can be employed to mitigate this issue. We did not claim Jacobian regularization as a novel technique but rather as a theoretically grounded approach to address the challenges we have identified. To avoid any possible misinterpretation, we revised the draft to make this distinction more explicit in the introduction (on Line 102-104, Line 106-109, highlighted in red). We hope this has addressed your concern.

Once again, we greatly appreciate your suggestions and feedback on our work. We are happy to address any further concerns.