Can one diffusion model tell where it is unsure?

Offline reinforcement learning with a diffusion world model. Work in progress.

The setting

In offline reinforcement learning, an agent has to learn from a fixed dataset and is never allowed to try things out. A popular approach is model-based: learn a model of the world from the data, then practice inside it. The danger is that the model is wrong wherever the data is thin, and a learning agent is good at finding those mistakes and exploiting them.

The remedy is to be pessimistic where the model is unsure, in effect practicing in a simulated worst-case world. That needs a measure of uncertainty. The usual measure is an ensemble: train several world models and look at how much they disagree. It works, but it costs several times the training.

The question

Can one diffusion model give a useful uncertainty signal? Here I measure output spread for a fixed conditioning state. The fixed-start SDE analysis in my diffusion theory post motivates studying that spread, but does not establish that it measures prediction error. Spread can reflect environmental randomness as well as model uncertainty; the empirical comparisons below test its usefulness as a proxy.

Does the spread rank uncertainty correctly?

To test it, I trained a five-model ensemble next to each diffusion model and compared the two: does the single model's spread rank states the same way the ensemble's disagreement does, as the states move further from the data?

Rank agreement (Spearman) between one diffusion model's spread and the disagreement of five independently trained models, averaged over 12 models (3 datasets, 4 seeds), on 2,000 states per distance. 1.0 means the two always rank states the same way.

On and near the data, the single model ranks states the way the ensemble does (0.89 on the data with the deterministic sampler). Far from the data it weakens, and at 3 standard deviations away the stochastic sampler is actually a little better than the deterministic one. That is where uncertainty matters most, and it is still the weak point.

Does it predict real errors?

Agreement with an ensemble is only useful if the ensemble itself is a good guide. So I also compared against the true one-step prediction error on the data.

How well the uncertainty predicts the model's actual one-step prediction error, on the data. Averaged over four seeds per dataset, with the deterministic sampler.

On the data, one model's spread predicts the real error about as well as five models do. Only the deterministic sampler is shown here. With the stochastic sampler, hopper was weaker for both (0.47 for one model, 0.43 for the ensemble).

What went wrong

The ranking can be right while the size is wrong. With the stochastic sampler, the spread grew much more slowly than the ensemble's disagreement as states moved off the data, so it under-reported uncertainty exactly where it should be highest. Switching to a deterministic sampler fixed most of this for single steps, but not for long rollouts, and an uncertainty that is right for one step is not yet enough to build a worst-case world for a whole trajectory.

Next tests

Policy-level runs. The next test is to train a policy with diffusion uncertainty and compare it with the same policy without it, with two standard offline RL algorithms, on three datasets with four seeds each. This is the test of whether a good one-step reading actually helps the agent.
Uncertainty over whole trajectories. Making the reading hold up over long rollouts, not only one step ahead. This is where the stochastic sampler failed, and where the policy has been exploiting imagined states.