DreamerV2ForWorldModeling
WorldModelingModelDreamerV2ForWorldModeling(config: DreamerV2Config)DreamerV2 with its world-model, actor and critic objectives.
Parameters
configDreamerV2ConfigNotes
Reference: Hafner, Lillicrap, Norouzi, and Ba, "Mastering Atari with Discrete World Models", ICLR, 2021 (arXiv:2010.02193).
Same three-optimiser contract as Dreamer, 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_target 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 follows the released implementation, which is worth
stating because it is not obvious. Imagination produces H + 1
states. Two are lost at the end — one is the bootstrap, and one has
an action that leads nowhere — and one target is lost at the start,
because the first state came from the replay buffer rather than from
the policy. So the actor is scored on states 0 .. H-2 against
targets 1 .. H-1, and the critic on states 0 .. H-1.
Examples
>>> import lucid
>>> from lucid.models import dreamer_v2_world_model
>>> model = dreamer_v2_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, 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
outputDreamerV2Outputforward, with behavior populated.Raises
ValueErroroutput carries no losses.Notes
Identical in shape to Dreamer's, 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_v2_world_model
>>> model = dreamer_v2_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, 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 target critic is no optimiser's to move — it holds no gradient,
and follows the learner through update_slow_target, called
once per step:
>>> target = model.dreamer_v2.target_value_head.parameters()
>>> any(p.grad is not None for p in target)
False
>>> for opt in opts:
... opt.step()
>>> model.update_slow_target()forward(observations: Tensor, actions: Tensor, rewards: Tensor, discounts: Tensor | None = None)Train the world model and the behaviour on one batch.
Parameters
Returns
DreamerV2Outputloss is the world-model loss; the behaviour losses are on
.behavior and take their own optimisers.
Raises
ValueErrorpcont is configured and discounts is omitted.Refresh the target critic, on the paper's schedule.
Notes
Called once per gradient step; it moves the target only every
slow_target_update of them. The very first call copies
outright whatever slow_target_fraction says, since a target
left at its initialisation is not a useful thing to regress onto.
The learning critic's parameters — not the target copy.
Returns
list of ParameterTrained by value_loss. The target copy is written by
update_slow_target and never by an optimiser, which is
what makes it a fixed point to regress onto.
Everything the world-model loss trains.
Returns
list of ParameterEncoder, RSSM, decoder, reward head, and the discount head when there is one.