[tab] diff.diff (diff)
diff --git a/genie.py b/genie.py
index 04ba383..b3d5022 100644
--- a/genie.py
+++ b/genie.py
@@ -565,11 +565,11 @@ class Genie(nnx.Module):
 
 # FIXME (f.srambical): add conversion script for old checkpoints
 def restore_genie_components(
-    optimizer: nnx.ModelAndOptimizer,
+    optimizer: nnx.Optimizer,
     sharding: jax.sharding.NamedSharding,
     rng: jax.Array,
     args,
-) -> nnx.ModelAndOptimizer:
+) -> nnx.Optimizer:
     """Restore pre-trained Genie components"""
     rng_tokenizer, rng_lam = jax.random.split(rng)
     rngs_tokenizer = nnx.Rngs(rng_tokenizer)
@@ -606,7 +606,7 @@ def restore_genie_components(
         use_flash_attention=args.use_flash_attention,
         rngs=rngs_tokenizer,
     )
-    dummy_tokenizer_optimizer = nnx.ModelAndOptimizer(dummy_tokenizer, tx)
+    dummy_tokenizer_optimizer = nnx.Optimizer(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
@@ -645,7 +645,7 @@ def restore_genie_components(
             use_flash_attention=args.use_flash_attention,
             rngs=rngs_lam,
         )
-        dummy_lam_optimizer = nnx.ModelAndOptimizer(dummy_lam, tx)
+        dummy_lam_optimizer = nnx.Optimizer(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
@@ -665,7 +665,7 @@ def restore_genie_components(
         lam_checkpoint_manager.close()
 
     # Reinitialize the optimizer states
-    optimizer = nnx.ModelAndOptimizer(model, tx)
+    optimizer = nnx.Optimizer(model, tx)
     return optimizer
 
 
diff --git a/models/dynamics.py b/models/dynamics.py
index 6f0288b..338631a 100644
--- a/models/dynamics.py
+++ b/models/dynamics.py
@@ -104,6 +104,11 @@ class DynamicsMaskGIT(nnx.Module):
             padded_act_embed_BT1M, vid_embed_BTNM.shape
         )
         vid_embed_BTNM += padded_act_embed_BTNM
+        
+        # add a little gaussian noise
+        noise = jax.random.normal(_rng_prob, vid_embed_BTNM.shape) * 0.01
+        vid_embed_BTNM += noise
+
         logits_BTNV = self.transformer(vid_embed_BTNM)
         return logits_BTNV, mask
 
diff --git a/requirements.txt b/requirements.txt
index b6a4550..13c75e6 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -14,5 +14,5 @@ ffmpeg-python==0.2.0
 hf-transfer==0.1.9
 huggingface-hub[cli]>=0.34.3
 pre-commit>=4.2.0
-procgen>=0.10.7
+# procgen>=0.10.7
 tqdm>=4.67.1
\ No newline at end of file
diff --git a/train_dynamics.py b/train_dynamics.py
index 657921a..deab9a2 100644
--- a/train_dynamics.py
+++ b/train_dynamics.py
@@ -153,7 +153,7 @@ def build_model(args: Args, rng: jax.Array) -> tuple[Genie, jax.Array]:
     return genie, rng
 
 
-def build_optimizer(genie: Genie, args: Args) -> nnx.ModelAndOptimizer:
+def build_optimizer(genie: Genie, args: Args) -> nnx.Optimizer:
     lr_schedule = get_lr_schedule(
         args.lr_schedule,
         args.init_lr,
@@ -170,7 +170,7 @@ def build_optimizer(genie: Genie, args: Args) -> nnx.ModelAndOptimizer:
         weight_decay=1e-4,
         mu_dtype=args.param_dtype,  # moments in full precision
     )
-    optimizer = nnx.ModelAndOptimizer(genie, tx)
+    optimizer = nnx.Optimizer(genie, tx)
     return optimizer
 
 
@@ -186,7 +186,7 @@ def build_mesh_and_sharding(
 
 
 def shard_optimizer_states(
-    optimizer: nnx.ModelAndOptimizer, replicated_sharding: NamedSharding
+    optimizer: nnx.Optimizer, replicated_sharding: NamedSharding
 ) -> None:
     model_state = nnx.state(optimizer.model)
     model_sharded_state = jax.lax.with_sharding_constraint(
@@ -277,7 +277,7 @@ def build_checkpoint_manager(args: Args) -> Optional[ocp.CheckpointManager]:
 def restore_or_initialize_components(
     args: Args,
     checkpoint_manager: Optional[ocp.CheckpointManager],
-    optimizer: nnx.ModelAndOptimizer,
+    optimizer: nnx.Optimizer,
     train_iterator: grain.DataLoaderIterator,
     rng: jax.Array,
     replicated_sharding: NamedSharding,
@@ -285,7 +285,7 @@ def restore_or_initialize_components(
     restore_step: Optional[int] = None,
 ) -> tuple[
     int,
-    nnx.ModelAndOptimizer,
+    nnx.Optimizer,
     grain.DataLoaderIterator,
     grain.DataLoaderIterator,
     jax.Array,
@@ -468,7 +468,7 @@ def main(args: Args) -> None:
 
     @nnx.jit(donate_argnums=0)
     def train_step(
-        optimizer: nnx.ModelAndOptimizer, inputs: dict
+        optimizer: nnx.Optimizer, inputs: dict
     ) -> tuple[jax.Array, jax.Array, dict]:
         def loss_fn(model: Genie) -> tuple[jax.Array, tuple[jax.Array, dict]]:
             model.train()
@@ -513,16 +513,17 @@ def main(args: Args) -> None:
                 args.val_sample_argmax,
                 args.val_maskgit_steps,
             )
+            # Calculate metrics for the last frame only
             step_outputs = {
-                "recon": recon_full_frame,
-                "token_logits": logits_full_frame,
-                "video_tokens": tokens_full_frame,
-                "mask": jnp.zeros_like(tokens_full_frame).at[:, -1].set(True),
+                "recon": recon_full_frame[:,-1],
+                "token_logits": logits_full_frame[:,-1],
+                "video_tokens": tokens_full_frame[:,-1],
+                "mask": jnp.ones_like(tokens_full_frame[:,-1]),
             }
             if lam_indices is not None:
                 step_outputs["lam_indices"] = lam_indices
             loss_full_frame, metrics_full_frame = _calculate_step_metrics(
-                step_outputs, gt, args.num_actions, args.num_patch_latents
+                step_outputs, gt[:,-1], args.num_actions, args.num_patch_latents
             )
             val_output.update(
                 {
diff --git a/train_lam.py b/train_lam.py
index 9d9874a..6b4ba92 100644
--- a/train_lam.py
+++ b/train_lam.py
@@ -109,7 +109,7 @@ def build_model(args: Args, rng: jax.Array) -> tuple[LatentActionModel, jax.Arra
     )
 
 
-def build_optimizer(model: LatentActionModel, args: Args) -> nnx.ModelAndOptimizer:
+def build_optimizer(model: LatentActionModel, args: Args) -> nnx.Optimizer:
     lr_schedule = get_lr_schedule(
         args.lr_schedule,
         args.init_lr,
@@ -126,7 +126,7 @@ def build_optimizer(model: LatentActionModel, args: Args) -> nnx.ModelAndOptimiz
         weight_decay=1e-4,
         mu_dtype=args.param_dtype,  # moments in full precision
     )
-    optimizer = nnx.ModelAndOptimizer(model, tx)
+    optimizer = nnx.Optimizer(model, tx)
     return optimizer
 
 
@@ -141,7 +141,7 @@ def build_mesh_and_sharding(
 
 
 def shard_optimizer_states(
-    optimizer: nnx.ModelAndOptimizer, replicated_sharding: NamedSharding
+    optimizer: nnx.Optimizer, replicated_sharding: NamedSharding
 ) -> None:
     model_state = nnx.state(optimizer.model)
     model_sharded_state = jax.lax.with_sharding_constraint(
@@ -232,12 +232,12 @@ def build_checkpoint_manager(args: Args) -> Optional[ocp.CheckpointManager]:
 def restore_checkpoint_if_needed(
     args: Args,
     checkpoint_manager: Optional[ocp.CheckpointManager],
-    optimizer: nnx.ModelAndOptimizer,
+    optimizer: nnx.Optimizer,
     train_iterator: grain.DataLoaderIterator,
     val_iterator: Optional[grain.DataLoaderIterator],
     restore_step: Optional[int] = None,
 ) -> tuple[
-    int, nnx.ModelAndOptimizer, grain.DataLoaderIterator, grain.DataLoaderIterator
+    int, nnx.Optimizer, grain.DataLoaderIterator, grain.DataLoaderIterator
 ]:
     step = 0
     if checkpoint_manager and restore_step is None:
@@ -380,7 +380,7 @@ def main(args: Args) -> None:
 
     @nnx.jit(donate_argnums=0)
     def train_step(
-        optimizer: nnx.ModelAndOptimizer,
+        optimizer: nnx.Optimizer,
         inputs: dict,
         action_last_active: jax.Array,
         rng: jax.Array,
diff --git a/train_tokenizer.py b/train_tokenizer.py
index 33206b4..1ce846f 100644
--- a/train_tokenizer.py
+++ b/train_tokenizer.py
@@ -109,7 +109,7 @@ def build_model(args: Args, rng: jax.Array) -> tuple[TokenizerVQVAE, jax.Array]:
     )
 
 
-def build_optimizer(model: TokenizerVQVAE, args: Args) -> nnx.ModelAndOptimizer:
+def build_optimizer(model: TokenizerVQVAE, args: Args) -> nnx.Optimizer:
     lr_schedule = get_lr_schedule(
         args.lr_schedule,
         args.init_lr,
@@ -126,7 +126,7 @@ def build_optimizer(model: TokenizerVQVAE, args: Args) -> nnx.ModelAndOptimizer:
         weight_decay=1e-4,
         mu_dtype=args.param_dtype,  # moments in full precision
     )
-    optimizer = nnx.ModelAndOptimizer(model, tx)
+    optimizer = nnx.Optimizer(model, tx)
     return optimizer
 
 
@@ -141,7 +141,7 @@ def build_mesh_and_sharding(
 
 
 def shard_optimizer_states(
-    optimizer: nnx.ModelAndOptimizer, replicated_sharding: NamedSharding
+    optimizer: nnx.Optimizer, replicated_sharding: NamedSharding
 ) -> None:
     model_state = nnx.state(optimizer.model)
     model_sharded_state = jax.lax.with_sharding_constraint(
@@ -232,12 +232,12 @@ def build_checkpoint_manager(args: Args) -> Optional[ocp.CheckpointManager]:
 def restore_checkpoint_if_needed(
     args: Args,
     checkpoint_manager: Optional[ocp.CheckpointManager],
-    optimizer: nnx.ModelAndOptimizer,
+    optimizer: nnx.Optimizer,
     train_iterator: grain.DataLoaderIterator,
     val_iterator: Optional[grain.DataLoaderIterator],
     restore_step: Optional[int] = None,
 ) -> tuple[
-    int, nnx.ModelAndOptimizer, grain.DataLoaderIterator, grain.DataLoaderIterator
+    int, nnx.Optimizer, grain.DataLoaderIterator, grain.DataLoaderIterator
 ]:
     step = 0
     if checkpoint_manager and restore_step is None:
@@ -372,7 +372,7 @@ def main(args: Args) -> None:
 
     @nnx.jit(donate_argnums=0)
     def train_step(
-        optimizer: nnx.ModelAndOptimizer, inputs: dict
+        optimizer: nnx.Optimizer, inputs: dict
     ) -> tuple[jax.Array, jax.Array, dict]:
         def loss_fn(model: TokenizerVQVAE) -> tuple[jax.Array, tuple[jax.Array, dict]]:
             model.train()

            if lam_indices_E is not None:
                lam_indices_BT = lam_indices_E.reshape((-1, args.seq_len - 1))
                lam_indices_F = lam_indices_BT[:, -1].reshape((-1))
                step_outputs["lam_indices"] = lam_indices_F