Density-Guided Bootstrap Boundaries for Latent World-Model Reinforcement Learning

Shashank Gupta; Pilhwa Lee

cs.LGOct 29, 2026

Density-Guided Bootstrap Boundaries for Latent World-Model Reinforcement Learning

Authors: Shashank Gupta, Pilhwa Lee

Organizations: Department of EECS, University of Michigan, Ann Arbor, MI 48109, USA · Department of Mathematics, Morgan State University, Baltimore, MD 21251, USA

Abstract

Latent world models allow reinforcement learning agents to learn policies from imagined trajectories, yet determining the optimal depth of model reliance remains a fundamental challenge. Fixed imagination horizons apply uniform prediction depths regardless of trajectory reliability, exposing target values to compounding model error or over-relying on critic approximations. We introduce density-guided bootstrap boundaries, a trajectory-dependent mechanism that adapts the depth of imagined returns using a score-informed relative density proxy. The proxy accumulates the estimated divergence of a score-corrected auxiliary field formed by a feature displacement network and a noise-conditioned score estimator along imagined paths. We propose two gating mechanisms: a hard boundary that halts return expansion at a discrete state, and a soft boundary that continuously shifts the return recursion toward critic bootstrapping. Both variants preserve boundary value estimates and apply corresponding survival weights to actor and critic losses. Formulated independently of specific dynamics models and instantiated within Dreamer V3, we prove that the resulting gated Bellman operator is a strict contraction mapping and that early bootstrapping monotonically reduces conditional target variance. We clarify the theoretical assumptions distinguishing our expansion proxy from exact path likelihoods. Across continuous-control tasks and Atari Pong, empirical evaluations show that dynamic boundary selection enhances sample efficiency and late-training performance over fixed-horizon and auxiliary-enabled baselines.

Figures & tables

Setting	Observations	Action repeat	Interaction budget
Control Suite	Proprioceptive vectors	2	500,000 environment steps
Atari Pong	$64\times 64$ RGB images	4	400,000 emulator frames

Table 1: Observation and interaction settings. Budgets follow the repeat-adjusted logging convention, corresponding nominally to 250,000 agent decisions for Control Suite and 100,000 for Pong.

Panel of Figure 1 — Figure 1: Evaluation returns for vanilla DreamerV3 and density-guided hard and soft bootstrap boundaries across five tasks. Curves summarize 10 independent training seeds; shaded regions indicate one standard deviation across seeds. TensorBoard-style exponential smoothing with weight 0.7 is used for visualization.

Environment	Method	Early	Middle	Late
Walker Walk	Vanilla	$544.2\pm 128.0$	$875.2\pm 75.1$	$945.6\pm 26.9$
	Hard	$735.9\pm 104.3$	$917.3\pm 69.0$	$944.2\pm 27.5$
	Soft	$663.5\pm 82.6$	$925.6\pm 79.8$	$951.3\pm 24.2$
Walker Run	Vanilla	$254.3\pm 18.7$	$495.6\pm 61.5$	$667.1\pm 45.7$
	Hard	$367.2\pm 45.3$	$651.9\pm 43.2$	$743.5\pm 46.6$
	Soft	$276.9\pm 52.8$	$562.6\pm 54.7$	$723.4\pm 29.9$

Table 2: Evaluation returns at early, middle, and late training stages. Stage targets are 100k, 250k, and 500k environment steps for DeepMind Control tasks, and 80k, 200k, and 400k for Pong. Values are mean

\pm

std.

Panel of Figure 2 — Figure 2: Comparison of vanilla DreamerV3, fixed bootstrap boundaries, and density-guided hard boundaries on Cheetah Run (left) and Walker Run (right). Fixed boundaries use six and five steps, respectively. The Cheetah fixed control retains the auxiliary components, whereas the Walker fixed control disables them.

Panel of Figure 3 — Figure 3: Comparison of vanilla DreamerV3, the auxiliary-enabled cutoff-off control, and density-guided hard boundaries on Cheetah Run (left) and Walker Run (right). The cutoff-off control retains auxiliary displacement and score learning without using the density signal to modify imagined returns. Curves use TensorBoard-style smoothing with coefficient 0.7.

Panel of Figure 4 — Figure 4: Density-threshold sensitivity on Cheetah Run for hard boundaries (left) and soft boundaries (right), using $\tau\in\{-500,-700,-900\}$ . Each threshold configuration uses seed 0, with boundary activation at 200k environment steps, increment cap $C=200$ , and minimum retained depth $m=2$ . Soft boundaries use temperature $T=100$ . The dashed curve shows the vanilla reference.

Panel of Figure 5 — Figure 5: Effective imagination length on Cheetah Run (left) and Walker Run (right), with one training run per variant. Vertical dotted lines mark boundary activation at 200k environment steps. Horizontal dashed lines indicate six- and five-transition references, respectively; these are not measured fixed-control curves. Hard effective length counts retained transitions, whereas soft effective length sums survival weights.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

\mathcal{L}_{\text{disp}}\leftarrow\frac{\sum_{i}m_{i}\|f_{\omega}(u_{i},\bar{x}_{i},b_{i})-(\bar{x}_{i+1}-\bar{x}_{i})\|_{2}^{2}}{d\max(1,\sum_{i}m_{i})}.

Appendix

Algorithm 2 Dreamer V3 World-Model and Auxiliary Update

Model-Based RL

Model-based reinforcement learning uses environment models to plan actions and learn policies with fewer real-world interactions.

See all 253 papers

Dreaming Smoothly and Sample Efficiently with Gradient Penalized Latent Dynamics
May 21, 2026Romil V. Sonigra, P. R. KumarLatent Dynamics Modeling Jacobian Regularization
Texas A&M University · Department of Electrical and Computer Engineering, Texas A&M University, College Station, TX 77843
GIRL: Generative Imagination Reinforcement Learning via Information-Theoretic Hallucination Control
Apr 8, 2026Prakul Sunil HiremathModel-Based RL World Models
Department of Computer Science and Engineering, Visvesvaraya Technological University (VTU), Belagavi, India · Aliens on Earth (AoE) Autonomous Research Group, Belagavi, India
Mind Dreamer: Untethering Imagination via Active Causal Intervention on Latent Manifolds
May 15, 2026Shaojun Xu, Xiaoling Zhou, Yihan Lin +4Reinforcement Learning Credit Assignment in RL
Center for Brain-Inspired Computing Research, Department of Precision Instrument, Tsinghua University, Beijing, China · College of Computer Science and Technology, Zhejiang University, Hangzhou, China · Pen-Tung Sah Institute of Micro-Nano Science and Technology, Xiamen University, Xiamen, China +1

May 21, 2026J/K move · Enter open · S save

Dreaming Smoothly and Sample Efficiently with Gradient Penalized Latent Dynamics

Romil V. Sonigra, P. R. Kumar

Texas A&M University · Department of Electrical and Computer Engineering, Texas A&M University, College Station, TX 77843

Latent Dynamics Modeling Jacobian Regularization Latent World Models Sample-Efficient RL Model-Based RL

Open paper AskPDF arXiv

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 .

Abstract

Model-Based RL

Explore similar work

Dreaming Smoothly and Sample Efficiently with Gradient Penalized Latent Dynamics

GIRL: Generative Imagination Reinforcement Learning via Information-Theoretic Hallucination Control

Mind Dreamer: Untethering Imagination via Active Causal Intervention on Latent Manifolds