[content] jasmine/genie.py (python)
            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,
                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
                ),
                rng=batch["rng"],
            )
            x_pred, t = self.dynamics(outputs)
            outputs["x_pred"] = x_pred
            outputs["x_gt"] = videos_BTHWC
            outputs["noise_level"] = t
            outputs["recon"] = x_pred  # TODO mihir: fix this
            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,
        diffusion_steps: int = 64,
        window_size: Optional[int] = None,
    ) -> 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)
        elif self.dyna_type == "diffusion":
            return self.sample_diffusion(batch, seq_len, diffusion_steps, window_size)
        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:]

            # --- Sample new tokens for final frame ---
            if sample_argmax:
                sampled_token_idxs_BSN = jnp.argmax(final_logits_BSNV, axis=-1)
            else:
                rng, _rng = jax.random.split(rng)
                sampled_token_idxs_BSN = jax.random.categorical(_rng, final_logits_BSNV)
            gather_fn = jax.vmap(jax.vmap(jax.vmap(lambda x, y: x[y])))
            final_token_probs_BSN = gather_fn(
                jax.nn.softmax(final_logits_BSNV), sampled_token_idxs_BSN
            )
            final_token_probs_BSN += ~mask_BSN
            # Update masked tokens and logits only
            token_idxs_BSN = jnp.where(mask_BSN, sampled_token_idxs_BSN, token_idxs_BSN)
            logits_BSNV = jnp.where(
                jnp.expand_dims(mask_BSN, -1), final_logits_BSNV, logits_BSNV
            )

            # --- Update mask ---
            num_unmasked_tokens = jnp.round(N * (1.0 - unmasked_ratio)).astype(int)
            final_token_probs_flat_BP = einops.rearrange(
                final_token_probs_BSN, "b s n -> b (s n)"
            )
            idx_mask_P = (
                jnp.arange(final_token_probs_flat_BP.shape[-1])
                <= N - num_unmasked_tokens
            )
            sorted_idxs_BP = jnp.argsort(final_token_probs_flat_BP, axis=-1)
            mask_update_fn = jax.vmap(lambda msk, ids: msk.at[ids].set(idx_mask_P))
            mask_flat_BP = einops.rearrange(mask_BSN, "b s n -> b (s n)")
            new_mask_flat_BP = mask_update_fn(mask_flat_BP, sorted_idxs_BP)
            new_mask_BSN = einops.rearrange(new_mask_flat_BP, "b (s n) -> b s n", n=N)

            new_carry = (
                rng,
                token_idxs_BSN,
                logits_BSNV,
                new_mask_BSN,
                action_tokens_EL,
            )
            return new_carry

        @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)
        def generation_step_fn(
            carry: tuple[jax.Array, jax.Array, jax.Array], step_t: jax.Array
        ) -> tuple[jax.Array, jax.Array, jax.Array]:
            rng, current_token_idxs_BSN, current_logits_BSNV = carry
            rng, step_rng = jax.random.split(rng)

            # Mask current frame (i.e., t == step_t)
            mask_S = jnp.arange(seq_len) == step_t
            mask_BSN = jnp.broadcast_to(mask_S[None, :, None], (B, seq_len, N)).astype(
                bool
            )
            masked_token_idxs_BSN = current_token_idxs_BSN * ~mask_BSN
            masked_logits_BSNV = current_logits_BSNV * jnp.expand_dims(~mask_BSN, -1)

            # --- Initialize and run MaskGIT loop ---
            init_carry_maskgit = (
                step_rng,
                masked_token_idxs_BSN,
                masked_logits_BSNV,
                mask_BSN,
                action_tokens_EL,
            )
            final_carry_maskgit = maskgit_step_fn(init_carry_maskgit, jnp.arange(steps))
            updated_token_idxs_BSN = final_carry_maskgit[1]
            updated_logits_BSNV = final_carry_maskgit[2]
            new_carry = (rng, updated_token_idxs_BSN, updated_logits_BSNV)
            return new_carry

        # --- Run the autoregressive generation using jax.lax.scan ---
        initial_carry = (batch["rng"], token_idxs_BSN, init_logits_BSNV)
        timesteps_to_scan = jnp.arange(T, seq_len)
        final_carry = generation_step_fn(initial_carry, timesteps_to_scan)
        final_token_idxs_BSN = final_carry[1]
        final_logits_BSNV = final_carry[2]

        # --- Decode all tokens at once at the end ---
        H, W = batch["videos"].shape[2:4]
        final_frames_BSHWC = self.tokenizer.decode(
            final_token_idxs_BSN,
            video_hw=(H, W),
        )
        return final_frames_BSHWC, final_logits_BSNV

    def sample_causal(
        self,
        batch: Dict[str, jax.Array],
        seq_len: int,
        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)
        """
        assert isinstance(self.dynamics, DynamicsCausal)
        # --- 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)
        logits_BSNV = jnp.zeros((*token_idxs_BSN.shape, self.num_patch_latents))
        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)

        @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)
        def causal_step_fn(
            carry: tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array],
            step_n: jax.Array,
        ) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]:
            rng, token_idxs_BSN, logits_BSNV, action_tokens_EL, step_t = 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_causal = 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=self.decode,
                rngs=nnx.Rngs(0),
            )
            nnx.update(dynamics_causal, dynamics_state)

            # --- Construct + encode video ---
            vid_embed_BSNM = dynamics_causal.patch_embed(token_idxs_BSN)

            # --- Predict transition ---
            action_tokens_BSm1L = jnp.reshape(action_tokens_EL, (B, S - 1, L))
            act_embed_BSm1M = dynamics_causal.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)
            final_logits_BTNp1V = (
                dynamics_causal.transformer(vid_embed_BSNp1M, (step_t, step_n))
                / temperature
            )
            final_logits_BV = final_logits_BTNp1V[:, step_t, step_n, :]

            # --- Sample new tokens for final frame ---
            if sample_argmax:
                sampled_token_idxs_B = jnp.argmax(final_logits_BV, axis=-1)
            else:
                rng, _rng = jax.random.split(rng)
                sampled_token_idxs_B = jax.random.categorical(_rng, final_logits_BV)
            # Update next tokens only
            token_idxs_BSN = token_idxs_BSN.at[:, step_t, step_n].set(
                sampled_token_idxs_B
            )
            logits_BSNV = logits_BSNV.at[:, step_t, step_n].set(final_logits_BV)

            new_carry = (rng, token_idxs_BSN, logits_BSNV, action_tokens_EL, step_t)
            return new_carry

        @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)
        def generation_step_fn(
            carry: tuple[jax.Array, jax.Array, jax.Array], step_t: jax.Array
        ) -> tuple[jax.Array, jax.Array, jax.Array]:
            rng, current_token_idxs_BSN, current_logits_BSNV = carry
            rng, step_rng = jax.random.split(rng)

            # --- Initialize and run causal loop ---
            init_carry_causal = (
                step_rng,
                current_token_idxs_BSN,
                current_logits_BSNV,
                action_tokens_EL,
                step_t,
            )
            final_carry_causal = causal_step_fn(init_carry_causal, jnp.arange(N))
            updated_token_idxs_BSN = final_carry_causal[1]
            updated_logits_BSNV = final_carry_causal[2]
            new_carry = (rng, updated_token_idxs_BSN, updated_logits_BSNV)
            return new_carry

        # --- Run the autoregressive generation using jax.lax.scan ---
        initial_carry = (batch["rng"], token_idxs_BSN, logits_BSNV)
        timesteps_to_scan = jnp.arange(T, seq_len)
        final_carry = generation_step_fn(initial_carry, timesteps_to_scan)
        final_token_idxs_BSN = final_carry[1]
        final_logits_BSNV = final_carry[2]

        # --- Decode all tokens at once at the end ---
        H, W = batch["videos"].shape[2:4]
        final_frames_BSHWC = self.tokenizer.decode(
            final_token_idxs_BSN,
            video_hw=(H, W),
        )
        return final_frames_BSHWC, final_logits_BSNV

    def sample_diffusion(
        self,
        batch: Dict[str, jax.Array],
        seq_len: int,
        diffusion_steps: int,
        window_size: Optional[int] = None,
    ) -> tuple[jax.Array, jax.Array]:
        assert isinstance(self.dynamics, DynamicsDiffusion)
        # --- Encode videos and actions ---
        videos_BTHWC = batch["videos"]
        window_size = window_size or seq_len
        B, T, H, W, C = videos_BTHWC.shape
        rng, _rng_noise_window, _rng_noise_full = jax.random.split(batch["rng"], 3)
        pad_shape_window = (B, min(window_size, seq_len) - T, H, W, C)
        pad_window = jax.random.normal(_rng_noise_window, pad_shape_window)
        videos_BSHWC = jnp.concatenate([videos_BTHWC, pad_window], axis=1)
        pad_shape_full = (B, seq_len - T, H, W, C)
        pad_full = jax.random.normal(_rng_noise_full, pad_shape_full)
        videos_BFHWC = jnp.concatenate([videos_BTHWC, pad_full], axis=1)

        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)

        delta_t = 1.0 / diffusion_steps

        @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)
        def denoise_step_fn(carry, denoise_step: jax.Array) -> jax.Array:
            vid_BSHWC, frame_i = carry

            # align actions with sliding window of frames
            act_embed_window_BSM, frame_i_window = nnx.cond(
                frame_i < window_size,
                lambda: (act_embed_BSM[:, :window_size], frame_i),
                lambda: (
                    jax.lax.dynamic_slice(
                        act_embed_BSM,
                        (0, frame_i - window_size + 1, 0),
                        (B, window_size, act_embed_BSM.shape[-1]),
                    ),
                    window_size - 1,
                ),
            )

            act_embed_window_BSM = act_embed_window_BSM.at[:, 0].set(0)

            t = denoise_step / diffusion_steps

            # TODO mihir: corrupt the context frames like in Dreamer 4 section 3.2
            t_vector = jnp.ones((B, min(window_size, seq_len)))
            t_vector = t_vector.at[:, frame_i_window].set(t)

            dt_flow = jnp.log2(diffusion_steps).astype(jnp.int32)
            dt_base = (
                jnp.ones((B, min(window_size, seq_len)), dtype=jnp.int32) * dt_flow
            )  # Smallest dt.

            v_pred_BSHWC = self.dynamics.diffusion_transformer(
                vid_BSHWC, t_vector, dt_base, act_embed_window_BSM
            )
            frame_BHWC = (
                vid_BSHWC[:, frame_i_window] + v_pred_BSHWC[:, frame_i_window] * delta_t
            )
            vid_BSHWC = vid_BSHWC.at[:, frame_i_window].set(frame_BHWC)
            new_carry = (vid_BSHWC, frame_i_window)
            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:
            vid_t_window_BSHWC, vid_t_full_BFHWC, rng = carry
            rng, _rng_noise = jax.random.split(rng)

            # slide the window forward by 1 frame if necessary
            vid_t_window_BSHWC, frame_id_window = nnx.cond(
                frame_t < window_size,
                lambda: (vid_t_window_BSHWC, frame_t),
                lambda: (
                    jnp.roll(vid_t_window_BSHWC, -1, axis=1)
                    .at[:, -1]
                    .set(jax.random.normal(_rng_noise, (B, H, W, C))),
                    window_size - 1,
                ),
            )

            carry_denoise = (vid_t_window_BSHWC, frame_id_window)
            final_carry_denoise = denoise_step_fn(
                carry_denoise, jnp.arange(diffusion_steps)
            )
            vid_t_window_BSHWC = final_carry_denoise[0]
            vid_t_full_BFHWC = vid_t_full_BFHWC.at[:, frame_t].set(
                vid_t_window_BSHWC[:, frame_id_window]
            )
            new_carry = (vid_t_window_BSHWC, vid_t_full_BFHWC, rng)
            return new_carry

        rng, _rng_sample = jax.random.split(rng)
        initial_carry = (videos_BSHWC, videos_BFHWC, _rng_sample)
        final_carry = autoregressive_step_fn(initial_carry, jnp.arange(T, seq_len))
        final_videos_BFHWC = final_carry[1]

        return final_videos_BFHWC, None  # TODO return something meaningful?

    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