Diffusion and Flow World Models › Worked Example: A Bimodal One-Dimensional World

To make this concrete, we will build a minimal one-dimensional conditional testbed in which the MSE-optimal point estimate lies between two modes. The state $s_t$ and action $a_t$ are scalars. Given $(s_t, a_t)$, the next state is

with $m$ and $\epsilon$ independent, $\Delta = 1.2$, and $\sigma = 0.25$. The hidden variable $m$ is never observed. Thus $p(s_{t+1}\mid s_t,a_t)=\tfrac12\mathcal N(s_t+a_t-\Delta,\sigma^2)+\tfrac12\mathcal N(s_t+a_t+\Delta,\sigma^2)$. A single point cannot represent both components. Under squared error, its population optimum is the conditional mean $s_t+a_t$, halfway between them in a region of extremely low, but nonzero, density. This mirrors the puck example from the introduction: the dynamics are simple conditional on the hidden mode, and the difficulty is that the mode is unknown to the model.

This testbed also separates marginal and conditional diagnostics. In the data generator below, $s_t\sim\mathcal N(0,1)$ and $a_t\sim\mathcal N(0,0.5^2)$ independently of $m$ and $\epsilon$. After integrating over $s_t+a_t$, the analytic marginal is only weakly bimodal: its center is about $0.63\%$ below its nearby maxima. A coarse marginal histogram can therefore look like one broad peak even though every fixed-context conditional has two well-separated components. Evaluating only the marginal can hide the conditional ambiguity that matters for prediction.

We begin with imports and dataset generation.

The dataset contains 20,000 one-step transitions, and each target is a scalar next state, so the generative models operate on a one-dimensional continuous variable. The conditioning vector is two-dimensional, stacking $s_t$ and $a_t$, and the hidden mode is discarded at training time so that no model can cheat by observing it. This compact implementation runs on CPU with float32 tensors. A device-portable version must place all tensors participating in an operation on compatible devices, align the floating dtype of model parameters, schedules, inputs, and temporary floating tensors, and preserve the required integer or Boolean dtypes of indices and masks. Samples must then move to CPU before NumPy, SciPy, or Matplotlib use. The seeds stabilize this run within a fixed environment, but exact losses and metrics can change across library versions or hardware.

We can see the bimodality directly by plotting the held-out samples for the fixed context.

The distribution peaks near $+1.2$ and $-1.2$ and has extremely low density near the origin. A point model optimized by MSE targets the conditional mean at the origin. We will compare that analytic consequence with the empirical behavior of learned diffusion and flow models.

We define a small MLP denoiser conditioned on $t$, $s_t$, and $a_t$. Fixed sinusoidal time features are a common diffusion baseline, although many architectures transform them further or use learned alternatives. Here they give the network a smooth, multi-scale encoding of the timestep so that nearby timesteps produce nearby embeddings, which helps the network share structure across the schedule.

The forward diffusion schedule uses a linear beta increase from $10^{-4}$ to $0.05$ over $T = 500$ steps, which gives $\bar\alpha_T \approx 3 \times 10^{-6}$, essentially pure noise.

We train the denoiser for 30 epochs. Each minibatch samples a random timestep per example and applies the denoising score matching objective.

The reported denoising loss is an optimization diagnostic, not a calibrated likelihood or a guarantee of sample quality. It need not reach zero because, after corruption, a particular noise realization is not uniquely determined by $(x_t,t,c)$; this ambiguity is especially relevant at low-noise timesteps where the observation remains strongly tied to the unknown clean sample.

Two samplers are useful. Full DDPM ancestral sampling visits every timestep and injects stochastic noise; DDIM sampling follows a selected timestep sequence and is deterministic given its initial noise when $\eta=0$. Both use the learned denoiser to approximate a reverse generative process. They do not literally invert the independently sampled forward corruptions used during training.

We now define and train a mean-square-error baseline for comparison. This is a plain MLP that maps $(s_t, a_t)$ to a point estimate of $s_{t+1}$.

The learned point estimate should land near the analytic conditional mean at the origin, in the low-density trough between the modes. This is not a bug: it is the population-optimal answer for squared error, even though it is unrepresentative of a typical draw from the conditional distribution.

Now we turn to flow matching. The model class is identical, but the corruption path, time sampling, and regression target change.

The flow-matching loss is another training diagnostic. Its numerical value should not be compared directly with the denoising loss because the objectives use different targets, scales, and timestep weighting. Here, separate instances of the same small MLP architecture optimize the two objectives well enough for the distributional diagnostics below.

Sampling from the flow model uses forward Euler integration from $t = 0$ to $t = 1$. The numerical path is deterministic conditional on the sampled Gaussian initial value.

We now generate fixed-context samples using three procedures: ancestral DDPM and DDIM sampling from the learned diffusion model, and Euler sampling from the learned flow model. The moment summary includes all three, while the histogram directly compares the synthetic reference, ancestral DDPM, and flow samples. All sampling happens before plotting.

The printed means and standard deviations provide a first check of symmetry and spread, but they cannot by themselves establish bimodality: a broad unimodal distribution can share both moments. The histogram below is the more relevant diagnostic because it shows whether probability mass concentrates near both component means.

We visualize three sampled distributions on the same axes together with the MSE point estimate. The DDIM sampler is used again for the rollout and step-sensitivity diagnostics below; the 40-step DDIM samples generated above are used only in the moment summary.

In this seeded run, DDPM and flow matching place samples near both synthetic-reference peaks, while the MSE baseline lies near zero in a very low-density region. The general lesson is narrower than a claim of universal superiority: a point prediction cannot represent a multimodal conditional, whereas a well-trained generative model can represent multiple alternatives.

Distribution recovery is only half the story. We next illustrate recursive rollout behavior by repeatedly calling each per-state model on its own output. The hidden mode in this simulator is resampled independently at every transition, so the six-step process has many possible mode sequences rather than one persistent left/right branch.

The shape is (8, 7) for each model: eight separately initialized rollout runs and seven states including the initial state. The generative rows use distinct seeds, while the deterministic MSE rows are identical.

We plot the trajectories from both generative models and the MSE baseline.

The generative rollouts diverge because each per-state call draws another conditional sample, while all eight deterministic MSE traces overlap. These paths are a behavior illustration, not evidence that the learned models reproduce the true joint rollout distribution. Indeed, because the simulator resamples $m$ at each transition, a rollout can switch branch sign; joint horizon consistency would need a separate trajectory-level evaluation.

We close the experiment with a sampling-step sensitivity check. For this implementation, both DDIM and Euler use one network evaluation per step, so step count equals NFE. We measure Wasserstein-1 distance to the finite synthetic-reference sample under one fixed seed; this is an illustrative diagnostic, not a benchmark.

The resulting quality-versus-steps curves show a practical difference between these trained DDIM and Euler-flow implementations.

Here the measured flow errors at 4, 8, 16, 32, and 64 evaluations are approximately $0.209$, $0.185$, $0.159$, $0.156$, and $0.158$; the DDIM errors are $1.139$, $0.313$, $0.150$, $0.104$, and $0.087$. Euler flow is markedly better at the two coarsest budgets, the curves cross around 16 evaluations, and DDIM is better at the larger budgets. This single seeded, one-dimensional experiment supports a local quality-versus-NFE comparison only; it does not establish a general ordering between flow and diffusion samplers. It is not a latency measurement: deployment requires end-to-end timing on the target control stack.