PlaNetForWorldModeling
WorldModelingModelPlaNetForWorldModeling(config: PlaNetConfig)PlaNet with the variational training objective.
Wraps PlaNetModel with the bound of Hafner et al., 2019:
Both likelihoods are unit-variance Gaussians, so their log-densities
reduce to one-half squared error up to constants that do not affect the
gradient — the observation term summed over pixels, the reward term a
scalar. The KL carries a free-nats floor; see rssm_kl.
Parameters
configPlaNetConfigfree_nats and kl_weight shape the KL term.Attributes
planetPlaNetModelNotes
Reference: Hafner, Lillicrap, Fischer, Villegas, Ha, Lee, and Davidson, "Learning Latent Dynamics for Planning from Pixels", ICML, 2019 (arXiv:1811.04551).
The reconstruction term never reaches the prior head. Nothing reconstructed is computed from the prior — the decoder and the reward head both read the posterior — so the KL is the prior's only teacher. Train without it and the dynamics never learn to predict, while every shape and every loss value still looks reasonable.
The released implementation scales the reward term by 10 and uses 3-layer, 300-unit heads. Neither appears in the paper, which specifies "two fully connected layers of size 200"; the paper's values are used here, and the discrepancy is recorded rather than silently split.
Examples
>>> import lucid
>>> from lucid.models.generative.planet import (
... PlaNetConfig, PlaNetForWorldModeling,
... )
>>> cfg = PlaNetConfig(action_dim=2, stoch_size=4, deter_size=8,
... hidden_size=8, cnn_depth=4, reward_hidden=8)
>>> model = PlaNetForWorldModeling(cfg).eval()
>>> obs, act = lucid.randn((1, 3, 3, 64, 64)), lucid.zeros(1, 3, 2)
>>> out = model(obs, act, rewards=lucid.zeros(1, 3))
>>> out.loss.shape, out.observation.shape
((), (1, 3, 3, 64, 64))Used by 2
Constructors
1Instance methods
2forward(observations: Tensor, actions: Tensor, rewards: Tensor | None = None)plan(state: RSSMState, horizon: int = 12, iterations: int = 10, candidates: int = 1000, elites: int = 100)Choose an action by searching imagined trajectories (CEM).
The paper's planner. It never touches an environment: candidate action sequences are rolled forward through the learned dynamics and scored by the learned reward head, so the entire search is a forward pass over this model's own parameters.
The cross-entropy method fits a diagonal Gaussian over action sequences to its own best samples, repeatedly:
and the first action of the final mean is returned. Searching in a compact latent rather than in pixels is what makes a thousand candidates per step affordable.
Parameters
stateRSSMState(B, ·) — typically one step of a
posterior produced by PlaNetModel.observe.horizonint= 12iterationsint= 10candidatesint= 1000elitesint= 100Returns
TensorThe first action of the planned sequence, (B, action_dim).
Notes
Defaults are the paper's: , , , .
The search is unbounded — the paper clips actions to the environment's range, which this model cannot know. Clip the result yourself if your action space is bounded.
Exploration noise is also the caller's: the paper adds Gaussian noise to the planned action when collecting episodes, which is a property of the data-collection loop rather than of the planner.
Examples
>>> import lucid
>>> from lucid.models.generative.planet import (
... PlaNetConfig, PlaNetForWorldModeling,
... )
>>> cfg = PlaNetConfig(action_dim=2, stoch_size=4, deter_size=8,
... hidden_size=8, cnn_depth=4, reward_hidden=8)
>>> model = PlaNetForWorldModeling(cfg).eval()
>>> start = model.planet.rssm.initial(1)
>>> model.plan(start, horizon=3, iterations=2, candidates=8,
... elites=2).shape
(1, 2)