Title: Towards Unraveling and Improving Generalization in World Models
Venue: NeurIPS 2024
Accepted: no
Keywords: world models

reinforcement learning

generalization
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.

Abstract

World models have recently emerged as a promising approach for reinforcement learning (RL), as evidenced by its stimulating successes that world model based agents achieve state-of-the-art performance on a wide range of tasks in empirical studies. The primary goal of this study is to obtain a deep understanding of the mysterious generalization capability of world models, based on which we devise new methods to enhance it further. Thus motivated, we develop a stochastic differential equation formulation by treating the world model learning as a stochastic dynamic system in the latent state space, and characterize the impact of latent representation errors on generalization, for both cases with zero-drift representation errors and with non-zero-drift representation errors. Our somewhat surprising findings, based on both theoretic and experimental studies, reveal that for the case with zero drift, modest latent representation errors can in fact function as implicit regularization and hence result in generalization gain. We further propose a Jacobian regularization scheme to mitigate the compounding error propagation effects of non-zero drift, thereby enhancing training stability and generalization. Our experimental results corroborate that this regularization approach not only stabilizes training but also accelerates convergence and improves performance on predictive rollouts.

1 Introduction

Model-based reinforcement learning (RL) has emerged as a promising learning paradigm to improve 19 sample efficiency by enabling agents to exploit a learned model for the physical environment. Notably, 20 in recent works [14, 13, 15, 16, 21, 10, 32, 22] on world models, an RL agent learns the latent 21 dynamics model of the environment, based on the observations and action signals, and then optimizes 22 the policy over the learned dynamics model. Different from conventional approaches, world-model 23 based RL takes an end-to-end learning approach, where the building blocks (such as dynamics model, 24 perception and action policy) are trained and optimized to achieve a single overarching goal, offering 25 significant potential to improve generalization capability. For example, DreamerV2 and DreamerV3 26 achieve great progress in mastering diverse tasks involving continuous and discrete actions, image27 based inputs, and both 2D and 3D environments, thereby facilitating robust learning across unseen 28 task domains [14, 13, 15]. Recent empirical studies have also demonstrated the capacity of world 29 models to generalize to unseen states in complex environments, such as autonomous driving [19]. 30 Nevertheless, it remains not well understood when and how world models can generalize well in 31 unseen environments. 32

In this work, we aim to first obtain a deep understanding of the generalization capability of world 33 models by examining the impact of latent representation errors , and then to devise new methods to 34 enhance its generalization. While one may expect that optimizing a latent dynamics model (LDM) 35 prior to training the task policy would minimize latent representation errors and hence can achieve 36 better world model training, our somewhat surprising findings, based on both theoretical and empirical 37

studies, reveal that modest latent representation errors in the training phase may in fact be beneficial. 38 In particular, the alternating training strategy for world model learning, which simultaneously refines 39 both the LDM and the action policy, could actually bring generalization gain, because the modest 40 latent representation errors (and the corresponding induced gradient estimation errors) could enable 41 the world model to visit unseen states and thus lead to improved generalization capacities. For 42 instance, as shown in Table 1, our experimental results suggest that moderate batch sizes (e.g., 16 or 43 32) appear to position the induced errors within a regime conferring notable generalization benefits, 44 leading to higher generalization improvement, when compared to the cases with very small (e.g., 8) 45 or large (e.g., 64) batch sizes. 46

47

48

49

50

51

52

53

54

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.

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

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

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

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.

2 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.

3 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):

In this work, we consider a popular class of world models, including Dreamer and PlaNet, where { z , ˜ z , ˜ s } have distributions parameterized by neural networks' outputs, and are Gaussian when the outputs are known. It is worth noting that { z , ˜ z , ˜ s } may not be Gaussian and are non-Gaussian in general. This is because while z is conditional Gaussian, its mean and variance are random variables which are learned by the encoder with s and h being the inputs, rendering that z is non-Gaussian due to the mixture effect. For this setting, we have a continuous-time formulation where the latent dynamics model can be interpreted as stochastic differential equations (SDEs) with coefficient functions of known inputs. Due to space limitation, we refer to Proposition B.1 in the Appendix for a more detailed treatment.

Consider a complete, filtered probability space (Ω , F , {F t } t ∈ [0 ,T ] , P ) where independent standard 132 Brownian motions B enc t , B pred t , B seq t , B dec t are defined such that F t is their augmented filtration, and 133 T ∈ R as the time length of the task environment. We interpret the stochastic dynamics of LDM 134 with latent representation errors through coupled SDEs representing continuous-time analogs of the 135 discrete components: 136

where π ( h, ˜ z ) is a policy function as a local maximizer of value function and the stochastic process s 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 functions ¯ q enc , ¯ p and ¯ q dec are bounded and Borel-measurable over the interval [0 , T ] , and of class C 3 with bounded Lipschitz continuous partial derivatives. The initial values z 0 , h 0 , ˜ z 0 , ˜ s 0 are squareintegrable random variables.

Assumption 3.2. σ and ¯ σ are bounded and Borel-measurable and are of class C 3 with bounded Lipschitz continuous partial derivatives over the interval [0 , T ] .

3.1 Latent Representation Errors in CNN Encoder-Decoder Networks

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.

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 . As high-dimensional state space in image-based tasks frequently exhibit intrinsic lower-dimensional geometric structure , we adopt the latent manifold assumption, formally stated as follows:

Assumption 3.3. (Latent manifold assumption) For a positive integer k , there exists a d M -dimensional C k,α submanifold M (with C k +3 ,α boundary) with Riemannian metric g and has positive reach and also isometrically embedded in the state space S ⊂ R d S and d M < < d S , where the state probability measure is supported on. In addition, M is a compact, orientable, connected manifold.

Assumption 3.4. (Smoothness of state probability measure) Q is a probability measure supported on M with its Radon-Nikodym derivative q ∈ 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 . In practice, it is usually an easyto-sample distribution such as uniform distribution which is determined by a specific encoder-decoder architecture choice.

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:

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

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 d i × d i -1 is a structured sparse Toeplitz matrix from the convolutional filter { w ( i ) j } s ( i ) j =0 with filter length s ( i ) ∈ N + , b c i ∈ R d i is a bias vector, and σ is the ReLU activation function. For downsampling layers, A i ( x ) = D i ( x ) = ( x jm i ) ⌊ d i -1 /m i ⌋ j =1 , where D i : R d i × d i -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:

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( dF ( 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.

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.

3.2 Latent Representation Errors as Implicit Regularization towards Generalization

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), as one stochastic system. Let x t = ( z t , h t , ˜ z t , ˜ s t ) and B t = ( B enc t , B seq t , B pred t , B dec t ) :

where g , and ¯ g 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 .

3.2.1 The Case with Zero-drift Representation Errors

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

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

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 , ˜ z t , ˜ s t . 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 ,

where the regularization term R is given by R := ε P + ε 2 ( Q + 1 2 S ) , with

Square matrix Φ t is the stochastic fundamental matrix of the corresponding homogeneous equation:

and ξ k t is the shorthand for ∫ t 0 Φ -1 s ¯ σ k ( x 0 s , s ) dB k t . Additionally, H k ( x 0 s , s ) is represented by for ∑ k 1 ,k 2 ∂ 2 ¯ g k ∂x i ∂x j ( x 0 s , s ) ( ξ k 1 s ) i ( ξ k 2 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 , ˜ z t , ˜ s t . In principle, the PDE formulation corresponding to the PDFs of the perturbed x ε 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.

3.2.2 The Case with Non-Zero-Drift Representation Errors

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 ¯ g 0 as g from Equation (9) for convenience.

Corollary 3.8. (Additional Bias Induced by Non-Zero Drift Representation Error)

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 :

where the additional bias term ˜ R is given by ˜ R := ε ˜ P + ε 2 ( ˜ Q + ˜ S ) , with

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

and ˜ ξ t being the shorthand for ∫ t 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 by

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.

4 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:

with random variables h 0 , ˜ z 0 + ε 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 -˜ z 0 t ∥ ∥ 2 . As δ → 0 ,

where C is a constant dependent on T. J 1 and J 2 are Jacobian-related terms, and H 1 and H 2 are Hessianrelated terms.

The Jacobian-related terms J 1 and J 2 are defined as J 0 := exp ( F h + F z + P h ) , J 1 := exp ( ¯ P h ) ; the Hessian-related terms H 0 and H 1 are defined as H 0 := F hh + F hz + F zh + F zz + P hh , H 1 := ¯ P hh , 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 ˜ z t 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 .

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

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:

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:

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 above theoretical results have established a close connection of input-output Jacobian matrices with the stabilized generalization capacity of world models (shown in 18 under non-zero drift form), and perturbation magnitude in predictive rollouts (indicated in the presence of Jacobian terms in Theorem 4.1 and Corollary 4.2.) Based on this, we propose a regularizer on input-output Jacobian norm ∥ ∂g k ∂x ∥ F that could modulate ˜ ξ ( and in addition ξ k ) for stabilized implicit regularization.

The regularized loss function for LDM is defined as follows:

where 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.

5 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 regularization in model trained against a vanilla model 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 , σ 2 1 ) ; (2) rotation; and (3) noise applied to a percentage of the image, N ( µ 2 , σ 2 2 ) . (In Walker task, µ 1 = µ 2 = 0 . 5 , σ 2 2 = 0 . 15 ; in Quadruped task, µ 1 = 0 , µ 2 = 0 . 05 , σ 2 2 = 0 . 2 .) In each case of perturbations, we examine a collection of noise levels: (1) variance σ 2 from 0 . 05 to 0 . 55 ; (2) rotation degree α 20 and 30 ; and (3) masked image percentage β % from 25 to 75 .

400

400

WithJac.

WithJac.

800

800

Baseline

350

Baseline

Reward

300

300

600

600

250

400

WithJac.

Baseline

WithJac.

Baseline

200

200

400

0.50.150.250.350.450.55

25%35%45%55%65%75%

0.15

0.250.35

0.45

0.55

25%35%45%55%65%75%

N(O.5, o²)

Noise(mask)

N(O, o2)

Noise(mask)

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

It can be seen from Table 3 and Figure 1 that thanks to the adoption of Jacobian regularization in training, the rewards (averaged over 5 trials) are higher compared to the baseline, indicating improved generalization to unseen image states in all cases. The experimental results corroborate the findings in Corollary 3.8 that the regularized Jacobian norm could stabilize the induced implicit regularization.

Table 2: Evaluation on unseen states by various perturbation (Clean means without perturbation). λ = 0 . 01 .

Robustness against encoder errors. Next, we focus on the effects of Jacobian regularization on controlling the error process to the latent states z during training. Since it is very challenging, if not impossible, to characterize the latent representation errors and hence the drift therein explicitly, we consider to evaluate the robustness against two exogenous error signals, namely (1) zero-drift error with µ t = 0 , σ 2 t ( σ 2 t = 5 in Walker, σ 2 t = 0 . 1 in Quadruped), and (2) non-zero-drift error with µ t ∼ [0 , 5] , σ 2 t ∼ [0 , 5] uniformly. Table 3 shows that the model with regularization can consistently learn policies with high returns and also converges faster, compared to the vanilla case. This corroborates our theoretical findings in Corollary 3.8 that the impacts of error to loss L can be controlled through the model's Jacobian norm.

Faster convergence on tasks with extended horizon. We further evaluate the efficacy of Jacobian regularization in tasks with extended horizon, particularly by extending the horizon length in MuJoCo Walker from 50 to 100 steps. Table 4 shows that the model with regularization converges significantly faster ( ∼ 100K steps) than the case without Jacobian regularization in training. This corroborates results in Theorem 4.1 that regularizing the Jacobian norm can reduce error propagation.

6 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.

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

References

Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research , 47:253-279, 2013.

Alexander Camuto, Matthew Willetts, Umut ¸ Sim¸ sekli, Stephen Roberts, and Chris Holmes. Explicit regularisation in gaussian noise injections, 2021.

Henri Cartan. Differential calculus on normed spaces . Createspace Independent Publishing Platform, North Charleston, SC, August 2017.

Bo Chang, Minmin Chen, Eldad Haber, and Ed H. Chi. Antisymmetricrnn: A dynamical system view on recurrent neural networks, 2019.

Zhengdao Chen, Jianyu Zhang, Martin Arjovsky, and Léon Bottou. Symplectic recurrent neural networks, 2020.

Bernard Dacorogna and Jürgen Moser. On a partial differential equation involving the jacobian determinant. Annales de l'I.H.P. Analyse non linéaire , 7(1):1-26, 1990.

Carl Doersch. Tutorial on variational autoencoders. arXiv preprint arXiv:1606.05908 , 2016.

Sean C Duncan. Minecraft, beyond construction and survival. 2011. 384

385

386

Lawrence Craig Evans and Ronald F Gariepy. Measure theory and fine properties of functions, revised edition . Textbooks in Mathematics. Apple Academic Press, Oakville, MO, April 2015.

C. Daniel Freeman, Luke Metz, and David Ha. Learning to predict without looking ahead: 387 World models without forward prediction. Thirty-third Conference on Neural Information 388 Processing Systems (NeurIPS 2019) , 2019. 389

Alison L. Gibbs and Francis Edward Su. On choosing and bounding probability metrics. 390 International Statistical Review / Revue Internationale de Statistique , 70(3):419-435, 2002. 391

David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122 , 2018. 392

Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: 393 Learning behaviors by latent imagination, 2020. 394

Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and 395 James Davidson. Learning latent dynamics for planning from pixels. In International conference 396 on machine learning , pages 2555-2565. PMLR, 2019. 397

Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with 398 discrete world models, 2022. 399

Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains 400 through world models. arXiv preprint arXiv:2301.04104 , 2023. 401

Paul Louis Hennequin, R. M. Dudley, H. Kunita, and F. Ledrappier. Ecole d'ete de Probabilites 402 de Saint-Flour XII-1982 . Springer-Verlag, 1984. 403

Judy Hoffman, Daniel A. Roberts, and Sho Yaida. Robust learning with jacobian regularization, 404 2019. 405

Anthony Hu, Lloyd Russell, Hudson Yeo, Zak Murez, George Fedoseev, Alex Kendall, Jamie 406 Shotton, and Gianluca Corrado. Gaia-1: A generative world model for autonomous driving. 407 arXiv preprint arXiv:submit/1234567 , Sep 2023. Submitted on 29 Sep 2023. 408

Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping 409 Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima, 410 2017. 411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

Samuel Kessler, Mateusz Ostaszewski, Michał Bortkiewicz, Mateusz ˙ Zarski, Maciej Wołczyk, Jack Parker-Holder, Stephen J. Roberts, and Piotr Miło´ s. The effectiveness of world models for continual reinforcement learning. CoLLAs 2023 , 2023.

Kuno Kim, Megumi Sano, Julian De Freitas, Nick Haber, and Daniel Yamins. Active world model learning with progress curiosity. In Proceedings of the 37th International Conference on Machine Learning (ICML) , 2020.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Yann LeCun, Bernhard Boser, John S Denker, Donnie Henderson, Richard E Howard, Wayne Hubbard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation , 1(4):541-551, 1989.

John M. Lee. Introduction to Riemannian Manifolds . Springer International Publishing, 2018.

Zhong Li, Jiequn Han, Weinan E, and Qianxiao Li. Approximation and optimization theory for linear continuous-time recurrent neural networks. Journal of Machine Learning Research , 23(42):1-85, 2022.

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.

Lynn Harold Loomis and Shlomo Sternberg. Advanced calculus (revised edition) . World Scientific Publishing, Singapore, Singapore, March 2014.

Guohao Shen, Yuling Jiao, Yuanyuan Lin, and Jian Huang. Approximation with cnns in sobolev space: with applications to classification. In NeurIPS , Oct 2022.

J. Michael Steele. Stochastic calculus and Financial Applications . Springer, 2001.

Colin Wei, Sham Kakade, and Tengyu Ma. The implicit and explicit regularization effects of dropout, 2020.

Philipp Wu, Alejandro Escontrela, Danijar Hafner, Pieter Abbeel, and Ken Goldberg. Daydreamer: World models for physical robot learning. In Proceedings of The 6th Conference on Robot Learning , volume 205 of PMLR , pages 2226-2240, 2023.

Yong Yu, Xiaosheng Si, Changhua Hu, and Jianxun Zhang. A review of recurrent neural 440 networks: Lstm cells and network architectures. Neural computation , 31(7):1235-1270, 2019. 441

442

443

444

445

446

447

448

449

450

451

452

Supplementary Materials

In this appendix, we provide the supplementary materials supporting the findings of the main paper on the latent representation of latent representations in world models. The organization is as follows:

In Section A, we provide proof on showing the approximation capacity of CNN encoderdecoder architecture in latent representation of world models.

In Section B, we provide proof on implicit regularization of zero-drift errors and additional effects of non-zero-drift errors by showing a proposition on the general form.

In Section C, we provide proof on showing the effects of non-zero-drift errors during predictive rollouts by again showing a result on the general form.

In Section D, we provide additional results and implementation details on our empirical studies.

453

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

A Approximation Power of Latent Representation with CNN Encoder and Decoder

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 by

Definition 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 , ψ α : U α → V α ⊆ R n such that

⋃ α ∈ A U α ⊇ M , meaning that the the open subsets form an open cover,

Each chart ψ α is a diffeomorphism that is a smooth map with smooth inverse (in the C k,α sense),

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 d vol M is the canonical volume form on X if for any point x ∈ X , for a chosen local coordinate chart ( x 1 , ..., x d ) , d vol M = √ det g ij dx 1 ∧ ... ∧ dx d , where g ij ( x ) := g ( ∂ ∂x i , ∂ ∂x j )( x ) .

Then the induced volume measure by the canonical volume form d vol X is denoted as µ X , defined by µ X : A ↦→ ∫ A d vol 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 .

Assumption A.3. (Latent manifold assumption) For a positive integer k , there exists a d M -dimensional C k,α submanifold M (with C k +3 ,α boundary) with Riemannian metric g and has positive reach and also isometrically embedded in the state space S ⊂ R d S and d M < < d S , where the state probability measure is supported on. In addition, M is a compact, orientable, connected manifold.

Assumption A.4. (Smoothness of state probability measure) Q is a probability measure supported on M with its Radon-Nikodym derivative q ∈ 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:

where A i 's are either convolutional or downsampling operators. For convolutional layers,

where W c i ∈ R d i × d i -1 is a structured sparse Toeplitz matrix from the convolutional filter { w ( i ) j } s ( i ) j =0 with filter length s ( i ) ∈ N + , b c i ∈ R d i is a bias vector, and σ is the ReLU activation function.

For downsampling layers,

495

496

497

498

499

500

501

502

503

504

505

506

507

508

509

510

511

512

513

514

515

516

517

518

519

520

521

522

523

524

RNs

IRNs

M

M

Encoder

Z

→Decoder

P

where D i : R d i × d i -1 is the downsampling operator with scaling parameter m i ≤ d i -1 in the i -th layer. The convolutional and downsampling operations are elaborated in Appendix [63]. We examine the class of functions represented by CNNs, denoted by F CNN, defined as:

For more details in the definitions of CNN functions, we refer to [29].

̸

Assumption A.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( dF ( x )) = 0 .

Theorem A.6. (Approximation Error of Latent Representation). Under Assumption A.3, A.4 and A.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.

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 satisfy

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 measure g # P = Q , that is for any measurable subset A ∈ B ( M ) , Q ( A ) = P ( g -1 ( A )) .

Proof. ( Proposition A.7 ) Let ω := p d vol 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 d vol N = ∫ N p dµ N = P ( N ) = 1 . Similarly, let η := q d vol 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 point x ∈ M , there exists an open neighborhood U x of x and a discrete set D x such that F -1 ( U ) = ⊔ α ∈ D V α ⊂ N and F | V α = V α → U is a diffeomorphism. Furthermore, | D x | = | D y | for any points x, 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 , 525

¯ η is well-defined as it is not dependent on the choice of neighborhoods and the sum and 1 | D x | are 526 always finite. Furthermore, ¯ η is a C k,α volume form on M , as p ◦ ( F ∣ ∣ V α -1 ) is C k,α . 527

528

529

530

531

532

533

534

535

536

537

538

539

540

541

542

543

544

545

546

547

Notice that F ∣ ∣ V α -1 is orientation-preserving as det dF ∣ ∣ V α -1 = 1 det dF ∣ ∣ 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 α -1 is 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 ,α . Let g := ψ ◦ F , then it holds that g ∗ η = ( ψ ◦ F ) ∗ η = F ∗ ◦ ψ ∗ η = F ∗ ¯ η = ω .

Then, for any measurable subset A on the manifold M , we verify that Q ( A ) = ∫ A η = ∫ g -1 ( A ) g ∗ η = ∫ g -1 ( A ) ω = ∫ g -1 ( A ) p d vol 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,

where the last inequality follows from the special case ρ = 0 of Theorem 2.4 in [29]. 548

Similarly, for decoder, from Proposition A.7, there exists an C k +1 ,α oracle map ¯ g : Z → M such 549 that the pushforward measure ¯ g # P = Q . 550

551

552

553

554

555

556

557

558

559

560

561

562

563

564

565

566

567

568

569

570

571

572

573

574

575

576

577

578

579

580

581

582

583

584

585

586

B Explicit Regularization of Latent Representation Error in World Model Learning

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 motions B enc t , B pred t , B seq t , B dec t 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:

where π ( h, ˜ z ) is a policy function as a local maximizer of value function and the stochastic process s 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,

with 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 ¯ q enc ( 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 functions ¯ q enc , ¯ p and ¯ q dec are bounded and Borel-measurable over the interval [0 , T ] , and of class C 3 with bounded Lipschitz continuous partial derivatives. The initial values z 0 , h 0 , ˜ z 0 , ˜ s 0 are squareintegrable random variables.

Assumption B.3. σ and ¯ σ are bounded and Borel-measurable and are of class C 3 with bounded Lipschitz continuous partial derivatives over the interval [0 , T ] .

One of our main results is the following:

Theorem B.4. (Explicit Regularization Induced by Zero-Drift Representation Error) Under Assumption B.2 and B.3 and considering a loss function L ∈ C 2 , the explicit effects of the zero-drift error can be marginalized out as follows:

as ε → 0 , where the regularization term R is given by R := ε P + ε 2 ( Q + 1 2 S ) . 587 Each term of R is as follows: 588

where square matrix Φ t is the stochastic fundamental matrix of the corresponding homogeneous 589 equation: 590

and ξ k t is as the shorthand for ∫ t 0 Φ -1 s ¯ σ k ( x 0 s , s ) dB k t . Additionally, H k ( x 0 s , s ) is represented by for ∑ k 1 ,k 2 ∂ 2 ¯ g k ∂x i ∂x j ( x 0 s , s ) ( ξ k 1 s ) i ( ξ k 2 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 by

with initial values x (0) = x 0 ,

591

592

593

594

595

Proposition B.5. Suppose that f is a real-valued function that is C 2 . Then it holds that, with 596 probability 1, as ε → 0 , for t ∈ [0 , T ] , 597

where the stochastic process x 0 t is the solution to SDE 32 with ε = 0 , with its first and second-order 598 derivatives w.r.t ε denoted as ∂ ε x 0 t , ∂ 2 ε x 0 t . 599

Furthermore, it holds that ∂ ε x 0 t , ∂ 2 ε x 0 t satisfy the following SDEs with probability 1, 600

with initial values ∂ ε x (0) = 0 , ∂ 2 ε x (0) = 0 , where 601

602

603

604

605

with 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 . 606 Our main tool is an important result on smoothness of solutions w.r.t. initial data from Theorem 3.1 607 from Section 2 in [17]. 608

For ∂ ε x , consider the SDEs 609

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 as

with initial values x (0) = x 0 , ε (0) = ε . From an application of Theorem 3.1 from Section 2 in [17] 610 on *, we have ∂ ε x that satisfies the following SDE with probability 1: 611

with initial value ∂ ε x 0 = 0 ∈ R n , with probability 1, where x t is the solution to Equation (35) and 612 the functions α ε k are given by 613

where k = 0 , ..., m . 614

To characterize ∂ 2 ε x t , consider the following SDEs 615

with initial value x (0) = x 0 , ∂ ε x (0) = 0 , ε (0) = ε . 616

From a similar application of Theorem 3.1 from Section 2 in [17], the second derivative ∂ 2 ε x satisfies 617 the following SDE with probability 1: 618

619

620

621

with 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

When ε = 0 in the obtained SDEs (35), (36) and (37), the corresponding solutions of which are 622 x 0 t , ∂ ε x 0 t , ∂ 2 ε x 0 t , we now have the following: 623

with initial values x (0) = x 0 , ∂ ε x (0) = 0 , ∂ 2 ε x (0) = 0 . In particular, Ψ k := β 0 k is given by 624

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 625 the following lemma B.6: 626

Lemma B.6.

Proof. To simplify the notations, we take the liberty to write constants as C and notice that C is not 627 necessarily identical in its each appearance. 628

(1) We first show that E sup t ∈ [0 ,T ] ∥ ∥ x 0 t ∥ ∥ 2 is bounded. 629

From Equation (38), we have that

By Jensen's inequality. it holds that 630

For the second term on the right hand side, it is a sum over k from 0 to m by Einstein notation. 631 For k = 0 , recall that we write t as B 0 t : 632

where we used Jensen's inequality, the assumption on the linear growth, the inequality property of 633 sup and Fubini's theorem, respectively. 634

For k is equal to 1 , . . . , m , 635

where (iv) holds from the Burkholder-Davis-Gundy inequality as ∫ t 0 g k ( x 0 τ , τ ) dB τ is a continuous 636 local martingale with respect to the filtration F t ; and then one can obtain (v) by following a similar 637 reasoning of (ii) and (iii). 638

Hence, now from the previous inequality (41),

By the Gronwall's lemma, it holds true that

As x 0 is square-integrable by assumption, therefore we have shown that E sup t ∈ [0 ,T ] ∥ ∥ x 0 t ∥ ∥ 2 is 639 bounded. 640

We then show that E sup t ∈ [0 ,T ] || ∂ ε x 0 t || 2 is also bounded. 641

From the SDE (39), as we have derived that

then we have

For k = 0 , we have 642

where to get to (vi), we used Jensen's inequality; for (vii), we used the linear growth assumption an 643 η 0 , then we obtain (viii) by as derivatives of function g 0 are bounded by assumption. 644

Similarly, for k = 1 , ..., m , 645

where we obtain (ix) by the Burkholder-Davis-Gundy inequality and (x) by following similar steps as 646 have shown in (vii) and (viii). 647

We are now ready to sum up each term to acquire a new inequality: 648

By Gronwall's lemma, we have that 649

As it is previously shown that E sup t ∈ [0 ,τ ] ∥ x ◦ ( t ) ∥ 2 is bounded, it is clear that E sup t ∈ [0 ,T ] ∥ ∥ ∂ ε x 0 t ∥ ∥ 2 650 is bounded too. 651

(3) From similar steps, one can also show that E sup t ∈ [0 ,T ] ∥ ∥ ∂ 2 ε x 0 t ∥ ∥ 2 is bounded. 652

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 653 proving the following lemma. 654

Lemma B.7. For a given ε ∈ R , let

where the stochastic process x ε t is the solution to Equation (32). Then it holds true that 655

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:

Hence, we have ε 3 R ε ( t ) = ∑ 1 k =0 θ k ( t ) + φ k ( t ) + σ k ( t ) . 656

For θ k ( t ) , we have 657

where to obtain (i) we used Jensen's inequality when k = 0 and by the Burkholder-Davis-Gundy 658 inequality when k = 1 , used the Lipschitz condition of g k to obtain (ii), and for (iii), it is because 659 ε 3 R ε ( t ) = ˜ x ε t -x ε t . 660

We note that from Taylor's theorem, for any s ∈ [0 , t ] , k = 0 , 1 , there exists some ε s ∈ (0 , ε ) s.t. 661

For φ k ( t ) , we have 662

where for (iv), we used Equation (42) and Jensen's inequality for k = 0 and the Burkholder-Davis663 Gundy inequality for k = 1 ; to obtain (v), we applied Jensen's equality; we then derived (vi) from 664 the Lipschitz conditions of g k and Ψ k ; and finally another application of Jensen's inequality gives 665 (vii) which is bounded as a result from the Lemma B.6. 666

667

For σ k ( t ) , 668

where we obtained (ix) by Jensen's inequality when k = 0 and by Burkholder-Davis-Gundy inequality 669 when k = 1 , and (x) by the linear growth assumption on η k ; one can see that (xi) is bounded by 670 recalling the Lemma B.6 and the assumption that η k has bounded derivatives. 671

Hence, by Jensen's inequality and Gronwall's lemma, we have 672

673

674

675

676

677

678

679

680

681

682

Therefore, E sup ∥R ε ( t ) ∥ 2 is bounded.

Finally, it is now straightforward to show Equation (33) by applying a second-order Taylor expansion on 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:

where n × n matrix Φ t is the fundamental matrix of the corresponding homogeneous equation: 683

with initial value 684

It is worthy to note that the fundamental matrix Φ t is non-deterministic and when ∂g i ∂x and ∂g j ∂x 685 commutes, Φ t has explicit solution 686

Having obtained the explicit solutions, one can plug in corresponding terms and obtain the results of 687 Theorem 3.7 ) after a Taylor expansion of the loss function L . 688

C Error Accumulation During the Inference Phase and its Effects to Value 689 Functions 690

691

Theorem C.1. (Error accumulation due to initial representation error )

692

where 693

694

where for brevity, when functions always have inputs (˜ z 0 t , h 0 t , t ) , we adopt the shorthand to write, for 695 example, f (˜ z 0 t , h 0 t , t ) as f . 696

Before proving the main result C.1, we first show the general case of perturbation in initial values. 697 Consider the following general system with noise at the initial value: 698

where 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 ) .

699

700

701

702

703

704

705

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 derivatives ∂ i x 0 t , ∂ 2 ij x 0 t , respectively.

Proposition C.2. Let δ := E ∥ ε ∥ , it holds that 706

as δ → 0 . 707

Proof. Similar to the previous section, for notational convenience, we write t as B 0 t and employs 708 Einstein summation notation. Hence, Equation (47) can be shorten as 709

with initial values x (0) = x 0 + ε . 710

711

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 712 following SDE with probability 1, 713

with initial value ∂ i x ε 0 to be the unit vector e i =(0 , 0 , . . . , 1 , . . . , 0) that is all zeros except one in 714 the i th coordinate. 715

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 716 satisfy the following SDE with probability 1, 717

with the initial value ∂ ij x ε (0) = e j , where 718

For the next step, we show that with probability 1, the following holds 719

720

721

722

723

724

Indeed, for E sup t ∈ [0 ,T ] ∥ ∥ ∂ i x 0 t ∥ ∥ 2 , it holds that 725

Similarly, for E sup t ∈ [0 ,T ] ∥ ∥ ∂ 2 ij x 0 t ∥ ∥ 2 , it holds that 726

Therefore, we could obtain the proposition by applying Jensen's inequality to Equation (54) and 727 plugging with 56 and 57. 728

as ∥ ε ∥ → 0 .

One can follow the similar steps of proofs for Lemma (B.6) and (B.7) in the previous section to show that E sup t ∈ [0 ,T ] ∥ ∥ x 0 t ∥ ∥ 2 , E sup t ∈ [0 ,T ] ∥ ∥ ∂ i x 0 t ∥ ∥ 2 , E sup t ∈ [0 ,T ] ∥ ∥ ∂ 2 ij x 0 t ∥ ∥ 2 and the remainder term are bounded. Hence, Equation (54) holds with probability 1.

Now we are ready to prove Theorem C.1. We note that one could then obtain Corollary 4.2 without 729 much more effort by a standard application of Taylor's theorem. 730

731

732

733

734

735

736

Proof. (Proof for Theorem C.1)

At ( h t , ˜ z t , π ( h t , ˜ z t )) , where the local optimal policy π ( h t , ˜ z t ) , denoted as a ∗ t , there exists an open neighborhood V ⊆ A of a ∗ t such that a ∗ t is the local maximizer for Q ( h t , ˜ z t , · ) by definition.

-

1

D Experimental Details 742

743

In this section, we provide additional details and results beyond thoese in the main paper.

744

D.1 Model Implementation and Training

Our baseline is based on the DreamerV2 Tensorflow implementation. Our theoretical and empirical 745 results should not matter on the choice of specific version; so we chose DreamerV2 as its codebase 746 implementation is simpler than V3. We incorporated a computationally efficient approximation of 747 the Jacobian norm for the sequence model, as detailed in [18], using a single projection. During our 748 experiments, all models were trained using the default hyperparameters (see Table 5) for the MuJoCo 749 tasks. The training was conducted on an NVIDIA A100 and a GTX 4090, with each session lasting 750 less than 15 hours. 751

752

753

754

755

756

757

758

759

760

D.2 Additional Results on Generalization on Perturbed States

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 , σ 2 1 ) ; (2) rotation; and (3) noise applied to a percentage of the image, N ( µ 2 , σ 2 2 ) . (In Walker task, µ 1 = µ 2 = 0 . 5 , σ 2 2 = 0 . 15 ; in Quadruped task, µ 1 = 0 , µ 2 = 0 . 05 , σ 2 2 = 0 . 2 .) In each case of perturbations, we examine a collection of noise levels: (1) variance σ 2 from 0 . 05 to 0 . 55 ; (2) rotation degree α 20 and 30 ; and (3) masked image percentage β % from 25 to 75 .

D.3 Walker Task

D.4 Quardruped Task 761

762

D.5 Additional Results on Robustness against Encoder Errors

In this experiment, we evaluate the robustness of model trained with Jacobian regularization against 763 two exogenous error signals (1) zero-drift error with µ t = 0 , σ 2 t ( σ 2 t = 5 in Walker, σ 2 t = 0 . 1 in 764 Quadruped), and (2) non-zero-drift error with µ t ∼ [0 , 5] , σ 2 t ∼ [0 , 5] uniformly. λ weight of Jacobian 765 regularization is 0 . 01 . In this section, we included plot results of both evaluation and training scores. 766

767

D.5.1 Walker Task

Under the Walker task, Figures 3 and 4 show that model with regularization is significantly less 768 sensitive to perturbations in latent state z t compared to the baseline model without regularization. 769 This empirical observation supports our theoretical findings in Corollary 3.8, which assert that the 770 impact of latent representation errors on the loss function L can be effectively controlled by regulating 771 the model's Jacobian norm.

1000

With Jacobian regularizer

1000

WithJacobianregularizer

No Jacobian regularizer

No Jacobian regularizer

800

800

eval score

600

train score

600

400

400

200

200

0

0

0

200000

400000

600000

800000

0

200000

400000

600000

800000

num steps num steps

.

1000

With Jacobian regularizer

WithJacobianregularizer

No Jacobian regularizer

No Jacobian regularizer

800

800

600

eval score

train score

600

400

400

200

200

0

100000

200000

300000

400000

500000

600000

0

100000

200000

300000

400000

500000

600000

num steps

num steps

772

D.5.2 Quadruped Task 773

Under the Quadruped task,we initially examined a smaller latent error process ( µ t = 0 , σ 2 t = 0 . 1 ) and 774 observed that the model with Jacobian regularization converged significantly faster, even though the 775 adversarial effects on the model without regularization were less severe (Figure 5). When considering 776 the more challenging latent error process ( µ t ∼ [0 , 5] , σ 2 t ∼ [0 , 5] ), we noted that the regularized 777 model remained significantly less sensitive to perturbations in latent state z t , whereas the baseline 778 model struggled to learn (Figure 6). These empirical observations reinforce our theoretical findings 779 in Corollary 3.8, demonstrating that regulating the model's Jacobian norm effectively controls the 780 impact of latent representation errors.

1000

With jac reg lambda=0.05

1000

With jac reg lambda=0.05

No Jac reg

No Jac reg

800

008

evalscore

600

train score

600

400

400

200

200

0

0

0.0

0.2

0.4

0.6

0.8

1.0

1.2

0.0

0.2

0.4

0.6

0.8

1.0

1.2

num steps

1e6

num steps

1e6

1000

1000

With jac reg lambda=0.05

No Jac reg

No Jac reg

800

800

eval score

600

train score

600

400

400

200

200

0

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

1e6

1e6

num steps num steps

D.6 Additional Results on Faster convergence on tasks with extended horizon. 782

In this experiment, we evaluate the efficacy of Jacobian regularization in extended horizon tasks, 783 specifically by increasing the horizon length in MuJoCo Walker from 50 to 100 steps. We tested two 784 regularization weights λ = 0 . 1 and λ = 0 . 05 . Figure 7 demonstrates that models with regularization 785 converge faster, with λ = 0 . 05 achieving convergence approximately 100,000 steps ahead of the 786 model without Jacobian regularization. This supports the findings in Theorem 4.1, indicating that 787 regularizing the Jacobian norm can reduce error propagation, especially over longer time horizons.

1000

1000

With Jacobian reg lambda=0.05

No Jac reg

WithJacobianreg lambda=0.1

800

800

eval score

600

train score

600

400

400

200

200

0

0

0

50000

100000

150000

200000

250000

0

50000

100000

150000

200000

250000

300000

num steps num steps

788

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

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.

2. Limitations

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.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: For each of our theoretical results, we state the assumptions required in both the main text and the provided appendix. We provide the full proofs of all of our theoretical results in Sections A, B and C in Appendix,

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The full source code required to reproduce the experimental results is included in the submission.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in

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.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Our task environments Walker and Quardruped are from open source package MuJoCo. Our baseline implementation is from open source codebase DreamerV2. Our implementation of Jacobian regularization has a full description in Section D.1.

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.

6. Experimental Setting/Details

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.

7. 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:

947

948

949

950

951

952

953

954

955

956

957

958

959

960

961

962

963

964

965

966

967

968

969

970

971

972

973

974

975

976

977

978

979

980

981

982

983

984

985

986

987

988

989

990

991

992

993

994

995

996

997

998

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.

8. 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).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: This work conforms with the NeurIPS Code of Ethics.

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

10. 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.

999

1000

1001

1002

1003

1004

1005

1006

1007

1008

1009

1010

1011

1012

1013

1014

1015

1016

1017

1018

1019

1020

1021

1022

1023

1024

1025

1026

1027

1028

1029

1030

1031

1032

1033

1034

1035

1036

1037

1038

1039

1040

1041

1042

1043

1044

1045

1046

1047

1048

1049

1050

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks as we do not have any released data or models. Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We cited the baseline implementation in Section D.1.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

1051

1052

1053

1054

1055

1056

1057

1058

1059

1060

1061

1062

1063

1064

1065

1066

1067

1068

1069

1070

1071

1072

1073

1074

1075

1076

1077

1078

1079

1080

1081

1082

1083

1084

1085

1086

1087

1088

1089

1090

1091

1092

1093

1094

1095

1096

1097

1098

1099

1100

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.

13. 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.

14. 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.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

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.

1101

1102

1103

1104

1105

1106

1107

1108

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- World models are a popular area of research in the RL community, but there is a lack of theoretical understanding. This paper takes one step towards theoretically analyzing the generalization capabilities of world models.\n- The analysis of the effect of latent representation error is a novel theoretical contribution, to the best of my knowledge.\n- The results in the paper seem mathematically sound and provide useful insights. The empirical results demonstrate that the Jacobian regularization, which naturally arises from the theoretical analysis, is helpful in improving robustness.\n- As a very theory-heavy paper, the authors structured the writing such that it makes it easy to follow each individual result (though there is some room for improvement here, see weaknesses).\n\nWEAKNESSES:\nWhile 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.\n\n- 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.\n- 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.\n- 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.\n- 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.\n- 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?\n\nQUESTIONS:\n- 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.\n- What do the drift and diffusion terms corresponding to the latent representation error signify?\n- 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)?\n- 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.\n- 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.\n- How much computational overhead is added by the calculation of the Jacobian regularization term?\n\n[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.\n\n[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*.\n\nLIMITATIONS:\nThere 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-04"}

comments: 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.

STRENGTHS:
+ A deep understanding of the generalization of world models via stochastic differential equation formulation;
+ A careful study of the different effects of zero drift and non-zero drift on gn

WEAKNESSES:
+ 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:
See the Weaknesses part and explain how to extend the analysis to more general cases.

LIMITATIONS:
NA
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. This work investigates an interesting aspect, the generalization of world models that learn the dynamics of the environment. Very limited work has been done in this facet of RL, thus it will share significant insights with the DRL research community.\n\n2. The paper followed a structured methodology to analyze the world model and its representation errors. They interpret the learned model function as stochastic differential equations (SDEs) and model the variation as Brownian motions. \n\n3. I liked the way they theoretically analyzed it case-by-case and established connections with prior findings. \n\n4. The paper articulately presents the findings of zero-drift error as a regularizer and the Jacobian correction term for non-zero-drift representation error.  It systematically proves its hypotheses and shows evidence against the claims. They presented corresponding formulas and interpretations.\n\nWEAKNESSES:\n1. 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. \n\n2. 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.\n\nQUESTIONS:\n1. 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?\n\n2. Do you observe any substantial relation between $\\lambda$ and the task horizon?\n\nLIMITATIONS:\nWhile 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}