[content] jasmine/genie.py (python)
from models.dynamics import DynamicsMaskGIT, DynamicsCausal, DynamicsDiffusion
from models.lam import LatentActionModel
from models.tokenizer import TokenizerVQVAE


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_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,
        decode: 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.decode = decode

        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,
        )
        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,
                decode=decode,
                rngs=rngs,
            )
        elif self.dyna_type == "diffusion":
            assert self.denoising_steps > 0, "denoising_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,
                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,
                denoise_steps=self.diffusion_denoise_steps,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                decode=decode,
                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"]
        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),
            )
        
        if self.dyna_type == "diffusion":
            outputs = dict(
                videos=videos_BTHWC,
                latent_actions=(
                    action_embeddings_BTm11L
                    if self.use_gt_actions
                    else latent_actions_BTm11L
                ),
            )
            v_pred, v_t = self.dynamics(outputs)
            outputs["v_pred"] = v_pred
            outputs["v_t"] = v_t
            # TODO add recons as well
            return outputs

        tokenizer_outputs = self.tokenizer.vq_encode(videos_BTHWC, training=False)
        token_indices_BTN = tokenizer_outputs["indices"]
        outputs = dict(
            video_tokens=jax.lax.stop_gradient(token_indices_BTN),
            latent_actions=(
                action_embeddings_BTm11L
                if self.use_gt_actions
                else latent_actions_BTm11L
            ),
        )
        outputs["mask_rng"] = batch["rng"]
        dyna_logits_BTNV, dyna_mask = self.dynamics(outputs)
        outputs["token_logits"] = dyna_logits_BTNV
        outputs["mask"] = dyna_mask
        mle_indices_BTN = jnp.argmax(outputs["token_logits"], axis=-1)
        H, W = batch["videos"].shape[2:4]
        outputs["recon"] = self.tokenizer.decode(mle_indices_BTN, (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,
        temperature: float = 1,
        sample_argmax: bool = False,
        maskgit_steps: int = 25,
    ) -> tuple[jax.Array, jax.Array]:
        if self.dyna_type == "maskgit":
            return self.sample_maskgit(
                batch, seq_len, maskgit_steps, temperature, sample_argmax
            )
        elif self.dyna_type == "causal":
            return self.sample_causal(batch, seq_len, temperature, sample_argmax)
        else:
            raise ValueError(f"Dynamics model type unknown: {self.dyna_type}")

    def sample_maskgit(
        self,
        batch: Dict[str, jax.Array],
        seq_len: int,
        steps: int = 25,
        temperature: float = 1,
        sample_argmax: bool = False,
    ) -> tuple[jax.Array, jax.Array]:
        """
        Autoregressively samples up to `seq_len` future frames, following Figure 8 of the paper.

        - Input frames are tokenized once.
        - Future frames are generated autoregressively in token space.
        - All frames are detokenized in a single pass.

        Note:
        - For interactive or step-wise sampling, detokenization should occur after each action.
        - To maintain consistent tensor shapes across timesteps, all current and future frames are decoded at every step.
        - Temporal causal structure is preserved by
            a) reapplying the mask before each decoding step.
            b) a temporal causal mask is applied within each ST-transformer block.

        Dimension keys:
            B: batch size
            T: number of input (conditioning) frames
            N: number of patches per frame
            M: model dimension
            S: sequence length
            H: height
            W: width
            E: B * (S - 1)
            P: S * N
        """
        assert isinstance(self.dynamics, DynamicsMaskGIT)
        # --- Encode videos and actions ---
        videos_BTHWC = batch["videos"]
        tokenizer_out = self.tokenizer.vq_encode(videos_BTHWC, training=False)
        token_idxs_BTN = tokenizer_out["indices"]
        B, T, N = token_idxs_BTN.shape
        pad_shape = (B, seq_len - T, N)
        pad = jnp.zeros(pad_shape, dtype=token_idxs_BTN.dtype)
        token_idxs_BSN = jnp.concatenate([token_idxs_BTN, pad], axis=1)
        init_logits_BSNV = jnp.zeros(
            shape=(*token_idxs_BSN.shape, self.num_patch_latents)
        )
        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)

        # --- Extract submodule state ---
        dynamics_state = nnx.state(self.dynamics)

        @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)
        def maskgit_step_fn(
            carry: tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array],
            step: jax.Array,
        ) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]:
            rng, token_idxs_BSN, logits_BSNV, mask_BSN, action_tokens_EL = carry
            S, N = token_idxs_BSN.shape[1:]
            L = action_tokens_EL.shape[-1]

            # We need to reconstruct the submodule inside scan body to prevent trace context mismatches
            dynamics_maskgit = 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=nnx.Rngs(0),
            )
            nnx.update(dynamics_maskgit, dynamics_state)

            # --- Construct + encode video ---
            vid_embed_BSNM = dynamics_maskgit.patch_embed(token_idxs_BSN)
            mask_token_111M = dynamics_maskgit.mask_token.value
            mask_expanded_BSN1 = mask_BSN[..., None]
            vid_embed_BSNM = jnp.where(
                mask_expanded_BSN1, mask_token_111M, vid_embed_BSNM
            )

            # --- Predict transition ---
            action_tokens_BSm1L = jnp.reshape(action_tokens_EL, (B, S - 1, L))
            act_embed_BSm1M = dynamics_maskgit.action_up(action_tokens_BSm1L)
            act_embed_BSM = jnp.pad(act_embed_BSm1M, ((0, 0), (1, 0), (0, 0)))
            act_embed_BS1M = jnp.reshape(
                act_embed_BSM, (B, S, 1, act_embed_BSM.shape[-1])
            )
            vid_embed_BSNp1M = jnp.concatenate([act_embed_BS1M, vid_embed_BSNM], axis=2)
            unmasked_ratio = jnp.cos(jnp.pi * (step + 1) / (steps * 2))
            step_temp = temperature * (1.0 - unmasked_ratio)
            final_logits_BSNp1V = (
                dynamics_maskgit.transformer(vid_embed_BSNp1M) / step_temp
            )
            final_logits_BSNV = final_logits_BSNp1V[:, :, 1:]