[tab] jasmine/genie_continous.py (python)
from typing import Dict

import einops
import jax
import jax.numpy as jnp
import flax.nnx as nnx
import orbax.checkpoint as ocp

from models.dynamics import DynamicsMaskGIT, DynamicsCausal, DynamicsDiffusion
from models.lam import LatentActionModel
from models.tokenizer import TokenizerVQVAE, TokenizerMAE


class Genie(nnx.Module):
    """Genie model"""

    def __init__(
        self,
        in_dim: int,
        tokenizer_dim: int,
        tokenizer_ffn_dim: int,
        latent_patch_dim: int,
        num_patch_latents: int,
        patch_size: int,
        tokenizer_type: str,
        tokenizer_num_blocks: int,
        tokenizer_num_heads: int,
        lam_dim: int,
        lam_ffn_dim: int,
        latent_action_dim: int,
        num_actions: int,
        lam_patch_size: int,
        lam_num_blocks: int,
        lam_num_heads: int,
        lam_co_train: bool,
        use_gt_actions: bool,
        dyna_type: str,
        dyna_dim: int,
        dyna_ffn_dim: int,
        dyna_num_blocks: int,
        dyna_num_heads: int,
        param_dtype: jnp.dtype,
        dtype: jnp.dtype,
        use_flash_attention: bool,
        rngs: nnx.Rngs,
        dropout: float = 0.0,
        mask_limit: float = 0.0,
        diffusion_denoise_steps: int = 0,
    ):
        # --- Tokenizer ---
        self.in_dim = in_dim
        self.tokenizer_dim = tokenizer_dim
        self.tokenizer_ffn_dim = tokenizer_ffn_dim
        self.latent_patch_dim = latent_patch_dim
        self.num_patch_latents = num_patch_latents
        self.patch_size = patch_size
        self.tokenizer_num_blocks = tokenizer_num_blocks
        self.tokenizer_num_heads = tokenizer_num_heads
        # --- LAM ---
        self.lam_dim = lam_dim
        self.lam_ffn_dim = lam_ffn_dim
        self.latent_action_dim = latent_action_dim
        self.num_actions = num_actions
        self.lam_patch_size = lam_patch_size
        self.lam_num_blocks = lam_num_blocks
        self.lam_num_heads = lam_num_heads
        self.lam_co_train = lam_co_train
        self.use_gt_actions = use_gt_actions
        # --- Dynamics ---
        self.dyna_type = dyna_type
        self.dyna_dim = dyna_dim
        self.dyna_ffn_dim = dyna_ffn_dim
        self.dyna_num_blocks = dyna_num_blocks
        self.dyna_num_heads = dyna_num_heads
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention
        self.dropout = dropout
        self.mask_limit = mask_limit
        self.diffusion_denoise_steps = diffusion_denoise_steps
        self.tokenizer_type = tokenizer_type
        if self.tokenizer_type == "vqvae":
            self.tokenizer = TokenizerVQVAE(
                in_dim=self.in_dim,
                model_dim=self.tokenizer_dim,
                ffn_dim=self.tokenizer_ffn_dim,
                latent_dim=self.latent_patch_dim,
                num_latents=self.num_patch_latents,
                patch_size=self.patch_size,
                num_blocks=self.tokenizer_num_blocks,
                num_heads=self.tokenizer_num_heads,
                dropout=0.0,
                codebook_dropout=0.0,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=rngs,
            )
        elif self.tokenizer_type == "mae":
            self.tokenizer = TokenizerMAE(
                in_dim=self.in_dim,
                model_dim=self.tokenizer_dim,
                ffn_dim=self.tokenizer_ffn_dim,
                latent_dim=self.latent_patch_dim,
                num_latents=self.num_patch_latents,
                patch_size=self.patch_size,
                num_blocks=self.tokenizer_num_blocks,
                num_heads=self.tokenizer_num_heads,
                dropout=0.0,
                max_mask_ratio=0.0,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=rngs,
            )
        else:
            raise ValueError(f"Invalid tokenizer type: {self.tokenizer_type}")
        if self.use_gt_actions:
            self.action_embed = nnx.Embed(
                self.num_actions, self.latent_action_dim, rngs=rngs
            )
            self.lam = None
        else:
            self.lam = LatentActionModel(
                in_dim=self.in_dim,
                model_dim=self.lam_dim,
                ffn_dim=self.lam_ffn_dim,
                latent_dim=self.latent_patch_dim,
                num_latents=self.num_actions,
                patch_size=self.lam_patch_size,
                num_blocks=self.lam_num_blocks,
                num_heads=self.lam_num_heads,
                dropout=0.0,
                codebook_dropout=0.0,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=rngs,
            )
            self.action_embed = None
        if self.dyna_type == "maskgit":
            self.dynamics = DynamicsMaskGIT(
                model_dim=self.dyna_dim,
                ffn_dim=self.dyna_ffn_dim,
                num_latents=self.num_patch_latents,
                latent_action_dim=self.latent_action_dim,
                num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
                dropout=self.dropout,
                mask_limit=self.mask_limit,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=rngs,
            )
        elif self.dyna_type == "causal":
            self.dynamics = DynamicsCausal(
                model_dim=self.dyna_dim,
                ffn_dim=self.dyna_ffn_dim,
                num_latents=self.num_patch_latents,
                latent_action_dim=self.latent_action_dim,
                num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=rngs,
            )
        elif self.dyna_type == "diffusion":
            assert (
                self.diffusion_denoise_steps > 0
            ), "diffusion_denoise_steps must be greater than 0 when using the diffusion backend"
            self.dynamics = DynamicsDiffusion(
                model_dim=self.dyna_dim,
                ffn_dim=self.dyna_ffn_dim,
                latent_patch_dim=self.latent_patch_dim,
                latent_action_dim=self.latent_action_dim,
                num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
                denoise_steps=self.diffusion_denoise_steps,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=rngs,
            )
        else:
            raise ValueError(f"Invalid dynamics type: {self.dyna_type}")

    def __call__(
        self,
        batch: Dict[str, jax.Array],
    ) -> Dict[str, jax.Array]:
        videos_BTHWC = batch["videos"]
        H, W = videos_BTHWC.shape[2:4]
        latent_actions_BTm11L = None
        action_embeddings_BTm11L = None
        if self.use_gt_actions:
            assert self.action_embed is not None
            action_indices_E = None
            action_embeddings_BT1L = self.action_embed(batch["actions"]).reshape(
                *batch["actions"].shape[:2], 1, self.latent_action_dim
            )
            action_embeddings_BTm11L = action_embeddings_BT1L[:, :-1]
        else:
            assert self.lam is not None
            lam_outputs = self.lam.vq_encode(videos_BTHWC, training=False)
            z_q_BTm11L = lam_outputs["z_q"]
            action_indices_E = lam_outputs["indices"]
            latent_actions_BTm11L = jax.lax.cond(
                self.lam_co_train,
                lambda: z_q_BTm11L,
                lambda: jax.lax.stop_gradient(z_q_BTm11L),
            )

        rng, _rng = jax.random.split(batch["rng"])
        tokenizer_outputs = self.tokenizer.mask_and_encode(
            videos_BTHWC, rng, training=False
        )
        token_latents_BTNL = tokenizer_outputs["z"]
        token_latents_BTNL = jax.lax.stop_gradient(token_latents_BTNL)
        outputs = dict(
            token_latents=token_latents_BTNL,
            latent_actions=(
                action_embeddings_BTm11L
                if self.use_gt_actions
                else latent_actions_BTm11L
            ),
            rng=rng,
        )
        pred_latents_BTNL, denoise_t = self.dynamics(outputs)
        outputs["x_pred"] = pred_latents_BTNL
        outputs["x_gt"] = token_latents_BTNL
        outputs["signal_level"] = denoise_t
        outputs["recon"] = self.tokenizer.decode(pred_latents_BTNL, (H, W))
        if action_indices_E is not None:
            outputs["lam_indices"] = action_indices_E

        return outputs

    def sample(
        self,
        batch: Dict[str, jax.Array],
        seq_len: int,
        diffusion_steps: int = 64,
        diffusion_corrupt_context_factor: float = 0.1,
    ) -> tuple[jax.Array, jax.Array]:
        return self.sample_diffusion(
            batch, seq_len, diffusion_steps, diffusion_corrupt_context_factor
        )

    def sample_diffusion(
        self,
        batch: Dict[str, jax.Array],
        seq_len: int,
        diffusion_steps: int,
        diffusion_corrupt_context_factor: float = 0.1,
    ) -> tuple[jax.Array, jax.Array]:
        """
        Autoregressively samples up to `seq_len` future frames.

        - Input frames are tokenized once.
        - Future frames are generated autoregressively in latent space.
        - The context frames are corrupted by noise like in Dreamer 4 section 3.2.
        - All frames are decoded in a single pass.

        Dimension keys:
            B: batch size
            T: number of input (conditioning) frames
            N: number of patches per frame
            L: latent dimension
            S: sequence length
            H: height
            W: width
        """
        assert isinstance(self.dynamics, DynamicsDiffusion)
        # --- Encode videos and actions ---
        videos_BTHWC = batch["videos"]
        H, W = videos_BTHWC.shape[2:4]
        rng, _rng_mask, _rng_noise_full = jax.random.split(batch["rng"], 3)
        tokenizer_outputs = self.tokenizer.mask_and_encode(
            videos_BTHWC, _rng_mask, training=False
        )
        token_latents_BTNL = tokenizer_outputs["z"]
        B, T, N, L = token_latents_BTNL.shape
        pad_shape = (B, seq_len - T, N, L)
        pad = jax.random.normal(_rng_noise_full, pad_shape)
        token_latents_BSNL = jnp.concatenate([token_latents_BTNL, pad], axis=1)
        dynamics_state = nnx.state(self.dynamics)

        if self.use_gt_actions:
            assert self.action_embed is not None
            latent_actions_BT1L = self.action_embed(batch["actions"]).reshape(
                *batch["actions"].shape[:2], 1, self.latent_action_dim
            )
            latent_actions_BTm11L = latent_actions_BT1L[:, :-1]
            action_tokens_EL = latent_actions_BTm11L.reshape(-1, self.latent_action_dim)
        else:
            assert self.lam is not None
            latent_actions_E = batch["latent_actions"]
            action_tokens_EL = self.lam.vq.get_codes(latent_actions_E)

        action_tokens_BSm1L = jnp.reshape(action_tokens_EL, (B, seq_len - 1, -1))
        act_embed_BSL = jnp.pad(action_tokens_BSm1L, ((0, 0), (1, 0), (0, 0)))
        act_embed_BSM = self.dynamics.action_up(act_embed_BSL)

        t_corrupt_context = 1 - diffusion_corrupt_context_factor
        t_corrupt_context = jnp.argmin(
            jnp.abs(jnp.arange(diffusion_steps) / diffusion_steps - t_corrupt_context)
        )
        t_corrupt_context = t_corrupt_context / diffusion_steps

        @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)
        def denoise_step_fn(carry, denoise_step: jax.Array) -> jax.Array:
            latents_BSNL, frame_i, rng = carry
            rng, _rng_noise_context = jax.random.split(rng)

            # We need to reconstruct the submodule inside scan body to prevent trace context mismatches
            dynamics_diffusion = DynamicsDiffusion(
                model_dim=self.dyna_dim,
                ffn_dim=self.dyna_ffn_dim,
                latent_patch_dim=self.latent_patch_dim,
                latent_action_dim=self.latent_action_dim,
                num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
                denoise_steps=self.diffusion_denoise_steps,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=nnx.Rngs(0),
            )
            nnx.update(dynamics_diffusion, dynamics_state)

            t = denoise_step / diffusion_steps

            # corrupt the context frames like in Dreamer 4 section 3.2
            t_vector = jnp.ones((B, seq_len)) * t_corrupt_context
            t_vector = t_vector.at[:, frame_i].set(t)

            noise_context_BSNL = jax.random.normal(
                _rng_noise_context, (B, seq_len, N, L)
            )
            corrupted_latents_BSNL = (
                latents_BSNL * t_corrupt_context
                + (1 - t_corrupt_context) * noise_context_BSNL
            )
            frame_mask = jnp.arange(seq_len) < frame_i
            frame_mask_1S11 = frame_mask.reshape(1, seq_len, 1, 1)
            corrupted_tok_latents_BSNL = jnp.where(
                frame_mask_1S11, corrupted_latents_BSNL, latents_BSNL
            )

            # --- Predict transition ---
            pred_latents_BSNL = dynamics_diffusion.diffusion_transformer(
                corrupted_tok_latents_BSNL, t_vector, act_embed_BSM
            )
            latents_BSNL = latents_BSNL.at[:, frame_i].set(
                pred_latents_BSNL[:, frame_i]
            )
            new_carry = (latents_BSNL, frame_i, rng)
            return new_carry

        @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)
        def autoregressive_step_fn(carry, frame_t: jax.Array) -> jax.Array:
            latents_BSNL, rng = carry
            rng, _rng_noise_context = jax.random.split(rng)

            carry_denoise = (latents_BSNL, frame_t, _rng_noise_context)
            final_carry_denoise = denoise_step_fn(
                carry_denoise, jnp.arange(diffusion_steps)
            )
            latents_BSNL = final_carry_denoise[0]
            new_carry = (latents_BSNL, rng)
            return new_carry

        rng, _rng_sample = jax.random.split(rng)
        initial_carry = (token_latents_BSNL, _rng_sample)
        final_carry = autoregressive_step_fn(initial_carry, jnp.arange(T, seq_len))
        final_latents_BSNL = final_carry[0]

        final_videos_BSHWC = self.tokenizer.decode(final_latents_BSNL, video_hw=(H, W))

        return final_videos_BSHWC, None

    def vq_encode(self, batch: Dict[str, jax.Array], training: bool) -> jax.Array:
        # --- Preprocess videos ---
        assert self.lam is not None
        video_BTHWC = batch["videos"]
        lam_output: Dict[str, jax.Array] = self.lam.vq_encode(
            video_BTHWC, training=training
        )
        lam_indices_E = lam_output["indices"]
        return lam_indices_E


# FIXME (f.srambical): add conversion script for old checkpoints
def restore_genie_components(
    optimizer: nnx.ModelAndOptimizer,
    sharding: jax.sharding.NamedSharding,
    rng: jax.Array,
    tokenizer_type: str,
    args,
) -> nnx.ModelAndOptimizer:
    """Restore pre-trained Genie components"""
    rng_tokenizer, rng_lam = jax.random.split(rng)
    rngs_tokenizer = nnx.Rngs(rng_tokenizer)
    rngs_lam = nnx.Rngs(rng_lam)

    tx = optimizer.tx
    model = optimizer.model
    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add(
        "model_state", ocp.args.PyTreeRestore, ocp.handlers.PyTreeCheckpointHandler
    )

    checkpoint_options = ocp.CheckpointManagerOptions(
        step_format_fixed_length=6,
    )
    tokenizer_checkpoint_manager = ocp.CheckpointManager(
        directory=args.tokenizer_checkpoint,
        options=checkpoint_options,
        handler_registry=handler_registry,
    )
    if tokenizer_type == "vqvae":
        dummy_tokenizer = TokenizerVQVAE(
            in_dim=args.image_channels,
            model_dim=args.tokenizer_dim,
            ffn_dim=args.tokenizer_ffn_dim,
            latent_dim=args.latent_patch_dim,
            num_latents=args.num_patch_latents,
            patch_size=args.patch_size,
            num_blocks=args.tokenizer_num_blocks,
            num_heads=args.tokenizer_num_heads,
            dropout=args.dropout,
            codebook_dropout=args.dropout,
            param_dtype=args.param_dtype,
            dtype=args.dtype,
            use_flash_attention=args.use_flash_attention,
            rngs=rngs_tokenizer,
        )
    elif tokenizer_type == "mae":
        dummy_tokenizer = TokenizerMAE(
            in_dim=args.image_channels,
            model_dim=args.tokenizer_dim,
            ffn_dim=args.tokenizer_ffn_dim,
            latent_dim=args.latent_patch_dim,
            num_latents=args.num_patch_latents,
            patch_size=args.patch_size,
            num_blocks=args.tokenizer_num_blocks,
            num_heads=args.tokenizer_num_heads,
            dropout=args.dropout,
            max_mask_ratio=0.0,
            param_dtype=args.param_dtype,
            dtype=args.dtype,
            use_flash_attention=args.use_flash_attention,
            rngs=rngs_tokenizer,
        )
    else:
        raise ValueError(f"Invalid tokenizer type: {tokenizer_type}")

    dummy_tokenizer_optimizer = nnx.ModelAndOptimizer(dummy_tokenizer, tx)
    dummy_tokenizer_optimizer_state = nnx.state(dummy_tokenizer_optimizer)
    abstract_sharded_tokenizer_optimizer_state = _create_abstract_sharded_pytree(
        dummy_tokenizer_optimizer_state, sharding
    )
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
        step=tokenizer_checkpoint_manager.latest_step(),
        args=ocp.args.Composite(
            model_state=ocp.args.PyTreeRestore(  # type: ignore
                abstract_sharded_tokenizer_optimizer_state  # type: ignore
            ),
        ),
    )["model_state"]
    nnx.update(dummy_tokenizer_optimizer.model, restored_tokenizer.model)
    model.tokenizer = dummy_tokenizer_optimizer.model
    tokenizer_checkpoint_manager.close()

    if args.lam_checkpoint:
        lam_checkpoint_manager = ocp.CheckpointManager(
            directory=args.lam_checkpoint,
            options=checkpoint_options,
            handler_registry=handler_registry,
        )
        dummy_lam = LatentActionModel(
            in_dim=args.image_channels,
            model_dim=args.lam_dim,
            ffn_dim=args.lam_ffn_dim,
            latent_dim=args.latent_patch_dim,
            num_latents=args.num_actions,
            patch_size=args.lam_patch_size,
            num_blocks=args.lam_num_blocks,
            num_heads=args.lam_num_heads,
            dropout=args.dropout,
            codebook_dropout=args.dropout,
            param_dtype=args.param_dtype,
            dtype=args.dtype,
            use_flash_attention=args.use_flash_attention,
            rngs=rngs_lam,
        )
        dummy_lam_optimizer = nnx.ModelAndOptimizer(dummy_lam, tx)
        dummy_lam_optimizer_state = nnx.state(dummy_lam_optimizer)
        abstract_sharded_lam_optimizer_state = _create_abstract_sharded_pytree(
            dummy_lam_optimizer_state, sharding
        )
        restored_lam_optimizer = lam_checkpoint_manager.restore(
            step=lam_checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.PyTreeRestore(  # type: ignore
                    abstract_sharded_lam_optimizer_state  # type: ignore
                ),
            ),
        )["model_state"]
        nnx.update(dummy_lam_optimizer.model, restored_lam_optimizer.model)
        model.lam = dummy_lam_optimizer.model
        # Remove the LAM decoder to save memory and avoid unnecessary computation.
        del model.lam.decoder
        lam_checkpoint_manager.close()

    # Reinitialize the optimizer states
    optimizer = nnx.ModelAndOptimizer(model, tx)
    return optimizer


def _create_abstract_sharded_pytree(
    pytree_template: nnx.GraphState, sharding_spec: jax.sharding.NamedSharding
) -> jax.Array:
    """Replaces arrays in a pytree with ShapeDtypeStructs having the given sharding."""

    def map_fn(leaf_template):
        if hasattr(leaf_template, "shape") and hasattr(leaf_template, "dtype"):
            return jax.ShapeDtypeStruct(
                leaf_template.shape, leaf_template.dtype, sharding=sharding_spec
            )
        return leaf_template

    return jax.tree_util.tree_map(map_fn, pytree_template)