[tab] dynamics_model_on_indices.patch (diff)
diff --git a/genie.py b/genie.py
index 2a7b868..7f35207 100644
--- a/genie.py
+++ b/genie.py
@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, cast
 
 import einops
 import jax
@@ -143,30 +143,38 @@ class Genie(nnx.Module):
     def __call__(
         self, batch: Dict[str, jax.Array], training: bool = True
     ) -> Dict[str, jax.Array]:
-        videos_BTHWC = batch["videos"]
-        tokenizer_outputs = self.tokenizer.vq_encode(videos_BTHWC, training=False)
-        token_indices_BTN = tokenizer_outputs["indices"]
-        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),
+        # --- Random index path (bypass tokenizer and LAM) ---
+        B, T, N = 12, 16, 64
+
+        rng = batch["mask_rng"]
+        rng, rng_tok, rng_act, rng_mask = jax.random.split(rng, 4)
+
+        # Create random token indices in [0, num_patch_latents)
+        token_indices_BTN = jax.random.randint(
+            rng_tok,
+            shape=(B, T, N),
+            minval=0,
+            maxval=self.num_patch_latents,
+            dtype=jnp.int32,
+        )
+
+        # Create random latent actions with shape (B, T-1, 1, latent_action_dim)
+        latent_actions_BTm11L = jax.random.normal(
+            rng_act,
+            shape=(B, max(T - 1, 1), 1, self.latent_action_dim),
+            dtype=self.dtype,
         )
+
         outputs = dict(
-            video_tokens=jax.lax.stop_gradient(token_indices_BTN),
+            video_tokens=token_indices_BTN,
             latent_actions=latent_actions_BTm11L,
+            mask_rng=rng_mask,
         )
-        outputs["mask_rng"] = batch["mask_rng"]
+
         dyna_logits_BTNV, dyna_mask = self.dynamics(outputs, training)
         outputs["token_logits"] = dyna_logits_BTNV
         if dyna_mask is not None:
             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))
-        outputs["lam_indices"] = action_indices_E
         return outputs
 
     def sample(
diff --git a/train_dynamics.py b/train_dynamics.py
index cd9f431..4bd339a 100644
--- a/train_dynamics.py
+++ b/train_dynamics.py
@@ -93,10 +93,8 @@ args = tyro.cli(Args)
 
 def dynamics_loss_fn(
     model: Genie, inputs: dict
-) -> tuple[jax.Array, tuple[jax.Array, dict]]:
+) -> tuple[jax.Array, tuple[jax.Array | None, dict]]:
     """Compute masked dynamics loss"""
-    gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
-    inputs["videos"] = gt.astype(args.dtype)
     model.train()
     outputs = model(inputs, training=True)
     mask = outputs["mask"]
@@ -108,17 +106,9 @@ def dynamics_loss_fn(
     acc = outputs["token_logits"].argmax(-1) == outputs["video_tokens"]
     acc = (mask * acc).sum() / mask.sum()
     select_probs = jax.nn.softmax(outputs["token_logits"])
-    gt = gt.clip(0, 1).reshape(-1, *gt.shape[2:])
-    recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
-    psnr = jnp.asarray(pix.psnr(gt, recon)).mean()
-    ssim = jnp.asarray(pix.ssim(gt, recon)).mean()
-    _, index_counts_lam = jnp.unique_counts(
-        jnp.ravel(outputs["lam_indices"]), size=args.num_latent_actions, fill_value=0
-    )
     _, index_counts_tokenizer = jnp.unique_counts(
         jnp.ravel(outputs["video_tokens"]), size=args.num_patch_latents, fill_value=0
     )
-    codebook_usage_lam = (index_counts_lam != 0).mean()
     codebook_usage_tokenizer = (index_counts_tokenizer != 0).mean()
     metrics = dict(
         cross_entropy_loss=ce_loss,
@@ -126,21 +116,18 @@ def dynamics_loss_fn(
         select_logit=outputs["token_logits"].max(-1).mean(),
         select_p=select_probs.max(-1).mean(),
         entropy=jax.scipy.special.entr(select_probs).sum(-1).mean(),
-        psnr=psnr,
-        ssim=ssim,
-        codebook_usage_lam=codebook_usage_lam,
         codebook_usage_tokenizer=codebook_usage_tokenizer,
     )
-    return ce_loss, (outputs["recon"], metrics)
+    return ce_loss, (None, metrics)
 
 
 @nnx.jit
 def train_step(
     model: Genie, optimizer: nnx.Optimizer, inputs: dict
-) -> tuple[jax.Array, jax.Array, dict]:
+) -> tuple[jax.Array, jax.Array | None, dict]:
     """Update state and compute metrics"""
 
-    def loss_fn(model: Genie) -> tuple[jax.Array, tuple[jax.Array, dict]]:
+    def loss_fn(model: Genie) -> tuple[jax.Array, tuple[jax.Array | None, dict]]:
         return dynamics_loss_fn(model, inputs)
 
     (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)