Valdi: Value Diffusion World Models

Abstract: World models can enable Model Predictive Control (MPC), but this requires dynamics prediction that is both fast enough for online use and expressive enough to represent uncertain futures. Diffusion models offer a natural mechanism for modeling uncertain dynamics, yet their iterative inference procedure makes them difficult to use for low-latency latent planning. We bridge this gap with Value Diffusion World Models (Valdi), combining end-to-end online training for MPC with a latent diffusion dynamics model. In preliminary experiments on the CarRacing environment, we show that Valdi, using a single diffusion step at both training and inference, matches a deterministic MLP baseline. Our experiments expose a trade-off between predictive multimodality and control performance in this setup. Code is available at https://github.com/Kit115/ValueDiffusionWorldModels.

• The paper introduces a hybrid model that embeds latent diffusion dynamics into a TD-MPC framework, bridging expressive generative modeling with efficient planning.
• By replacing deterministic MLP dynamics with a transformer-based diffusion model, the approach enables diverse trajectory predictions while maintaining control performance.
• Empirical results in a CarRacing environment show that a single diffusion denoising step achieves competitive performance, highlighting a trade-off between generative expressivity and planning accuracy.

Value Diffusion World Models: A Technical Analysis

Introduction and Motivation

The paper "Valdi: Value Diffusion World Models" (2607.00917) introduces an approach for model-based RL that synergizes the strengths of diffusion models with the efficiency requirements of value-based Model Predictive Control (MPC). Standard approaches to planning in latent space, such as TD-MPC, use deterministic MLP dynamics models, which are efficient but limited in representing multimodal or uncertain environment transitions. Conversely, diffusion models have established state-of-the-art results for generative modeling of complex distributions in high-dimensional spaces, but their iterative denoising processes are computationally prohibitive for online MPC.

Valdi proposes to bridge this gap by embedding a latent-space diffusion dynamics model within a TD-MPC-style value-based planning framework. The system is trained end-to-end and online, with objectives for representation learning, reward prediction, and value estimation, providing practical insight into the trade-offs between generative expressivity and control performance.

Methodology

Model Architecture and Training

Valdi retains the core TOLD architecture of TD-MPC—consisting of a state encoder $E_\theta$, dynamics model $D_\theta$, reward model $R_\theta$, and state value model $V_\theta$—but substitutes the deterministic MLP dynamics with a bidirectional, encoder-only transformer-based diffusion model. The diffusion model, parameterized in latent space rather than pixel space, enables efficient online rollout for planning while preserving the property of modeling multimodal future trajectories.

Training is conducted online, using a trajectory replay buffer. The loss is a weighted sum of the standard velocity-based latent diffusion objective ($\mathcal{L}_{\rm diff}$), temporal-difference value prediction error ($\mathcal{L}_{\rm val}$), and reward prediction error ($\mathcal{L}_{\rm rew}$), with additional regularization to guard against latent collapse. Unlike most latent diffusion models—which freeze the encoder—Valdi jointly trains the encoder, ensuring latent representations are optimized for the downstream planning task.

Inference and Planning

At inference, Valdi uses CEM to optimize over action sequences by maximizing cumulative discounted reward and value predictions along imagined latent trajectories. A critical empirical finding is that a single diffusion denoising step suffices both at training and inference to match the MLP baseline in control, avoiding the prohibitive latency costs of multi-step diffusion processes.

Experimental Evaluation

All experiments were conducted in a modified CarRacing environment, with results compared to a TD-MPC-style MLP baseline under equivalent training and evaluation conditions.

Control Performance

Valdi, with a single diffusion step, matches the control performance of the MLP baseline across both training and evaluation, with differences well within run-to-run variance. This result holds despite the capacity for richer trajectory modeling inherent to diffusion models and the transformer architecture being slightly smaller than the baseline MLP.

Figure 1: Training (left) and evaluation (right) returns for Valdi and MLP baseline on the CarRacing environment.

Predictive Diversity

The diffusion dynamics allow Valdi to generate diverse, plausible trajectory continuations from identical initial states—behavior unattainable for deterministic models, which are forced to output either a mean or a single mode.

Figure 2: Example of diverse trajectories from the same start state using Valdi with 8 diffusion steps; frames 1–5 are generated by the model in latent space and then decoded to pixels.

Quantitative analysis using LPIPS demonstrates that increasing the number of inference diffusion steps leads to substantial increases in predictive diversity. However, this multimodality comes at the cost of degraded control, particularly because the planning algorithm (CEM) cannot efficiently exploit the increased uncertainty.

Value Function Behavior

Diagnostics of the value function—focusing on self-consistency (one-step TD residuals along imagined rollouts) and grounding (agreement of value estimates between imagined and real trajectories)—revealed that Valdi, while less accurate at short rollout horizons, accumulates lower errors relative to the MLP as rollout depth increases. This suggests that the predictions most critical for control (i.e., at the bootstrapping depth of planning) are more robust in Valdi, despite short-horizon inaccuracies.

Implications and Future Directions

Theoretical and Practical Takeaways

Valdi demonstrates that latent diffusion-based world models can be practically deployed in online model-based RL frameworks without sacrificing control efficiency, provided denoising is constrained to a single step. The system uniquely exposes the trade-off between generative expressivity (multimodality) and planning performance, a property not previously studied in this context.

The approach highlights a critical insight for future model-based RL research: increased predictive expressivity—often tacitly assumed to support superior control—may introduce variance that planning algorithms are not robust to, ultimately degrading policy effectiveness.

Prospective Developments

Further research is warranted on jointly optimizing training and inference for multimodal predictions—possibly through knowledge distillation or scheduled denoising curricula—to enable flexible test-time scaling of generative expressivity without compromising control. There is also an opportunity to design planning algorithms that can reliably exploit uncertain, multimodal dynamics models, particularly in more complex, stochastic environments where deterministic rollout is insufficient.

Conclusion

Valdi advances model-based RL by embedding a diffusion-based latent world model within a value-optimized control framework, empirically verifying that such stochastic dynamics models can achieve parity with traditional MLPs in online planning tasks when configured for minimal diffusion steps. The work exposes a nuanced but important tension between modeling capacity and planning efficacy, informing both theoretical understanding and practical implementation of world models in RL. The demonstrated capacity for diverse trajectory generation—without incurring computational bottlenecks—establishes a new point of reference for scalable, uncertainty-aware planning in future value-driven reinforcement learning systems.