DreamerV2's architecture with distributional heads and a slow critic.
Parameters
configDreamerV3ConfigNotes
Reference: Hafner, Pasukonis, Ba, and Norouzi, "Mastering Diverse Domains through World Models", Nature 640 (2025), 647-653 (arXiv:2301.04104).
Two architectural differences from DreamerV2 are visible here.
The reward and value heads are .TwoHotHead\ s — they predict
a distribution over exponentially spaced bins and are read back
through symexp, which is what decouples the gradients from the
reward's magnitude.
The critic is duplicated, but the copy plays the opposite role to DreamerV2's. There the frozen copy produced the targets; here the returns bootstrap from the live critic and the slow copy appears only as a regulariser pulling the live one toward its own moving average. The copy is therefore updated every step by a small fraction rather than replaced wholesale on a schedule.
Examples
>>> import lucid
>>> from lucid.models import dreamer_v3_12m
>>> model = dreamer_v3_12m(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)
>>> _, posteriors = model.observe(lucid.randn((1, 3, 3, 64, 64)),
... lucid.randn((1, 3, 2)))
>>> posteriors.is_discrete, model.act(posteriors, sample=False).shape
(True, (1, 3, 2))Used by 2
Constructors
1Instance methods
9Propose actions — (B, T, action_dim) or (B, action_dim).
Parameters
stateRSSMStatesamplebool= True, keyword-onlyReturns
TensorActions inside (-1, 1), or one-hot when discrete.
Examples
>>> import lucid
>>> from lucid.models import dreamer_v3_12m
>>> kwargs = dict(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)
>>> model = dreamer_v3_12m(**kwargs)
>>> obs, act = lucid.randn((1, 3, 3, 64, 64)), lucid.randn((1, 3, 2))
>>> _, posteriors = model.observe(obs, act)
>>> last = posteriors.map(lambda t: t[:, -1])
>>> model.act(posteriors).shape, model.act(last).shape
((1, 3, 2), (1, 2))
>>> bool((model.act(last).abs() < 1).all())
True
A discrete action space picks one alternative instead, as a one-hot:
>>> game = dreamer_v3_12m(action_space="discrete", **kwargs)
>>> _, beliefs = game.observe(obs, act)
>>> choice = game.act(beliefs.map(lambda t: t[:, -1]))
>>> choice.shape, choice.sum(dim=-1).tolist()
((1, 2), [1.0])Reconstruct frames from a state — (B, T, C, 64, 64).
Notes
Frames are not passed through symlog. The transform is for
quantities whose range differs between domains; a pixel is always
in the same interval, and the reference applies it to vector
observations and to reward and value only.
Embed a frame sequence — (B, T, C, 64, 64) -> (B, T, embed_size).
forward(observations: Tensor, actions: Tensor)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.
Roll the dynamics forward under the actor's own policy.
Parameters
stateRSSMState(N, ·).horizonintsample(bool or None, optional, keyword - only)= Nonemean_only.Returns
RSSMStateThe imagined states including the start, (N, horizon + 1, ·).
Notes
The state the actor reads is detached. Unlike the earlier families this costs nothing: DreamerV3's actor is trained purely by the score function, so no gradient was ever going to travel back through the dynamics from the return.
Examples
>>> import lucid
>>> from lucid.models import dreamer_v3_12m
>>> model = dreamer_v3_12m(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)
>>> _, posteriors = model.observe(lucid.randn((1, 3, 3, 64, 64)),
... lucid.randn((1, 3, 2)))
Every filtered step becomes an imagination start:
>>> start = posteriors.map(lambda t: t.reshape(3, *t.shape[2:]))
>>> states, actions = model.imagine(start, horizon=5)
>>> states.stoch.shape, states.logits.shape, actions.shape
((3, 6, 12), (3, 6, 3, 4), (3, 5, 2))
Reward along the trajectory is decoded from its bins in reward
units — one per state, the start included:
>>> model.predict_reward(states).shape
(3, 6)observe
→priors, posteriors : RSSMStateobserve(observations: Tensor, actions: Tensor, state: RSSMState | None = None, sample: bool | None = None)Filter a trajectory into posterior states.
Parameters
Returns
priors, posteriors : RSSMState(B, T, ·) each, carrying categorical logits already mixed
with unimix.
Examples
>>> import lucid
>>> from lucid.models import dreamer_v3_12m
>>> model = dreamer_v3_12m(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)
>>> obs, act = lucid.randn((1, 3, 3, 64, 64)), lucid.randn((1, 3, 2))
>>> priors, posteriors = model.observe(obs, act)
>>> posteriors.stoch.shape, posteriors.logits.shape
((1, 3, 12), (1, 3, 3, 4))
unimix is already folded in: every class keeps at least
unimix / discrete of the mass, so no outcome is ever ruled out:
>>> floor = model.config.unimix / model.config.discrete
>>> bool((lucid.softmax(posteriors.logits, dim=-1) >= floor).all())
TruePredict the discount at a state — logits, (B, T).
Parameters
stateRSSMStateReturns
TensorBernoulli logits; apply sigmoid for the probability.
Raises
ValueErrorpcont.Examples
>>> import lucid
>>> from lucid.models import dreamer_v3_12m
>>> model = dreamer_v3_12m(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)
>>> _, posteriors = model.observe(lucid.randn((1, 3, 3, 64, 64)),
... lucid.randn((1, 3, 2)))
>>> logits = model.predict_pcont(posteriors)
>>> logits.shape
(1, 3)
The head is on by default in this family; sigmoid gives the
probability of continuing that imagination discounts by:
>>> model.config.pcont
True
>>> keep = lucid.sigmoid(logits)
>>> bool(((keep > 0) & (keep < 1)).all())
TruePredict reward from a state — (B, T), in reward units.
Estimate a state's value — (B, T).
Parameters
stateRSSMStateslowbool= False, keyword-onlyReturns
TensorValue estimates, in return units.
Examples
>>> import lucid
>>> from lucid.models import dreamer_v3_12m
>>> model = dreamer_v3_12m(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)
>>> _, posteriors = model.observe(lucid.randn((1, 3, 3, 64, 64)),
... lucid.randn((1, 3, 2)))
>>> value = model.predict_value(posteriors)
>>> value.shape
(1, 3)
Both critics are zero-initialised, so an untrained critic asserts no
return for the actor to chase, and the slow copy agrees with it:
>>> slow = model.predict_value(posteriors, slow=True)
>>> float(value.abs().max().item()), float(slow.abs().max().item())
(0.0, 0.0)