Diffusion and Flow World Models › Diffusion over States and Trajectories

When the goal is to represent multimodal conditional futures, a distributional predictor or sampler is more suitable than a single MSE point estimate. A design question then appears: diffuse over what? Temporal horizon is one axis, with one-step successor diffusion and joint finite-horizon diffusion as common endpoints and multi-step chunks between them. Separate axes specify which variables are generated, such as states or actions, and whether the target is observable, compressed, or hierarchical and latent.

The first option, which we will call one-step or per-state diffusion, generates one next state or latent at a time. In the Markov version, conditioning is $c=(s_t,a_t)$ and each rollout step draws $s_{t+1}\sim p_\theta(\cdot\mid s_t,a_t)$; histories or belief states can be added when the current state is not sufficient. Relative to generating an $H$-step trajectory jointly, the target omits the multiplicative horizon dimension. That makes it a natural probabilistic transition interface, but not automatically cheap: latency still depends on state dimension, architecture, and sampler.

Per-state diffusion inherits the rollout-shift problem flagged in Part IV, Ch 5: Stochasticity and Uncertainty. Transition uncertainty and model error propagate through later conditioning; they do not simply add as independent variances. If training uses only dataset states while deployment conditions recursively on model samples, errors can move the rollout into contexts where the learned conditional has not been validated. The risk generally grows with horizon, though its rate depends on stability, calibration, and data coverage.

The second option is trajectory diffusion. Here $x_0$ is an entire horizon, typically a matrix $\tau\in\mathbb{R}^{H\times d}$ stacking states and sometimes actions. Diffuser (Janner et al., 2022), for example, uses a temporal U-Net over joint state-action trajectories. Both designs can represent cross-time dependence. Joint trajectory models parameterize it directly and can represent dependencies missed by a one-step model when that model's state or history is not Markov-sufficient. Recursively composed one-step kernels usually induce it through later transitions' dependence on earlier states. Given an initial-state distribution and either a fixed action sequence or a policy, an exact sufficient-state Markov kernel determines the corresponding joint trajectory law. Joint generation also exposes the full horizon to one generative computation, but neither formulation by itself guarantees dynamically feasible or constraint-satisfying plans.

Trajectory diffusion makes horizon-wide inpainting direct because every constrained time index belongs to one generated variable. We can fix $s_0=\bar{s}_0$, reimpose selected states or actions during each denoising update, and generate the remaining entries. A standard causal one-step model needs an additional planner, smoother, bridge-conditioning scheme, or other coordination mechanism to impose comparable future constraints. Decision Diffuser (Ajay et al., 2023) conditions state-sequence generation on returns, skills, or constraints and uses inverse dynamics to recover actions. Diffusion Policy (Chi et al., 2023) is instead an observation-conditioned diffusion policy over action sequences. These are adjacent uses of diffusion, not interchangeable definitions of a predictive world model.

The tradeoffs between the two designs are real:

• Per-state generation omits the horizon factor in its output and can accept either Markov or history-based context, but it must be invoked recursively for a rollout.
• Trajectory generation represents a joint distribution over the horizon and supports direct inpainting, but its output dimension, memory, and per-evaluation cost grow with horizon length.
• Per-state composes naturally with model-based RL pipelines that alternate between sampling transitions and planning; trajectory models equipped with return, goal, or constraint control are closer to generative planners and can replace a separate classical trajectory optimizer by turning planning into guided or constrained sampling, as in Diffuser.
• Both designs depend on training support. A trajectory constraint far outside the dataset can produce an implausible completion even when its masked entries are enforced exactly.

A useful mental model is that per-state diffusion parameterizes a stochastic transition kernel, whereas trajectory diffusion acts more like a generative trajectory optimizer when sampling is guided or constrained by a planning objective. Without such control, it remains a joint trajectory generator. Either design can be placed in a receding-horizon loop: propose a short plan, execute its first action, observe the environment, and replan. Combining joint proposals with a separate reactive transition model is possible, but it is a system design choice rather than an intrinsic property of either method.

Actions can appear either inside the generated variable or in the conditioning. Joint state-action trajectory generation returns a plan and executable actions, as in the original Diffuser formulation. A predictive transition model instead conditions on candidate actions supplied by an external policy or planner. A state-only trajectory generator may use inverse dynamics afterward to infer actions, as in Decision Diffuser. Diffusion Policy generates action sequences directly from observation context and is a policy rather than a state predictor. We return to these roles in Part VI, Ch 2: Actions, Inverse Dynamics, and Latent Actions.