PlaNet's latent dynamics model (Hafner et al., 2019).
A convolutional encoder embeds each frame, a recurrent state-space model carries the belief forward, and a decoder plus a reward head read predictions back off that belief. The point of the architecture is that everything downstream of the encoder happens in a compact latent — which is what makes searching over thousands of imagined action sequences affordable.
This class computes no loss. Use PlaNetForWorldModeling for
the training objective, or read the distributions off the returned
PlaNetOutput and build your own.
Parameters
configPlaNetConfigPlaNetConfig.Attributes
Notes
Reference: Hafner, Lillicrap, Fischer, Villegas, Ha, Lee, and Davidson, "Learning Latent Dynamics for Planning from Pixels", ICML, 2019 (arXiv:1811.04551).
Action alignment. actions[:, t] is the action taken into
step t — the one that produced observations[:, t]. A caller
holding shifts by one and passes a
zero action at index 0.
Planning is not implemented. The paper pairs this model with a
cross-entropy-method planner, which is a control algorithm rather than
a network; imagine is the piece of it that belongs to the
model, and the search on top is left to the caller.
Examples
>>> import lucid
>>> from lucid.models.generative.planet import PlaNetConfig, PlaNetModel
>>> cfg = PlaNetConfig(action_dim=2, stoch_size=4, deter_size=8,
... hidden_size=8, cnn_depth=4, reward_hidden=8)
>>> model = PlaNetModel(cfg).eval()
>>> obs, act = lucid.randn((1, 3, 3, 64, 64)), lucid.zeros(1, 3, 2)
>>> priors, posteriors = model.observe(obs, act)
>>> posteriors.stoch.shape
(1, 3, 4)
>>> model.imagine(model.rssm.initial(1), act).stoch.shape
(1, 3, 4)Used by 2
Constructors
1Instance methods
6Reconstruct 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 with no observations at all.
Parameters
stateRSSMState(B, ·).actionsTensor(B, T, action_dim).sample(bool or None, optional, keyword - only)= NoneTrue) or take its mean (False).
None follows the config's mean_only setting.Returns
RSSMStateThe imagined prior states, (B, T, ·).
Examples
>>> import lucid
>>> from lucid.models.generative.planet import PlaNetConfig, PlaNetModel
>>> cfg = PlaNetConfig(action_dim=2, stoch_size=4, deter_size=8,
... hidden_size=8, cnn_depth=4, reward_hidden=8)
>>> model = PlaNetModel(cfg).eval()
>>> obs, act = lucid.randn((1, 3, 3, 64, 64)), lucid.zeros(1, 3, 2)
>>> _, posteriors = model.observe(obs, act)
>>> belief = posteriors.map(lambda t: t[:, -1])
Score four five-step plans from that one belief — the part of a
planner that belongs to the model:
>>> plans = lucid.randn((4, 5, 2)).clip(-1, 1)
>>> starts = belief.map(lambda t: lucid.cat([t] * 4))
>>> imagined = model.imagine(starts, plans, sample=False)
>>> imagined.stoch.shape
(4, 5, 4)
>>> model.predict_reward(imagined).sum(dim=1).shape
(4,)
Planning takes the mean: a search that re-sampled each candidate
would rank its own noise instead of the actions.
>>> again = model.imagine(starts, plans, sample=False)
>>> bool((imagined.stoch == again.stoch).all())
Trueobserve
→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) — see the
class docstring on alignment.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.planet import PlaNetConfig, PlaNetModel
>>> cfg = PlaNetConfig(action_dim=2, stoch_size=4, deter_size=8,
... hidden_size=8, cnn_depth=4, reward_hidden=8)
>>> model = PlaNetModel(cfg).eval()
>>> obs, act = lucid.randn((1, 3, 3, 64, 64)), lucid.zeros(1, 3, 2)
>>> priors, posteriors = model.observe(obs, act)
>>> posteriors.mean.shape, posteriors.std.shape
((1, 3, 4), (1, 3, 4))
The frame refines the latent, not the path that led to it, so prior
and posterior share one deterministic state — and min_std
floors every scale:
>>> bool((priors.deter == posteriors.deter).all())
True
>>> bool((posteriors.std >= cfg.min_std).all())
True
A long episode is filtered in chunks by carrying the last belief in:
>>> last = posteriors.map(lambda t: t[:, -1])
>>> model.observe(obs, act, last)[1].deter.shape
(1, 3, 8)Predict reward from a state — (B, T).