Dreaming Smoothly and Sample Efficiently with Gradient Penalized Latent Dynamics

# Dreaming Smoothly and Sample Efficiently with Gradient Penalized Latent Dynamics
Source: [https://arxiv.org/html/2605.23089](https://arxiv.org/html/2605.23089)
Romil V\. Sonigra Department of Electrical and Computer Engineering Texas A&M University College Station, TX 77843 romils@tamu\.edu &P\. R\. Kumar Department of Electrical and Computer Engineering Texas A&M University College Station, TX 77843 prk@tamu\.edu

###### Abstract

Model\-based reinforcement learning improves sample efficiency by learning a world model\. However, existing latent world models such as DreamerV3 do not explicitly enforce local smoothness in their learned transition dynamics, leaving a useful inductive bias for transition dynamics learning unexploited\. We propose GPLD, a gradient\-penalized latent dynamics regularizer for DreamerV3 that applies a row\-wise Jacobian penalty to the posterior latent distribution to encourage locally smooth transition learning\. We show that this penalty can be interpreted as the continuous\-latent analog of finite\-difference smoothing of transition laws in discrete embedded\-state MDPs, and estimate it efficiently using Hutchinson\-style stochastic probes\. Empirically, across DeepMind Control proprioceptive tasks, GPLD improves aggregate sample efficiency, with particularly strong gains on higher\-complexity locomotion environments\. On more challenging quadruped tasks, GPLD reaches high\-return behavior earlier and exhibits more consistent late\-stage learning over longer horizons\. Explicit local smoothness regularization is a simple and effective way to improve latent world models for smooth continuous control environments\. Code for GPLD is available at[github\.com/romils9/gpld\-mbrl](https://github.com/romils9/gpld-mbrl)\.

## 1Introduction

Model\-based reinforcement learning \(MBRL\) improves sample efficiency by learning a predictive world model that can be used for planning and policy optimization\(Ha and Schmidhuber,[2018](https://arxiv.org/html/2605.23089#bib.bib12)\)\. DreamerV3 has demonstrated that latent world models can scale to diverse continuous\-control tasks by learning compact recurrent state\-space representations and training policies through imagined rollouts\(Hafneret al\.,[2025](https://arxiv.org/html/2605.23089#bib.bib13)\)\. Despite this success, standard latent world models do not explicitly exploit a basic structural prior of many continuous\-control systems: nearby states induce similar short\-horizon transition behavior\. Encoding this prior directly in the learned latent dynamics can help the world model share information across nearby states, reducing the amount of interaction data needed to learn useful transition structure\.

We study how to exploit this local smoothness prior through*gradient\-penalized latent dynamics*\(GPLD\), a differentiable regularizer for latent world\-model dynamics\. This regularizer is motivated through a discrete\-to\-continuous smoothness argument\. In a finite embedded\-state MDP, local smoothness can be imposed by penalizing differences between transition laws at neighboring states\. When the state representation is continuous and the transition model is differentiable, these normalized finite differences become directional derivatives of the learned dynamics\. Averaging over local directions yields a Frobenius Jacobian penalty\. This provides a principled route from neighborhood smoothing in tabular MDPs to the gradient penalty used by GPLD\.

We implement GPLD in DreamerV3 and evaluate it on DeepMind Control tasks\. Empirically, GPLD improves aggregate sample efficiency on proprioceptive tasks, with the largest gains attained on higher\-complexity locomotion environments\. On difficult quadruped tasks, the benefit becomes clearer over longer horizons, where GPLD reaches high\-return behavior earlier and exhibits more consistent late\-stage learning\. Pixel\-observation results show milder aggregate gains, suggesting that the effect of latent smoothness regularization is weaker when dynamics learning is coupled to high\-dimensional visual encoding\.

Our contributions are:

1. 1\.A discrete\-to\-continuous justification showing how finite\-difference smoothing of tabular transition laws leads to Frobenius Jacobian regularization in differentiable latent models\.
2. 2\.GPLD, a posterior Jacobian regularizer for DreamerV3 that explicitly encourages local smoothness in latent probability maps, yielding a normalized aggregate gain of17\.7%17\.7\\%across the DeepMind Control proprioceptive benchmark\.
3. 3\.Evaluation of GPLD on DeepMind Control, showing a34\.6%34\.6\\%normalized aggregate gain on higher\-complexity proprioceptive locomotion tasks, clearer long\-horizon benefits on difficult quadruped tasks, and ablations over the main design choices\.

Code for GPLD is available at[github\.com/romils9/gpld\-mbrl](https://github.com/romils9/gpld-mbrl)\.

## 2Related Work

#### Smoothness in reinforcement learning\.

Smoothness has long been recognized as important for generalization, stable prediction, and control in reinforcement learning\(Boyan and Moore,[1994](https://arxiv.org/html/2605.23089#bib.bib3); Munos and Szepesvári,[2008](https://arxiv.org/html/2605.23089#bib.bib15); Asadiet al\.,[2018](https://arxiv.org/html/2605.23089#bib.bib1); Christmannet al\.,[2024](https://arxiv.org/html/2605.23089#bib.bib11)\)\. Unconstrained neural networks can produce high\-frequency oscillations that degrade stability and generalization\(Boyan and Moore,[1994](https://arxiv.org/html/2605.23089#bib.bib3); Roscaet al\.,[2020](https://arxiv.org/html/2605.23089#bib.bib17); Christmannet al\.,[2024](https://arxiv.org/html/2605.23089#bib.bib11)\), and such errors can compound in multi\-step value or dynamics prediction\(Thrun and Schwartz,[1993](https://arxiv.org/html/2605.23089#bib.bib22); Venkatramanet al\.,[2015](https://arxiv.org/html/2605.23089#bib.bib23)\)\. In model\-based RL,Asadiet al\.\([2018](https://arxiv.org/html/2605.23089#bib.bib1)\)show that Lipschitz constraints on learned transition models yield tighter multi\-step estimation error bounds, directly linking transition smoothness to world\-model accuracy\. These works motivate smoothness as a useful inductive bias for learned dynamics, but do not provide a local Jacobian regularizer for stochastic latent world models\.

#### Gradient penalties and control smoothness\.

Gradient penalties provide a differentiable way to encourage local smoothness and have been widely used in representation and generative modeling\(Gulrajaniet al\.,[2017](https://arxiv.org/html/2605.23089#bib.bib7);terjék2020adversariallipschitzregularization; Goodfellowet al\.,[2015](https://arxiv.org/html/2605.23089#bib.bib8)\)\. In reinforcement learning, related methods often regularize value functions, policies, or action trajectories\. For example,Wanget al\.\([2025](https://arxiv.org/html/2605.23089#bib.bib24)\)apply gradient penalties to Q\-functions in model\-free RL, whileChenet al\.\([2024](https://arxiv.org/html/2605.23089#bib.bib4)\)introduce Lipschitz\-constrained policies as a differentiable alternative to task\-specific reward\-smoothing heuristics\. Other control approaches encourage smooth behavior through penalties on mechanical energy or low\-pass filtering\(Fuet al\.,[2021](https://arxiv.org/html/2605.23089#bib.bib6); Penget al\.,[2020](https://arxiv.org/html/2605.23089#bib.bib16)\)\. In contrast, GPLD targets the learned world model itself: it applies a row\-wise Jacobian penalty to the posterior latent probability map in DreamerV3, improving the latent dynamics used for inference and imagination\.

#### Global versus local smoothness constraints\.

Global smoothness methods such as spectral normalization constrain the Lipschitz constant of neural networks by bounding layer\-wise spectral norms\(Gogianuet al\.,[2021](https://arxiv.org/html/2605.23089#bib.bib10); Bjorcket al\.,[2022](https://arxiv.org/html/2605.23089#bib.bib2)\)\. However, global constraints can be overly restrictive because the global Lipschitz constant may arise from rare, unseen, or irrelevant regions of the input space\(Dherinet al\.,[2022](https://arxiv.org/html/2605.23089#bib.bib5)\)\. GPLD instead imposes a local, data\-dependent smoothness penalty on sampled latent inputs during world\-model training\. This distinguishes GPLD from global normalization methods and from policy\-smoothing approaches: GPLD regularizes the posterior transition representation used by a model\-based agent, while leaving the rest of the world model unconstrained except through the training objective\.

## 3A discrete\-to\-continuous justification for GPLD

We justify GPLD by connecting a discrete local\-smoothness prior on transition laws to the continuous Jacobian penalty used in our latent world\-model setting\. We begin with a finite embedded\-state MDP, where smoothness can be imposed by penalizing finite differences of neighboring transition distributions\. We then consider a differentiable transition model over continuous state representations and show that these finite differences become directional Jacobian norms, whose isotropic average yields the Frobenius Jacobian norm\. This produces the row\-wise posterior regularizer used in GPLD\.

### 3\.1Discrete local smoothness prior

We begin with a finite state and finite action MDP\. Let𝒮\\mathcal\{S\}denote the state space,𝒜\\mathcal\{A\}the action space, and lete:𝒮→ℝde:\\mathcal\{S\}\\to\\mathbb\{R\}^\{d\}be an embedding of states into a metric space\. For each actiona∈𝒜a\\in\\mathcal\{A\}, letP\(⋅∣s,a\)∈Δ\|𝒮\|−1P\(\\cdot\\mid s,a\)\\in\\Delta^\{\|\\mathcal\{S\}\|\-1\}denote the transition probability vector, and define theε\\varepsilon\-neighborhood ofssby

𝒩ε​\(s\):=\{k∈𝒮:‖e​\(k\)−e​\(s\)‖2≤ε\}\.\\mathcal\{N\}\_\{\\varepsilon\}\(s\):=\\\{k\\in\\mathcal\{S\}:\\\|e\(k\)\-e\(s\)\\\|\_\{2\}\\leq\\varepsilon\\\}\.
Given a dataset of transitions𝒟=\{\(s,a,s′\)\}t=1T\\mathcal\{D\}=\\\{\(s,a,s^\{\\prime\}\)\\\}\_\{t=1\}^\{T\}, letN​\(s,a,s′\)N\(s,a,s^\{\\prime\}\)denote the number of observed transitions from\(s,a\)\(s,a\)tos′s^\{\\prime\}\. The standard maximum\-likelihood estimator minimizes the negative log\-likelihood

ℒMLE​\(P;𝒟\)=−∑s,a,s′N​\(s,a,s′\)​log⁡P​\(s′∣s,a\)\.\\mathcal\{L\}\_\{\\mathrm\{MLE\}\}\(P;\\mathcal\{D\}\)=\-\\sum\_\{s,a,s^\{\\prime\}\}N\(s,a,s^\{\\prime\}\)\\log P\(s^\{\\prime\}\\mid s,a\)\.
A natural structural prior is that nearby states should induce similar transition laws\. We encode this by coupling neighboring transition vectors through a finite\-difference regularizer:

ℒFD​\(P\)=∑a∑s∑k∈𝒩ε​\(s\)∑s′\(P​\(s′∣k,a\)−P​\(s′∣s,a\)‖e​\(k\)−e​\(s\)‖2\)2\.\\mathcal\{L\}\_\{\\mathrm\{FD\}\}\(P\)=\\sum\_\{a\}\\sum\_\{s\}\\sum\_\{k\\in\\mathcal\{N\}\_\{\\varepsilon\}\(s\)\}\\sum\_\{s^\{\\prime\}\}\\left\(\\frac\{P\(s^\{\\prime\}\\mid k,a\)\-P\(s^\{\\prime\}\\mid s,a\)\}\{\\\|e\(k\)\-e\(s\)\\\|\_\{2\}\}\\right\)^\{2\}\.\(1\)This yields the regularized objective

minP⁡ℒMLE​\(P;𝒟\)\+λ​ℒFD​\(P\)\.\\min\_\{P\}\\;\\mathcal\{L\}\_\{\\mathrm\{MLE\}\}\(P;\\mathcal\{D\}\)\+\\lambda\\mathcal\{L\}\_\{\\mathrm\{FD\}\}\(P\)\.
This is the discrete precursor of GPLD: it penalizes normalized local differences of transition laws across neighboring states and therefore explicitly favors locally smooth transition estimates\.

### 3\.2Continuous\-state limit

The discrete finite\-difference regularizer in \([1](https://arxiv.org/html/2605.23089#S3.E1)\) leads to a Jacobian penalty when the transition law is represented by a differentiable function over continuous state embeddings\. Consider a neighboring statek∈𝒩ε​\(s\)k\\in\\mathcal\{N\}\_\{\\varepsilon\}\(s\)\. In the embedding space, its displacement fromsscan be written as

e​\(k\)−e​\(s\)=h​u,e\(k\)\-e\(s\)=hu,\(2\)whereh=‖e​\(k\)−e​\(s\)‖2h=\\\|e\(k\)\-e\(s\)\\\|\_\{2\}andu∈ℝdu\\in\\mathbb\{R\}^\{d\}is a unit vector\. Substituting \([2](https://arxiv.org/html/2605.23089#S3.E2)\) into a summand of \([1](https://arxiv.org/html/2605.23089#S3.E1)\) yields

‖P\(⋅∣k,a\)−P\(⋅∣s,a\)‖e​\(k\)−e​\(s\)‖2‖22↝‖fθ​\(x\+h​u,a\)−fθ​\(x,a\)h‖22,\\left\\\|\\frac\{P\(\\cdot\\mid k,a\)\-P\(\\cdot\\mid s,a\)\}\{\\\|e\(k\)\-e\(s\)\\\|\_\{2\}\}\\right\\\|\_\{2\}^\{2\}\\;\\leadsto\\;\\left\\\|\\frac\{f\_\{\\theta\}\(x\+hu,a\)\-f\_\{\\theta\}\(x,a\)\}\{h\}\\right\\\|\_\{2\}^\{2\},wherex=e​\(s\)x=e\(s\), andfθf\_\{\\theta\}denotes a differentiable continuous\-state transition model\. Thus, the discrete regularizer naturally gives rise to normalized local finite differences offθf\_\{\\theta\}\.

The small\-neighborhood limith→0h\\to 0then converts these finite differences into directional derivatives\.

Letfθ:ℝd→ℝmf\_\{\\theta\}:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{m\}be locally Fréchet differentiable atx∈ℝdx\\in\\mathbb\{R\}^\{d\}\. Then for any unit vectoru∈ℝdu\\in\\mathbb\{R\}^\{d\},

limh→0‖fθ​\(x\+h​u\)−fθ​\(x\)h‖22=‖Jfθ​\(x\)​u‖22\.\\lim\_\{h\\to 0\}\\left\\\|\\frac\{f\_\{\\theta\}\(x\+hu\)\-f\_\{\\theta\}\(x\)\}\{h\}\\right\\\|\_\{2\}^\{2\}=\\\|J\_\{f\_\{\\theta\}\}\(x\)u\\\|\_\{2\}^\{2\}\.\(3\)
This shows that each normalized finite\-difference term converges to‖Jfθ​\(x\)​u‖22\\\|J\_\{f\_\{\\theta\}\}\(x\)u\\\|\_\{2\}^\{2\}\. Since the discrete regularizer aggregates neighboring states around each reference state, its small\-neighborhood continuous limit induces an average over all unit directions aroundxx\.

To see how isotropic averaging yields Frobenius Jacobian energy, letuube uniformly distributed on the unit sphere𝕊d−1\\mathbb\{S\}^\{d\-1\}\. Then

𝔼u​\[‖Jfθ​\(x\)​u‖22\]=1d​‖Jfθ​\(x\)‖F2\.\\mathbb\{E\}\_\{u\}\\\!\\left\[\\\|J\_\{f\_\{\\theta\}\}\(x\)u\\\|\_\{2\}^\{2\}\\right\]=\\frac\{1\}\{d\}\\\|J\_\{f\_\{\\theta\}\}\(x\)\\\|\_\{F\}^\{2\}\.\(4\)
Thus, under a local isotropy assumption, the continuous counterpart of the discrete finite\-difference regularizer is, up to a constant factor, a squared Frobenius Jacobian penalty\. This establishes the theoretical bridge from discrete neighborhood smoothing to the continuous Jacobian regularization used by GPLD\. In the next section, we instantiate this penalty in DreamerV3 by applying it row\-wise to the posterior latent probability map\. Proofs of \([3](https://arxiv.org/html/2605.23089#S3.E3)\) and \([4](https://arxiv.org/html/2605.23089#S3.E4)\) are deferred to Appendix[A](https://arxiv.org/html/2605.23089#A1)\.

## 4Gradient\-Penalized Latent Dynamics in DreamerV3

We now instantiate the Frobenius Jacobian penalty motivated in Section[3](https://arxiv.org/html/2605.23089#S3)within DreamerV3\. GPLD regularizes the posterior latent probability map rather than the raw environment transition map\. This choice matches the object used for latent\-state inference during world\-model learning and yields a direct implementation through reverse\-mode automatic differentiation\.

DreamerV3 trains a recurrent state\-space world model with deterministic latent statehth\_\{t\}, stochastic latent stateztz\_\{t\}, encoder outputete\_\{t\}, and actionata\_\{t\}, all at timett\. Its world\-model objective can be written as

ℒDv3:=𝔼qϕ​\[∑t=1T\(βpred​ℒpred,t\+βdyn​ℒdyn,t\+βrep​ℒrep,t\)\],\\mathcal\{L\}\_\{\\mathrm\{Dv3\}\}:=\\mathbb\{E\}\_\{q\_\{\\phi\}\}\\left\[\\sum\_\{t=1\}^\{T\}\\left\(\\beta\_\{\\mathrm\{pred\}\}\\mathcal\{L\}\_\{\\mathrm\{pred\},t\}\+\\beta\_\{\\mathrm\{dyn\}\}\\mathcal\{L\}\_\{\\mathrm\{dyn\},t\}\+\\beta\_\{\\mathrm\{rep\}\}\\mathcal\{L\}\_\{\\mathrm\{rep\},t\}\\right\)\\right\],whereℒpred,t\\mathcal\{L\}\_\{\\mathrm\{pred\},t\}combines reconstruction, reward, and continuation prediction losses at timett\. The dynamics and representation KL terms are

ℒdyn,t=max\(1,DKL\(sg\[qϕ\(zt∣ht,et\)\]∥pϕ\(zt∣ht\)\)\),\\mathcal\{L\}\_\{\\mathrm\{dyn\},t\}=\\max\\left\(1,D\_\{\\mathrm\{KL\}\}\\\!\\left\(\\mathrm\{sg\}\\\!\\left\[q\_\{\\phi\}\(z\_\{t\}\\mid h\_\{t\},e\_\{t\}\)\\right\]\\,\\\|\\,p\_\{\\phi\}\(z\_\{t\}\\mid h\_\{t\}\)\\right\)\\right\),and

ℒrep,t=max\(1,DKL\(qϕ\(zt∣ht,et\)∥sg\[pϕ\(zt∣ht\)\]\)\),\\mathcal\{L\}\_\{\\mathrm\{rep\},t\}=\\max\\left\(1,D\_\{\\mathrm\{KL\}\}\\\!\\left\(q\_\{\\phi\}\(z\_\{t\}\\mid h\_\{t\},e\_\{t\}\)\\,\\\|\\,\\mathrm\{sg\}\\\!\\left\[p\_\{\\phi\}\(z\_\{t\}\\mid h\_\{t\}\)\\right\]\\right\)\\right\),wheresg​\[⋅\]\\mathrm\{sg\}\[\\cdot\]denotes the stop\-gradient operator, which prevents gradients from flowing through its argument\.

Letut:=\[ht,et\]u\_\{t\}:=\[h\_\{t\},e\_\{t\}\]denote the posterior input\. After row\-wise softmax, the posterior distributionqϕ​\(zt∣ht,et\)q\_\{\\phi\}\(z\_\{t\}\\mid h\_\{t\},e\_\{t\}\)is represented as aK×CK\\times Cprobability table

qϕ​\(ut\)∈\[0,1\]K×C,q\_\{\\phi\}\(u\_\{t\}\)\\in\[0,1\]^\{K\\times C\},whereKKis the number of categorical variables andCCis the number of classes per variable\. We refer to the maput↦qϕ​\(ut\)u\_\{t\}\\mapsto q\_\{\\phi\}\(u\_\{t\}\)as the posterior latent probability map\. Letqϕ\(i\)​\(ut\)∈ΔC−1q\_\{\\phi\}^\{\(i\)\}\(u\_\{t\}\)\\in\\Delta^\{C\-1\}denote theii\-th row, whereΔC−1\\Delta^\{C\-1\}is the probability simplex overCCclasses\.

GPLD penalizes the average row\-wise Frobenius Jacobian norm of this posterior latent probability map with respect toutu\_\{t\}:

ℛGPLD​\(ut\)=1K​∑i=1K‖Jqϕ\(i\)​\(ut\)‖F2\.\\mathcal\{R\}\_\{\\mathrm\{GPLD\}\}\(u\_\{t\}\)=\\frac\{1\}\{K\}\\sum\_\{i=1\}^\{K\}\\left\\\|J\_\{q\_\{\\phi\}^\{\(i\)\}\}\(u\_\{t\}\)\\right\\\|\_\{F\}^\{2\}\.\(5\)The resulting world\-model objective is

ℒGPLD:=𝔼qϕ​\[∑t=1T\(βpred​ℒpred,t\+βdyn​ℒdyn,t\+βrep​ℒrep,t\+λtpost​ℛGPLD​\(ut\)\)\]\.\\mathcal\{L\}\_\{\\mathrm\{GPLD\}\}:=\\mathbb\{E\}\_\{q\_\{\\phi\}\}\\left\[\\sum\_\{t=1\}^\{T\}\\left\(\\beta\_\{\\mathrm\{pred\}\}\\mathcal\{L\}\_\{\\mathrm\{pred\},t\}\+\\beta\_\{\\mathrm\{dyn\}\}\\mathcal\{L\}\_\{\\mathrm\{dyn\},t\}\+\\beta\_\{\\mathrm\{rep\}\}\\mathcal\{L\}\_\{\\mathrm\{rep\},t\}\+\\lambda^\{\\mathrm\{post\}\}\_\{t\}\\mathcal\{R\}\_\{\\mathrm\{GPLD\}\}\(u\_\{t\}\)\\right\)\\right\]\.
Computing the full Jacobian in \([5](https://arxiv.org/html/2605.23089#S4.E5)\) explicitly is expensive\. We therefore estimate each row\-wise Frobenius term using a Hutchinson\-style output\-space probe\. For each posterior rowii, sampleϵi∼RadC\\epsilon\_\{i\}\\sim\\mathrm\{Rad\}^\{C\}\. Then

gi​\(ut\):=∇ut⟨ϵi,qϕ\(i\)​\(ut\)⟩=Jqϕ\(i\)​\(ut\)⊤​ϵi\.g\_\{i\}\(u\_\{t\}\):=\\nabla\_\{u\_\{t\}\}\\left\\langle\\epsilon\_\{i\},q\_\{\\phi\}^\{\(i\)\}\(u\_\{t\}\)\\right\\rangle=J\_\{q\_\{\\phi\}^\{\(i\)\}\}\(u\_\{t\}\)^\{\\top\}\\epsilon\_\{i\}\.Since𝔼​\[ϵi​ϵi⊤\]=IC\\mathbb\{E\}\[\\epsilon\_\{i\}\\epsilon\_\{i\}^\{\\top\}\]=I\_\{C\}, we have

𝔼ϵi​\[‖Jqϕ\(i\)​\(ut\)⊤​ϵi‖22\]=‖Jqϕ\(i\)​\(ut\)‖F2\.\\mathbb\{E\}\_\{\\epsilon\_\{i\}\}\\left\[\\left\\\|J\_\{q\_\{\\phi\}^\{\(i\)\}\}\(u\_\{t\}\)^\{\\top\}\\epsilon\_\{i\}\\right\\\|\_\{2\}^\{2\}\\right\]=\\left\\\|J\_\{q\_\{\\phi\}^\{\(i\)\}\}\(u\_\{t\}\)\\right\\\|\_\{F\}^\{2\}\.Thus, the implemented estimator is

ℛ^GPLD​\(ut\)=1K​∑i=1K‖Jqϕ\(i\)​\(ut\)⊤​ϵi‖22\.\\widehat\{\\mathcal\{R\}\}\_\{\\mathrm\{GPLD\}\}\(u\_\{t\}\)=\\frac\{1\}\{K\}\\sum\_\{i=1\}^\{K\}\\left\\\|J\_\{q\_\{\\phi\}^\{\(i\)\}\}\(u\_\{t\}\)^\{\\top\}\\epsilon\_\{i\}\\right\\\|\_\{2\}^\{2\}\.Algorithm[1](https://arxiv.org/html/2605.23089#alg1)summarizes the minibatch implementation\.

Algorithm 1Gradient penalty on the posterior latent probability map using Hutchinson’s estimator1:Input:batch of latent states

hh, encoded observations

ee, posterior model

q​\(h,e\)q\(h,e\), penalty coefficient

λT\\lambda\_\{T\}, sampling fraction

ρ\\rho
2:Output:posterior gradient penalty

λT​Lgp\\lambda\_\{T\}\\,L\_\{\\mathrm\{gp\}\}
3:

B←B\\leftarrowtotal batch size

4:

N←⌊ρ​B⌋N\\leftarrow\\lfloor\\rho B\\rfloor// number of sampled states

5:sample index set

I⊂\{1,…,B\}I\\subset\\\{1,\\dots,B\\\}with

\|I\|=N\|I\|=N
6:

ℒGPLD←0\\mathcal\{L\}\_\{\\mathrm\{GPLD\}\}\\leftarrow 0
7:foreach index

k∈Ik\\in Ido

8:extract

hk,ekh\_\{k\},e\_\{k\}from the batch

9:

xk←\[hk,ek\]x\_\{k\}\\leftarrow\[h\_\{k\},e\_\{k\}\]
10:

𝐪k←q​\(hk,ek\)∈\[0,1\]32×16\\mathbf\{q\}\_\{k\}\\leftarrow q\(h\_\{k\},e\_\{k\}\)\\in\[0,1\]^\{32\\times 16\}// row\-wise probabilities

11:for

i=1i=1to

KKdo

12:sample Rademacher noise

ϵk,i∼\{−1,\+1\}C\\epsilon\_\{k,i\}\\sim\\\{\-1,\+1\\\}^\{C\}
13:

sk,i←⟨ϵk,i,𝐪k,i⟩s\_\{k,i\}\\leftarrow\\langle\\epsilon\_\{k,i\},\\mathbf\{q\}\_\{k,i\}\\rangle
14:

gk,i←∇xksk,ig\_\{k,i\}\\leftarrow\\nabla\_\{x\_\{k\}\}s\_\{k,i\}
15:

Lgp←Lgp\+‖gk,i‖22L\_\{\\mathrm\{gp\}\}\\leftarrow L\_\{\\mathrm\{gp\}\}\+\\\|g\_\{k,i\}\\\|\_\{2\}^\{2\}
16:endfor

17:endfor

18:

Lgp←1K⋅N​LgpL\_\{\\mathrm\{gp\}\}\\leftarrow\\frac\{1\}\{K\\cdot N\}L\_\{\\mathrm\{gp\}\}
19:return

λT​Lgp\\lambda\_\{T\}\\,L\_\{\\mathrm\{gp\}\}

### 4\.1Why regularize only the posterior?

GPLD regularizes the posterior latent probability map rather than both the posterior and prior\. This targets the observation\-conditioned distributionqϕ​\(zt∣ht,et\)q\_\{\\phi\}\(z\_\{t\}\\mid h\_\{t\},e\_\{t\}\), which is used to infer latent states from real trajectories during world\-model training\. Since the posterior depends directly on the encoded observationete\_\{t\}, regularizing this map constrains the latent representation learned from environment data\.

The priorpϕ​\(zt∣ht\)p\_\{\\phi\}\(z\_\{t\}\\mid h\_\{t\}\)is not directly penalized by GPLD, but it is trained against the posterior through the dynamics KL termℒdyn,t\\mathcal\{L\}\_\{\\mathrm\{dyn\},t\}\. Thus, posterior\-only regularization directly smooths the observation\-conditioned latent map while still influencing the target that the prior learns to predict\. Ablations in Section[5\.4](https://arxiv.org/html/2605.23089#S5.SS4)show that this choice provides the best performance–cost trade\-off among posterior\-only, prior\-only, and joint regularization\.

### 4\.2Time\-decayed gradient penalty coefficient

A fixed smoothness penalty can be useful early in training, when the world model is poorly estimated, but may become overly restrictive later as the model fit improves\. We therefore decay the posterior penalty coefficient over training\. Our default schedule uses square\-root decay with a minimum threshold:

λTpost=max⁡\(λ0post1\+Tupdates/c,λmin\),\\lambda^\{\\mathrm\{post\}\}\_\{T\}=\\max\\left\(\\frac\{\\lambda^\{\\mathrm\{post\}\}\_\{0\}\}\{\\sqrt\{1\+T\_\{\\mathrm\{updates\}\}/c\}\},\\lambda\_\{\\min\}\\right\),whereTupdatesT\_\{\\mathrm\{updates\}\}is the cumulative number of optimizer updates andc\>0c\>0is a decay\-scale constant\. The square\-root form is motivated by the finite\-state intuition that explicit smoothing should matter most when transition estimates are data\-limited, and should weaken as more data are collected\. Algorithm[2](https://arxiv.org/html/2605.23089#alg2)summarizes the schedule\.

Algorithm 2Square\-root decay of the gradient penalty coefficient1:Input:initial coefficient

λ0\\lambda\_\{0\}, current update step

TupdatesT\_\{\\mathrm\{updates\}\}, decay scale

cc, minimum coefficient

λmin\\lambda\_\{\\min\}
2:Output:decayed coefficient

λT\\lambda\_\{T\}
3:

s←1\+Tupdates/cs\\leftarrow 1\+T\_\{\\mathrm\{updates\}\}/c
4:

λT←max⁡\(λ0/s,λmin\)\\lambda\_\{T\}\\leftarrow\\max\\left\(\\lambda\_\{0\}/\\sqrt\{s\},\\lambda\_\{\\min\}\\right\)
5:return

λT\\lambda\_\{T\}

## 5Experiments

We evaluate GPLD on DeepMind Control \(DMC\) tasks\(Tunyasuvunakoolet al\.,[2020](https://arxiv.org/html/2605.23089#bib.bib20)\)to test whether local smoothness regularization improves latent world\-model learning\. Our main evaluation uses proprioceptive observations, where the agent receives low\-dimensional physical state features\. We additionally evaluate pixel observations, where the agent receives image frames and must learn visual representations jointly with latent dynamics\.

### 5\.1Experimentation

DMC locomotion tasks provide a natural testbed for GPLD because many of them evolve through locally smooth body motions over short time intervals, as qualitatively illustrated in Fig\.[1](https://arxiv.org/html/2605.23089#S5.F1)\. We compare GPLD\-DreamerV3 against DreamerV3 using aggregate learning curves, representative task\-level curves, and ablations over the main design choices\. Experimental details, including hyperparameters, seed counts, evaluation protocol, and aggregate normalization, are provided in Appendix[B](https://arxiv.org/html/2605.23089#A2)\.

Cheetah run ![Refer to caption](https://arxiv.org/html/2605.23089v1/figures/montage_cheetah_run_4.png)

Walker walk ![Refer to caption](https://arxiv.org/html/2605.23089v1/figures/montage_walker_walk_4.png)

Figure 1:Representative frames from DMC locomotion tasks\. Many proprioceptive locomotion environments evolve through smooth body motions over short time intervals, making them a natural testbed for evaluating local smoothness regularization in latent dynamics\.
### 5\.2DMC Proprioceptive Tasks

GPLD provides the strongest evidence for its effectiveness in proprioceptive control, where the world model learns from low\-dimensional physical state features and the local\-smoothness prior acts directly on latent dynamics learning\. Across the proprioceptive DMC benchmark, GPLD\-DreamerV3 improves sample efficiency over DreamerV3, as shown in Fig\.[2](https://arxiv.org/html/2605.23089#S5.F2)\. Importantly, these aggregate gains are obtained using a single GPLD configuration across tasks, without per\-environment tuning\. The improvement appears in both raw aggregate return and normalized aggregate return, indicating that GPLD does not merely improve a single high\-return environment, but produces broader gains across the benchmark\.

The largest gains appear on tasks that demand coordinated locomotion and richer transition modeling\. On the higher\-complexity proprioceptive locomotion subset in Table[1](https://arxiv.org/html/2605.23089#S5.T1), GPLD improves the normalized aggregate score by34\.6%34\.6\\%at 1M environment steps\. The improvement is especially large on hopper\-hop and walker\-run, while gains are smaller on tasks where DreamerV3 already approaches high return\. The 1M\-step comparison across all proprioceptive tasks is reported in Appendix[C](https://arxiv.org/html/2605.23089#A3), Table[3](https://arxiv.org/html/2605.23089#A3.T3), where GPLD improves the normalized aggregate mean by17\.7%17\.7\\%\. Individual 2M\-step learning curves for the proprioceptive suite are shown in Appendix[C](https://arxiv.org/html/2605.23089#A3), Fig\.[6](https://arxiv.org/html/2605.23089#A3.F6)\.

![Refer to caption](https://arxiv.org/html/2605.23089v1/x1.png)

\(a\) Raw Aggregate Mean Score

![Refer to caption](https://arxiv.org/html/2605.23089v1/x2.png)

\(b\) Normalized Aggregate Mean Score

Figure 2:Aggregate proprioceptive\-control performance on DMC tasks\. GPLD\-DreamerV3 improves sample efficiency over DreamerV3 in both raw aggregate return and normalized aggregate return\.Table 1:Performance at 1M environment steps on higher\-complexity proprioceptive DMC tasks\. Individual task scores are raw episodic returns\. For the aggregate summary reported in the text, scores are normalized by the final DreamerV3 baseline score for each task before aggregation\.Difficult quadruped tasks reveal a complementary long\-horizon effect\. Unlike the 1M\-step locomotion subset, quadruped\-run and quadruped\-walk remain noisy early in training for both methods, suggesting that reliable latent dynamics for these environments require substantially more interaction\. Over 4M environment steps, however, GPLD\-DreamerV3 reaches the high\-return regime earlier and maintains stronger late\-stage performance than DreamerV3, as shown in Fig\.[3](https://arxiv.org/html/2605.23089#S5.F3)\. This indicates that the benefit of local smoothness regularization is not limited to early learning: on harder locomotion tasks, it can also improve the quality and consistency of learning over longer horizons\.

![Refer to caption](https://arxiv.org/html/2605.23089v1/x3.png)

\(a\) Quadruped run

![Refer to caption](https://arxiv.org/html/2605.23089v1/x4.png)

\(b\) Quadruped Walk

Figure 3:Long\-horizon quadruped performance\. Mean episodic return is reported across seeds, with shaded regions indicating one standard deviation around the mean\. On quadruped\-run and quadruped\-walk, GPLD\-DreamerV3 reaches high\-return behavior earlier and maintains stronger late\-stage performance than DreamerV3\.
### 5\.3DMC Pixel Tasks

We also evaluate GPLD with pixel observations, where the agent receives image frames rather than low\-dimensional proprioceptive state features\. This setting is harder because the world model must learn visual encodings and latent dynamics jointly\. As shown in Fig\.[4](https://arxiv.org/html/2605.23089#S5.F4), GPLD\-DreamerV3 remains competitive with DreamerV3 in the normalized aggregate, but the gain is smaller than in proprioceptive control\.

The task\-level curves help explain this weaker aggregate effect\. GPLD shows clearer late\-stage gains on walker\-run and quadruped\-run, but hopper\-hop is mixed: it learns more slowly early in training and catches up later\. Thus, pixel results support a more cautious conclusion than the proprioceptive results\. GPLD can still help under visual observations, but its effect is less direct when the same world model must simultaneously learn visual representations and smooth latent dynamics\. Appendix[D\.1](https://arxiv.org/html/2605.23089#A4.SS1)provides individual pixel\-observation learning curves for all evaluated tasks, while Appendix[D\.2](https://arxiv.org/html/2605.23089#A4.SS2)further studies an encoder\-decoder warm\-start diagnostic to separate the role of visual representation learning from latent\-dynamics regularization\.

![Refer to caption](https://arxiv.org/html/2605.23089v1/x5.png)

\(a\) Normalized aggregate

![Refer to caption](https://arxiv.org/html/2605.23089v1/x6.png)

\(b\)Walker\-run

![Refer to caption](https://arxiv.org/html/2605.23089v1/x7.png)

\(c\) Quadruped\-run

![Refer to caption](https://arxiv.org/html/2605.23089v1/x8.png)

\(d\)Hopper\-hop

Figure 4:Pixel\-observation DMC results\. With image\-frame observations, GPLD\-DreamerV3 remains competitive in normalized aggregate performance but shows smaller gains than in proprioceptive control\. Representative curves show clearer late\-stage improvements on walker\-run and quadruped\-run, while hopper\-hop exhibits slower early learning and a smaller late\-stage advantage\.
### 5\.4Ablations

The ablations in Fig\.[5](https://arxiv.org/html/2605.23089#S5.F5)show that GPLD works best as a targeted posterior regularizer rather than as a broad smoothness constraint on all latent distributions\. Posterior\-only regularization performs competitively with or better than prior\-only and joint prior–posterior regularization, while avoiding the extra derivative cost of explicitly penalizing both distributions\. This supports the design choice in Section[4\.1](https://arxiv.org/html/2605.23089#S4.SS1): GPLD directly regularizes the observation\-conditioned latent map, while the prior remains coupled through the DreamerV3 dynamics loss\.

The sampling fractionρ\\rhocontrols the fraction of batch states on which the gradient penalty is evaluated, while the penalty coefficient controls the strength of the smoothness penalty on those states\. Increasingρ\\rhoapplies GPLD to more batch states, but the gains do not increase monotonically enough to justify always using the largest value\. The decay ablation further shows that maintaining a strong penalty throughout training is less effective than square\-root decay, supporting the schedule in Section[4\.2](https://arxiv.org/html/2605.23089#S4.SS2)\. Overall, the ablations indicate that GPLD is most effective when applied selectively to the posterior with moderate, time\-decayed regularization\.

![Refer to caption](https://arxiv.org/html/2605.23089v1/x9.png)

\(a\) Varyρ\\rhoatλ0post=0\.5\\lambda\_\{0\}^\{\\mathrm\{post\}\}=0\.5

![Refer to caption](https://arxiv.org/html/2605.23089v1/x10.png)

\(b\) Varyρ\\rhowithλ0prior=0\.5\\lambda\_\{0\}^\{\\mathrm\{prior\}\}=0\.5

![Refer to caption](https://arxiv.org/html/2605.23089v1/x11.png)

\(c\) Varyλ0post\\lambda\_\{0\}^\{\\mathrm\{post\}\}andλ0prior\\lambda\_\{0\}^\{\\mathrm\{prior\}\}

![Refer to caption](https://arxiv.org/html/2605.23089v1/x12.png)

\(d\) \(d\) Sqrt decay vs no decay

Figure 5:GPLD ablations over sampling fraction, posterior/prior regularization, and penalty scheduling\. The results support using GPLD as a targeted posterior regularizer with time\-decayed smoothness rather than a broad or fixed smoothness constraint\.
### 5\.5Sensitivity diagnostics and computational cost

We include two additional diagnostics to check whether GPLD behaves as intended\. To assess the effect of the gradient penalty, we measure how much the learned posterior and prior distributions change under small perturbations of their inputs\. A smoother latent map should change less under such perturbations\. GPLD reduces this measured local sensitivity on representative locomotion tasks, supporting the intended effect of the regularizer\. Full details and plots are provided in Appendix[F](https://arxiv.org/html/2605.23089#A6)\.

To assess computational overhead, we measure wall\-clock runtime relative to DreamerV3\. The additional cost comes from estimating the gradient penalty through vector\-Jacobian products, and scales with the number of batch states selected by the sampling fractionρ\\rho\. Timing results are reported in Appendix[E](https://arxiv.org/html/2605.23089#A5)\.

## 6Conclusion

We have introduced GPLD, a local smoothness regularizer for latent world models, and instantiated it in DreamerV3 through a row\-wise Jacobian penalty on the posterior latent probability map\. This penalty is motivated by a discrete\-to\-continuous argument showing how finite\-difference smoothing of neighboring transition laws leads to Frobenius Jacobian regularization in differentiable latent models\. Empirically, GPLD improves aggregate sample efficiency on proprioceptive DMC tasks, with the strongest gains on higher\-complexity locomotion environments\. The quadruped results further show that local smoothness regularization improves the quality and consistency of learning over longer horizons\. These gains are obtained with a single GPLD configuration across tasks, without per\-environment tuning\.

There are some important limitations\. GPLD is not uniformly beneficial across all environments, and its gains are weaker in pixel\-observation settings, where latent dynamics learning is coupled to high\-dimensional visual encoding\. More broadly, GPLD encodes a local smoothness prior, so it may be less appropriate in environments where discontinuities dominate the relevant dynamics\.

## Impact Statement

This work aims to improve the sample efficiency of model\-based reinforcement learning by regularizing learned latent dynamics\. More sample\-efficient world models can reduce the number of interactions needed to train continuous\-control agents, potentially lowering simulation cost and making model\-based RL more practical\. As with other reinforcement\-learning methods, downstream impact depends on the systems in which the trained agents are deployed\. If such methods are later used in physical\-control systems, deployment would require the same safety validation expected for other reinforcement\-learning agents\. Our experiments are limited to simulated DeepMind Control tasks and do not involve deployed robotic systems or human subjects\.

## Acknowledgements

Portions of this research were conducted with the advanced computing resources provided by Texas A&M High Performance Research Computing\. This material is based upon work partially supported by the U\.S\. Army Contracting Command under Contract Numbers W911NF2120064, W911NF2520046, and W911NF2210151, the Office of Naval Research under Contracts N000142412615, and N00014\-21\-1\-2385, and the National Science Foundation under Contract Numbers CNS\-2328395 and CMMI\-2038625\. The views expressed herein and conclusions contained in this document are those of the authors and should not be interpreted as representing the views or official policies, either expressed or implied, of the U\.S\. Army, Navy, or the United States Government\. The U\.S\. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein\.

## References

- Lipschitz continuity in model\-based reinforcement learning\.InProceedings of the 35th International Conference on Machine Learning,J\. Dy and A\. Krause \(Eds\.\),Proceedings of Machine Learning Research, Vol\.80,pp\. 264–273\.External Links:[Link](https://proceedings.mlr.press/v80/asadi18a.html)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px1.p1.1)\.
- J\. Bjorck, C\. P\. Gomes, and K\. Q\. Weinberger \(2022\)Towards deeper deep reinforcement learning with spectral normalization\.External Links:2106\.01151,[Link](https://arxiv.org/abs/2106.01151)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px3.p1.1)\.
- J\. Boyan and A\. Moore \(1994\)Generalization in reinforcement learning: safely approximating the value function\.InAdvances in Neural Information Processing Systems,G\. Tesauro, D\. Touretzky, and T\. Leen \(Eds\.\),Vol\.7,pp\.\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/1994/file/ef50c335cca9f340bde656363ebd02fd-Paper.pdf)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px1.p1.1)\.
- Z\. Chen, X\. He, Y\. Wang, Q\. Liao, Y\. Ze, Z\. Li, S\. S\. Sastry, J\. Wu, K\. Sreenath, S\. Gupta, and X\. B\. Peng \(2024\)Learning smooth humanoid locomotion through lipschitz\-constrained policies\.External Links:2410\.11825,[Link](https://arxiv.org/abs/2410.11825)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px2.p1.1)\.
- G\. Christmann, Y\. Luo, H\. Mandala, and W\. Chen \(2024\)Benchmarking smoothness and reducing high\-frequency oscillations in continuous control policies\.External Links:2410\.16632,[Link](https://arxiv.org/abs/2410.16632)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px1.p1.1)\.
- B\. Dherin, M\. Munn, M\. Rosca, and D\. G\. T\. Barrett \(2022\)Why neural networks find simple solutions: the many regularizers of geometric complexity\.External Links:2209\.13083,[Link](https://arxiv.org/abs/2209.13083)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px3.p1.1)\.
- Z\. Fu, A\. Kumar, J\. Malik, and D\. Pathak \(2021\)Minimizing energy consumption leads to the emergence of gaits in legged robots\.External Links:2111\.01674,[Link](https://arxiv.org/abs/2111.01674)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px2.p1.1)\.
- F\. Gogianu, T\. Berariu, M\. Rosca, C\. Clopath, L\. Busoniu, and R\. Pascanu \(2021\)Spectral normalisation for deep reinforcement learning: an optimisation perspective\.External Links:2105\.05246,[Link](https://arxiv.org/abs/2105.05246)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px3.p1.1)\.
- I\. J\. Goodfellow, J\. Shlens, and C\. Szegedy \(2015\)Explaining and harnessing adversarial examples\.External Links:1412\.6572,[Link](https://arxiv.org/abs/1412.6572)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px2.p1.1)\.
- I\. Gulrajani, F\. Ahmed, M\. Arjovsky, V\. Dumoulin, and A\. Courville \(2017\)Improved training of wasserstein gans\.External Links:1704\.00028,[Link](https://arxiv.org/abs/1704.00028)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px2.p1.1)\.
- D\. Ha and J\. Schmidhuber \(2018\)World models\.Zenodo\.External Links:[Document](https://dx.doi.org/10.5281/ZENODO.1207631),[Link](https://zenodo.org/record/1207631)Cited by:[§1](https://arxiv.org/html/2605.23089#S1.p1.1)\.
- D\. Hafner, J\. Pasukonis, J\. Ba, and T\. Lillicrap \(2025\)Mastering diverse control tasks through world models\.Nature640,pp\. 647–653\.External Links:[Document](https://dx.doi.org/10.1038/s41586-025-08744-2)Cited by:[§1](https://arxiv.org/html/2605.23089#S1.p1.1)\.
- R\. Munos and C\. Szepesvári \(2008\)Finite\-time bounds for fitted value iteration\.Journal of Machine Learning Research9\(27\),pp\. 815–857\.External Links:[Link](https://www.jmlr.org/papers/v9/munos08a.html)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px1.p1.1)\.
- X\. B\. Peng, E\. Coumans, T\. Zhang, T\. Lee, J\. Tan, and S\. Levine \(2020\)Learning agile robotic locomotion skills by imitating animals\.External Links:2004\.00784,[Link](https://arxiv.org/abs/2004.00784)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px2.p1.1)\.
- M\. Rosca, T\. Weber, A\. Gretton, and S\. Mohamed \(2020\)A case for new neural network smoothness constraints\.InProceedings on "I Can’t Believe It’s Not Better\!" at NeurIPS Workshops,J\. Zosa Forde, F\. Ruiz, M\. F\. Pradier, and A\. Schein \(Eds\.\),Proceedings of Machine Learning Research, Vol\.137,pp\. 21–32\.External Links:[Link](https://proceedings.mlr.press/v137/rosca20a.html)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px1.p1.1)\.
- S\. Thrun and A\. Schwartz \(1993\)Issues in using function approximation for reinforcement learning\.InProceedings of the 1993 Connectionist Models Summer School,M\. Mozer, P\. Smolensky, D\. Touretzky, J\. Elman, and A\. Weigend \(Eds\.\),pp\. 255–263\.External Links:[Link](http://www.ri.cmu.edu/pub_files/pub1/thrun_sebastian_1993_1/thrun_sebastian_1993_1.pdf)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px1.p1.1)\.
- S\. Tunyasuvunakool, A\. Muldal, Y\. Doron, S\. Liu, S\. Bohez, J\. Merel, T\. Erez, T\. Lillicrap, N\. Heess, and Y\. Tassa \(2020\)Dm\-control: software and tasks for continuous control\.Software Impacts6,pp\. 100022\.External Links:ISSN 2665\-9638,[Document](https://dx.doi.org/https%3A//doi.org/10.1016/j.simpa.2020.100022),[Link](https://www.sciencedirect.com/science/article/pii/S2665963820300099)Cited by:[§5](https://arxiv.org/html/2605.23089#S5.p1.1)\.
- A\. Venkatraman, M\. Hebert, and J\. Bagnell \(2015\)Improving multi\-step prediction of learned time series models\.Proceedings of the AAAI Conference on Artificial Intelligence29,pp\.\.External Links:[Document](https://dx.doi.org/10.1609/aaai.v29i1.9590)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px1.p1.1)\.
- H\. Wang, Y\. Sun, Z\. Tang, H\. Shi, and C\. Jiao \(2025\)HG2P: hippocampus\-inspired high\-reward graph and model\-free q\-gradient penalty for path planning and motion control\.Neural Networks191,pp\. 107781\.External Links:ISSN 0893\-6080,[Document](https://dx.doi.org/https%3A//doi.org/10.1016/j.neunet.2025.107781),[Link](https://www.sciencedirect.com/science/article/pii/S0893608025006616)Cited by:[§2](https://arxiv.org/html/2605.23089#S2.SS0.SSS0.Px2.p1.1)\.

## Appendix AProofs for the discrete\-to\-continuous justification of GPLD

1. 1\.Proof of Finite Differences converging to directional Jacobian sensitivities\([3](https://arxiv.org/html/2605.23089#S3.E3)\) ###### Proof\. Sincefθ:ℝd→ℝmf\_\{\\theta\}:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{m\}is locally Fréchet differentiable atxx, there exists a linear mapJfθ​\(x\)∈ℝm×dJ\_\{f\_\{\\theta\}\}\(x\)\\in\\mathbb\{R\}^\{m\\times d\}such that fθ​\(x\+h​u\)=fθ​\(x\)\+Jfθ​\(x\)​\(h​u\)\+r​\(h\),f\_\{\\theta\}\(x\+hu\)=f\_\{\\theta\}\(x\)\+J\_\{f\_\{\\theta\}\}\(x\)\(hu\)\+r\(h\),\(6\)where the remainder satisfies ‖r​\(h\)‖2\|h\|→0as​h→0\.\\frac\{\\\|r\(h\)\\\|\_\{2\}\}\{\|h\|\}\\to 0\\qquad\\text\{as \}h\\to 0\.\(7\)Dividing \([6](https://arxiv.org/html/2605.23089#A1.E6)\) byhhgives fθ​\(x\+h​u\)−fθ​\(x\)h=Jfθ​\(x\)​u\+r​\(h\)h\.\\frac\{f\_\{\\theta\}\(x\+hu\)\-f\_\{\\theta\}\(x\)\}\{h\}=J\_\{f\_\{\\theta\}\}\(x\)u\+\\frac\{r\(h\)\}\{h\}\.\(8\)By \([7](https://arxiv.org/html/2605.23089#A1.E7)\), we have ‖r​\(h\)h‖2→0as​h→0\.\\left\\\|\\frac\{r\(h\)\}\{h\}\\right\\\|\_\{2\}\\to 0\\qquad\\text\{as \}h\\to 0\.Therefore, fθ​\(x\+h​u\)−fθ​\(x\)h→Jfθ​\(x\)​uas​h→0\.\\frac\{f\_\{\\theta\}\(x\+hu\)\-f\_\{\\theta\}\(x\)\}\{h\}\\to J\_\{f\_\{\\theta\}\}\(x\)u\\qquad\\text\{as \}h\\to 0\.Taking squaredℓ2\\ell\_\{2\}\-norms on both sides and using continuity of the norm yields limh→0‖fθ​\(x\+h​u\)−fθ​\(x\)h‖22=‖Jfθ​\(x\)​u‖22\.\\lim\_\{h\\to 0\}\\left\\\|\\frac\{f\_\{\\theta\}\(x\+hu\)\-f\_\{\\theta\}\(x\)\}\{h\}\\right\\\|\_\{2\}^\{2\}=\\\|J\_\{f\_\{\\theta\}\}\(x\)u\\\|\_\{2\}^\{2\}\.This proves the claim\. ∎
2. 2\.Isotropic averaging of Directional Jacobian norm yields Frobenius Jacobian energy\([4](https://arxiv.org/html/2605.23089#S3.E4)\) ###### Proof\. LetJ:=Jfθ​\(x\)∈ℝm×dJ:=J\_\{f\_\{\\theta\}\}\(x\)\\in\\mathbb\{R\}^\{m\\times d\}\. Then ‖J​u‖22=u⊤​J⊤​J​u\.\\\|Ju\\\|\_\{2\}^\{2\}=u^\{\\top\}J^\{\\top\}Ju\.\(9\)Taking expectation overu∼Unif​\(𝕊d−1\)u\\sim\\mathrm\{Unif\}\(\\mathbb\{S\}^\{d\-1\}\), we obtain 𝔼u​\[‖J​u‖22\]=𝔼u​\[u⊤​J⊤​J​u\]\.\\mathbb\{E\}\_\{u\}\[\\\|Ju\\\|\_\{2\}^\{2\}\]=\\mathbb\{E\}\_\{u\}\[u^\{\\top\}J^\{\\top\}Ju\]\.\(10\)Using the trace identityv⊤​A​v=Tr⁡\(A​v​v⊤\)v^\{\\top\}Av=\\operatorname\{Tr\}\(Avv^\{\\top\}\), we have 𝔼u​\[u⊤​J⊤​J​u\]=𝔼u​\[Tr⁡\(J⊤​J​u​u⊤\)\]=Tr⁡\(J⊤​J​𝔼u​\[u​u⊤\]\)\.\\mathbb\{E\}\_\{u\}\[u^\{\\top\}J^\{\\top\}Ju\]=\\mathbb\{E\}\_\{u\}\\\!\\left\[\\operatorname\{Tr\}\(J^\{\\top\}J\\,uu^\{\\top\}\)\\right\]=\\operatorname\{Tr\}\\\!\\left\(J^\{\\top\}J\\,\\mathbb\{E\}\_\{u\}\[uu^\{\\top\}\]\\right\)\.\(11\)Sinceuuis uniformly distributed on the unit sphere, its distribution is isotropic, and hence 𝔼u​\[u​u⊤\]=1d​Id\.\\mathbb\{E\}\_\{u\}\[uu^\{\\top\}\]=\\frac\{1\}\{d\}I\_\{d\}\.\(12\)Substituting \([12](https://arxiv.org/html/2605.23089#A1.E12)\) into \([11](https://arxiv.org/html/2605.23089#A1.E11)\) yields 𝔼u​\[‖J​u‖22\]=Tr⁡\(J⊤​J⋅1d​Id\)=1d​Tr⁡\(J⊤​J\)\.\\mathbb\{E\}\_\{u\}\[\\\|Ju\\\|\_\{2\}^\{2\}\]=\\operatorname\{Tr\}\\\!\\left\(J^\{\\top\}J\\cdot\\frac\{1\}\{d\}I\_\{d\}\\right\)=\\frac\{1\}\{d\}\\operatorname\{Tr\}\(J^\{\\top\}J\)\.Finally, by the definition of the Frobenius norm, Tr⁡\(J⊤​J\)=‖J‖F2\.\\operatorname\{Tr\}\(J^\{\\top\}J\)=\\\|J\\\|\_\{F\}^\{2\}\.Therefore, 𝔼u​\[‖Jfθ​\(x\)​u‖22\]=1d​‖Jfθ​\(x\)‖F2,\\mathbb\{E\}\_\{u\}\[\\\|J\_\{f\_\{\\theta\}\}\(x\)u\\\|\_\{2\}^\{2\}\]=\\frac\{1\}\{d\}\\\|J\_\{f\_\{\\theta\}\}\(x\)\\\|\_\{F\}^\{2\},as claimed\. ∎

## Appendix BExperimental Details

We use the open\-source DreamerV3 implementation, which is released under the MIT license, and the DeepMind Control Suite, which is released under the Apache\-2\.0 license\.

Unless otherwise stated, GPLD uses posterior regularization withλ0post=0\.5\\lambda^\{\\mathrm\{post\}\}\_\{0\}=0\.5, no prior regularization, the square\-root decay schedule from Algorithm[2](https://arxiv.org/html/2605.23089#alg2), imagination horizonH=25H=25, sampling fractionρ=0\.5\\rho=0\.5, decay scalec=1000c=1000, and minimum coefficientλmin=0\.001\\lambda\_\{\\min\}=0\.001\. We use the 12M\-parameter DreamerV3 configuration and keep all non\-GPLD hyperparameters fixed between the baseline and GPLD variants\.

Our primary metric is sample efficiency, measured by evaluation return as a function of environment steps\. Learning curves report mean episodic return across available seeds, with shaded regions indicating one standard deviation around the mean; unless otherwise noted, experiments use five seeds\. For aggregate curves, we report both raw and normalized aggregate return\. In the normalized aggregate, each task is scaled by its final DreamerV3 baseline score before averaging, so that tasks contribute equally relative to their baseline performance\.

Table 2:Training and evaluation hyperparameters\. All non\-GPLD hyperparameters are shared between DreamerV3 and GPLD\-DreamerV3\. Values are taken from the 12M DMC configuration, with GPLD\-specific overrides listed separately\.Table[2](https://arxiv.org/html/2605.23089#A2.T2)lists the main training and evaluation hyperparameters used in our experiments\. DreamerV3 and GPLD\-DreamerV3 share all non\-GPLD hyperparameters; the only differences are the posterior gradient penalty, sampling fraction, and decay schedule\.

## Appendix CDMC Proprioceptive Results

Fig[6](https://arxiv.org/html/2605.23089#A3.F6)shows that the early GPLD gains generally persist or remain competitive over longer horizons, while quadruped tasks are separated in the main text because their benefits emerge most clearly only beyond the 1M\-step regime\.

![Refer to caption](https://arxiv.org/html/2605.23089v1/x13.png)Figure 6:DMC Proprioceptive Individual TasksTable 3:Full 1M\-step proprioceptive DMC results\. Scores are reported as mean±\\pmstandard deviation across seeds, using checkpoint evaluation with 20 episodes per seed\. GPLD\-DreamerV3 improves the normalized aggregate mean by17\.7%17\.7\\%\.
## Appendix DDMC Pixel Results

### D\.1Individual learning plots

Fig\.[7](https://arxiv.org/html/2605.23089#A4.F7)shows individual plots for the DMC Vision tasks\.

![Refer to caption](https://arxiv.org/html/2605.23089v1/x14.png)Figure 7:DMC Pixel Individual Tasks
### D\.2Encoder\-decoder warm\-start diagnostic for pixel observations

To better understand why GPLD yields weaker aggregate gains under pixel observations, we run an encoder\-decoder warm\-start diagnostic\. In this setting, the encoder and decoder are initialized from a previously trained DreamerV3 baseline run, while all other components are reinitialized\. The encoder and decoder are not frozen; they continue to train during the new run\. Both DreamerV3 and GPLD\-DreamerV3 use the same encoder\-decoder initialization, so differences in performance reflect how the methods use the improved visual initialization rather than a different starting representation\.

Figure[8](https://arxiv.org/html/2605.23089#A4.F8)shows results on hopper\-hop and quadruped\-walk\. Under the same warm start, GPLD improves performance over DreamerV3 on these tasks\. This supports the interpretation that the weaker fully end\-to\-end pixel gains may partly arise because visual representation learning and latent\-dynamics learning must be solved simultaneously\. We treat this experiment as a diagnostic rather than a main empirical claim\.

![Refer to caption](https://arxiv.org/html/2605.23089v1/x15.png)\(a\)Hopper\-hop
![Refer to caption](https://arxiv.org/html/2605.23089v1/x16.png)\(b\)Quadruped\-walk

Figure 8:Encoder\-decoder warm\-start diagnostic for pixel observations\. The encoder and decoder are initialized from a trained DreamerV3 baseline run and continue training, while all other components are reinitialized\. With the same visual warm start, GPLD\-DreamerV3 improves performance on hopper\-hop and quadruped\-walk, suggesting that visual representation learning may partially mask the effect of latent\-dynamics regularization in fully end\-to\-end pixel training\.

## Appendix EComputational Cost Analysis

Computation cost analysis: \(Avg Baseline time = 2:01:29/seed for 500k env steps\.\)

Table 4:Computation\-time overhead of GPLD relative to vanilla DreamerV3 across different\(ρ,λp​o​s​t,λp​r​i​o​r\)\(\\rho,\\lambda^\{post\},\\lambda^\{prior\}\)settings\. Runtime is reported as the mean ratio over runs, with standard deviation\.GPLD introduces additional computation because the gradient penalty requires vector\-Jacobian products\. The main algorithmic cost driver is the sampling fractionρ\\rho, which determines the fraction of batch states on which the penalty is evaluated\. In contrast, changing the penalty coefficientλpost\\lambda^\{\\mathrm\{post\}\}does not substantially change the number of derivative evaluations and should not, by itself, determine the algorithmic overhead\.

Table[4](https://arxiv.org/html/2605.23089#A5.T4)reports wall\-clock overhead relative to vanilla DreamerV3\. For the default setting used in our main experiments,λ0post=0\.5\\lambda^\{\\mathrm\{post\}\}\_\{0\}=0\.5andρ=0\.5\\rho=0\.5, the measured runtime ratio is1\.511\.51\. However, this value should be interpreted cautiously: the standard deviation is relatively large, and a nearby posterior\-only run with the same sampling fraction,λ0post=0\.25\\lambda^\{\\mathrm\{post\}\}\_\{0\}=0\.25andρ=0\.5\\rho=0\.5, has a substantially lower runtime ratio of1\.231\.23\. Since these two settings require the same type of gradient\-penalty computation, the difference suggests that the1\.511\.51measurement likely reflects wall\-clock variability in the execution environment in addition to the algorithmic overhead of GPLD\.

Overall, the timing results indicate that GPLD adds measurable overhead relative to DreamerV3, but the1\.51×1\.51\\timesvalue should be viewed as a conservative high\-end wall\-clock estimate rather than the inherent cost of the method\. Runs with the same sampling fraction suggest a more typical overhead closer to the1\.21\.2–1\.3×1\.3\\timesrange, although precise runtime depends on hardware utilization and system\-level variability\.

## Appendix FLocal sensitivity analysis

To verify that GPLD changes the learned world model in the intended direction, we estimate the local sensitivity of the posterior and prior distributions during training\. For each saved checkpoint, we generate trajectories and perturb the corresponding latent inputs\. For the posterior, we perturb the encoder\-dependent inputete\_\{t\}; for the prior, we perturb the deterministic statehth\_\{t\}\. We then measure the change in the output distribution using KL divergence and normalize by the perturbation norm\.

Letrϕ\(⋅∣ut\)r\_\{\\phi\}\(\\cdot\\mid u\_\{t\}\)denote either the posterior or prior distribution evaluated at its corresponding inpututu\_\{t\}\. For a perturbationδ\\delta, we estimate

DKL\(rϕ\(⋅∣ut\)∥rϕ\(⋅∣ut\+δ\)\)‖δ‖2,\\frac\{D\_\{\\mathrm\{KL\}\}\\\!\\left\(r\_\{\\phi\}\(\\cdot\\mid u\_\{t\}\)\\,\\\|\\,r\_\{\\phi\}\(\\cdot\\mid u\_\{t\}\+\\delta\)\\right\)\}\{\\\|\\delta\\\|\_\{2\}\},with perturbation magnitudes ranging from0\.1%0\.1\\%to10%10\\%of the average input norm\. We average this quantity across time steps, perturbation magnitudes, and seeds\. Lower values indicate that the learned latent distribution changes less under small input perturbations\.

Figure[9](https://arxiv.org/html/2605.23089#A6.F9)shows aggregate local sensitivity across tasks\. GPLD reduces the sensitivity of the learned posterior distribution relative to DreamerV3, consistent with the intended effect of the gradient penalty\. The prior sensitivity is also reduced in several cases, even though GPLD is applied directly only to the posterior, suggesting that the posterior regularizer influences the prior through the DreamerV3 dynamics loss\.

![Refer to caption](https://arxiv.org/html/2605.23089v1/x17.png)

\(a\) Walker walk: Posterior aggregate

![Refer to caption](https://arxiv.org/html/2605.23089v1/x18.png)

\(b\) Walker walk: Prior aggregate

![Refer to caption](https://arxiv.org/html/2605.23089v1/x19.png)

\(c\)Hopper stand: posterior aggregate

![Refer to caption](https://arxiv.org/html/2605.23089v1/x20.png)

\(d\) Hopper stand: prior aggregate

Figure 9:Aggregate local sensitivity of the learned posterior and prior distributions for Walker walk and Hopper stand environments\. We perturb the corresponding latent inputs and measure the KL change in the output distribution normalized by the perturbation norm\. GPLD reduces posterior sensitivity relative to DreamerV3, consistent with the intended effect of the gradient penalty\.