raodreamer: deep & broad model imagination

Contrary to popular belief, machines can dream. For many years, the technologies around you have been improved by machines that imagine paths into the future and learn to act based on this imagination. This machine-learning paradigm uses what is known as a world model. I ask: should agents dream broadly or deeply? Should models imagine paths far into the future, many shallow paths, or some mixture of the two? Is there an optimal way to derive decisions from a world model given a fixed budget?

Dreamer is a model-based reinforcement-learning algorithm that learns latent environment dynamics and trains an actor-critic using trajectories imagined inside the learned world model. For a discrete actor π_θ(a∣z), the action-dependent component of the imagined policy gradient at latent state z has the familiar actor-gradient theorem form:

g(z) = Σ_a∈𝒜 π_θ(a∣z) Y(z,a) ∇_θ log π_θ(a∣z)

where Y(z,a) is Dreamer’s stop-gradient long-horizon advantage coefficient. Standard discrete Dreamer estimates this expectation using a single sampled action, which is inexpensive but may exhibit substantial action-selection variance.

We introduce RaoDreamer, a compute-budgeted estimator that evaluates a weighted bank of counterfactual actions from the same latent state. The bank

ℬ(z) = {(a_i,w_i)}_i=1^B

is constructed so that its weighted expectation reproduces the original policy expectation. One branch is selected according to I∼Categorical(w) and continued as an ordinary Dreamer trajectory, preserving the marginal distribution of the long-horizon imagined rollout.

The remaining branches define a short-horizon control variate:

ĝ_Rao = Y_Is_I + β [Σ_i=1^B w_iC_is_i − C_Is_I]

Because

𝔼_I[Σ_iw_iC_is_i − C_Is_I ∣ ℬ] = 0

RaoDreamer preserves the expected pre-optimizer Dreamer actor gradient while potentially reducing its variance and improving training stability. The control C_i need not be an unbiased return estimate; its accuracy affects variance rather than expectation. When B=1, the correction cancels exactly and the method reduces to standard Dreamer. RaoDreamer therefore allows allocating limited world-model computation between action breadth, trajectory depth, and additional imagined roots.

Results

12 paired training seeds; 100 evaluation episodes per policy at 1M frames.
metric	Dreamer	Rao B=4	paired / relative
mean score	8.97	11.41	+2.44
paired median	n/a	n/a	+1.52
paired wins	3/12	9/12	n/a
paired-effect SD	n/a	n/a	4.80
95% interval	n/a	n/a	[-0.62, +5.49]
preregistered two-sided p	n/a	n/a	0.107
logical transitions	1.00×	1.20×	+20%
wall time / policy	1.73 h	1.85 h	+6.9%

Paired Dreamer and Rao B=4 Pong scores for twelve training seeds, alongside ordered paired effects — Fresh 1M-frame policies by paired training seed.

Mean online Pong learning curves for Dreamer and Rao B=4 across twelve paired training seeds — Fresh 12-seed learning curves, mean ± SEM.

I get the intuition that this method will generalize to common world-model actor-critic methods. RaoDreamer’s shallow imagination improved its long-term outcomes here, even at the cost of additional model transitions. Perhaps androids do dream of electric sheep, but when trying to predict how many will jump the fence, it may be better to imagine many immediate scenarios than one deep one. That choice is not yet made by the model itself; RaoDreamer performs the meta-thinking on the model’s behalf, because we hand-roll the rule that decides how its imagination should be spent. Humans do this intrinsically; we learn to learn more effectively, we consider how we think, we consider why we think. I believe very deeply that if we can teach models to make these choices for themselves, rather than hand-rolling each meta-optimization, there is a path to something greater for classical actor-critic methods.

The figures are regenerated from retained per-seed evaluations and training time series.

Lev Petersen

Results