Genie's three networks and their objectives.
Parameters
configGenieConfigAttributes
Notes
Reference: Bruce et al., ICML 2024, Section 2.
The paper trains the tokenizer first and the other two after it, on
its codes. forward computes all three objectives in one pass;
staging is a matter of which parameter group an optimiser is given —
tokenizer_parameters, latent_action_parameters and
dynamics_parameters.
Videos are (B, T, C, H, W) in [0, 1] with
2 <= T <= num_frames.
Examples
>>> import lucid
>>> from lucid.models.generative.genie import GenieConfig, GenieModel
>>> 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 = GenieModel(config)
>>> out = model(lucid.rand(2, 4, 3, 8, 8))
>>> out.tokens.shape, out.actions.shape, out.logits.shape
((2, 4, 4), (2, 3), (2, 4, 4, 16))Used by 2
Constructors
1Instance methods
11The dynamics model's MaskGIT objective, on frozen codes and actions.
The video is tokenized and its actions inferred without gradient — the stop-gradient of Section 2.1 — so this trains the dynamics model and nothing else.
Parameters
videoTensor(B, T, C, H, W) in [0, 1], with T >= 2.Returns
TensorScalar cross-entropy over the masked tokens.
The dynamics model's parameters, trained by dynamics_loss.
Returns
list of ParameterEmbeddings, transformer and output head.
forward(video: Tensor)Compute all three objectives on one video batch.
Parameters
videoTensor(B, T, C, H, W) in [0, 1], with 2 <= T <= num_frames.Returns
GenieOutputThe summed loss, each objective, and what they were computed from.
Name each transition of a video with a latent action.
This is how a real video is replayed through the model: its first frame and the actions inferred here reproduce it (Section 2.2).
Parameters
videoTensor(B, T, C, H, W) in [0, 1], with 2 <= T <= num_frames.Returns
TensorLatent actions (B, T - 1), each in [0, num_latent_actions).
The latent action model's parameters, trained by latent_action_loss.
Returns
list of ParameterEncoder, action codebook and decoder.
Move unused codes of both codebooks back onto the data.
Call it after the optimiser step. Genie's action codebook holds
eight entries and collapses to a handful without this; the loss
does not show it, because the codes that survive take up the
slack. It writes the codebooks in place, which is why it is a
call of its own rather than part of forward: doing it while a
graph that read them is alive severs the two, and their gradients
stop arriving silently.
Does nothing when config.code_reset_threshold is 0.
Examples
>>> import lucid
>>> from lucid.models import genie_coinrun
>>> model = genie_coinrun(
... sample_size=(8, 8), num_frames=2, num_codes=8, code_dim=4,
... tokenizer_encoder_layers=1, tokenizer_encoder_dim=8,
... tokenizer_encoder_heads=1, tokenizer_decoder_layers=1,
... tokenizer_decoder_dim=8, tokenizer_decoder_heads=1,
... action_patch_size=4, action_dim=4, action_encoder_layers=1,
... action_encoder_dim=8, action_encoder_heads=1,
... action_decoder_layers=1, action_decoder_dim=8,
... action_decoder_heads=1, dynamics_layers=1, dynamics_dim=8,
... dynamics_heads=1)
>>> optimizer = lucid.optim.Adam(model.parameters(), lr=1e-3)
>>> model(lucid.rand(1, 2, 3, 8, 8)).loss.backward()
>>> optimizer.step()
>>> model.revive_codes()The video tokenizer's parameters, trained by tokenizer_loss.
Returns
list of ParameterEncoder, codebook and decoder.