Source Paper: https://arxiv.org/pdf/2509.24527
Main Ideas
I picked this up because I wanted to really see how RL in imagination itself works, and this is quite a seminal paper in the domain of world models.
As a broad overview, the governing principles of Dreamer 4 to do imagination RL is:
- learn a rich latent representation
- be stable while generating latent space rollouts
- establish a good baseline policy and reward estimator
After these three principles are established, it’s easy to do RL on countless imagined scenarios simply because you can make an agent that looks at the latent space directly, a reasonably good reward estimator to supply to the agent and finally a starting policy that makes sure the agent isn’t fumbling around in its own imagination.
The question sort of arose within me… why is imagination RL important? The immediate answer I could think of was that you can create very usable policies without having to deploy the agent while it is still fumbling around in potentially harmful scenarios and be a liability to other humans. Another reason is that, a rich latent representation and being able to generate good stable rollouts for that representation comes together to provide a very nice RL environment. So this is more scalable, cheaper and more efficient than a live environment. The issue is that imagination is an approximation. The model learns some dynamics but we can never truly be sure that the learned dynamics really represents the information we’d need for the agent to do the task.
Background Things I Needed To Learn
Here’s a bunch of notes on things I added to study this paper, anything that seems unique to this paper I will inline into the notes for this paper directly:
- Flow Matching
- Shortcut Models
- Diffusion Forcing
- LPIPS Loss (Learned Perceptual Image Patch Similarity)
- Temporal Difference Learning
Symexp two-hot
Break this down into two parts,
Define:
And its inverse:
You use symlog when you want to compress the number down to a small space near zero.
symexp on the other hand is the function inverse of that. Get your original values back by running symexp on your symlogged values.
In two-hot encoding instead of taking the top probabiity, we first bin the values into a set of evenly spaced points. A separate classifier network outputs logits which we use to weight the bins to find the value. We train this classifier using cross-entropy.
PMPO (Preference-based Maximimum a Posteriori Policy Optimization)
Retain only the sign of the advantage rather than the value of the advantage while running policy-gradient RL.
Grouped Query Attention
Instead of maintaining separate keys and values for each Q matrix, we maintain a common set for Q matrices.
This significantly reduces KV cache size.
This does come at a risk of degrading model quality.
Methodology
Broadly, we have 3 phases:
- World Model Pre-Training
- Agent Fine-tuning
- Imagination Training
World Model Pre-Training
Causal Tokenizer
For a given time step, create image patch embeddings and a learned latent token.
A transformer model converts images into patch tokens along with a learned latent token. The encoder converts them into learned latent token embeddings. The output of this encoder is projected into a lower dimensional embedding with a tanh activation. The decoder projects this back to the model embedding dimension and continues.
Trick about attention: During the encoding step, the latents are allowed to attend to all modalities and the tokens from each modality are only allowed to attend to other tokens from that same modality. Vice versa with the decoder, the modalities attend to itself and the learned latents. The learned latents can only attend to itself.
As far as I understand, this seems to align with the objective of each component, we’re trying to force the model to really learn a good representation that carries enough information for downstream tasks. The latents are trying to “get good” in the encoder step, so they must learn from all the modalities. During the decoder step, the latents must be “good enough” to reconstruct the original image now.
The goal of this tokenizer is to learn a compressed representation that makes the original image reconstructable. Hence there’s a lower dimension projection made in between the encoder and decoder.
Masked Autoencoding - The tokenizer has an image reconstruction objective, a combination of mean squared error and LPIPS Loss.
There’s also a masking step, some input patches are replaced with a learned mask embedding sampled to a probabiity distribution. The hope is that this will force the model to actually learn spatial relationships instead of simply copying pixel values. During inference, just stop the random dropouts.
Interactive Dynamics
Once we have the tokenizer trained. The latent constructed between the encoder and decoder should now be good enough to contain enough information to reconstruct the image, or in other words contain information about the game state and history and other useful details.
Building on top of Diffusion Forcing and Shortcut Models, we have shortcut forcing. The two main things to really keep in mind is that instead of predicting the velocity, as seen in shortcut models, we predict the latent itself and perform the euler integration, the loss is computed by converting the latent back to velocity.
Say we have a network that is parameterized on distance and action ,
If we know,
If we break it down, b’ and b” are the same velocities as seen in Shortcut Models.
The new loss function becomes:
I struggled at first with the second case, so here’s the breakdown:
The term,
is nothing but the predicted latent being mapped back to velocity.
Subtracting that from gives us the error compared to the “smarter” velocities which take two subsequent jumps
Scaling the resultant norm by is to negate the effects of the division by
After doing this, we have a ramp loss to supply weight to the that meaningfully contributes the most learning to the network.
The intuition for this is that, near zero is practically meaningless. There’s zero information for the model to learn, but we still want the model to predict something near 0, because that’s what it will be doing during inference. So we can’t omit those samples from the training inputs, so we just downweight them in a linear manner as seen above.
Imagination Training
After pre-training, add another modality. Agent tokens. The agent tokens get a task in embedding form, and we predict the policy and reward using MLP heads.
We’re at the part where the model is fine-tuned for taking task inputs, using these two additional MLP heads.
Behaviour Cloning
To bootstrap the model with a policy, we can train it to imitate the players in the dataset already. Predict multiple action and reward tokens in one go from timestep , and calculate the Negative Log Loss of your distribution against what the players in your dataset did. We have be the action of the player in the footage and be the reward sampled from your policy and reward heads. Sample tokens in one loss update.
These reward heads are trained using the symexp two heads mentioned in an earlier section.
Note of course that during this step, the tokenizer is frozen and the dynamics model is finetuned.
RL
Imitating behaviours can only take us so far. What happens when the model sees something it hasn’t encountered in the training data? What happens when the player performs a mistake in the data and it learns to imitate that? We need RL to take things further.
The special thing about this RL is that it performs no actions on a live environment. During RL training, the model never actually plays Minecraft. The RL runs are done completely on the imagination of the model.
We initialize a value head and take the trained policy head from the behaviour cloning step.
Get the reward using the reward head, and the value head is trained using Temporal Difference Learning.
Annotate the resulting trajectories using the reward head,
Breaking down each term,
This is a negative log loss to train the value head.
This is a TD-learning objective. is a binary term, if the episode hasn’t terminated, is 1, else 0. is a reward output from the frozen reward head.
We then compute the advantage,
Revising what PMPO does, we only take the sign of the advantage rather than the magnitude. Why do this? My hypothesis is that the model can only see that it’s doing the right thing and doesn’t try to reward hack its way to optimize for specific high advantage actions.
Conclusion
I think my main takeaway from this paper is just that behaviour cloning is pretty good to get a baseline policy. Imagination RL is good to keep in mind, the framework that they’ve setup around doing it seems very transferable to other domains. I didn’t get a chance to go deep on the efficient transformer architecture they have proposed but my main interest was in their pipeline around the model architecture itself.