DOCSEARCHbyAlgolia
Submit Docs

Genie2 PyTorch

repository·main·Indexed 41 days ago

https://github.com/lucidrains/genie2-pytorch

A PyTorch implementation of a world model for generating video trajectories conditioned on discrete actions. It features a transformer architecture, vector quantization (VQ) for latent representation, and support for multiple simultaneous actions. Version 0.0.21 includes the Genie2 class for autoregressive sampling, state encoding via encode_state(), and guided sampling using Classifier-Free Guidance (CFG).

Tokens
1.5K
Snippets
5
Records
7
Agent score
50%

What's inside Genie2

  1. How temporal awareness is added to VQ

    main

    If you are using a standard image encoder/decoder (not a native video encoder/decoder), you can enable temporal awareness in the latent space by setting add_temporal_convs = True during Genie2 initialization.

    This applies a CausalConv3d (3D causal convolution) both before and after the Vector Quantization (VQ) step. This allows the quantization process to consider temporal context, making the resulting latent codes more temporally coherent. Note that this requires setting is_video_enc_dec = False (the default).

  2. How Genie2 handles multiple actions

    main

    Genie2 supports conditioning on multiple simultaneous actions (e.g., pressing multiple keys at once).

    To use this feature:

    1. Initialize Genie2 with allow_multiple_actions = True and set max_num_actions.
    2. Provide actions as a tensor of shape [b, t, a], where a is the number of actions being performed at that timestep.
    3. The model uses an internal action_transformer to process these multiple actions before projecting them into the latent space.
    4. During training, the model can also learn to predict these multiple actions autoregressively if action_autoregressive_loss_weight > 0.

  4. Perform forward pass with Classifier-Free Guidance (CFG)

    main

    The forward_with_cfg method allows for guided sampling by interpolating between conditioned and unconditioned (null action) logits.

    Parameters:

    • actions: The action sequence to condition on.
    • cond_scale: The guidance scale. A value $> 1$ amplifies the effect of the actions.
    • parallel_keep_frac: A fraction used to control how much of the parallel component of the update is kept (useful for specific guidance techniques).
    # logits will be conditioned on actions using the specified scale
logits = model.forward_with_cfg(
    state_codes = current_codes, 
    actions = current_actions, 
    cond_scale = 2.0
)

  5. Generate video trajectories with generate()

    main

    Use the generate method to autoregressively sample video frames conditioned on an initial image and a sequence of actions.

    Parameters:

    • image: The starting frame of shape [b, c, h, w].
    • num_frames: Number of future frames to generate.
    • init_action: An integer representing the starting action (required if interactive=True).
    • interactive: If True, the model enters an interactive loop where it prompts the user for the next action via CLI input at every frame.
    • cond_scale: Scale for classifier-free guidance.
    • temperature: Sampling temperature for the Gumbel-softmax sampling.
    • filter_kwargs: Arguments for min_p_filter to prune low-probability logits.

    Returns:

    • If interactive=False: Returns the generated video tensor.
    • If interactive=True: Returns a tuple of (video, actions).
    # Non-interactive generation
video = model.generate(
    image = initial_frame, 
    num_frames = 16, 
    temperature = 0.9,
    cond_scale = 1.5
)

# Interactive generation (requires CLI input)
video, actions = model.generate(
    image = initial_frame, 
    num_frames = 30, 
    init_action = 5, 
    interactive = True
)

  6. Encode video states with encode_state()

    main

    The encode_state method converts a video sequence into quantized latent codes.

    Parameters:

    • state: The input video tensor of shape [b, c, t, h, w].
    • return_recon_loss: If True, returns the reconstruction loss (MSE between original latents and reconstructed latents) if add_temporal_convs was enabled during initialization.

    Returns:

    • quantized_out: The output from the VQ module (containing quantized latents and indices).
    • recon_loss (optional): The MSE reconstruction loss.
    # Returns (quantized_out, recon_loss) if return_recon_loss=True
quantized_out, recon_loss = model.encode_state(video_tensor, return_recon_loss = True)

  7. Initialize the Genie2 model

    main

    The Genie2 class is the main entry point for the world model. It is a transformer-based model designed for video generation conditioned on actions. It uses a Vector Quantized (VQ) latent space and can optionally support multiple simultaneous actions.

    Key initialization parameters:

    • dim: Dimension of the transformer tokens.
    • dim_latent: Dimension of the latent space.
    • num_actions: Total number of possible discrete actions. If provided, the model will learn to condition on actions.
    • allow_multiple_actions: If True, allows the model to handle multiple actions per timestep (e.g., multiple keys pressed).
    • max_num_actions: Maximum number of simultaneous actions allowed if allow_multiple_actions is True.
    • vq_kwargs: Dictionary of arguments passed to the VectorQuantize module.
    • encoder / decoder: Custom modules for encoding/decoding states. Defaults to nn.Identity().
    • action_autoregressive_loss_weight: Weight for the loss used to predict future actions.
    from genie2_pytorch import Genie
import torch
from torch import nn

model = Genie(
    dim = 1024,
    dim_latent = 512,
    num_actions = 100,
    allow_multiple_actions = True,
    max_num_actions = 5,
    vq_kwargs = dict(
        codebook_size = 4096,
    )
)