DreamerV3ForWorldModeling
WorldModelingModelDreamerV3ForWorldModeling(config: DreamerV3Config)DreamerV3 with its world-model, actor and critic objectives.
Parameters
configDreamerV3ConfigNotes
Reference: Hafner, Pasukonis, Ba, and Norouzi, "Mastering Diverse Domains through World Models", Nature 640 (2025), 647-653 (arXiv:2301.04104).
Same three-optimiser contract as the earlier families, and the same
reason for it. Use backward — spending the losses by hand
either contaminates the world model or raises, depending on the order.
update_slow_critic must be called once per gradient step. It
is not folded into backward because it counts optimiser
steps, not backward passes, and only the caller knows when one has
happened.
The indexing is simpler than DreamerV2's, and the simplification is
the point. Imagination produces H + 1 states and H actions;
action[t] is taken from state[t], the return R[t] starts
at state[t], and the critic's estimate v[t] scores the same
state. So the advantage is R[t] - v(s_t) paired with
log pi(a_t | s_t) at every t in 0 .. H-1 — the textbook
alignment, available here because the actor is score-function only
and does not need the one-step shift dynamics backpropagation forced
on DreamerV2.
Examples
>>> import lucid
>>> from lucid.models import dreamer_v3_12m_world_model
>>> model = dreamer_v3_12m_world_model(action_dim=2, cnn_depth=2, stoch_size=3,
... discrete=4, deter_size=8, hidden_size=8, actor_hidden=8,
... value_hidden=8, reward_hidden=8, num_bins=5, horizon=4, pcont=False)
>>> 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
6The actor's parameters.
Returns
list of ParameterTrained by actor_loss alone.
Give every parameter group the gradient of its own loss.
Parameters
outputDreamerV3Outputforward, with behavior populated.Raises
ValueErroroutput carries no losses.Notes
Identical in shape to the earlier families', and identical in motivation: the three losses share one graph, so backpropagating them by hand either lets the actor's gradient descend the world model or raises on a parameter the imagination's graph still needed.
Examples
>>> import lucid
>>> import lucid.optim as optim
>>> from lucid.models import dreamer_v3_12m_world_model
>>> model = dreamer_v3_12m_world_model(action_dim=2, cnn_depth=2,
... stoch_size=3, discrete=4, deter_size=8, hidden_size=8,
... actor_hidden=8, value_hidden=8, reward_hidden=8, num_bins=5,
... horizon=3, pcont=False)
>>> groups = (model.world_parameters(), model.actor_parameters(),
... model.value_parameters())
>>> opts = [optim.Adam(g, lr=1e-4) for g in groups]
>>> out = model(lucid.randn((1, 3, 3, 64, 64)),
... lucid.randn((1, 3, 2)), lucid.randn((1, 3)))
>>> model.backward(out)
>>> [all(p.grad is not None for p in g) for g in groups]
[True, True, True]
The slow critic is no optimiser's to move — it holds no gradient,
and is pulled toward the learner by update_slow_critic,
called once per step:
>>> slow = model.dreamer_v3.slow_value_head.parameters()
>>> any(p.grad is not None for p in slow)
False
>>> for opt in opts:
... opt.step()
>>> model.update_slow_critic()forward(observations: Tensor, actions: Tensor, rewards: Tensor, discounts: Tensor | None = None)Train the world model and the behaviour on one batch.
Parameters
Returns
DreamerV3Outputloss is the world-model loss; the behaviour losses are on
.behavior and take their own optimisers.
Raises
ValueErrorpcont is configured and discounts is omitted.Notes
The reward term is a two-hot cross-entropy rather than a squared error, so a domain whose rewards are in the thousands produces gradients of the same size as one whose rewards are in hundredths. The reconstruction stays a squared error: pixels are already in a fixed range and have nothing to normalise away.
Move the slow critic a little toward the live one.
Notes
Every gradient step, by critic_ema — 2% at the paper's value.
DreamerV2 copied its target outright every hundred steps; a
continuous average has no schedule to tune and no discontinuity
for the critic's loss to step over. The first call copies
outright, since a slow critic left at its initialisation is not
something worth regressing toward.
The learning critic's parameters — not the slow copy.
Returns
list of ParameterTrained by value_loss. The slow copy is written by
update_slow_critic and never by an optimiser, which is
what makes it a fixed point to regress toward.
Everything the world-model loss trains.
Returns
list of ParameterEncoder, RSSM, decoder, reward head, and the discount head when there is one.