data
GenieConfig
extends
GenerativeModelConfigGenieConfig(sample_size: int | tuple[int, int] = (90, 160), in_channels: int = 3, out_channels: int = 3, act_fn: GenerativeActivation = 'gelu', num_frames: int = 16, mlp_ratio: float = 4.0, commitment_cost: float = 0.25, code_reset_threshold: float = 0.1, tokenizer_patch_size: int = 4, num_codes: int = 1024, code_dim: int = 32, tokenizer_encoder_layers: int = 12, tokenizer_encoder_dim: int = 512, tokenizer_encoder_heads: int = 8, tokenizer_encoder_head_dim: int | None = 64, tokenizer_decoder_layers: int = 20, tokenizer_decoder_dim: int = 1024, tokenizer_decoder_heads: int = 16, tokenizer_decoder_head_dim: int | None = 64, action_patch_size: int = 16, num_latent_actions: int = 8, action_dim: int = 32, action_encoder_layers: int = 20, action_encoder_dim: int = 1024, action_encoder_heads: int = 16, action_encoder_head_dim: int | None = None, action_decoder_layers: int = 20, action_decoder_dim: int = 1024, action_decoder_heads: int = 16, action_decoder_head_dim: int | None = None, dynamics_layers: int = 48, dynamics_dim: int = 5120, dynamics_heads: int = 36, dynamics_head_dim: int | None = 128, dynamics_qk_norm: bool = True, mask_ratio_min: float = 0.5, mask_ratio_max: float = 1.0, maskgit_steps: int = 25, temperature: float = 2.0)Frozen configuration for the Genie family.
Defaults are the paper's Platformers model (Tables 5, 7 and 12). Fields marked not stated are absent from the paper; the value and the reason for it are given in each entry.
Parameters
sample_sizeint or tuple of int= (90, 160)Frame height and width. Platformers videos are 160x90.
in_channelsint= 3Channels per frame.
out_channelsint= 3Channels each decoder emits — one frame, so
in_channels.act_fn(gelu, relu, silu, swish, elu)= "gelu"Feed-forward activation. Not stated; GELU, as in the ViT
encoders the tokenizer is built from.
num_framesint= 16Frames every component sees at once, and so the dynamics model's
memory. Section 3: sequence length 16 at 10 FPS.
mlp_ratiofloat= 4.0Feed-forward width as a multiple of the model width. Not
stated; 4, the transformer convention.
commitment_costfloat= 0.25 of both vector quantisers. Not stated; 0.25,
the VQ-VAE paper's value.
code_reset_thresholdfloat= 0.1A code used less than this fraction of its fair share, averaged
over batches, is moved onto an encoder output from the current
batch while training.
0 switches that off. Not stated; a
codebook of eight entries collapses readily, and a collapsed one
is invisible in the loss.tokenizer_patch_sizeint= 4Tokenizer patch side (Table 7). A frame whose side is not a
multiple is padded at the bottom and right — not stated; 90
becomes 92, which is the reading under which the paper's token
counts are exact (942B tokens is 40x23 tokens per frame).
num_codesint= 1024Tokenizer codebook size (Table 7).
code_dimint= 32Tokenizer code width (Table 7).
tokenizer_encoder_layersint= 1212, 512 and 8 (Table 7).
tokenizer_encoder_dimint= 1212, 512 and 8 (Table 7).
tokenizer_encoder_headsint= 1212, 512 and 8 (Table 7).
tokenizer_encoder_head_dimint or None= 64Query/key size per head (Table 7).
None means
dim // heads.tokenizer_decoder_layersint= 2020, 1024 and 16 (Table 7) — the paper found scaling the decoder
worth more than scaling the encoder.
tokenizer_decoder_dimint= 2020, 1024 and 16 (Table 7) — the paper found scaling the decoder
worth more than scaling the encoder.
tokenizer_decoder_headsint= 2020, 1024 and 16 (Table 7) — the paper found scaling the decoder
worth more than scaling the encoder.
tokenizer_decoder_head_dimint or None= 64Query/key size per head (Table 7).
action_patch_sizeint= 16Latent action model patch side (Table 5).
num_latent_actionsint= 8, the latent action codebook (Table 5).
action_dimint= 32Latent action width (Table 5).
action_encoder_layersint= 2020, 1024 and 16 (Table 5).
action_encoder_dimint= 2020, 1024 and 16 (Table 5).
action_encoder_headsint= 2020, 1024 and 16 (Table 5).
action_encoder_head_dimint or None= NoneNot stated — Table 5 gives no query/key size, so
dim // heads.action_decoder_layersint= 2020, 1024 and 16 (Table 5).
action_decoder_dimint= 2020, 1024 and 16 (Table 5).
action_decoder_headsint= 2020, 1024 and 16 (Table 5).
action_decoder_head_dimint or None= NoneNot stated, as for the encoder.
dynamics_layersint= 4848, 5120 and 36 (Table 12).
dynamics_dimint= 4848, 5120 and 36 (Table 12).
dynamics_headsint= 4848, 5120 and 36 (Table 12).
dynamics_head_dimint or None= 128Query/key size per head (Table 12).
36 x 128 is 4608, not
5120: the attention's inner width is decoupled from the model's.dynamics_qk_normbool= TrueNormalise queries and keys per head. Section 3 uses it for the
dynamics model, for stability in bfloat16.
mask_ratio_minfloat= 0.5, 1.0The training mask rate is drawn uniformly from this range
(Section 2.1).
mask_ratio_maxfloat= 0.5, 1.0The training mask rate is drawn uniformly from this range
(Section 2.1).
maskgit_stepsint= 25Refinement steps per generated frame (Section 3).
temperaturefloat= 2.0Sampling temperature at play time (Section 3).
Notes
Reference: Bruce, Jake, et al., "Genie: Generative Interactive
Environments", ICML 2024 (arXiv:2402.15391). The CoinRun case
study of Appendix F is lucid.models.genie_coinrun.
Examples
>>> from lucid.models.generative.genie import GenieConfig
>>> config = GenieConfig()
>>> config.frame_shape, config.token_grid
((90, 160), (23, 40))
>>> config.num_latent_actions, config.maskgit_steps, config.temperature
(8, 25, 2.0)
Every component's attention may be narrower than its model width:
>>> config.dynamics_heads * config.attention_head_dim("dynamics")
4608Used by 3
Constructors
1dunder
__init__
→None__init__(sample_size: int | tuple[int, int] = (90, 160), in_channels: int = 3, out_channels: int = 3, act_fn: GenerativeActivation = 'gelu', num_frames: int = 16, mlp_ratio: float = 4.0, commitment_cost: float = 0.25, code_reset_threshold: float = 0.1, tokenizer_patch_size: int = 4, num_codes: int = 1024, code_dim: int = 32, tokenizer_encoder_layers: int = 12, tokenizer_encoder_dim: int = 512, tokenizer_encoder_heads: int = 8, tokenizer_encoder_head_dim: int | None = 64, tokenizer_decoder_layers: int = 20, tokenizer_decoder_dim: int = 1024, tokenizer_decoder_heads: int = 16, tokenizer_decoder_head_dim: int | None = 64, action_patch_size: int = 16, num_latent_actions: int = 8, action_dim: int = 32, action_encoder_layers: int = 20, action_encoder_dim: int = 1024, action_encoder_heads: int = 16, action_encoder_head_dim: int | None = None, action_decoder_layers: int = 20, action_decoder_dim: int = 1024, action_decoder_heads: int = 16, action_decoder_head_dim: int | None = None, dynamics_layers: int = 48, dynamics_dim: int = 5120, dynamics_heads: int = 36, dynamics_head_dim: int | None = 128, dynamics_qk_norm: bool = True, mask_ratio_min: float = 0.5, mask_ratio_max: float = 1.0, maskgit_steps: int = 25, temperature: float = 2.0)Initialise the three networks. See the class docstring for parameters.
Properties
3Rows and columns of latent-action-model patches, after padding.
Frame height and width, whichever way sample_size was written.
Rows and columns of tokenizer patches, after padding.