Valdi: Value Diffusion World Models

Christopher Lindenberg; Kashyap Chitta

Source-linked AI summary

Valdi: Value Diffusion World Models

Christopher Lindenberg, Kashyap Chitta

arXiv:2607.00917v1Jul 1, 2026cs.LGcs.AI

TL;DR

MPC needs world models that are fast enough for online planning while representing uncertain dynamics. Valdi combines online-trained latent diffusion dynamics with value learning, matching deterministic MLP control performance while producing multimodal trajectory predictions.

Problem
Online MPC needs latent dynamics models that combine rapid inference with expressive representations of uncertain, multimodal state transitions.
Method
Valdi jointly trains a latent diffusion dynamics model and value function end-to-end within a TD-MPC-style online control loop.
Results
Valdi matches MLP control performance while producing multimodal trajectory predictions.
Takeaways & Limitations
In preliminary CarRacing experiments, Valdi shows that diffusion-based latent dynamics can match deterministic control while retaining multimodal predictions.
Takeaways & Limitations
The one-step training and inference scheme may not scale to substantially more complex dynamics requiring multi-step denoising.

Abstract

from arXiv · show

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.

1 Introduction

MPC needs fast latent world-model inference, while diffusion models capture complex long-horizon dynamics but rely on iterative inference. Valdi addresses this tension by training an online latent diffusion dynamics model to predict value functions for MPC.

Motivation: MPC requires rapid inference, motivating lightweight latent world models that predict compressed observations rather than raw pixels.Latent world models can nevertheless suffer representation collapse from trivial optimization solutions.
Motivation: Diffusion models capture complex long-horizon distributions, but their iterative inference conflicts with MPC’s low-latency requirements.They are typically trained in observation space to avoid representation collapse.
Valdi: Valdi trains a latent diffusion dynamics model to predict value functions end-to-end within an online control loop, following TD-MPC.The approach is designed to bridge diffusion modeling with low-latency MPC.
Valdi: Valdi predicts action-sequence-conditioned rewards and values, enabling Model Predictive Control in the CarRacing environment.Figure 1 illustrates predictions for good and bad action sequences.

2 Related Work

This section situates Valdi within diffusion-based world modeling for model-based reinforcement learning and latent-space MPC approaches such as TD-MPC. Diffusion models have become a state-of-the-art world-modeling approach, while TD-MPC introduced task-oriented latent dynamics for online planning.

Diffusion World Models: Diffusion models have emerged as a state-of-the-art approach for both pixel-space and latent world modeling in model-based reinforcement learning.These models have been used to train RL agents inside learned world models.
Temporal Difference Learning for Model Predictive Control (TD-MPC): TD-MPC pioneered a Task-Oriented Latent Dynamics Model for online planning in a latent space.Given an observation, TD-MPC encodes it into a latent representation and models task-relevant dynamics for planning.
Temporal Difference Learning for Model Predictive Control (TD-MPC): TD-MPC uses representation, dynamics, reward, and value models as analogous high-level components to Valdi.The approach operates from an observation, its latent representation, and the corresponding action.

3 Value Diffusion World Models

Valdi combines joint latent-trajectory diffusion with value-based MPC, using one-step denoising during both training and inference. Its end-to-end training jointly learns the encoder, dynamics, reward, and value components for low-latency control.

Model: Valdi generates the entire H-step latent trajectory jointly with a diffusion model instead of rolling out latent states autoregressively.The model uses the velocity parameterization of Salimans & Ho (2022).
Model: Valdi predicts state values rather than action values and uses no policy prior; its CNN encoder, MLP reward/value models, and bidirectional transformer total ∼5M parameters.This simpler value formulation is sufficient because Valdi lacks TD-MPC’s policy prior.
Training: Training jointly optimizes latent diffusion, reward, and value losses while updating the encoder rather than freezing it.Reward and value losses use one-step denoised predictions from noisy state embeddings.
Training: One-step denoising keeps training tractable and matches single-step inference for low-latency MPC, with reward error and temporal-difference losses applied to denoised trajectories.The value loss is a temporal difference loss, and the reward loss follows TD-MPC’s reward error.
Inference: At inference, a Cross Entropy Method solver maximizes discounted returns over action sequences using predicted rewards and the terminal value.The solver executes the first action from the optimized sequence.

4 Experiments

Experiments in modified CarRacing show that single-step Valdi matches a deterministic MLP baseline in control, while additional diffusion steps increase predictive variety but slightly degrade control. Across imagined-rollout diagnostics, Valdi is initially less accurate but becomes more self-consistent and better grounded than the MLP near the planner’s bootstrap depth.

Performance and Multimodality: Valdi matches the deterministic MLP baseline in control at one diffusion step, within run-to-run variance.Training and default evaluation both use a single diffusion step.
Performance and Multimodality: More diffusion steps substantially increase trajectory variety but slightly degrade control, both within run-to-run variance.The authors attribute degradation to training-inference mismatch and increased trajectory variance affecting CEM selection.
Performance and Multimodality: Valdi’s predictive distribution is narrow at one diffusion step, typically committing to a single track continuation despite CarRacing’s partial observability.Predictive variety is measured using mean pairwise LPIPS across decoded final states from N = 100 rollouts.
Value Function: At short rollout depths, Valdi is less accurate than the MLP in both self-consistency and grounding.Self-consistency measures agreement across adjacent imagined steps, while grounding compares imagined and real latents at matching environment timesteps.
Value Function: Near bootstrap depth H, Valdi is more self-consistent at h = 4 and better grounded at h = 5 than the baseline, although gaps remain within run-to-run variance.The trend is consistent across both diagnostics.

5 Conclusion

The paper presents Value Diffusion World Models, combining latent diffusion dynamics with value-based online planning. The method matches MLP control performance while producing multimodal trajectory predictions, but its single-step design may not scale to more complex dynamics.

5 Conclusion: Valdi combines diffusion and latent world models for inference-time planning.The latent diffusion dynamics model is trained jointly with a value function in a TD-MPC-style online loop.
5 Conclusion: Valdi matches MLP control performance while producing multi-modal trajectory predictions.This result is obtained with joint training of the latent diffusion dynamics model and value function inside the online loop.
5 Conclusion: The one-step estimate used during training and inference may not scale to substantially more complex dynamics.Such environments may require multi-step denoising, motivating distillation from multi-step teachers to single-step students to preserve inference-time feasibility.

Supplementary Materials The following content was not necessarily subject to peer review. … A.2 Diffusion Models

The supplementary materials describe TD-MPC’s latent dynamics architecture, joint training and MPC procedure, then introduce diffusion-based dynamics modeling with a specified velocity objective and noise schedule. They also explain how Valdi differs from TD-MPC, especially by replacing Qθ and πθ with a state-value function Vθ and omitting policy bootstrapping.

A.1 TD-MPC: TD-MPC combines an encoder, deterministic latent dynamics model, reward function, action-value function, and prior policy for latent-space planning.The encoder maps observations to latent states, while the dynamics model predicts the next latent state from the current state and action.
A.1 TD-MPC: The TOLD components are jointly trained using reward prediction, temporal-difference learning, and multi-step latent consistency objectives.These objectives encourage accurate rewards, value learning, and self-consistent latent rollouts.
A.1 TD-MPC: TD-MPC’s deterministic transition model is efficient but does not explicitly represent ambiguous futures caused by partial observability or stochastic multimodal transitions.This limitation motivates diffusion-based dynamics modeling for uncertain futures.
A.1 TD-MPC: At inference, TD-MPC uses the learned latent model in an MPC loop with MPPI to search action sequences, execute only the first action, and replan.Rollouts satisfy the latent dynamics constraint z_t+1 = D_θ(z_t, a_t).
A.1 TD-MPC: Valdi omits TD-MPC’s prior-policy bootstrapping because its dynamics model predicts all future states jointly from the full action sequence.TD-MPC interleaves prior-policy actions with model predictions to improve planning efficiency and possibly asymptotic performance.
A.1 TD-MPC: Valdi replaces TD-MPC’s action-value function Qθ and prior policy πθ with a state-value function Vθ, reducing code complexity and training moving parts.The paper does not formally ablate this choice because its focus is the dynamics model and its effects.
A.2 Diffusion Models: The diffusion dynamics model is trained with a velocity-prediction objective using a linear variance schedule over T = 1000 steps, from β1 = 1e−4 to βT = 2e−2.The schedule uses the cumulative product of (1−β_i).

B Architecture Details

Valdi and the MLP baseline share the same TOLD architecture and differ only in their dynamics models, isolating dynamics as the studied variable. Valdi uses a latent diffusion transformer alongside shared encoder, reward, and value components.

Shared TOLD layout: Both systems use TOLD components Eθ, Dθ, Rθ, and Vθ, with identical non-dynamics components.Only the dynamics model differs between Valdi and the MLP baseline.
Parameterization: Valdi has 5,390,230 trainable parameters, versus 5,795,363 for the MLP baseline.The baseline is slightly larger, so performance parity does not result from a Valdi capacity advantage.
Encoder: The encoder separately processes masked frames and seven-dimensional proprioception, concatenates two 256-dimensional embeddings, and projects them to a 64-dimensional latent.The frame uses a convolutional stack, while proprioception uses a small MLP and the projection head is LayerNorm-prefixed.
Dynamics: Valdi’s dynamics model is a bidirectional encoder-only transformer with 6 pre-norm blocks at width 256 and one token per world-model step.A learned Embedding(1000, 32) maps diffusion timestep τ to 32 dimensions, and action plus timestep conditioning is concatenated before the first block.
Reward and Value: The reward model is a residual MLP over a 73-dimensional latent-action input, while the value model ensembles two identical MLP heads and takes their elementwise minimum for TD targets.The reward input is 64+9 = 73, and the minimum reduces value overestimation.

C Training Details … F Hyperparameters

The paper trains Valdi with a weighted four-term objective and performs latent-space MPC using CEM, action chunking, and single-step diffusion for low-latency inference. The environment separates visual and proprioceptive inputs, while hyperparameters are provided in Table 3.

C Training Details: Valdi optimizes a weighted sum of diffusion, reward, value, and SIGReg losses, with reward and value losses applied to one-step denoised latent rollouts.The TD target uses the minimum of two value heads to reduce value overestimation.
C Training Details: SIGReg regularizes only the first timestep of each sampled trajectory and is included chiefly to match the reported implementation rather than as a principled method component.It encourages an isotropic Gaussian latent distribution.
C Training Details: Training collects MPC-driven episodes, stores transitions in replay, samples trajectory segments, encodes observations, predicts denoising directions, and updates online and target networks.The target network is updated using an exponential moving average.
D Planning and Inference: At inference, Valdi plans in latent space with a Cross Entropy Method that samples action sequences, predicts latent trajectories, scores discounted returns, and replans after executing the first action.The sampling distribution is refit to elite candidates after each iteration.
D Planning and Inference: Each world-model step represents 3 environment steps, so horizon H = 5 corresponds to a 15-step environment horizon and conditions on 9 scalar action dimensions.The planner reasons over five world-model steps while action chunks contain three environment actions.
D Planning and Inference: At one diffusion step, the planner runs at over 10 Hz on a single RTX 4080 and the agent acts at over 30 Hz, while the MLP baseline is slightly faster.Inference cost grows linearly with the number of diffusion steps; neither approach received significant performance optimization.
E Environment: The modified CarRacing environment separates visual and proprioceptive inputs by blacking out the rendered proprioceptive bar and exposing the information as a seven-dimensional state vector.The modification was primarily intended to explore multimodality.
F Hyperparameters: All hyperparameters are listed in Table 3, with the action-chunking entry explained in the planning and inference supplement.The provided passage does not report individual hyperparameter values.

G Additional Results · G.1 Training and Evaluation Returns · G.2 Multimodal Trajectory Predictions

Valdi matches the MLP baseline in control returns within run-to-run variance, while additional visualizations demonstrate diverse imagined futures from a shared starting state. These trajectories are generated through latent dynamics and post hoc decoding.

G.1 Training and Evaluation Returns: Valdi matches the MLP baseline within run-to-run variance in both training-time and evaluation-time returns.The reported comparison finds neither significant improvement nor degradation in control performance.
G.1 Training and Evaluation Returns: Figure 4 reports training-time returns for Valdi and the MLP baseline.
G.1 Training and Evaluation Returns: Figure 4 also reports evaluation-time performance for both runs of both systems.
G.2 Multimodal Trajectory Predictions: The dynamics model generates diverse futures from the same starting state by selecting four minimally similar trajectories from 100 imagined futures.Pairwise visual similarity is computed between decoded final states before selecting the subset.
G.2 Multimodal Trajectory Predictions: Each selected trajectory state is decoded independently to visualize the resulting imagined futures.
G.2 Multimodal Trajectory Predictions: The post hoc Vision Transformer decoder reconstructs image patches from frozen encoder latents and is never used during training or planning.One decoder is trained per fully trained encoder on an on-policy dataset collected from the corresponding model.
G.2 Multimodal Trajectory Predictions: Figure 5 shows diverse trajectories generated from one starting state using 8 diffusion steps, with an identical initial frame across trajectories.Frames one through five are imagined in latent space and then decoded.