Dreamer-v2

DreamerV2 keeps its predecessor's shape — a recurrent state-space model, an actor and a critic trained on imagined trajectories — and changes what the stochastic latent is. Where Dreamer draws $s_t$ from a diagonal Gaussian, this draws a grid of categorical variables: 32 of them, 32 classes each, sampled as one-hots. The state is then a sparse binary vector of length 1024 with exactly 32 bits set.

A one-hot sample has no gradient, so the paper takes the straight-through estimator — the same device that makes vector quantisation trainable:

which passes the sample forward and the class probabilities' gradient backward.

The motivation the paper gives is representational rather than computational. A Gaussian latent must describe a multi-modal future — a ball that will bounce left or right — with a single mean, and so describes neither; a categorical one can put mass on both outcomes and keep them apart. Atari is full of such futures, which is where the Gaussian model had been losing.

The second change is to the divergence. A single KL term makes the prior chase the posterior and the posterior chase the prior at the same rate, and the paper observes the posterior giving way — it collapses toward a prior that has not learned anything yet. KL balancing splits the term and weights the two directions separately:

with $\alpha = 0.8$. Most of the pressure is on the prior to explain what the posterior saw, and only a fifth on the posterior to stay predictable — which is the asymmetry that keeps the latent carrying the observation.

Two smaller changes complete it. The critic regresses onto a target network refreshed every hundred steps, so its own moving estimate does not chase itself. And the actor's gradient can come either through the dynamics, as before, or from a score-function estimator; the paper uses the analytic path for continuous control and REINFORCE for Atari, where the actions are discrete and there is no path to differentiate.