Towards Unraveling and Improving Generalization in World Models

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 05:16:52

Abstract: World model has recently emerged as a promising approach to reinforcement learning (RL), as evidenced by its great successes that world model based agents exhibit state-of-the-art performance on a wide range visual control tasks. In this study, we aim to first obtain a clear understanding of the generalization capability of world models by examining the impact of _latent representation error_, and then devise new methods to enhance its generalization.  We hypothesize that latent representation errors  may paradoxically bring generalization to the model. We develop a continuous-time stochastic dynamics framework to quantify the impact of these errors, by examining the regularization effects for both cases with zero-drift representation errors and  non-zero-drift representation  errors. We propose a Jacobian regularization scheme to mitigate the "destabilizing'' effects of non-zero drift errors, thereby enhancing training stability and model generalization. Our empirical results confirm that this regularization approach not only stabilizes training but also accelerates convergence and improves performance on long-horizon prediction.

Main paper: 
 introduction: 
 Model-based reinforcement learning (RL) has emerged as a promising learning paradigm to improve sample efficiency by enabling agents to exploit a learned model for the physical environment. Notably, in recent works [14, 13, 15, 16, 21, 10, 32, 22] on world models, an RL agent learns the latent dynamics model of the environment, based on the observations and action signals, and then optimizes the policy over the learned dynamics model. Different from conventional approaches, world-model based RL takes an end-to-end learning approach, where the building blocks (such as dynamics model, perception and action policy) are trained and optimized to achieve a single overarching goal, offering significant potential to improve generalization capability. For example, DreamerV2 and DreamerV3 achieve great progress in mastering diverse tasks involving continuous and discrete actions, imagebased inputs, and both 2D and 3D environments, thereby facilitating robust learning across unseen task domains [14, 13, 15] . Recent empirical studies have also demonstrated the capacity of world models to generalize to unseen states in complex environments, such as autonomous driving [19] .Nevertheless, it remains not well understood when and how world models can generalize well in unseen environments.In this work, we aim to first obtain a deep understanding of the generalization capability of world models by examining the impact of latent representation errors, and then to devise new methods to enhance its generalization. While one may expect that optimizing a latent dynamics model (LDM) prior to training the task policy would minimize latent representation errors and hence can achieve 1 : Reward values on unseen perturbed states by rotation (α) or mask (β%) with N (0.15, 0.5).studies, reveal that modest latent representation errors in the training phase may in fact be beneficial.In particular, the alternating training strategy for world model learning, which simultaneously refines both the LDM and the action policy, could actually bring generalization gain, because the modest latent representation errors (and the corresponding induced gradient estimation errors) could enable the world model to visit unseen states and thus lead to improved generalization capacities. For instance, as shown in Table 1 , our experimental results suggest that moderate batch sizes (e.g., 16 or 32) appear to position the induced errors within a regime conferring notable generalization benefits, leading to higher generalization improvement, when compared to the cases with very small (e.g., 8) or large (e.g., 64) batch sizes.In a nutshell, latent representation errors incurred by latent encoders, if designed properly, may actually facilitate world model training and enhance generalization. This insight aligns with recent advances in deep learning, where noise injection schemes have been studied as a form of implicit regularization to enhance models' robustness. For instance, recent study [2] analyzes the effects of introducing isotropic Gaussian noise at each layer of neural networks, identifying it as a form of implicit regularization. Another recent work [27] explores the addition of zero-drift Brownian motion to RNN architectures, demonstrating its regularizing effects in improving network's stability against noise perturbations.We caution that latent representation errors in world models differ from the above noise injection schemes ( [27, 2] ), in the following aspects: 1) Unlike the artificially injected noise only added in training, these errors are inherent in world models, leading to error propagation in the rollouts; 2) Unlike the controlled conditions of isotropic or zero-drift noise examined in prior studies, the errors in world models may not exhibit such well-behaved properties in the sense that the drift may be non-zero and hence biased; 3) additionally, in the iterative training of world models and agents, the error originating from the encoder affects the policy learning and agent exploration. In light of these observations, we develop a continuous-time stochastic differential equation (SDE) formulation by treating the world model learning as a stochastic dynamic system with stochastic latent states. This approach offers an insightful view on model errors as stochastic perturbation, enabling us to obtain an explicit characterization to quantify the impacts of the errors on world models' generalization capability. Our main contributions can be summarized as follows.• Latent representation errors as implicit regularization: Aiming to understand the generalization capability of world models and improve it further, we develop a continuous-time SDE formulation by treating the world model learning as a stochastic dynamic system in latent state space.Leveraging tools in stochastic calculus and differential geometry, we characterize the impact of latent representation errors on world models' generalization. Our findings reveal that under some technical conditions, modest latent representation errors can in fact function as implicit regularization and hence result in generalization gain.• Improving generalization in non-zero drift cases via Jacobian regularization: For the case where latent representation errors exhibit non-zero drifts, we show that the additional bias term would degrade the implicit regulation and hence may make the learning unstable. We propose to add Jacobian regularization to mitigate the effects of non-zero-drift errors in training. Experimental studies are carried out to evaluate the efficacy of Jacobian regularization.• Reducing error propagation in predictive rollouts: We explicitly characterize the effect of latent representation errors on predictive rollouts. Our experimental results corroborate that Jacobian regularization can reduce the impact of error propagation on rollouts, leading to enhanced prediction performance and accelerated convergence in tasks with longer time horizons.• Bounding Latent Representation Error: We establish a novel bound on the latent representation error within CNN encoder-decoder architectures. To our knowledge, this is the first quantifiable bound applied to a learned latent representation model, and the analysis carries over to other architectures (e.g., ReLU) along the same line.Notation. We use Einstein summation convention for succinctness, where a i b i denotes i a i b i . We denote functions in C k,α as being k-times differentiable with α-Hölder continuity. The Euclidean norm of a vector is represented by ∥ • ∥, and the Frobenius norm of a matrix by | • | F ; this notation may occasionally extend to tensors. The notation x i indicates the i th coordinate of the vector x, and A ij the (i, j)-entry of the matrix A. Function composition is denoted by f • g, implying f (g). For a differentiable function f : R n → R m , its Jacobian matrix is denoted by ∂f ∂x ∈ R m×n . Its gradient, following conventional definitions, is denoted by ∇f . The constant C may represent different values in distinct contexts. 
related work: 
 World model based RL. World models have demonstrated remarkable efficacy in visual control tasks across various platforms, including Atari [1] and Minecraft [8] , as detailed in the studies by Hafner et al. [14, 13, 15] . These models typically integrate encoders and memory-augmented neural networks, such as RNNs [33] , to manage the latent dynamics. The use of variational autoencoders (VAE) [7, 23] to map sensory inputs to a compact latent space was pioneered by Ha et al. [12] .Furthermore, the Dreamer algorithm [13, 16] employs convolutional neural networks (CNNs) [24] to enhance the processing of both hidden states and image embeddings, yielding models with improved predictive capabilities in dynamic environments.Continuous-time RNNs. The continuous-time assumption is standard for theoretical formulations of RNN models. Li et al. [26] study the optimization dynamics of linear RNNs on memory decay.Chang et al. [4] propose AntisymmetricRNN, which captures long-term dependencies through the control of eigenvalues in its underlying ODE. Chen et al. [5] propose the symplectic RNN to model Hamiltonians. As continuous-time formulations can be discretized with Euler methods [4, 5] (or with Euler-Maruyama methods if stochastic in [27] ) and yield similar insights, this step is often eliminated for brevity.Implicit regularization by noise injection in RNN. Studies on noise injection as a form of implicit regularization have gained traction, with Lim et al. [27] deriving an explicit regularizer under small noise conditions, demonstrating bias towards models with larger margins and more stable dynamics.Camuto et al. [2] examine Gaussian noise injections at each layer of neural networks. Similarly, Wei et al. [31] provide analytic insights into the dual effects of dropout techniques.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: For each of our theoretical results, we state the required assumptions and provide relevant discussions that compare our assumptions to the practical implementations which involves certain limitations for theoretical simplifications. For empirical results, we also state the experiment settings and the number of trials run. We also discuss the possible future research to extend our work in the conclusion section.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.• The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
demystifying world model: a stochastic differential equation approach: 
 As pointed out in [14, 13, 15, 16] , critical to the effectiveness of the world model representation is the stochastic design of its latent dynamics model. The model can be outlined by the following key components: an encoder that compresses high dimensional observations s t into a low-dimensional latent state z t (Eq.1), a sequence model that captures temporal dependencies in the environment (Eq.2), a transition predictor that estimates the next latent state (Eq.3), and a latent decoder that reconstructs observed information from the posterior (Eq.4):EQUATIONEQUATIONIn this work, we consider a popular class of world models, including Dreamer and PlaNet, where {z, EQUATIONSequence Model:EQUATIONTransition Predictor:EQUATIONLatent Decoder:EQUATIONwhere π(h, z) is a policy function as a local maximizer of value function and the stochastic processs t is F t -adapted.Notice that f is often a zero function indicating that Equation ( 6) is an ODE, as the sequence model is generally designed as deterministic. Generally, the coefficient functions in dt and dB t terms in SDEs are referred to as the drift and diffusion coefficients. Intuitively, the diffusion coefficients here represent the stochastic model components. In Equation ( 5), σ(•, •) and σ(•, •) denotes the drift and diffusion coefficients of the latent representation errors, respectively.Both are assumed to be functions of hidden states h t and task states s t . In addition, ε indicates the magnitude of the error.Next, we impose standard assumptions on these SDEs ( 5) -( 8) to guarantee the well-definedness of the solution to SDEs. For further technical details, we refer readers to fundamental works on SDEs in the literature (e.g., [30, 17] ).Assumption 3.1. The drift coefficient functions q enc , f, p and q dec and the diffusion coefficient As shown in the empirical studies with different batch sizes (Table 1 ), the latent representation error would also enrich generalization when it is within a moderate regime. In this section, we show that the latent representation error, in the form of approximation error corresponding to widely used CNN encoder-decoder, could be made sufficiently small by finding appropriate CNN network configuration.In particular, this result provides theoretical justification to interpreting latent representation error as stochastic perturbation in the dynamical system defined in Equations (5 -8) , as the error magnitude ε can be made sufficiently small by CNN network configuration. Latent Representation Learning. We define the latent representation learning as to find encoder g enc : M → Z and decoder g dec : Z → M as maps that optimize the following objectives:min genc∈G W1 genc # Q, P ; min g dec ∈G W1 Q, gdec # P .Here, g enc # Q and g dec # P represent the pushforward measures of Q and P through the encoder map g enc and decoder map g dec , respectively. The latent representation error is understood as the "difference" of pushforward measure by the encoder/decoder and target measure. Here, to understand the "scale" of the error ε in Equation ( 5), we use W 1 for the discrepancy between probability measures. In particular, for Dreamer-type loss function that uses KL-divergence, we note that squared W 1 distance between two probability measures can be upper bounded by their KL-divergence up to a constant [11] , implying that one could reasonably expect the W 1 distance to also decrease when KL-divergence is used in the model.CNN configuration. As a popular choice choice in encoder-decoder architecture is CNN, we consider a general CNN function f CNN : X → R. Let f CNN have L hidden layers, represented as:for x ∈ X , f CNN (x) := A L+1 • A L • • • • • A 2 • A 1 (x), where A i 's are either convolutional or downsampling operators. For convolutional layers,A i (x) = σ(W c i x + b c i ), where W c i ∈ R di×di-1is a structured sparse Toeplitz matrix from the convolutional filter {w(i) j } s(i)j=0 with filter lengths(i) ∈ N + , b c i ∈ R diis a bias vector, and σ is the ReLU activation function. For downsampling layers,A i (x) = D i (x) = (x jmi ) ⌊di-1/mi⌋ j=1, where D i : R di×di-1 is the downsampling operator with scaling parameter m i ≤ d i-1 in the i-th layer. We examine the class of functions represented by CNNs, denoted by F CNN , defined as:F CNN = {f CNNas in defined above with any choice of A i , i = 1, . . . , L + 1}.For the specific definition of F CNN , we refer to [29] 's (4), ( 5) and ( 6).Assumption 3.5. Assume that M and Z are locally diffeomorphic, that is there exists a map F : M → Z such that at every point x on M, det(d F (x)) ̸ = 0.Theorem 3.6. (Approximation Error of Latent Representation). Under Assumption 3.3, 3.4 and 3.5,for θ ∈ (0, 1), let d θ := O(d M θ -2 log d θ ).For positive integers M and N , there exists an encoder g enc and decoder g dec ∈ F CNN (L, S, W ) s.t.EQUATION.Theorem 3.6 indicates that with an appropriate CNN configuration, the W 1 approximation error can be made to reside in a small region, as the best candidate within the function class is indeed capable of approximating the oracle encoder/decoder. In particular, this result indicates that the error magnitude ε in SDE (5) can be assumed to be small. This allows us to apply the perturbation analysis of the dynamical system defined in Equations (5 -8) in the following sections.In this section, we investigate the impact of latent representation errors on generalization, for the two cases with zero drift and non-zero drift, respectively. We show that under mild conditions, the zero-drift errors can function as a natural form of implicit regularization, promoting wider landscapes for improved robustness. Nevertheless, we caution that when latent representation errors have non-zero drift, it could lead to poor regularization with unstable bias and degrade world model's generalization, calling for explicit regularization.To simplify the notation here, we consider the system equations, specifically Equations ( 5), ( 6) -( 8 EQUATIONwhere g, and ḡi are structured accordingly for the respective components, employing the Einstein summation convention for concise representation. For abuse of notation, σ = (σ, 0, 0, 0), σ = (σ, 0, 0, 0). For a given error magnitude ε, we denote the solution to SDE (9) as x ε t . Intuitively, x ε t is the perturbed trajectory of the latent dynamics model. In particular, when ε = 0, indicating that the absence of latent representation error in the model, the solution is denoted as x 0 t .When the drift coefficient σ = 0, the latent representation errors correspond to a class of well-behaved stochastic processes. The following result translates the induced perturbation on the stochastic latent dynamics model's loss function L to a form of explicit regularization. We assume that L ∈ C 2 and depends on z t , h t , zt , st . Loss functions used in practical implementation, e.g. in DreamerV3, reconstruction loss J O , reward loss J R , consistency loss J D , all satisfy this condition.Theorem 3.7. (Explicit Effect Induced by Zero-Drift Representation Error) Under Assumptions 3.1 and 3.2 and considering a loss function L ∈ C 2 , the explicit effects of the zero-drift error can be marginalized out as follows: as ε → 0,EQUATION)where the regularization term R is given by R := ε P + ε 2 Q + 1 2 S , withEQUATIONEQUATIONEQUATIONSquare matrix Φ t is the stochastic fundamental matrix of the corresponding homogeneous equation:dΦt = ∂ḡ k ∂x (x 0 t , t) Φt dB k t , Φ(0) = I,and ξ k t is the shorthand fort 0 Φ -1 s σk (x 0 s , s)dB k t . Additionally, H k (x 0 s , s) is represented by for k1,k2 ∂ 2 ḡk ∂x i ∂x j (x 0 s , s) ξ k1 s i ξ k2 s j .The proof is relegated to Appendix B in the Supplementary Materials.When the loss L is convex, then its Hessian, ∇ 2 L, is positive semi-definite, which ensures that the term S is non-negative. The presence of this Hessian-dependent term S, under latent representation error, implies a tendency towards wider minima in the loss landscape. Empirical results from [20] indicates that wider minima correlate with improved robustness of implicit regularization during training. This observation also aligns with the theoretical insights in [27] that the introduction of Brownian motion, which is indeed zero-drift by definition, in training RNN models promotes robustness. We note that in addition, when the error σt (•) is too small, the effect of term S as implicit regularization would not be as significant as desired. Intuitively, this insight resonates with the empirical results in Table 1 that model's robustness gain is not significant when the error induced by small batch sizes is too small.We remark that the exact loss form treated here is simplified compared to that in the practical implementation of world models, which frequently depends on the probability density functions (PDFs) of z t , h t , zt , st . In principle, the PDE formulation corresponding to the PDFs of the perturbedx ε t can be derived from the Kolmogorov equation of the SDE (9) , and the technicality is more involved but can offer more direct insight. We will study this in future work.In practice, latent representation errors may not always exhibit zero drift as in idealized noise-injection schemes for deep learning ([27], [2] ). When the drift coefficient σ is non-zero or a function of input data h t and s t in general, the explicit regularization terms induced by the latent representation error may lead to unstable bias in addition to the regularization term R in Theorem 3.7. With a slight abuse of notation, we denote ḡ0 as g from Equation ( 9) for convenience.Under Assumptions 3.1 and 3.2 and considering a loss function L ∈ C 2 , the explicit effects of the general form error can be marginalized out as follows as ε → 0:EQUATION)where the additional bias term R is given by R := ε P + ε 2 Q + S , withEQUATIONEQUATIONEQUATIONand ξt being the shorthand fort 0 Φ -1 s σ k (x 0 s , s)dt.The presence of the new bias term R implies that regularization effects of latent representation error could be unstable. The presence of ξ in P, Q and S induces a bias to the loss function with its magnitude dependent on the error level ε, since ξ is a non-zero term influenced on the drift term σ. This contrasts with the scenarios described in [27] and [2] , where the noise injected for implicit regularization follows a zero-mean Gaussian distribution. To modulate the regularization and bias terms R and R respectively, we note that a common factor, the fundamental matrix Φ, can be bounded byEQUATION)which can be shown by using the Burkholder-Davis-Gundy Inequality and Gronwall's Lemma.Based on this observation, we next propose a regularizer on input-output Jacobian norm ∥ ∂g k ∂x ∥ F that could modulate the new bias term R for stabilized implicit regularization.some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details.• While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).• The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.• The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.• The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.• At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).• Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
enhancing predictive rollouts via jacobian regularization: 
 In this section, we study the effects of latent representation errors on predictive rollouts using latent state transitions, which happen in the inference phase in world models. We then propose to use Jacobian regularization to enhance the quality of rollouts. In particular, we first obtain an upper bound of state trajectory divergence in the rollout due to the representation error. We show that the error effects on task policy's Q function can be controlled through model's input-output Jacobian norm.In world model learning, the task policy is optimized over the rollouts of dynamics model with the initial latent state z 0 . Recall that latent representation error is introduced to z 0 when latent encoder encodes the initial state s 0 from task environment. Intuitively, the latent representation error would propagate under the sequence model and impact the policy learning, which would then affect the generalization capacity through increased exploration.Recall that the sequence model and the transition predictor are given as follows:EQUATIONwith random variables h 0 , z0 + ε as the initial values, respectively. In particular, ε is a random variable of proper dimension, representing the error from encoder introduced at the initial step. We impose the standard assumption on the error to ensure the well-definedness of the SDEs.Under Assumption 3.1, there exists a unique solution to the SDEs (for Equations 19 with squareintegrable ε), denoted as (h ε t , z ε t ). In the case of no error introduced, i.e., ε = 0, we denote the solution of the SDEs as (h 0 t , z 0 t ) understood as the rollout under the absence of latent representation error. To understand how to modulate impacts of the error in rollouts, our following result gives an upper bound on the expected divergence between the perturbed rollout trajectory (h ε t , z ε t ) and the original (h 0 t , z 0 t ) over the interval [0, T ].Theorem 4.1. (Bounding trajectory divergence) For a square-integrable random variable ε, letδ := E ∥ε∥ and dε := E sup t∈[0,T ] h ε t -h 0 t 2 + zε t -z0 t 2 . As δ → 0, dε ≤ δ C (J0 + J1) + δ 2 C exp ( H0 (J0 + J1)) + δ 2 C exp ( H1 (J0 + J1)) + O(δ 3 ),where C is a constant dependent on T. J1 and J2 are Jacobian-related terms, and H1 and H2 are Hessianrelated terms.The Jacobian-related terms J 1 and J 2 are defined asJ 0 := exp (F h + F z + P h ) , J 1 := exp Ph ;the Hessian-related terms H 0 and H 1 are defined asH 0 := F hh +F hz +F zh +F zz +P hh , H 1 := Phh ,where F h , F z are the expected sup Frobenius norm of Jacobians of f w.r.t h, z, respectively, and F hh , F hz , F zh , F zz are the corresponding expected sup Frobenius norm of second-order derivatives.Other terms are similarly defined. A detailed description of all terms, can be found in Appendix C.1.Theorem 4.1 correlates with the empirical findings in [14] regarding the diminished predictive accuracy of latent states zt over the extended horizons. In particular, Theorem 4.1 suggests that 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 .Our next result reveals how initial latent representation error influences the value function Q during the prediction rollouts, which again verifies that the perturbation is dependent on expected error magnitude, the model's Jacobian norms and the horizon length T :Corollary 4.2. For a square-integrable ε, let x t := (h t , z t ). Then, for any action a ∈ A, the following holds for value function Q almost surely:Q(x ε t , a) = Q(x 0 t , a) + ∂ ∂x Q(x 0 t , a) ε i ∂i x 0 t + 1 2 ε i ε j ∂ 2 ij x 0 t + 1 2 (ε i ∂i x 0 t ) ⊤ ∂ 2 ∂x 2 Q(x 0 t , a) (ε i ∂i x 0 t ) + O(δ 3 ),as δ → 0, where stochastic processes ∂ i x 0 t , ∂ 2 ij x 0 t are the first and second derivatives of x 0 t w.r.t ε and are bounded as follows:E sup t∈[0,T ] ∂i x 0 t ≤ C (J0 + J1) , E sup t∈[0,T ] ∂ 2 ij x 0 t ≤ C exp ( H0 (J0 + J1)) + C exp ( H1 (J0 + J1)) .This corollary reveals that latent representation errors implicitly encourage exploration of unseen states by inducing a stochastic perturbation in the value function, which again can be regularized through a controlled Jacobian norm.Jacobian Regularization against Non-Zero Drift. The regularized loss function for LDM is defined as follows:EQUATIONwhere L dyn is the original loss function for dynamics model, J θ denotes the data-dependent Jacobian matrix associated with the θ-parameterized dynamics model, and λ is the regularization weight.Our empirical results in 5 with an emphasis on sequential case align with the experimental findings from [18] that Jacobian regularization can enhance robustness against random and adversarial input perturbation in machine learning models. 
experimental studies: 
 In this section, experiments are carried out over a number of tasks in Mujoco environments. Due to space limitation, implementation details and additional results, including the standard deviation of the trials, are relegated to Section D in the Appendix.Enhanced generalization to unseen noisy states. We investigated the effectiveness of Jacobian It can be seen from Table 3 and Figure 1  
conclusion: 
 In this study, we investigate the impacts of latent representation errors on the generalization capacity of world models. We utilize a stochastic differential equation formulation to characterize the effects of latent representation errors as implicit regularization, for both cases with zero-drift errors and with non-zero drift errors. We develop a Jacobian regularization scheme to address the compounding effects of non-zero drift, thereby enhancing training stability and generalization. Our empirical findings validate that Jacobian regularization improves the generalization performance, expanding the applicability of world models in complex, real-world scenarios. Future research is needed to investigate how stabilizing latent errors can enhance generalization across more sophisticated tasks for general non-zero drift cases.The broader social impact of our work resides in its potential to enhance the robustness and reliability of RL agents deployed in real-world applications. By improving the generalization capacities of world models, our work could contribute to the development of RL agents that perform consistently across diverse and unseen environments. This is particularly relevant in safety-critical domains such as autonomous driving, where reliable agents can provide intelligent and trustworthy decision-making.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We state the hyperparameters used in Table 5 . The perturbations we considered is fully described in the experiment section form the main text.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.• The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: While our work is predominantly theoretical, we conducted 5 random trials for each perturbation degree and type. For additional results including standard deviation of trials, see Section D.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: Relevant computing information is provided in Section D.1.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: The work is of theoretical nature and has no societal impact of the work performed.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: Our work is mostly of theoretical nature and does not release new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.• The paper should discuss whether and how consent was obtained from people whose asset is used.• At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: This work does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.• For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: To mathematically describe this intrinsic lower-dimensional geometric structure, for an integer k > 0 and α ∈ (0, 1], we consider the notion of smooth manifold (in the C k,α sense), formally defined byDefinition A.1 (C k,α manifold). A C k,α manifold M of dimension n is a topological manifold (i.e.a topological space that is locally Euclidean, with countable basis, and Hausdorff) that has a C k,α structure Ξ that is a collection of coordinate charts {U α , ψ α } α∈A where U α is an open subset of M, • Each chart ψ α is a diffeomorphism that is a smooth map with smooth inverse (in the C k,α sense),ψ α : U α → V α ⊆ R n such that • α∈A U α ⊇ M,• Any two charts are C k,α -compatible with each other, that is for allα 1 , α 2 ∈ A, ψ α1 • ψ -1 α2 : ψ α2 (U α1 ∩ U α2 ) → ψ α1 (U α1 ∩ U α2 ) is C k,α .Intuitively, a C k,α manifold is a generalization of Euclidean space by allowing additional spaces with nontrivial global structures through a collection of charts that are diffeomorphisms mapping open subsets from the manifold to open subsets of euclidean space. For technical utility, the defined charts allow to transfer most familiar real analysis tools to the manifold space. For more references, see [25] .Definition A.2 (Riemannian volume form). Let X be a smooth, oriented d-dimensional manifold with Riemannian metric g. A volume form dvol M is the canonical volume form on X if for any pointx ∈ X , for a chosen local coordinate chart (x 1 , ..., x d ), dvol M = det g ij dx 1 ∧ ... ∧ dx d , whereg ij (x) := g ( ∂ ∂xi , ∂ ∂xj )(x).Then the induced volume measure by the canonical volume form dvol X is denoted as µ X , defined by µ X : A → A dvol X , for any Borel-measurable subset A on the space X . For more references, see [9] .We recall the latent representation problem defined in the main paper.Consider the state space S ⊂ R d S and the latent space Z. Consider a state probability measure Q on the state space S and a probability measure P on the latent space Z. ∈ C k,α (M, R) w.r.t µ M . Let Z be a closed ball in R d M , that is {x ∈ R d M : ∥x∥ ≤ 1 }. P is a probability measure supported on Z with its Radon-Nikodym derivative p ∈ C k,α (Z, R) w.r.t µ Z .We consider a general CNN function f CNN : X → R. Let f CNN have L hidden layers, represented as:f CNN (x) = A L+1 • A L • • • • • A 2 • A 1 (x), x ∈ X ,where A i 's are either convolutional or downsampling operators. For convolutional layers,A i (x) = σ(W c i x + b c i ),where W c i ∈ R di×di-1 is a structured sparse Toeplitz matrix from the convolutional filter {w(i) j } s(i) j=0with filter length s(i) ∈ N + , b c i ∈ R di is a bias vector, and σ is the ReLU activation function.For downsampling layers, For more details in the definitions of CNN functions, we refer to [29] .A i (x) = D i (x) = (x jmi ) ⌊di-1/mi⌋ j=1 ,Assumption A.5. Assume that M and Z are locally diffeomorphic, that is there exists a mapF : M → Z such that at every point x on M, det(d F (x)) ̸ = 0.Theorem A.6. (Approximation Error of Latent Representation). Under Assumption A.3, A.4 andA.5, for θ ∈ (0, 1), let d θ = O(d M θ -2 log d θ ).For positive integers M and N , there exists an encoder g enc and decoder g dec ∈ F CNN (L, S, W ) s.t.W 1 (g enc # Q, P ) ≤ d M C(N M ) - 2(k+1) d θ , W 1 (g dec # P, Q) ≤ d M C(N M ) - 2(k+1) d θ.The primary challenge to show Theorem A.6 is in demonstrating the existence of oracle encoder and decoder maps. These maps, denoted as g * enc : M → Z and g * dec : Z → M respectively, must satisfyEQUATION)and importantly they have the proper smoothness guarantee, namely g * enc ∈ C k+1,α (M, Z) and g * dec ∈ C k+1,α (Z, M). Proposition A.7 shows the existence of such oracle map(s).Proposition A.7 (C k,α , compact). Let M, N be compact, oriented d-dimensional Riemannian manifolds with C k+3,α boundary with the volume measure µ M and µ N respectively. Let Q, P be distributions supported on M, N respectively with their C k,α density functions q, p, that is Q, P are probability measures supported on M, N with their Radon-Nikodym derivatives q ∈ C k,α (M, R)w.r.t µ M and p ∈ C k,α (N , R) w.r.t µ N .Then, there exists a C k+1,α map g : N → M such that the pushforward measureg # P = Q, that is for any measurable subset A ∈ B(M), Q(A) = P (g -1 (A)).Proof. (Proposition A.7) Let ω := p dvol N , then ω is a C k,α volume form on N , as p ∈ C k,α and for any point x ∈ N , we have p(x) > 0. In addition, N ω = N p dvol N = N p dµ N = P (N ) = 1.Similarly, let η := q dvol M a C k,α volume form on M and M η = 1.Let F : N → M be an orientation-preserving local diffeomorphism, we then have det(dF ) > 0 everywhere on N .As N is compact and M is connected by assumption, F is a covering map, that is for every pointx ∈ M, there exists an open neighborhood U x of x and a discrete setD x such that F -1 (U ) = ⊔ α∈D V α ⊂ N and F | Vα = V α → U is a diffeomorphism. Furthermore, |D x | = |D y | for any pointsx, y ∈ M. In addition, |D x | is finite from the compactness of N .Let η be the pushforward of ω via F , defined by for any point x ∈ M and a neighborhood U x ,EQUATIONη is well-defined as it is not dependent on the choice of neighborhoods and the sum and 1 |Dx| are always finite. Furthermore, η is a C k,α volume form on M, as p• F Vα -1 is C k,α . Notice that F Vα -1 is orientation-preserving as det d F Vα -1 = 1 det d F Vα > 0 everywhere on V α .In addition, F Vα -1 is proper: as for any compact subset K of N , K is closed; and as F Vα -1is continuous, the preimage of K via F Vα -1 a closed subset of M which is compact, then the preimage of K must also be compact. Hence, F Vα -1 is proper. As every F Vα -1 is proper, orientation-preserving and surjective, then c := deg(F Vα -1 ) = 1.Then, M η = c N ω = 1.As we have shown that η and η ∈ C k,α and M η = M η, by [6] , there exists a diffeomorphism ψ : M → M fixing on the boundary such that ψ * η = η, where ψ, ψ -1 ∈ C k+1,α .Letg := ψ • F , then it holds that g * η = (ψ • F ) * η = F * • ψ * η = F * η = ω.Then, for any measurable subset A on the manifold M, we verify thatQ(A) = A η = g -1 (A) g * η = g -1 (A) ω = g -1 (A) p dvol N = g -1 (A) p dµ N = P (g -1 (A)).Hence, we have shown the existence by an explicit construction. As ψ ∈ C k+1,α , and F ∈ C ∞ , then we have g ∈ C k+1,α .We are now ready to show Theorem A.6 with the existence of oracle map and the low-dimensional approximation results from [29] .Proof. (Theorem A.6) For encoder, from Proposition A.7, there exists an C k+1,α oracle map g :M → Z such that the pushforward measure g # Q = P . Then,W 1 ((g enc ) # Q , P ) =W 1 ((g enc ) # Q , g # Q) = sup f ∈Lip 1 (Z) Z f (y) d((g enc ) # Q) - Z f (y) d(g # Q) ≤ sup f ∈Lip 1 (Z) M |f • g enc (x) -f • g(x)| dQ ≤ M ∥g enc (x) -g(x)∥ dQ ≤ d M C(N M ) - 2(k+1) d θ,where the last inequality follows from the special case ρ = 0 of Theorem 2.4 in [29] .Similarly, for decoder, from Proposition A.7, there exists an C k+1,α oracle map ḡ : Z → M such that the pushforward measureḡ# P = Q. W 1 ((g dec ) # P , Q) =W 1 ((g dec ) # P , ḡ# P ) ≤ Z ∥g dec (y) -ḡ(y)∥ dP ≤ d M C(N M ) - 2(k+1) d θ.We recall the SDEs for latent dynamics model defined in the main paper. Consider a complete, filtered probability space (Ω, F, {F t } t∈[0,T ] , P ) where independent standard Brownian motionsB enc t , B pred t , B seq t , B dect are defined such that F t is their augmented filtration, and T ∈ R as the time length of the task environment. We consider the stochastic dynamics of LDM through the following coupled SDEs after error perturbation:EQUATIONEQUATION)EQUATIONEQUATIONwhere π(h, z) is a policy function as a local maximizer of value function and the stochastic processs t is F t -adapted.As discussed in the main paper, our analysis applies to a common class of world models that uses Gaussian distributions parameterized by neural networks' outputs for z, z, s. Their distributions are not non-Gaussian in general.For example, as z is conditional Gaussian and its mean and variance are random variables which are learned by the encoder from r.v.s s and h as inputs, thus rendering z non-Gaussian. However, z is indeed Gaussian when the inputs are known. Under this conditional Gaussian class of world models, to see that the continuous formulation of latent dynamics model can be interrupted as SDEs, one notices that SDEs with coefficient functions of known inputs are indeed Gaussian, matching to this class of world models. Formally, in the context of z without latent representation error:Proposition B.1. (Latent states SDE with known inputs is Gaussian)For the latent state process z t∈[0,T ] without error,EQUATIONwith zero initial value. Given known h t∈[0,T ] and s t∈[0,T ] , the process z t is a Gaussian process.Furthermore, for any t ∈ [0, T ], z t follows a Gaussian distribution with mean µ t = t 0 q enc (h s , s s )ds and variance σ 2 t = t 0 qenc (h s , s s ) 2 ds.Proof. Proof follows from Proposition 7.6 in [30] .Next, we recall our assumptions from the main text:Assumption B.2. The drift coefficient functions q enc , f, p and q dec and the diffusion coefficient EQUATIONas ε → 0, where the regularization term R is given by R:= ε P + ε 2 Q + 1 2 S .Each term of R is as follows:EQUATIONEQUATIONEQUATIONwhere square matrix Φ t is the stochastic fundamental matrix of the corresponding homogeneous equation:dΦ t = ∂ḡ k ∂x (x 0 t , t) Φ t dB k t , Φ(0) = I,and ξ k t is as the shorthand fort 0 Φ -1 s σk (x 0 s , s)dB k t . Additionally, H k (x 0 s , s) is represented by for k1,k2 ∂ 2 ḡk ∂x i ∂x j (x 0 s , s) ξ k1 s i ξ k2 s j .Before proving Theorem B.4, we first show Proposition B.5 on the general case of perturbation to the stochastic system. Consider the following perturbed system given byEQUATION)with initial values x(0) = x 0 , Proposition B.5. Suppose that f is a real-valued function that is C 2 . Then it holds that, with probability 1, as ε → 0, for t ∈ [0, T ],EQUATION) where the stochastic process x 0 t is the solution to SDE 32 with ε = 0, with its first and second-order derivatives w.r.t ε denoted as∂ ε x 0 t , ∂ 2 ε x 0 t .Furthermore, it holds that ∂ ε x 0 t , ∂ 2 ε x 0 t satisfy the following SDEs with probability 1,EQUATIONwith initial values ∂ ε x(0) = 0, ∂ 2 ε x(0) = 0, whereΨ k : (∂ ε x, x, t) → ∂ ε x i ∂g k ∂x i ∂x j (x, t)∂ ε x j , for k = 0, 1, ..., m.Proof. We first apply the stochastic version of perturbation theory to SDE 32. For brevity, we will write t as B 0 t and use Einstein summation convention. Hence, SDE 32 is rewritten asEQUATIONwith initial value x(0) = x 0 .Step 1: We begin with the corresponding systems to derive the SDEs that characterize ∂ ε x ε t and ∂ 2 ε x ε t .Our main tool is an important result on smoothness of solutions w.r.t. initial data from Theorem 3.1 from Section 2 in [17] .For ∂ ε x, consider the SDEsd x t = γ ε k (x t , t) dB k t , (*) d ε t = 0,with initial values x (0) = x 0 , ε(0) = ε. From an application of Theorem 3.1 from Section 2 in [17] on *, we have ∂ ε x that satisfies the following SDE with probability 1:EQUATIONwith initial value ∂ ε x 0 = 0 ∈ R n , with probability 1, where x t is the solution to Equation (35) and the functions α ε k are given byα ε k : (x, t) → ∂g k ∂x j (x, t) + ε ∂η k ∂x j (x, t) ,where k = 0, ..., m.To characterize ∂ 2 ε x t , consider the following SDEsd x t = γ ε k (x t , t) dB k t , (**) d ∂ ε x t = (α ε k (x t , t) ∂ ε x t + η k (x t , t)) dB k t , d ε t = 0, with initial value x(0) = x 0 , ∂ ε x(0) = 0, ε(0) = ε.From a similar application of Theorem 3.1 from Section 2 in [17] , the second derivative ∂ 2 ε x satisfies the following SDE with probability 1:EQUATIONwith initial value ∂ 2 ε x(0) = 0 ∈ R n , where ∂ ε x t is the solution to Equation(36), x(t) is the solution to Equation ( 35), and the functionsβ ε k : (∂ ε x, x, t) → ∂ ε x j ∂g i k ∂x l ∂x j (x, t) + ε ∂η i k ∂x l ∂x j (x, t) ∂ ε x l ,where k = 0, ..., m.When ε = 0 in the obtained SDEs (35), ( 36) and (37), the corresponding solutions of which arex 0 t , ∂ ε x 0 t , ∂ 2 εx 0 t , we now have the following:EQUATIONEQUATIONEQUATIONwith initial values x(0) = x 0 , ∂ ε x(0) = 0, ∂ 2 ε x(0) = 0. In particular, Ψ k := β 0 k is given by(∂ ε x, x, t) → ∂ ε x i ∂g k ∂x i ∂x i (x, t)∂ ε x j .Step 2: For the next step, we show that the solutions x 0 t , ∂ s x 0 t , ∂ 2 ε x 0 t are indeed bounded by proving the following lemma B.6:Lemma B.6. E sup t∈[0,T ] x 0 t 2 , E sup t∈[0,T ] ∂ ε x 0 t 2 , and E sup t∈[0,T ] ∂ 2 ε x 0 t 2 are bounded.Proof. To simplify the notations, we take the liberty to write constants as C and notice that C is not necessarily identical in its each appearance.(1) We first show that E sup t∈[0,T ] x 0 t 2 is bounded.From Equation (38), we have thatx 0 t = x 0 + t 0 g k (x τ , τ ) dB k τ .By Jensen's inequality. it holds thatEQUATION)For the second term on the right hand side, it is a sum over k from 0 to m by Einstein notation.For k = 0, recall that we write t as B 0 t :EQUATIONEQUATIONwhere we used Jensen's inequality, the assumption on the linear growth, the inequality property of sup and Fubini's theorem, respectively.For k is equal to 1, . . . , m,EQUATIONEQUATIONwhere (iv) holds from the Burkholder-Davis-Gundy inequality as t 0 g k x 0 τ , τ dB τ is a continuous local martingale with respect to the filtration F t ; and then one can obtain (v) by following a similar reasoning of (ii) and (iii).Hence, now from the previous inequality (41),E sup t∈[0,T ] x 0 t 2 ≤ E ∥x 0 ∥ 2 + C + C T 0 E sup s∈[0,τ ]x 0 s dτ.By the Gronwall's lemma, it holds true thatE sup t∈[0,T ] x 0 t 2 ≤ C E ∥x 0 ∥ 2 + C exp(C).As x 0 is square-integrable by assumption, therefore we have shown that E sup t∈[0,T ] x 0 t 2 is bounded.(2) We then show that E supt∈[0,T ] ||∂ ε x 0 t || 2 is also bounded.From the SDE (39), as we have derived that∂ ε x 0 t = t 0 ∂g k ∂x x 0 τ , τ ∂ ε x 0 τ + η k x 0 τ , τ dB k τ ,then we haveE sup t∈[0,τ ] ∂ ε x 0 t 2 ≤ C E sup t∈[0,τ ] t 0 ∂g k ∂x x 0 τ , τ ∂ ε x 0 τ dB k τ 2 + C E sup t∈[0,T ] t 0 η k x 0 τ , τ dB k τ 2 .For k = 0, we haveEQUATIONEQUATION≤ C + C T 0 E sup s∈[0,τ ] ∂ ε x 0 s 2 dτ + C E sup t∈[0,T ] x 0 t 2 ,where to get to (vi), we used Jensen's inequality; for (vii), we used the linear growth assumption an η 0 , then we obtain (viii) by as derivatives of function g 0 are bounded by assumption.Similarly, for k = 1, ..., m,EQUATIONEQUATIONwhere we obtain (ix) by the Burkholder-Davis-Gundy inequality and (x) by following similar steps as have shown in (vii) and (viii).We are now ready to sum up each term to acquire a new inequality:E sup t∈[0,T ] ∂ ε x 0 t 2 ≤ C + C E sup t∈[0,T ] x 0 t 2 + C T 0 E sup s∈[0,τ ] ∂ ε x 0 s 2 dτ.By Gronwall's lemma, we have thatE sup t∈[0,T ] ∂ ε x 0 t 2 ≤ C + C E sup t∈[0,T ] x 0 t 2 exp(C).As it is previously shown thatE sup t∈[0,τ ] ∥x • (t)∥ 2 is bounded, it is clear that E sup t∈[0,T ] ∂ ε x 0 t 2is bounded too.(3) From similar steps, one can also show that E supt∈[0,T ] ∂ 2 ε x 0 t 2 is bounded.Step 3: Having shown that x 0 t , ∂ ε x 0 t , ∂ 2 ε x 0 t are bounded, we proceed to bound the remainder term by proving the following lemma.Lemma B.7. For a given ε ∈ R, letR ε := (t, ω) → 1 ε 3 x ε (t, ω) -x 0 (t, ω) -ε∂ ε x 0 (t, ω) -ε 2 ∂ 2 ε x 0 (t, ω) ,where the stochastic process x ε t is the solution to Equation (32) . Then it holds true thatE sup t∈[0,T ] ∥R ε (t)∥ 2 is bounded.Proof. The main strategy of this proof is to first rewrite ε 3 R ε as the sum of some simpler terms and then to bound each term. To simplify the notation, we denote xε t as x 0 t + ε∂ ε x 0 t + ε 2 ∂ 2 ε x 0 t . For k = 0, .., n, we define the following terms:θ k (t) := t 0 g k (x ε τ , τ ) -g k (x ε τ , τ ) dB k τ , φ k (t) := t 0 g k (x ε τ , τ ) -g k x 0 τ , τ -ε ∂g k ∂x x 0 τ , τ ∂ ε x 0 τ -ε 2 Ψ k ∂ ε x 0 τ , x 0 τ , τ -ε 2 ∂g k ∂x i x 0 τ , τ ∂ 2 ε x 0 τ dB k τ , σ k (t) := -ε t 0 η k x 0 τ , τ + 2ε ∂η ∂x x 0 τ , τ ∂ ε x 0 τ dB k τ .Hence, we haveε 3 R ε (t) = 1 k=0 θ k (t) + φ k (t) + σ k (t).For θ k (t), we haveEQUATIONEQUATIONEQUATIONwhere to obtain (i) we used Jensen's inequality when k = 0 and by the Burkholder-Davis-Gundy inequality when k = 1, used the Lipschitz condition of g k to obtain (ii), and for (iii), it is becauseε 3 R ε (t) = xε t -x ε t .We note that from Taylor's theorem, for any s ∈ [0, t], k = 0, 1, there exists some ε s ∈ (0, ε) s.t.EQUATIONFor φ k (t), we haveEQUATIONEQUATIONEQUATIONEQUATIONwhere for (iv), we used Equation (42) and Jensen's inequality for k = 0 and the Burkholder-Davis-Gundy inequality for k = 1; to obtain (v), we applied Jensen's equality; we then derived (vi) from the Lipschitz conditions of g k and Ψ k ; and finally another application of Jensen's inequality gives (vii) which is bounded as a result from the Lemma B.6.EQUATIONEQUATIONEQUATIONwhere we obtained (ix) by Jensen's inequality when k = 0 and by Burkholder-Davis-Gundy inequality when k = 1, and (x) by the linear growth assumption on η k ; one can see that (xi) is bounded by recalling the Lemma B.6 and the assumption that η k has bounded derivatives.Hence, by Jensen's inequality and Gronwall's lemma, we haveE sup t∈[0,T ] ∥R ε (t)∥ 2 ≤ C K k=0 E sup t∈[0,T ] ∥θ k (t)∥ 2 + E sup t∈[0,T ] ∥φ k (t)∥ 2 + E sup t∈[0,T ] ∥σ k (t)∥ 2 , ≤ C + C T 0 E sup t∈[0,τ ] ∥R ε (t)∥ 2 dτ, ≤ C exp (C) . Therefore, E sup ∥R ε (t)∥ 2 is bounded.Finally, it is now straightforward to show Equation ( 33) by applying a second-order Taylor expansionon f x 0 t + ε∂ ε x 0 t + ε 2 ∂ 2 ε x 0 t +ε 3 R ε (t) .We are now ready to show Theorem 3.7. One notes that Corollary 3.8 directly follows from the result too.Proof. (Theorem 3.7) From Proposition B.5, it is noteworthy to point out that the derived SDEs (34)for ∂ ε x 0 t and ∂ 2 ε x 0 t are vector-valued general linear SDEs. With some steps of derivations, one can express the solutions as:EQUATIONwhere n × n matrix Φ t is the fundamental matrix of the corresponding homogeneous equation:EQUATIONwith initial valueEQUATIONIt is worthy to note that the fundamental matrix Φ t is non-deterministic and when ∂gi ∂x and ∂gj ∂x commutes, Φ t has explicit solutionEQUATIONHaving obtained the explicit solutions, one can plug in corresponding terms and obtain the results of ε := E sup t∈[0,T ] h ε t -h 0 t 2 + zε t -z0 t 2 .It holds that as δ → 0,EQUATIONJ 0 = exp (F h + F z + P h ) , J 1 = exp Ph , H 0 =F hh + F hz + F zh + F zz + P hh , H 1 = Phh F h =C E sup t∈[0,T ] ∂f ∂h + ∂f ∂a ∂ h ρ 2 F , F z = C E sup t∈[0,T ] ∂f ∂z + ∂f ∂a ∂ z ρ 2 F , P h =C E sup t∈[0,T ] ∂p ∂h 2 F , Ph = C E sup t∈[0,T ] ∂ p ∂h 2 F , F hh =C E sup t∈[0,T ] ∂ 2 f ∂h 2 + ∂ 2 f ∂h∂a ∂ h ρ + ∂f ∂a ∂ 2 hh ρ 2 F , F hz =C E sup t∈[0,T ] ∂ 2 f ∂h∂z + ∂ 2 f ∂z∂a ∂ h ρ + ∂f ∂a ∂ 2 zh ρ 2 F F zh =C E sup t∈[0,T ] ∂ 2 f ∂h∂z + ∂ 2 f ∂h∂a ∂ z ρ + ∂f ∂a ∂ 2 hz ρ 2 F F zz =C E sup t∈[0,T ] ∂ 2 f ∂z 2 + ∂ 2 f ∂z∂a ∂ z ρ + ∂f ∂a ∂ 2 zz ρ 2 F , P hh =C E sup t∈[0,T ] ∂ 2 p ∂h 2 2 F , Phh = C E sup t∈[0,T ] ∂ 2 p ∂h 2 2 F, where for brevity, when functions always have inputs (z 0 t , h 0 t , t), we adopt the shorthand to write, for example, f (z 0 t , h 0 t , t) as f .Before proving the main result C.1, we first show the general case of perturbation in initial values.Consider the following general system with noise at the initial value:EQUATIONEQUATIONwhere the initial perturbation ε ∈ R n × Ω. As g k are C 2,α g functions, by the classical result on the existence and the uniqueness of solution to SDE, there exists a unique solution to Equation (47), denoted as x ε t or x ε (t).To simplify the notation, we write∂ i x ε t := ∂x ε (t) ∂x i , ∂ 2 ij x ε t = ∂ 2 x ε t ∂x i ∂x j , for i, j = 1, . . . , n that are,respectively, the first and second-order derivatives of the solution x ε (t) w.r.t. the changes in the corresponding coordinates of the initial value. When ε = 0 ∈ R n , we denote the solutions to Equation (47) as x 0 t with its first and second derivativesEQUATIONas δ → 0.Proof. Similar to the previous section, for notational convenience, we write t as B 0 t and employs Einstein summation notation. Hence, Equation (47) can be shorten asEQUATIONwith initial values x(0) = x 0 + ε.To begin, we find the SDEs that characterize ∂ i x ε t and ∂ 2 ij x ε t , for i, j = 1, ..., n.For ∂ i x ε t , we apply Theorem 3.1 from Section 2 in [17] on Equation (51) and ∂ i x ε t satisfy the following SDE with probability 1,EQUATION)with initial value ∂ i x ε 0 to be the unit vector e i = (0, 0, . . . , 1, . . . , 0) that is all zeros except one in the i th coordinate.For ∂ 2 ij x ε t , we again apply Theorem 3.1 from Section 2 in [17] on the SDE (52) and obtain that ∂ 2 ij x ε b satisfy the following SDE with probability 1,EQUATIONwith the initial value ∂ ij x ε (0) = e j , whereΨ k : R d × R d × [0, T ] → R d×d , (x, ∂ i x, t) → ∂ 2 g l k ∂x u ∂x v (x ε t , t) l,u,v ∂ i x v .For the next step, we show that with probability 1, the following holdsEQUATIONas ∥ε∥ → 0.One can follow the similar steps of proofs for Lemma (B. EQUATIONSimilarly, for such that there exists a C 2 map ρ : U → V such that ∂Q ∂a (h, z, ρ(h, z)) = 0 and ρ(h, z) is the local maximizer of Q(h, z, •) for any h, z ∈ U . Furthermore, we have thatEQUATION∂ h ρ = -∂ 2 Q ∂a 2 -1 ∂ 2 Q ∂a∂h .Similarly, other first-terms and second-order terms ∂ z ρ, ∂ 2 zz ρ, ∂ 2 zh ρ, ∂ 2 hz ρ, ∂ 2 hh ρ can be explicitly expressed without much additional effort (e.g., in [28] , [3] ).The rest of the proof is easy to see after plugging in the corresponding terms from Proposition C.2.In this experiment, we investigated the effectiveness of Jacobian regularization in model trained against a baseline during the inference phase with perturbed state images. We consider three types of perturbations: (1) Gaussian noise across the full image, denoted as N (µ 1 , σ Under the Quadruped task,we initially examined a smaller latent error process (µ t = 0, σ 2 t = 0.1) and observed that the model with Jacobian regularization converged significantly faster, even though the adversarial effects on the model without regularization were less severe (Figure 5 ). When considering the more challenging latent error process (µ t ∼ [0, 5], σ 2 t ∼ [0, 5]), we noted that the regularized model remained significantly less sensitive to perturbations in latent state z t , whereas the baseline model struggled to learn (Figure 6 ). These empirical observations reinforce our theoretical findings in Corollary 3.8, demonstrating that regulating the model's Jacobian norm effectively controls the impact of latent representation errors. In this experiment, we evaluate the efficacy of Jacobian regularization in extended horizon tasks, specifically by increasing the horizon length in MuJoCo Walker from 50 to 100 steps. We tested two regularization weights λ = 0.1 and λ = 0.05. Justification: In the abstract and the contribution section 1 in introduction, we provide a clear list of statements outlining the paper's contributions.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper.• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: The paper takes an SDE-based perspective on understanding errors and generalization in learned latent dynamics models, and ultimately propose to regularize a learned dynamics model with the Jacobian w.r.t. the current state, leading to eq. (20). This adds generalization against unseen states and encoder errors on walker and quadruped tasks, and also seems to help improve convergence on a walker task. Generalization remains an active and open topic, and advancing our understanding of it is important to the community.

We thoroughly discussed this paper and came to the conclusion that there are still a few remaining concerns that limit publication in the current form. These need another reviewing cycle to address:

1. The most significant concern, raised by reviewers pwHm and waKG, is that the experimental results are not extensive and do not evaluate on any established generalization setting. Despite connecting to the models used the Dreamer and TD-MPC variants, the experimental results are only on walker and quadruped systems with artificially created noisy states. Reviewer waKG raised the concern that there are no comparisons to works that perform image augmentations, and the authors responded with an experiment in the walker environment. It would have been convincing to compare to the standard Dreamer/TD-MPC models trained on the more complex environments, and also to use established benchmarks for robustness/generalization, e.g., as discussed in this next point:
2. Generalization in world models (and MDPs) can go significantly beyond the formulation they consider, e.g., as described in A Survey of Zero-shot Generalisation in Deep Reinforcement Learning. That paper also summarizes a number of existing generalization benchmarks (e.g., the distracting control suite) and methods that experimentally this paper could have compared with. However, none of them are used (or referenced). For this reason, the paper seems too disconnected from the existing research on generalization in RL --- neither experimentally nor their related work/other discussions discuss other literature of generalization in RL. This makes the claim of  "unraveling and improving generalization" unsupported.
3. The idea of a Jacobian regularization is straightforward and can be added without going through the SDE formulation or bounds. There could be some theoretical interest and contribution from the SDE bounds and relating them to generalization and the Jacobian, but they are also little disconnected from the experimental setting because the Jacobian regularization could also be interpreted from other perspectives without needing the SDE bounds. While the latent dynamics SDE with an error model is interesting to study, and bound against the reference 0-error SDE, we feel the claims to generalizing to non-trivial modeling errors are not well-supported. We do not find the additional theory in Section 3 to be experimentally supported.
4. We assume the finalized Crafter results will at least be consistent with Table 2, where they have a similar experiment for the walker and quadruped environments. While it would improve the generalization, it is not a significant kind of generalization as the perturbations they use in the environment are synthetic and not surprising given the gradient penalty they added for training the model. Furthermore, the table with Crafter results they posted seems a little misleading as 1) it is unlikely that they would surpass the performance of the baseline method in the original, unnoised environment (but it would be a positive result if they do!), and 2) it only uses a single seed (actually, all of the experiments seem to use a single seed --- we do assume the main trends would continue holding over multiple seeds, but it is scientifically important to show this, especially if they are comparing on the standard benchmarks where other methods are reporting averaged performance)

Reviews: review: paper_summary: paper_summary: This paper explores the generalization capability of world models in reinforcement learning. In particular, they investigate the latent representation error in world models. They show that zero-drift representation error is inherently a regularizer for the learned model functions. On the other hand, they show that the non-zero-drift representation error accumulates errors and Jacobian regularization can be used to alleviate the issue. They demonstrate their proposed approach improves stability, convergence, and performance.
strength_weakness: strength_weakness: 1. The paper is very thorough in terms of theoretical derivation. However, in my opinion, the experimental section of the paper is somewhat lacking. It utilizes only two tasks from Mujoco to prove the efficacy of the approach. More diverse tasks from other benchmarks and robust perturbations will certainly improve the paper. 

2. The experimental evaluation is limited to reward comparison. However, it would be interesting to see some visualization of how the trajectories unfold in the case of both types of errors and with Jacobian regularization.
questions: questions: 1. It appears in different tables different values of $\lambda$ (the regularization weight in eq. 20) have been used. How sensitive the models are to the value of this hyperparameter? Do you have any suggested range for better performance?

2. Do you observe any substantial relation between $\lambda$ and the task horizon?
limitations: limitations: While the paper discusses the potential social impact of the work, it doesn’t discuss any limitations. I believe the characterization of the models as SDE and the use of Brownian motion as variation have certain contributions to the identified claims. Other interpretations may alter the findings.
paper_summary: paper_summary: The paper studies the generalization capability of world models via a stochastic differential equation formulation. They try to understand latent representation errors on generalization, with both zero-drift representation errors and non-zero-drift representation errors. They found that zero drift latent representation errors are implicit regularization and thus bring generalization gain. Jacobian regularization is proposed to enhance training stability and generalization.
strength_weakness: strength_weakness: + 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;
questions: questions: See the Weaknesses part and explain how to extend the analysis to more general cases.
limitations: limitations: NA
paper_summary: paper_summary: This paper investigates the generalization capabilities of world models in RL, particularly with respect to latent representation errors, which arise when observations are encoded into a low-dimensional latent space. The authors provide a bound on latent representation error when using CNN encoder-decoder architectures. The world model is framed as a stochastic differential equation to characterize the impact of latent representation errors on generalization in terms of either zero or non-zero drift. The authors provide theoretical analysis which shows that these errors can result in implicit regularization in the zero drift case, and propose a Jacobian regularization scheme to tackle the unwanted bias term in the non-zero drift case. Finally, when performing model rollouts for learning a policy, the authors study the effect of these errors on the value function. Experiments on Mujoco tasks demonstrate that the proposed Jacobian regularization enhances robustness to noisy states, reduces the detrimental impact of latent representation errors, and improves convergence speed for longer horizon tasks.
strength_weakness: strength_weakness: While the paper studies a previously unexplored problem, there are some questions about the significance of these findings and the use of drift and diffusion terms to represent the error. Other areas for improvement include explaining the insights from the theoretical analysis more clearly, describing the experimental settings in more detail, and supporting certain claims with more evidence.

- Studying the effect of latent representation error is certainly useful, however, with recent advances in representation learning approaches, one can learn reasonably good representations such that the reconstruction error is negligible. When it comes to model-based RL, a much bigger issue is the compounding model error, which is a result of error in the latent/state dynamics model predictions. A comment from the authors on this aspect would be helpful.
- The decomposition of latent error into drift and diffusion terms seems a bit contrived. It is not clear how the error can be expressed in this form, and what defines the scenarios of zero versus non-zero drift.
- The interpretation that propagation of latent error leads to the model exploring novel states seems somewhat questionable. My understanding is that the erroneous states improve robustness similar to noise injection, but will most likely not be valid states belonging to the state space of the MDP. Some reasonable evidence is required to support this statement.
- The paper presents several results and including some intuitive or low-level explanation for each of those results would greatly improve readability. Additionally, due to the large amount of mathematical notation used throughout the paper, it would be helpful to include a notation table in the appendix for easy reference.
- The experimental setting is not sufficiently clear, especially in the introduction when the authors refer to Table 1. With regards to the perturbations - are they applied to every state in the trajectory? For masking, is the same mask used for every state, or is the mask also sampled randomly? With regards to injecting encoder error - how to interpret the $\mu_t$ and $\sigma_t$ values?
questions: questions: - I am not sure I fully understand the relation of batch size with the latent representation error in Table 1. Since one would take the mean over the batch size, it should not affect the error magnitude. Is the variation in performance due to stochasticity of the SGD updates? If so, how does it relate to latent representation error? Also, as mentioned in weaknesses, the experimental setting should be clearly explained at this point in the paper.
- What do the drift and diffusion terms corresponding to the latent representation error signify?
- Theorem 3.7 suggests that the regularizing effect of the latent representation error can be attributed to the Hessian of the loss function, which encourages wider minima. As noted by the authors, this term is non-negative only if the loss function is convex. Could the authors comment on how to interpret this result for non-convex loss functions (which is usually the case in deep learning)?
- The analysis in this paper focuses on Dreamer style models which use an RNN to represent the latent dynamics. A parallel line of work, TD-MPC [1], uses an MLP model to predict the next state which is applied recursively. I am interested to know if the authors have any thoughts on the applicability of their analysis to such methods, and if there are any major differences.
- An alternate method to improve robustness to perturbations is by training the model/value function on different augmentations of the states [2]. I understand the limitations of time during the short rebuttal period, but an empirical comparison would significantly enhance the paper. I am also curious to know the authors’ thoughts on the pros/cons of using the Jacobian regularization term over such data augmentation methods.
- How much computational overhead is added by the calculation of the Jacobian regularization term?

[1] Hansen, N.A., Su, H. and Wang, X., 2022, June. Temporal Difference Learning for Model Predictive Control. In *International Conference on Machine Learning* (pp. 8387-8406). PMLR.

[2] Yarats, D., Kostrikov, I. and Fergus, R., 2021, May. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In *International conference on learning representations*.
limitations: limitations: There is little discussion on the limitations of the analysis. Some points worth discussing could be the impact of various assumptions when deriving the results, the fact that the analysis is mostly focused on a specific setting - learning from pixels using a CNN encoder and an RNN latent dynamics model, and further investigation of the compounding model error problem.

Comments: title: Rebuttal 4/4 for Reviewer waKG
Rebuttal 3/4 for Reviewer waKG
Rebuttal 2/4 for Reviewer waKG
comment: We are glad that the additional experiment and clarification was helpful in addressing your concerns. Thank you very much for your updated score and your valuable insights.
The authors clarified my major concern about whether the tackled aspect is robustness or generalization. The added new experiment is very helpful.
Thank you for your positive feedback. We have carefully addressed your comments in our responses.

We kindly remind you to review our responses, and we are happy to address any further questions you may have.

Thank you once again for your thoughtful insights and for contributing to the improvement of our work.
Thank you for your valuable feedback and for taking the time to review our work. We have carefully addressed your comments in our responses.

We kindly remind you to review our responses, and we are happy to address any further questions you may have.

Thank you once again for your thoughtful insights and for contributing to the improvement of our work.
Thank you for your thoughtful feedback and for considering our responses. We are pleased that our clarifications and the comparison with data augmentation methods were helpful. We greatly appreciate your updated score and the opportunity to improve our work.
I thank the authors for their detailed responses, which clarify several questions regarding the interpretation of the drift and diffusion coefficients, the exploration of novel states, and the experiment details.
Additionally, the comparison with data augmentation methods enhances the empirical results and provides a more well-rounded understanding of the proposed Jacobian regularization method. 
I increase my score to reflect the changes.
### **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.
### 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. 


|   Gravity($m/s^2$) g                             |  g = 9.8 (default)| g = 6                   | g = 3                   | g = 1
|---------------------------------|--------------------------|-------------------------|-------------------------|-------------------------|
| Aug w. $\mathcal{N}(0.15, 0.1)$ | 847.19 ± 131.85          | 771.34 ± 88.112         | 550.4 ± 75.8            |  390.7 ± 94.28          |
| Jac Reg ($\lambda = 0.01$)      | **920.24 ± 39.952**      | **906.42 ± 42.664**     | **798.02 ± 95.936**     | **603.88 ± 162.224**    |

| Rotation ($^\circ$) $\alpha$     |  $\alpha = 20$           | $\alpha = 25$             | $\alpha = 30$
|---------------------------------|--------------------------|-------------------------|-------------------------|
| Aug w. $\mathcal{N}(0.15, 0.1)$ | 286.63 ± 81.678          | 284.09 ± 59.801         | 213.93 ± 42.44          | 
| Jac Reg ($\lambda = 0.01$)      | **423.81 ± 12.9**        | **301.84 ± 20.26**     | **226.04 ± 23.00**       | 

| $\beta\\%$ masked $\mathcal{N}(0.5 ,0.15)$  |  $\beta = 40$            | $\beta = 50$             | $\beta = 60$           | $\beta = 70$
|-------------------------------------------|--------------------------|-------------------------|-------------------------|-------------------------|
| Aug w. $\mathcal{N}(0.15, 0.1)$           | **846.76 ± 46.928**      | **767.92 ± 78.256**     | 373.08 ± 64.056     | 247.68 ± 54.576     |
| Jac Reg ($\lambda = 0.01$)                | 804.21 ± 80.369      | 725.81 ± 50.714     | **730.87 ± 65.263**     | **687.35 ± 63.222**    |

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).
### **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.