[tab] diff.diff (diff)
diff --git a/genie.py b/genie.py
index 50729ab..dfa352c 100644
--- a/genie.py
+++ b/genie.py
@@ -4,6 +4,7 @@ from orbax.checkpoint import PyTreeCheckpointer
 import jax
 import jax.numpy as jnp
 import flax.linen as nn
+import einops
 
 from models.dynamics import DynamicsMaskGIT
 from models.lam import LatentActionModel
@@ -87,25 +88,42 @@ class Genie(nn.Module):
     def sample(
         self,
         batch: Dict[str, Any],
+        seq_len: int,
         steps: int = 25,
-        temperature: int = 1,
+        temperature: float = 1,
         sample_argmax: bool = False,
     ) -> Any:
+        """
+        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: patches per frame  
+            S: sequence length  
+            A: action space  
+            D: model latent dimension
+        """
         # --- Encode videos and actions ---
         tokenizer_out = self.tokenizer.vq_encode(batch["videos"], training=False)
-        token_idxs = tokenizer_out["indices"]
-        new_frame_idxs = jnp.zeros_like(token_idxs)[:, 0]
+        token_idxs = tokenizer_out["indices"] # (B, T, N)
+        B, T, N = token_idxs.shape
+        pad_shape = (B, seq_len - T, N)
+        pad = jnp.zeros(pad_shape, dtype=token_idxs.dtype)
+        token_idxs = jnp.concatenate([token_idxs, pad], axis=1) # (B, S, N)
         action_tokens = self.lam.vq.get_codes(batch["latent_actions"])
 
-        # --- Initialize MaskGIT ---
-        init_mask = jnp.ones_like(token_idxs, dtype=bool)[:, 0]
-        init_carry = (
-            batch["rng"],
-            new_frame_idxs,
-            init_mask,
-            token_idxs,
-            action_tokens,
-        )
         MaskGITLoop = nn.scan(
             MaskGITStep,
             variable_broadcast="params",
@@ -123,13 +141,45 @@ class Genie(nn.Module):
             sample_argmax=sample_argmax,
             steps=steps,
         )
-        final_carry, _ = loop_fn(init_carry, jnp.arange(steps))
-        new_frame_idxs = final_carry[1]
-        new_frame_pixels = self.tokenizer.decode(
-            jnp.expand_dims(new_frame_idxs, 1),
+
+        def generation_step_fn(carry, step_t):
+            rng, current_token_idxs = carry
+            rng, step_rng = jax.random.split(rng)
+
+            # Mask current frame (future frames are masked by default using causal mask in ST-transformer)
+            mask = jnp.arange(seq_len) == step_t # (S,)
+            mask = jnp.broadcast_to(mask[None, :, None], (B, seq_len, N)) # (B, S, N)
+            mask = mask.astype(bool)
+            masked_token_idxs = current_token_idxs * ~mask
+
+            # --- Initialize and run MaskGIT loop ---
+            init_carry_maskgit = (
+                step_rng,
+                masked_token_idxs,
+                mask,
+                action_tokens,
+            )
+            final_carry_maskgit, _ = loop_fn(init_carry_maskgit, jnp.arange(steps))
+            updated_token_idxs = final_carry_maskgit[1]
+            new_carry = (rng, updated_token_idxs)
+            return new_carry, None
+
+        # --- Run the autoregressive generation using scan ---
+        initial_carry = (batch["rng"], token_idxs)
+        timesteps_to_scan = jnp.arange(T, seq_len)
+        final_carry, _ = jax.lax.scan(
+            generation_step_fn,
+            initial_carry,
+            timesteps_to_scan
+        )
+        final_token_idxs = final_carry[1]
+
+        # --- Decode all tokens at once at the end ---
+        final_frames = self.tokenizer.decode(
+            final_token_idxs,
             video_hw=batch["videos"].shape[2:4],
         )
-        return new_frame_pixels
+        return final_frames
 
     def vq_encode(self, batch, training) -> Dict[str, Any]:
         # --- Preprocess videos ---
@@ -181,10 +231,13 @@ class MaskGITStep(nn.Module):
 
         # --- Update mask ---
         num_unmasked_tokens = jnp.round(N * (1.0 - unmasked_ratio)).astype(int)
-        idx_mask = jnp.arange(final_token_probs.shape[-1]) > num_unmasked_tokens
-        sorted_idxs = jnp.argsort(final_token_probs, axis=-1, descending=True)
+        idx_mask = jnp.arange(final_token_probs.shape[-1]) <= N - num_unmasked_tokens
+        final_token_probs_flat = einops.rearrange(final_token_probs, "b s n -> b (s n)")
+        sorted_idxs = jnp.argsort(final_token_probs_flat, axis=-1)
         mask_update_fn = jax.vmap(lambda msk, ids: msk.at[ids].set(idx_mask))
-        new_mask = mask_update_fn(mask, sorted_idxs)
+        mask_flat = einops.rearrange(mask, "b s n -> b (s n)")
+        new_mask_flat = mask_update_fn(mask_flat, sorted_idxs)
+        new_mask = einops.rearrange(new_mask_flat, "b (s n) -> b s n", n=N)
 
         new_carry = (rng, token_idxs, new_mask, action_tokens)
         return new_carry, None