DreamerForWorldModeling
WorldModelingModelDreamerForWorldModeling(config: DreamerConfig)Dreamer with its world-model, actor and critic objectives.
Three losses come out of one forward pass:
- the world-model loss, identical in form to PlaNet's — pixel and reward reconstruction against a free-nats-clamped KL;
- the actor loss, over trajectories imagined under the policy itself;
- the critic loss, regressing onto those same targets held fixed.
Parameters
configDreamerConfigNotes
Reference: Hafner, Lillicrap, Ba, and Norouzi, "Dream to Control: Learning Behaviors by Latent Imagination", ICLR, 2020.
These three must be optimised separately. Summing them and taking
one step is a different algorithm: the actor's gradient would descend
the world model too, and the critic would chase a moving target. The
parameter groups are exposed as world_parameters,
actor_parameters and value_parameters, which partition
the model exactly.
Use backward to fill those groups' gradients. The losses share
a graph, so backpropagating them by hand either contaminates the world
model or raises, depending on the order chosen — backward is
the whole training step bar the step calls.
Imagination starts from detached posterior states, so no behaviour gradient reaches the encoder or the filtering that produced them. Within the horizon, though, the actor's gradient deliberately flows through the learned dynamics — through the RSSM transition, the reward head and the critic — because that analytic path is the paper's entire contribution. Both the latent draw and the action draw are reparameterised, which is what lets a return computed 15 steps out reach the policy that chose step 1. Only the actor's parameters are updated by it; that is what the parameter grouping is for.
The critic is the exception: it reads detached states, so value_loss
reaches nothing but the critic itself.
Reconstruction is a Gaussian log-likelihood with unit variance, so it reduces to a squared error summed over pixels and averaged over the batch — the same convention as PlaNet.
Examples
>>> import lucid
>>> from lucid.models.generative.dreamer import (
... DreamerConfig, DreamerForWorldModeling)
>>> cfg = DreamerConfig(action_dim=2, cnn_depth=2, stoch_size=4,
... deter_size=8, hidden_size=8, actor_hidden=8,
... value_hidden=8, reward_hidden=8, horizon=3)
>>> model = DreamerForWorldModeling(cfg)
>>> out = model(lucid.randn((1, 3, 3, 64, 64)), lucid.randn((1, 3, 2)),
... lucid.randn((1, 3)))
>>> bool(out.loss.ndim == 0), bool(out.behavior.actor_loss.ndim == 0)
(True, True)Used by 2
Constructors
1Instance methods
5The actor's parameters.
Returns
list of ParameterTrained by actor_loss alone.
Give every parameter group the gradient of its own loss.
The three losses share one graph, and that makes the obvious ways of using them both wrong:
- Backward all three and then step — every world-model parameter ends up carrying the actor's gradient as well as its own, so the world model quietly descends the policy's objective. Nothing errors; the algorithm is simply no longer Dreamer's.
- Backward-then-step each in turn — the first
stepmutates parameters that the imagination's deeper graph still needs, and the nextbackwardraises. It happens to work if the actor goes first, which is not a property anyone should have to know.
So this does it once, correctly: each loss is backpropagated in
isolation, its group's gradients are kept, and the rest are
discarded. Afterwards every group holds exactly its own gradient
and the three optimisers may step in any order.
Parameters
outputDreamerOutputforward, with behavior populated.Raises
ValueErroroutput carries no losses — it came from
DreamerModel rather than from this wrapper.Examples
>>> import lucid
>>> import lucid.optim as optim
>>> from lucid.models import dreamer_world_model
>>> model = dreamer_world_model(action_dim=2, cnn_depth=2,
... stoch_size=4, deter_size=8, hidden_size=8, actor_hidden=8,
... value_hidden=8, reward_hidden=8, horizon=3)
>>> opts = [optim.Adam(g, lr=1e-4) for g in (model.world_parameters(),
... model.actor_parameters(), model.value_parameters())]
>>> out = model(lucid.randn((1, 3, 3, 64, 64)),
... lucid.randn((1, 3, 2)), lucid.randn((1, 3)))
>>> model.backward(out)
>>> for opt in opts:
... opt.step()forward(observations: Tensor, actions: Tensor, rewards: Tensor, discounts: Tensor | None = None)Train the world model and the behaviour on one batch of trajectories.
Parameters
observationsTensor(B, T, C, 64, 64).actionsTensor(B, T, action_dim).rewardsTensor(B, T).(B, T), 1 where the episode continued past that step
and 0 where it ended. Required when the config asks for a
discount head, ignored otherwise.Returns
DreamerOutputloss is the world-model loss; the behaviour losses are on
.behavior and take their own optimisers.
Raises
ValueErrorpcont is configured and discounts is omitted — the
head has nothing to learn from without it, and defaulting to
"never terminates" would silently train it to a constant.The critic's parameters.
Returns
list of ParameterTrained by value_loss alone.
Everything the world-model loss trains — encoder, RSSM, decoder, reward.
Returns
list of ParameterThe world model's parameters, disjoint from
actor_parameters and value_parameters.