GenieForWorldModeling
WorldModelingModelGenieForWorldModeling(config: GenieConfig)Genie played as an environment: a prompt frame, then latent actions.
Parameters
configGenieConfigAttributes
genieGenieModelNotes
Reference: Bruce et al., ICML 2024, Section 2.2.
A user's action is an index into the latent action codebook. The rest
of the latent action model is needed only to fill in the actions
between prompt frames, which a prompt of one frame does not have —
Section 2.2 replays a real video that way, from its first frame and
the actions the model infers from it. Each frame is revealed over
maskgit_steps steps at temperature: at every step each still
masked token is sampled, and the most confident samples are kept, so
many remain masked under a cosine schedule until the last step
reveals them all. The schedule and the confidence rule are MaskGIT's
— not stated in the paper.
The model remembers num_frames frames, as the paper's does. A
longer rollout keeps the most recent ones — not stated; Section 5
names the 16-frame memory as a limitation and says nothing more.
Examples
>>> import lucid
>>> from lucid.models.generative.genie import GenieConfig, GenieForWorldModeling
>>> config = GenieConfig(
... sample_size=(8, 8), num_frames=4, num_codes=16, code_dim=4,
... tokenizer_encoder_layers=1, tokenizer_encoder_dim=16,
... tokenizer_encoder_heads=2, tokenizer_encoder_head_dim=8,
... tokenizer_decoder_layers=1, tokenizer_decoder_dim=16,
... tokenizer_decoder_heads=2, tokenizer_decoder_head_dim=8,
... action_patch_size=4, action_dim=4,
... action_encoder_layers=1, action_encoder_dim=16, action_encoder_heads=2,
... action_decoder_layers=1, action_decoder_dim=16, action_decoder_heads=2,
... dynamics_layers=1, dynamics_dim=16, dynamics_heads=2,
... dynamics_head_dim=8, maskgit_steps=2)
>>> model = GenieForWorldModeling(config).eval()
>>> prompt = lucid.rand(1, 1, 3, 8, 8)
>>> actions = lucid.tensor([[0, 5, 7, 2, 1]], dtype=lucid.int64)
>>> out = model(prompt, actions)
>>> out.frames.shape, out.tokens.shape
((1, 5, 3, 8, 8), (1, 6, 4))Used by 2
Constructors
1Instance methods
2forward(prompt: Tensor, actions: Tensor)Play the environment from a prompt, one latent action per frame.
Parameters
Returns
GenieRolloutOutputS generated frames and the codes of the whole sequence.
Computed without gradient, and with gradients off nothing is
counted either: a rollout leaves the codebooks and their usage
exactly as it found them, in training mode as in eval.
Generate the next frame's codes by MaskGIT decoding.
Parameters
Returns
TensorCodes of the next frame, (B, N).