PlaNet's world model plus a learned actor and critic.
The world model is unchanged from PlaNet — the same encoder, the same RSSM, the same decoder and reward head. Two heads are added: an actor that proposes actions from a latent state and a critic that scores one.
Parameters
configDreamerConfigNotes
Reference: Hafner, Lillicrap, Ba, and Norouzi, "Dream to Control: Learning Behaviors by Latent Imagination", ICLR, 2020.
actions[:, t] is the action taken into step t, matching
lucid.models.generative.planet.PlaNetModel.
This class holds no environment and never steps one. Everything it does — filtering, imagining, acting — happens in latent space, which is what makes the whole family trainable from a replay buffer alone.
Examples
>>> import lucid
>>> from lucid.models.generative.dreamer import DreamerConfig, DreamerModel
>>> 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)
>>> model = DreamerModel(cfg)
>>> _, posteriors = model.observe(lucid.randn((1, 4, 3, 64, 64)),
... lucid.randn((1, 4, 2)))
>>> model.act(posteriors, sample=False).shape
(1, 4, 2)Used by 2
Constructors
1Instance methods
9Propose actions for a state — in (-1, 1), or one-hot if discrete.
Parameters
stateRSSMState(B, T, ·) or a
single step (B, ·).samplebool= True, keyword-onlyTrue) or take its squashed mean
(False).Returns
Tensor(B, T, action_dim) for a sequence, (B, action_dim) for
a single step — the rank that went in. A discrete action
space gives one-hot rows rather than values in (-1, 1).
Notes
Both ranks are accepted because acting is inherently a single-step operation: an agent choosing its next move holds one belief, not a sequence of them. Demanding a length-1 time axis at the call site would be an artifact of how the heads are batched, and every caller would strip it again immediately.
Examples
>>> import lucid
>>> from lucid.models.generative.dreamer import DreamerConfig, DreamerModel
>>> 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)
>>> model = DreamerModel(cfg)
>>> _, posteriors = model.observe(lucid.randn((1, 3, 3, 64, 64)),
... lucid.randn((1, 3, 2)))
>>> model.act(posteriors).shape
(1, 3, 2)
One belief per batch element gives one action each, squashed inside
(-1, 1):
>>> last = posteriors.map(lambda t: t[:, -1])
>>> action = model.act(last)
>>> action.shape, bool((action.abs() < 1).all())
((1, 2), True)
sample=False is the squashed mean — the same action every time:
>>> mode = model.act(last, sample=False)
>>> bool((mode == model.act(last, sample=False)).all())
TrueReconstruct frames from a state — (B, T, C, 64, 64).
Embed a frame sequence — (B, T, C, 64, 64) -> (B, T, embed_size).
Propose an action for a state — (B, T, action_dim) in (-1, 1).
Parameters
featureTensor(B, T, latent_size).samplebool= True, keyword-onlyTrue) or take the squashed mean
(False, which is how a trained policy should act).Returns
TensorActions bounded to (-1, 1), or one-hot rows when discrete.
Roll the dynamics forward under the actor's own policy.
This is the loop the whole method rests on. Unlike PlaNet's
lucid.models.generative.planet.PlaNetModel.imagine, which
is handed a fixed action sequence to evaluate, here the action at
each step is produced by the actor from the state the model just
imagined — so the trajectory and the policy are coupled, and a
gradient taken at the end reaches the policy at every step along
the way.
Parameters
stateRSSMState(N, ·).horizonintsample(bool or None, optional, keyword - only)= NoneTrue) or take their
means (False). None follows the config's mean_only
setting, matching observe — a model configured
deterministic must imagine deterministically too.Returns
RSSMStateThe imagined states including the start, (N, horizon + 1, ·).
Notes
The start state is not detached here — that is the caller's
decision, and DreamerForWorldModeling does detach it so
the actor's gradient cannot reach the world model.
The state the actor reads is detached when the config says so, which is the released implementation's behaviour and the default. It does not stop the actor learning: the gradient still arrives through each action it produced. What it drops are the terms in which a return depends on the policy through the state it read.
Examples
>>> import lucid
>>> from lucid.models.generative.dreamer import DreamerConfig, DreamerModel
>>> 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)
>>> model = DreamerModel(cfg)
>>> _, posteriors = model.observe(lucid.randn((1, 3, 3, 64, 64)),
... lucid.randn((1, 3, 2)))
Every filtered step becomes an independent start, as the training
objective uses them:
>>> start = posteriors.map(lambda t: t.reshape(3, -1))
>>> states, actions = model.imagine(start, horizon=5)
>>> states.deter.shape, actions.shape
((3, 6, 8), (3, 5, 2))
The start is kept as the first imagined state — one more state than
action — and the heads score every one of them:
>>> bool((states.deter[:, 0] == start.deter).all())
True
>>> model.predict_reward(states).shape
(3, 6)observe
→RSSMStateobserve(observations: Tensor, actions: Tensor, state: RSSMState | None = None, sample: bool | None = None)Filter a trajectory into posterior states.
Parameters
observationsTensor(B, T, C, 64, 64).actionsTensor(B, T, action_dim).stateRSSMState or None= NoneNone starts from zeros.sample(bool or None, optional, keyword - only)= NoneTrue) or take its mean (False).
None follows the config's mean_only setting.Returns
RSSMStateWhat the dynamics predicted, (B, T, ·).
Examples
>>> import lucid
>>> from lucid.models.generative.dreamer import DreamerConfig, DreamerModel
>>> 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)
>>> model = DreamerModel(cfg)
>>> obs, act = lucid.randn((1, 3, 3, 64, 64)), lucid.randn((1, 3, 2))
>>> priors, posteriors = model.observe(obs, act)
>>> posteriors.deter.shape, posteriors.stoch.shape, posteriors.std.shape
((1, 3, 8), (1, 3, 4), (1, 3, 4))
Prior and posterior share one deterministic path — the frame refines
the belief about the latent, not the history that led to it:
>>> bool((priors.deter == posteriors.deter).all())
True
sample=False takes each latent's mean, so filtering repeats
exactly:
>>> first = model.observe(obs, act, sample=False)[1].stoch
>>> bool((first == model.observe(obs, act, sample=False)[1].stoch).all())
TruePredict the discount at a state — logits, (B, T).
The head is Bernoulli: its probability is how likely the episode is to continue past this state, and the discount used downstream is that probability rather than a constant. A state the agent will not survive therefore discounts everything after it to nothing, which is the point — a constant would have the planner keep collecting rewards past the end of the episode.
Parameters
stateRSSMStateReturns
TensorLogits, (B, T). Apply sigmoid for the probability;
the loss consumes the logits directly.
Raises
ValueErrorpcont.Examples
>>> import dataclasses
>>> import lucid
>>> from lucid.models.generative.dreamer import DreamerConfig, DreamerModel
>>> 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, pcont=True)
>>> model = DreamerModel(cfg)
>>> _, posteriors = model.observe(lucid.randn((1, 3, 3, 64, 64)),
... lucid.randn((1, 3, 2)))
>>> logits = model.predict_pcont(posteriors)
>>> logits.shape
(1, 3)
sigmoid turns the logits into the probability of continuing,
which is what imagination then discounts by:
>>> keep = lucid.sigmoid(logits)
>>> bool(((keep > 0) & (keep < 1)).all())
True
The head exists only when the configuration asks for it:
>>> plain = DreamerModel(dataclasses.replace(cfg, pcont=False))
>>> plain.predict_pcont(posteriors)
Traceback (most recent call last):
...
ValueError: this model has no discount head...Predict reward from a state — (B, T).
Estimate the value of a state — (B, T).