class

DreamerV2Model

extendsPretrainedModel

DreamerV2Model(config: DreamerV2Config)

source edit

Dreamer's architecture with a categorical latent and a target critic.

Parameters

configDreamerV2Config

Frozen configuration.

Notes

Reference: Hafner, Lillicrap, Norouzi, and Ba, "Mastering Atari with Discrete World Models", ICLR, 2021 (arXiv:2010.02193).

The critic is duplicated. value_head is what learns; the frozen target_value_head is what the returns are computed from, and it is refreshed on a schedule. Without it the critic regresses onto targets built from itself, and a small error compounds every step it is fed back through.

Examples

>>> import lucid
>>> from lucid.models import dreamer_v2
>>> model = dreamer_v2(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)
>>> _, 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

dunder

init

→None

__init__(config: DreamerV2Config)

source edit

Instance methods

act

→Tensor

act(state: RSSMState, sample: bool = True)

source edit

Propose actions — (B, T, action_dim) or (B, action_dim).

Parameters

stateRSSMState

A sequence or a single step; the rank is preserved.

samplebool= True, keyword-only

Draw, or take the policy's mode.

Returns

Tensor

Actions inside (-1, 1), or one-hot when discrete.

Examples

>>> import lucid
>>> from lucid.models import dreamer_v2
>>> 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)
>>> model = dreamer_v2(**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_v2(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])

decode

→Tensor

decode(state: RSSMState)

source edit

Reconstruct frames from a state — (B, T, C, 64, 64).

encode

→Tensor

encode(observations: Tensor)

source edit

Embed a frame sequence — (B, T, C, 64, 64) -> (B, T, embed_size).

forward

→DreamerV2Output

forward(observations: Tensor, actions: Tensor)

source edit

Train the world model and the behaviour on one batch.

Parameters

observationsTensor

Frames, (B, T, C, 64, 64).

actionsTensor

Actions taken into each step, (B, T, action_dim).

rewardsTensor

Observed reward, (B, T).

discountsTensor or None

(B, T), 0 where the episode terminated. Required when the config asks for a discount head.

Returns

DreamerV2Output

loss is the world-model loss; the behaviour losses are on .behavior and take their own optimisers.

Raises

ValueError

If pcont is configured and discounts is omitted.

imagine

→RSSMState

imagine(state: RSSMState, horizon: int, sample: bool | None = None)

source edit

Roll the dynamics forward under the actor's own policy.

Parameters

stateRSSMState

Flat starting beliefs, (N, ·).

horizonint

Steps to imagine.

sample(bool or None, optional, keyword - only)= None

Draw both latent and action, take their modes, or follow mean_only.

Returns

RSSMState

The imagined states including the start, (N, horizon + 1, ·).

Notes

The state the actor reads is detached, as the released implementation does — the gradient still reaches the policy through each action, and still travels back through the dynamics from the return.

Examples

>>> import lucid
>>> from lucid.models import dreamer_v2
>>> model = dreamer_v2(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)
>>> _, 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))
The start is kept as the first state, which is why there is one more
state than action:
>>> bool((states.stoch[:, 0] == start.stoch).all())
True

observe

→priors, posteriors : RSSMState

observe(observations: Tensor, actions: Tensor, state: RSSMState | None = None, sample: bool | None = None)

source edit

Filter a trajectory into posterior states.

Parameters

observationsTensor

Frames, (B, T, C, 64, 64).

actionsTensor

Actions taken into each step, (B, T, action_dim).

stateRSSMState or None= None

Starting belief; None starts from a uniform prior.

sample(bool or None, optional, keyword - only)= None

Draw the latent, take its mode, or follow mean_only.

Returns

priors, posteriors : RSSMState

(B, T, ·) each, carrying categorical logits.

Examples

>>> import lucid
>>> from lucid.models import dreamer_v2
>>> model = dreamer_v2(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)
>>> 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, posteriors.mean
((1, 3, 12), (1, 3, 3, 4), None)
The latent is stoch_size one-hot draws over discrete classes,
flattened, so each group of four holds a single 1:
>>> posteriors.stoch.reshape(1, 3, 3, 4).sum(dim=-1).tolist()
[[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]]

predict_pcont

→Tensor

predict_pcont(state: RSSMState)

source edit

Predict the discount at a state — logits, (B, T).

Parameters

stateRSSMState

The states to score.

Returns

Tensor

Bernoulli logits; apply sigmoid for the probability.

Raises

ValueError

If the model was configured without pcont.

Examples

>>> import lucid
>>> from lucid.models import dreamer_v2
>>> model = dreamer_v2(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)
>>> _, 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())
True

predict_reward

→Tensor

predict_reward(state: RSSMState)

source edit

Predict reward from a state — (B, T).

predict_value

→Tensor

predict_value(state: RSSMState, target: bool = False)

source edit

Estimate a state's value — (B, T).

Parameters

stateRSSMState

The states to score.

targetbool= False, keyword-only

Read the frozen target copy instead of the learning critic. The returns are built from the target; the critic's own loss is measured against the learner.

Returns

Tensor

Value estimates.

Examples

>>> import lucid
>>> from lucid.models import dreamer_v2_world_model
>>> wrapper = 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)
>>> model = wrapper.dreamer_v2
>>> _, posteriors = model.observe(lucid.randn((1, 3, 3, 64, 64)),
...                               lucid.randn((1, 3, 2)))
>>> value = model.predict_value(posteriors)
>>> value.shape
(1, 3)
The target copy starts from its own initialisation, and only
DreamerV2ForWorldModeling.update_slow_target moves it —
the first call copies the learner outright:
>>> target = model.predict_value(posteriors, target=True)
>>> bool((value == target).all())
False
>>> wrapper.update_slow_target()
>>> bool((value == model.predict_value(posteriors, target=True)).all())
True

>>> import lucid >>> from lucid.models import dreamer_v2 >>> model = dreamer_v2(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) >>> _, 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))

>>> import lucid >>> from lucid.models import dreamer_v2 >>> 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) >>> model = dreamer_v2(**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_v2(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])

>>> import lucid >>> from lucid.models import dreamer_v2 >>> model = dreamer_v2(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) >>> _, 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)) The start is kept as the first state, which is why there is one more state than action: >>> bool((states.stoch[:, 0] == start.stoch).all()) True

>>> import lucid >>> from lucid.models import dreamer_v2 >>> model = dreamer_v2(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) >>> 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, posteriors.mean ((1, 3, 12), (1, 3, 3, 4), None) The latent is stoch_size one-hot draws over discrete classes, flattened, so each group of four holds a single 1: >>> posteriors.stoch.reshape(1, 3, 3, 4).sum(dim=-1).tolist() [[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]]

>>> import lucid >>> from lucid.models import dreamer_v2 >>> model = dreamer_v2(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) >>> _, 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()) True

>>> import lucid >>> from lucid.models import dreamer_v2_world_model >>> wrapper = 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) >>> model = wrapper.dreamer_v2 >>> _, posteriors = model.observe(lucid.randn((1, 3, 3, 64, 64)), ... lucid.randn((1, 3, 2))) >>> value = model.predict_value(posteriors) >>> value.shape (1, 3) The target copy starts from its own initialisation, and only DreamerV2ForWorldModeling.update_slow_target moves it — the first call copies the learner outright: >>> target = model.predict_value(posteriors, target=True) >>> bool((value == target).all()) False >>> wrapper.update_slow_target() >>> bool((value == model.predict_value(posteriors, target=True)).all()) True