LatentPolicy
LatentPolicy(encode: Callable[[Tensor], Tensor], rssm: RSSM, act: Callable[[RSSMState], Tensor], action_dim: int, noise: float = 0.0)Carries an RSSM belief through an episode and asks the model to act.
A world model does not act from the frame in front of it; it acts from a belief it has been updating since the episode began. So this holds that belief, folds each new observation into it, and hands it to whatever the model uses to choose — which is not the same method across the family, and is therefore passed in rather than assumed: Dreamer has a learned actor, PlaNet searches with a planner.
Parameters
encodecallable(B, T, C, H, W) -> (B, T, embed), normally model.encode.rssmRSSMactcallableRSSMState -> (1, action_dim), given the single-step posterior.
For Dreamer this is lambda s: model.act(s, sample=False). For
PlaNet it is the task wrapper's planner — plan lives on
lucid.models.PlaNetForWorldModeling, not on the trunk,
so the encoder and RSSM come from wrapper.planet while the
planner comes from wrapper.action_dimintnoisefloat= 0.00.3 and evaluate with
0.Notes
The action is clipped into (-1, 1) whether or not noise was added,
because that is what Policy promises an environment. Dreamer
emits a tanh-squashed action and the clip does nothing; PlaNet's
planner searches an unbounded Gaussian and relies on it, which is the
clipping the paper describes.
Runs under lucid.no_grad — acting is not a thing to
differentiate, and holding a graph across a whole episode would keep
every frame alive.
Examples
>>> import lucid
>>> from lucid.models import dreamer
>>> from lucid.utils.rollout import LatentPolicy
>>> model = dreamer(action_dim=2, cnn_depth=2, stoch_size=4, deter_size=8,
... hidden_size=8, actor_hidden=8, value_hidden=8,
... reward_hidden=8).eval()
>>> policy = LatentPolicy(model.encode, model.rssm,
... lambda s: model.act(s, sample=False), 2)
>>> policy.reset()
>>> policy(lucid.zeros((3, 64, 64))).shape
(2,)Used by 1
Constructors
2__init__
→None__init__(encode: Callable[[Tensor], Tensor], rssm: RSSM, act: Callable[[RSSMState], Tensor], action_dim: int, noise: float = 0.0)Initialise the policy. See the class docstring for parameters.
Fold in the observation, then choose.