[tab] diff.diff (diff)
diff --git a/generate_dataset.py b/generate_dataset.py
index a67c424..4530f66 100644
--- a/generate_dataset.py
+++ b/generate_dataset.py
@@ -7,49 +7,95 @@ from dataclasses import dataclass
 from pathlib import Path
 
 from gym3 import types_np
+import os
 import numpy as np
 from procgen import ProcgenGym3Env
 import tyro
+import pickle
+import json
+from array_record.python.array_record_module import ArrayRecordWriter 
+
 
 
 @dataclass
 class Args:
-    num_episodes: int = 10000
+    num_episodes_train: int = 10000
+    num_episodes_val: int = 500
+    num_episodes_test: int = 500
     output_dir: str = "data/coinrun_episodes"
     min_episode_length: int = 50
+    seed: int = 0
 
 
 args = tyro.cli(Args)
-output_dir = Path(args.output_dir)
-output_dir.mkdir(parents=True, exist_ok=True)
-
-# --- Generate episodes ---
-i = 0
-metadata = []
-while i < args.num_episodes:
-    seed = np.random.randint(0, 10000)
-    env = ProcgenGym3Env(num=1, env_name="coinrun", start_level=seed)
-    dataseq = []
-
-    # --- Run episode ---
-    for j in range(1000):
-        env.act(types_np.sample(env.ac_space, bshape=(env.num,)))
-        rew, obs, first = env.observe()
-        dataseq.append(obs["rgb"])
-        if first:
-            break
-
-    # --- Save episode ---
-    if len(dataseq) >= args.min_episode_length:
-        episode_data = np.concatenate(dataseq, axis=0)
-        episode_path = output_dir / f"episode_{i}.npy"
-        np.save(episode_path, episode_data.astype(np.uint8))
-        metadata.append({"path": str(episode_path), "length": len(dataseq)})
-        print(f"Episode {i} completed, length: {len(dataseq)}")
-        i += 1
-    else:
-        print(f"Episode too short ({len(dataseq)}), resampling...")
-
-# --- Save metadata ---
-np.save(output_dir / "metadata.npy", metadata)
-print(f"Dataset generated with {len(metadata)} valid episodes")
+
+def generate_episodes(num_episodes, split):
+    i = 0
+    episode_metadata = []
+    while i < num_episodes:
+        seed = np.random.randint(0, 10000)
+        env = ProcgenGym3Env(num=1, env_name="coinrun", start_level=seed)
+        observations_seq = []
+
+        # --- Run episode ---
+        for _ in range(1000):
+            env.act(types_np.sample(env.ac_space, bshape=(env.num,)))
+            _ , obs, first = env.observe()
+            observations_seq.append(obs["rgb"])
+            if first:
+                break
+
+        # --- Save episode ---
+        if len(observations_seq) >= args.min_episode_length:
+            observations_data = np.concatenate(observations_seq, axis=0).astype(np.uint8)
+            episode_path = os.path.join(args.output_dir, split, f"episode_{i}.array_record")  
+
+            # --- Save as ArrayRecord ---
+            writer = ArrayRecordWriter(str(episode_path), "group_size:1")
+            record = {"raw_video": observations_data.tobytes(), "sequence_length": len(observations_seq)}
+            writer.write(pickle.dumps(record))
+            writer.close()
+
+            episode_metadata.append({"path": str(episode_path), "length": len(observations_seq)})
+            print(f"Episode {i} completed, length: {len(observations_seq)}")
+            i += 1
+        else:
+            print(f"Episode too short ({len(observations_seq)}), resampling...")
+    print(f"Done generating {split} split")
+    return episode_metadata
+
+
+def main():
+    # Set random seed and create dataset directories
+    np.random.seed(args.seed)
+    output_dir = Path(args.output_dir)
+    (output_dir / "train").mkdir(parents=True, exist_ok=True)
+    (output_dir / "val").mkdir(parents=True, exist_ok=True)
+    (output_dir / "test").mkdir(parents=True, exist_ok=True)
+
+    # --- Generate episodes ---
+    train_episode_metadata = generate_episodes(args.num_episodes_train, "train")
+    val_episode_metadata = generate_episodes(args.num_episodes_val, "val")
+    test_episode_metadata = generate_episodes(args.num_episodes_test, "test")
+
+    # --- Save metadata ---
+    metadata = {
+        "env": "coinrun",
+        "num_episodes_train": args.num_episodes_train,
+        "num_episodes_val": args.num_episodes_val,
+        "num_episodes_test": args.num_episodes_test,
+        "avg_episode_len_train": np.mean([ep["length"] for ep in train_episode_metadata]),
+        "avg_episode_len_val": np.mean([ep["length"] for ep in val_episode_metadata]),
+        "avg_episode_len_test": np.mean([ep["length"] for ep in test_episode_metadata]),
+        "episode_metadata_train": train_episode_metadata,
+        "episode_metadata_val": val_episode_metadata,
+        "episode_metadata_test": test_episode_metadata,
+
+    }
+    with open(output_dir / "metadata.json", "w") as f:
+        json.dump(metadata, f)
+
+    print(f"Done generating dataset.")
+
+if __name__ == "__main__":
+    main()
diff --git a/genie.py b/genie.py
index f6aaf28..2a7b868 100644
--- a/genie.py
+++ b/genie.py
@@ -224,17 +224,13 @@ class Genie(nnx.Module):
             vid_embed_BSNM = self.dynamics.patch_embed(token_idxs_BSN)
             mask_token_111M = self.dynamics.mask_token.value
             mask_expanded_BSN1 = mask_BSN[..., None]
-            vid_embed_BSNM = jnp.where(
-                mask_expanded_BSN1, mask_token_111M, vid_embed_BSNM
-            )
+            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 = self.dynamics.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])
-            )
+            act_embed_BS1M = jnp.reshape(act_embed_BSM, (B, S, 1, act_embed_BSM.shape[-1]))
             vid_embed_BSNM += act_embed_BS1M
             unmasked_ratio = jnp.cos(jnp.pi * (step + 1) / (steps * 2))
             step_temp = temperature * (1.0 - unmasked_ratio)
@@ -256,13 +252,8 @@ class Genie(nnx.Module):
 
             # --- 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
-            )
+            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)")
@@ -316,12 +307,12 @@ class Genie(nnx.Module):
         return final_frames_BSHWC
 
     def sample_causal(
-        self,
-        batch: Dict[str, jax.Array],
-        seq_len: int,
-        temperature: float = 1,
-        sample_argmax: bool = False,
-    ) -> jax.Array:
+            self,
+            batch: Dict[str, jax.Array],
+            seq_len: int,
+            temperature: float = 1,
+            sample_argmax: bool = False,
+        ) -> jax.Array:
         """
         Autoregressively samples up to `seq_len` future frames, following Figure 8 of the paper.
 
@@ -373,14 +364,9 @@ class Genie(nnx.Module):
             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])
-            )
+            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_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 ---
@@ -390,9 +376,7 @@ class Genie(nnx.Module):
                 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
-            )
+            token_idxs_BSN = token_idxs_BSN.at[:, step_t, step_n].set(sampled_token_idxs_B)
 
             new_carry = (rng, token_idxs_BSN, action_tokens_EL, step_t)
             return new_carry, None
@@ -440,7 +424,6 @@ class Genie(nnx.Module):
         lam_indices_E = lam_output["indices"]
         return lam_indices_E
 
-
 # FIXME (f.srambical): add conversion script for old checkpoints
 def restore_genie_components(
     optimizer: nnx.Optimizer,
@@ -449,9 +432,7 @@ def restore_genie_components(
     args,
 ) -> nnx.Optimizer:
     """Restore pre-trained Genie components"""
-    rng_tokenizer, rng_lam = jax.random.split(rng)
-    rngs_tokenizer = nnx.Rngs(rng_tokenizer)
-    rngs_lam = nnx.Rngs(rng_lam)
+    rngs = nnx.Rngs(rng)
 
     tx = optimizer.tx
     model = optimizer.model
@@ -482,7 +463,7 @@ def restore_genie_components(
         param_dtype=args.param_dtype,
         dtype=args.dtype,
         use_flash_attention=args.use_flash_attention,
-        rngs=rngs_tokenizer,
+        rngs=rngs,
     )
     dummy_tokenizer_optimizer = nnx.Optimizer(dummy_tokenizer, tx)
     dummy_tokenizer_optimizer_state = nnx.state(dummy_tokenizer_optimizer)
@@ -521,7 +502,7 @@ def restore_genie_components(
             param_dtype=args.param_dtype,
             dtype=args.dtype,
             use_flash_attention=args.use_flash_attention,
-            rngs=rngs_lam,
+            rngs=rngs,
         )
         dummy_lam_optimizer = nnx.Optimizer(dummy_lam, tx)
         dummy_lam_optimizer_state = nnx.state(dummy_lam_optimizer)
@@ -541,7 +522,7 @@ def restore_genie_components(
         # Remove the LAM decoder to save memory and avoid unnecessary computation.
         del model.lam.decoder
         lam_checkpoint_manager.close()
-
+    
     # Reinitialize the optimizer states
     optimizer = nnx.Optimizer(model, tx)
     return optimizer
diff --git a/input_pipeline/preprocess/npy_to_tfrecord.py b/input_pipeline/preprocess/npy_to_tfrecord.py
deleted file mode 100644
index d97ea43..0000000
--- a/input_pipeline/preprocess/npy_to_tfrecord.py
+++ /dev/null
@@ -1,125 +0,0 @@
-from dataclasses import dataclass
-
-import tensorflow as tf
-import concurrent.futures
-import numpy as np
-import logging
-import tyro
-from pathlib import Path
-from tqdm import tqdm
-
-logging.basicConfig(level=logging.INFO)
-
-
-@dataclass
-class Args:
-    source_data_dir: str = "data/coinrun_episodes"
-    output_tfrecords_dir: str = "data/coinrun_tfrecords"
-    num_shards: int = 50
-
-
-args = tyro.cli(Args)
-
-
-def _bytes_feature(value):
-    if isinstance(value, type(tf.constant(0))):
-        value = value.numpy()
-    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
-
-
-def _int64_feature(value):
-    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
-
-
-def create_tfrecord_example(episode_numpy_array):
-    feature = {
-        "height": _int64_feature(episode_numpy_array.shape[1]),
-        "width": _int64_feature(episode_numpy_array.shape[2]),
-        "channels": _int64_feature(episode_numpy_array.shape[3]),
-        "sequence_length": _int64_feature(episode_numpy_array.shape[0]),
-        "raw_video": _bytes_feature(episode_numpy_array.tobytes()),
-    }
-    return tf.train.Example(features=tf.train.Features(feature=feature))
-
-
-def process_shard(shard_idx, episode_paths, output_filename):
-    """Process a single shard: load episodes, write to one TFRecord file."""
-    with tf.io.TFRecordWriter(output_filename) as writer:
-        for npy_path in tqdm(
-            episode_paths,
-            desc=f"Shard {shard_idx:03d}",
-            leave=False,
-        ):
-            try:
-                episode_data = np.load(npy_path)
-                tf_example = create_tfrecord_example(episode_data)
-                writer.write(tf_example.SerializeToString())
-            except Exception as e:
-                logging.error(
-                    f"Shard {shard_idx}: Skipping {npy_path} due to error: {e}"
-                )
-
-
-def main_preprocess(data_dir_str, output_dir_str, num_shards):
-    data_dir = Path(data_dir_str)
-    output_dir = Path(output_dir_str)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    metadata = np.load(data_dir / "metadata.npy", allow_pickle=True)
-    episode_source_paths = [Path(item["path"]) for item in metadata]
-    num_total_episodes = len(episode_source_paths)
-
-    if num_shards <= 0:
-        raise ValueError("num_shards must be positive.")
-    if num_shards > num_total_episodes:
-        logging.warning(
-            f"Warning: num_shards ({num_shards}) is greater than total episodes ({num_total_episodes}). "
-            f"Setting num_shards to {num_total_episodes}."
-        )
-        num_shards = num_total_episodes
-
-    logging.info(
-        f"Preparing to write {num_total_episodes} episodes to {num_shards} TFRecord shards in {output_dir}..."
-    )
-
-    output_filenames = [
-        str(output_dir / f"shard-{i:05d}-of-{num_shards:05d}.tfrecord")
-        for i in range(num_shards)
-    ]
-
-    # Split episode paths into shards
-    shards = [[] for _ in range(num_shards)]
-    for idx, npy_path in enumerate(episode_source_paths):
-        shards[idx % num_shards].append(npy_path)
-
-    with concurrent.futures.ProcessPoolExecutor(max_workers=num_shards) as executor:
-        futures = []
-        for shard_idx, (shard_paths, out_fname) in enumerate(
-            zip(shards, output_filenames)
-        ):
-            futures.append(
-                executor.submit(process_shard, shard_idx, shard_paths, out_fname)
-            )
-        for f in tqdm(
-            concurrent.futures.as_completed(futures), total=num_shards, desc="Shards"
-        ):
-            f.result()  # Propagate exceptions
-
-    logging.info(
-        f"TFRecord sharding complete. {num_shards} shards written to {output_dir}."
-    )
-    logging.info("Generated shard files:")
-    for fname in output_filenames:
-        logging.info(f"  {fname}")
-
-
-if __name__ == "__main__":
-    if (
-        not Path(args.source_data_dir).exists()
-        or not (Path(args.source_data_dir) / "metadata.npy").exists()
-    ):
-        logging.error(f"Please generate data in '{args.source_data_dir}' first.")
-    else:
-        main_preprocess(
-            args.source_data_dir, args.output_tfrecords_dir, args.num_shards
-        )
diff --git a/input_pipeline/preprocess/pngs_to_array_records.py b/input_pipeline/preprocess/pngs_to_array_records.py
new file mode 100644
index 0000000..8f371e4
--- /dev/null
+++ b/input_pipeline/preprocess/pngs_to_array_records.py
@@ -0,0 +1,130 @@
+import os
+import numpy as np
+from PIL import Image
+import tyro
+from dataclasses import dataclass
+import pickle
+import json
+import multiprocessing as mp
+from array_record.python.array_record_module import ArrayRecordWriter
+
+
+@dataclass
+class Args:
+    input_path: str
+    output_path: str
+    env_name: str
+    original_fps: int = 60
+    target_fps: int = 10
+    target_width: int = 64 
+
+def preprocess_pngs(input_dir, output_path, original_fps, target_fps, target_width=None):
+    print(f"Processing PNGs in {input_dir}")
+    try:
+        png_files = sorted([
+            f for f in os.listdir(input_dir)
+            if f.lower().endswith('.png')
+        ], key=lambda x: int(os.path.splitext(x)[0]))
+
+        if not png_files:
+            print(f"No PNG files found in {input_dir}")
+            return {"path": input_dir, "length": 0} 
+
+        # Downsample indices
+        n_total = len(png_files)
+        if original_fps == target_fps:
+            selected_indices = np.arange(n_total)
+        else:
+            n_target = int(np.floor(n_total * target_fps / original_fps))
+            selected_indices = np.linspace(0, n_total-1, n_target, dtype=int)
+
+        selected_files = [png_files[i] for i in selected_indices]
+
+        # Load images
+        frames = []
+        for fname in selected_files:
+            img = Image.open(os.path.join(input_dir, fname)).convert("RGB")
+            if target_width is not None:
+                w, h = img.size  # PIL gives (width, height)
+                if w != target_width:
+                    target_height = int(round(h * (target_width / float(w))))
+                    resample_filter = Image.LANCZOS
+                    img = img.resize((target_width, target_height), resample=resample_filter)
+            frames.append(np.array(img))
+
+        frames = np.stack(frames, axis=0)  # (n_frames, H, W, 3)
+        environment = os.path.basename(os.path.dirname(input_dir)) 
+        episode_id = os.path.basename(input_dir)
+        # Write to array_record
+        os.makedirs(output_path, exist_ok=True)
+        out_file = os.path.join(
+            output_path,
+            f"{environment}_{episode_id}.array_record"
+        )
+        writer = ArrayRecordWriter(str(out_file), "group_size:1")
+        record = {"raw_video": frames.tobytes(), 
+                  "environment": environment,
+                  "sequence_length": frames.shape[0]}
+        writer.write(pickle.dumps(record))
+        writer.close()
+        print(f"Saved {frames.shape[0]} frames to {out_file}")
+        return {"path": input_dir, "length": frames.shape[0]}
+    except Exception as e:
+        print(f"Error processing {input_dir}: {e}")
+        return {"path": input_dir, "length": 0}
+
+def main():
+    args = tyro.cli(Args)
+    os.makedirs(args.output_path, exist_ok=True)
+    print(f"Output path: {args.output_path}")
+
+    games = [
+        os.path.join(args.input_path, d)
+        for d in os.listdir(args.input_path)
+        if os.path.isdir(os.path.join(args.input_path, d))
+    ]
+    episodes = [
+        os.path.join(game, d)
+        for game in games
+        for d in os.listdir(game)
+    ]
+
+    results = []
+    num_processes = mp.cpu_count()
+    print(f"Number of processes: {num_processes}")
+    pool_args = [
+        (episode, args.output_path, args.original_fps, args.target_fps, args.target_width)
+        for episode in episodes
+    ]
+    with mp.Pool(processes=num_processes) as pool:
+        for result in pool.starmap(preprocess_pngs, pool_args):
+            results.append(result)
+
+    print("Done converting png to array_record files")
+
+    # count the number of failed videos
+    failed_videos = [result for result in results if result["length"] == 0]
+    short_videos = [result for result in results if result["length"] < 1600]
+    num_successful_videos = len(results) - len(failed_videos) - len(short_videos)
+    print(f"Number of failed videos: {len(failed_videos)}")
+    print(f"Number of short videos: {len(short_videos)}")
+    print(f"Number of successful videos: {num_successful_videos}")
+    print(f"Number of total videos: {len(results)}")
+
+    metadata = {
+        "env": args.env_name,
+        "total_videos": len(results),
+        "num_successful_videos": len(results) - len(failed_videos) - len(short_videos),
+        "num_failed_videos": len(failed_videos),
+        "num_short_videos": len(short_videos),
+        "avg_episode_len": np.mean([ep["length"] for ep in results]),
+        "episode_metadata": results,
+    }
+
+    with open(os.path.join(args.output_path, "metadata.json"), "w") as f:
+        json.dump(metadata, f)
+
+    print("Done.")
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/input_pipeline/preprocess/video_to_array_records.py b/input_pipeline/preprocess/video_to_array_records.py
index 9cee29d..9a61f2e 100644
--- a/input_pipeline/preprocess/video_to_array_records.py
+++ b/input_pipeline/preprocess/video_to_array_records.py
@@ -11,10 +11,12 @@ from array_record.python.array_record_module import ArrayRecordWriter
 
 @dataclass
 class Args:
-    target_width, target_height = 160, 90
-    target_fps = 10
-    input_path: str = "data/minecraft_videos"
-    output_path: str = "data/minecraft_arrayrecords"
+    input_path: str
+    output_path: str
+    env_name: str
+    target_width: int = 160
+    target_height: int = 90
+    target_fps: int = 10
 
 
 def preprocess_video(
@@ -48,10 +50,10 @@ def preprocess_video(
         writer.write(pickle.dumps(record))
         writer.close()
 
-        return in_filename, n_frames
+        return {"path": in_filename, "length": n_frames}
     except Exception as e:
         print(f"Error processing video {idx} ({in_filename}): {e}")
-        return in_filename, 0
+        return {"path": in_filename, "length": 0}
 
 
 def main():
@@ -63,38 +65,48 @@ def main():
     num_processes = mp.cpu_count()
     print(f"Number of processes: {num_processes}")
 
-    print("Converting mp4 to npy files...")
-    pool_args = [
-        (
-            idx,
-            os.path.join(args.input_path, in_filename),
-            args.output_path,
-            args.target_width,
-            args.target_height,
-            args.target_fps,
-        )
-        for idx, in_filename in enumerate(os.listdir(args.input_path))
-        if in_filename.endswith(".mp4") or in_filename.endswith(".webm")
-    ]
+    print("Converting video to array_record files...")
+    pool_args = []
+    for idx, in_filename in enumerate(os.listdir(args.input_path)):
+        if in_filename.endswith(".mp4") or in_filename.endswith(".webm"):
+            pool_args.append((
+                idx,
+                os.path.join(args.input_path, in_filename),
+                args.output_path,
+                args.target_width,
+                args.target_height,
+                args.target_fps,
+            ))
+        else:
+            print(f"Warning: {in_filename} is not a supported video format. Skipping...")
 
     results = []
     with mp.Pool(processes=num_processes) as pool:
         for result in pool.starmap(preprocess_video, pool_args):
             results.append(result)
-    print("Done converting mp4 to npy files")
+    print("Done converting video to array_record files")
 
     # count the number of failed videos
-    failed_videos = [result for result in results if result[1] == 0]
-    short_episodes = [result for result in results if result[1] < 1600]
+    failed_videos = [result for result in results if result["length"] == 0]
+    short_videos = [result for result in results if result["length"] < 1600]
+    num_successful_videos = len(results) - len(failed_videos) - len(short_videos)
     print(f"Number of failed videos: {len(failed_videos)}")
-    print(f"Number of short episodes: {len(short_episodes)}")
-    print(
-        f"Number of successful videos: {len(results) - len(failed_videos) - len(short_episodes)}"
-    )
+    print(f"Number of short videos: {len(short_videos)}")
+    print(f"Number of successful videos: {num_successful_videos}")
     print(f"Number of total videos: {len(results)}")
 
-    with open(os.path.join(args.output_path, "meta_data.json"), "w") as f:
-        json.dump(results, f)
+    metadata = {
+        "env": args.env_name,
+        "total_videos": len(results),
+        "num_successful_videos": len(results) - len(failed_videos) - len(short_videos),
+        "num_failed_videos": len(failed_videos),
+        "num_short_videos": len(short_videos),
+        "avg_episode_len": np.mean([ep["length"] for ep in results]),
+        "episode_metadata": results,
+    }
+
+    with open(os.path.join(args.output_path, "metadata.json"), "w") as f:
+        json.dump(metadata, f)
 
 
 if __name__ == "__main__":
diff --git a/input_pipeline/preprocess/video_to_npy.py b/input_pipeline/preprocess/video_to_npy.py
deleted file mode 100644
index 0f196f7..0000000
--- a/input_pipeline/preprocess/video_to_npy.py
+++ /dev/null
@@ -1,117 +0,0 @@
-import ffmpeg
-import numpy as np
-import os
-import tyro
-import multiprocessing as mp
-from dataclasses import dataclass
-import json
-
-
-@dataclass
-class Args:
-    target_width, target_height = 160, 90
-    target_fps = 10
-    input_path: str = "data/minecraft_videos"
-    output_path: str = "data/minecraft_npy"
-
-
-def preprocess_video(
-    idx, in_filename, output_path, target_width, target_height, target_fps
-):
-    print(f"Processing video {idx}, Filename: {in_filename}")
-    try:
-        out, _ = (
-            ffmpeg.input(in_filename)
-            .filter("fps", fps=target_fps, round="up")
-            .filter("scale", target_width, target_height)
-            .output("pipe:", format="rawvideo", pix_fmt="rgb24")
-            .run(capture_stdout=True, quiet=True)
-        )
-
-        frame_size = target_height * target_width * 3
-        n_frames = len(out) // frame_size
-
-        frames = np.frombuffer(out, np.uint8).reshape(
-            n_frames, target_height, target_width, 3
-        )
-
-        output_file = os.path.join(
-            output_path, os.path.splitext(os.path.basename(in_filename))[0] + ".npy"
-        )
-
-        if not os.path.exists(os.path.dirname(output_file)):
-            os.makedirs(os.path.dirname(output_file))
-
-        np.save(output_file, frames)
-        print(f"Saved {n_frames} frames to {output_file} with shape {frames.shape}")
-        return in_filename, True
-    except Exception as e:
-        print(f"Error processing video {idx} ({in_filename}): {e}")
-        return in_filename, False
-
-
-def get_meta_data(filename, directory):
-    filepath = os.path.join(directory, filename)
-    arr = np.load(filepath, mmap_mode="r")
-    return filepath, arr.shape[0]
-
-
-def main():
-    args = tyro.cli(Args)
-
-    output_path = os.path.join(
-        args.output_path,
-        f"{args.target_fps}fps_{args.target_width}x{args.target_height}",
-    )
-    print(f"Output path: {output_path}")
-
-    num_processes = mp.cpu_count()
-    print(f"Number of processes: {num_processes}")
-
-    print("Converting mp4 to npy files...")
-    pool_args = [
-        (
-            idx,
-            os.path.join(args.input_path, in_filename),
-            output_path,
-            args.target_width,
-            args.target_height,
-            args.target_fps,
-        )
-        for idx, in_filename in enumerate(os.listdir(args.input_path))
-        if in_filename.endswith(".mp4") or in_filename.endswith(".webm")
-    ]
-
-    results = []
-    with mp.Pool(processes=num_processes) as pool:
-        for result in pool.starmap(preprocess_video, pool_args):
-            results.append(result)
-    print("Done converting mp4 to npy files")
-
-    # count the number of failed videos
-    failed_videos = [result for result in results if not result[1]]
-    print(f"Number of failed videos: {len(failed_videos)}")
-    print(f"Number of successful videos: {len(results) - len(failed_videos)}")
-    print(f"Number of total videos: {len(results)}")
-
-    with open(os.path.join(output_path, "failed_videos.json"), "w") as f:
-        json.dump(failed_videos, f)
-
-    print("Creating metadata file...")
-    metadata = []
-    filenames = [
-        filename
-        for filename in os.listdir(output_path)
-        if filename.endswith(".npy") and filename != "metadata.npy"
-    ]
-    pool_args = [(filename, output_path) for filename in filenames]
-
-    with mp.Pool(processes=num_processes) as pool:
-        results = list(pool.starmap(get_meta_data, pool_args))
-        metadata = [{"path": path, "length": length} for path, length in results]
-    np.save(os.path.join(output_path, "metadata.npy"), metadata)
-    print(f"Saved {len(metadata)} videos to {output_path}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/models/dynamics.py b/models/dynamics.py
index 827dce8..13381e3 100644
--- a/models/dynamics.py
+++ b/models/dynamics.py
@@ -11,7 +11,7 @@ from utils.nn import STTransformer, Transformer
 class DynamicsMaskGIT(nnx.Module):
     """
     MaskGIT dynamics model
-
+    
     Dimension keys:
         B: batch size
         T: sequence length
@@ -99,17 +99,12 @@ class DynamicsMaskGIT(nnx.Module):
 
         # --- Predict transition ---
         act_embed_BTm11M = self.action_up(latent_actions_BTm11L)
-        padded_act_embed_BT1M = jnp.pad(
-            act_embed_BTm11M, ((0, 0), (1, 0), (0, 0), (0, 0))
-        )
-        padded_act_embed_BTNM = jnp.broadcast_to(
-            padded_act_embed_BT1M, vid_embed_BTNM.shape
-        )
+        padded_act_embed_BT1M = jnp.pad(act_embed_BTm11M, ((0, 0), (1, 0), (0, 0), (0, 0)))
+        padded_act_embed_BTNM = jnp.broadcast_to(padded_act_embed_BT1M, vid_embed_BTNM.shape)
         vid_embed_BTNM += padded_act_embed_BTNM
         logits_BTNV = self.transformer(vid_embed_BTNM)
         return logits_BTNV, mask
 
-
 class DynamicsCausal(nnx.Module):
     """Causal dynamics model"""
 
@@ -170,12 +165,8 @@ class DynamicsCausal(nnx.Module):
 
         vid_embed_BTNM = self.patch_embed(video_tokens_BTN)
         act_embed_BTm11M = self.action_up(latent_actions_BTm11L)
-        padded_act_embed_BT1M = jnp.pad(
-            act_embed_BTm11M, ((0, 0), (1, 0), (0, 0), (0, 0))
-        )
-        vid_embed_BTNp1M = jnp.concatenate(
-            [padded_act_embed_BT1M, vid_embed_BTNM], axis=2
-        )
+        padded_act_embed_BT1M = jnp.pad(act_embed_BTm11M, ((0, 0), (1, 0), (0, 0), (0, 0)))
+        vid_embed_BTNp1M = jnp.concatenate([padded_act_embed_BT1M, vid_embed_BTNM], axis=2)
 
         logits_BTNp1V = self.transformer(vid_embed_BTNp1M)
         logits_BTNV = logits_BTNp1V[:, :, :-1]
diff --git a/models/lam.py b/models/lam.py
index b68a2ef..25f80e7 100644
--- a/models/lam.py
+++ b/models/lam.py
@@ -10,7 +10,7 @@ from utils.nn import STTransformer, VectorQuantizer
 
 class LatentActionModel(nnx.Module):
     """Latent Action ST-ViVit VQ-VAE
-
+    
     Dimension keys:
         B: batch size
         T: sequence length
@@ -81,7 +81,6 @@ class LatentActionModel(nnx.Module):
             self.latent_dim,
             self.num_latents,
             self.codebook_dropout,
-            self.dtype,
             rngs=rngs,
         )
         self.patch_up = nnx.Linear(
@@ -156,6 +155,4 @@ class LatentActionModel(nnx.Module):
         z_EL = z_BTm1L.reshape(B * (T - 1), self.latent_dim)
         z_q_EL, z_EL, emb_EL, indices_E = self.vq(z_EL, training)
         z_q_BTm11L = z_q_EL.reshape(B, T - 1, 1, self.latent_dim)
-        return dict(
-            patches=patch_BTNP, z_q=z_q_BTm11L, z=z_EL, emb=emb_EL, indices=indices_E
-        )
+        return dict(patches=patch_BTNP, z_q=z_q_BTm11L, z=z_EL, emb=emb_EL, indices=indices_E)
diff --git a/models/tokenizer.py b/models/tokenizer.py
index a07097b..625038f 100644
--- a/models/tokenizer.py
+++ b/models/tokenizer.py
@@ -72,7 +72,6 @@ class TokenizerVQVAE(nnx.Module):
             self.latent_dim,
             self.num_latents,
             self.codebook_dropout,
-            self.dtype,
             rngs=rngs,
         )
         self.out_dim = self.in_dim * self.patch_size**2
diff --git a/sample.py b/sample.py
index 524354a..4d5e327 100644
--- a/sample.py
+++ b/sample.py
@@ -33,7 +33,7 @@ class Args:
     maskgit_steps: int = 25
     temperature: float = 1.0
     sample_argmax: bool = True
-    start_frame: int = 1
+    start_frame: int = 0
     # Tokenizer checkpoint
     tokenizer_dim: int = 512
     tokenizer_ffn_dim: int = 2048
@@ -207,14 +207,10 @@ if __name__ == "__main__":
     # --- Sample + evaluate video ---
     recon_video_BSHWC = _autoreg_sample(genie, rng, video_batch_BSHWC, action_batch_E)
     recon_video_BSHWC = recon_video_BSHWC.astype(jnp.float32)
-    gt = (
-        gt_video[:, : recon_video_BSHWC.shape[1]]
-        .clip(0, 1)
-        .reshape(-1, *gt_video.shape[2:])
-    )
+    gt = gt_video[:, : recon_video_BSHWC.shape[1]].clip(0, 1).reshape(-1, *gt_video.shape[2:])
     recon = recon_video_BSHWC.clip(0, 1).reshape(-1, *recon_video_BSHWC.shape[2:])
     ssim = jnp.asarray(
-        pix.ssim(gt[:, args.start_frame :], recon[:, args.start_frame :])
+        pix.ssim(gt[:, args.start_frame:], recon[:, args.start_frame:])
     ).mean()
     print(f"SSIM: {ssim}")
 
@@ -230,7 +226,7 @@ if __name__ == "__main__":
     imgs = [Image.fromarray(img) for img in frames]
     # Write actions on each frame, on each row (i.e., for each video in the batch, on the GT row)
     B, S, _, _, _ = video_batch_BSHWC.shape
-    action_batch_BSm11 = jnp.reshape(action_batch_E, (B, S - 1, 1))
+    action_batch_BSm11 = jnp.reshape(action_batch_E, (B, S-1, 1))
     for t, img in enumerate(imgs[1:]):
         d = ImageDraw.Draw(img)
         for row in range(action_batch_BSm11.shape[0]):
diff --git a/train_dynamics.py b/train_dynamics.py
index f75238c..c9707f0 100644
--- a/train_dynamics.py
+++ b/train_dynamics.py
@@ -1,10 +1,6 @@
-import os
-
-os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.98")
-
 from dataclasses import dataclass, field
-import itertools
-from typing import cast, Optional
+import os
+from typing import cast
 
 import einops
 from jax.sharding import Mesh, PartitionSpec, NamedSharding
@@ -22,13 +18,8 @@ import flax.nnx as nnx
 
 from genie import Genie, restore_genie_components
 from utils.dataloader import get_dataloader
-from utils.train_utils import (
-    get_lr_schedule,
-    count_parameters_by_component,
-    print_mem_stats,
-    print_compiled_memory_stats,
-    print_compiled_cost_analysis,
-)
+from utils.lr_utils import get_lr_schedule
+from utils.parameter_utils import count_parameters_by_component
 
 
 @dataclass
@@ -97,7 +88,88 @@ class Args:
     wandb_id: str = ""
 
 
-def build_model(args: Args, rng: jax.Array) -> tuple[Genie, jax.Array]:
+args = tyro.cli(Args)
+
+
+def dynamics_loss_fn(
+    model: Genie, inputs: dict
+) -> tuple[jax.Array, tuple[jax.Array, 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"]
+    outputs["token_logits"] = outputs["token_logits"].astype(jnp.float32)
+    ce_loss = optax.softmax_cross_entropy_with_integer_labels(
+        outputs["token_logits"], outputs["video_tokens"]
+    )
+    ce_loss = (mask * ce_loss).sum() / mask.sum()
+    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,
+        masked_token_accuracy=acc,
+        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)
+
+
+@nnx.jit
+def train_step(
+    model: Genie, optimizer: nnx.Optimizer, inputs: dict
+) -> tuple[jax.Array, jax.Array, dict]:
+    """Update state and compute metrics"""
+
+    def loss_fn(model: Genie) -> tuple[jax.Array, tuple[jax.Array, dict]]:
+        return dynamics_loss_fn(model, inputs)
+
+    (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)
+    optimizer.update(grads)
+    if args.log_gradients:
+        metrics["gradients_std/"] = jax.tree.map(
+            lambda x: x.std(), grads["params"]["dynamics"]
+        )
+    return loss, recon, metrics
+
+
+if __name__ == "__main__":
+    jax.distributed.initialize()
+    num_devices = jax.device_count()
+    if num_devices == 0:
+        raise ValueError("No JAX devices found.")
+    print(f"Running on {num_devices} devices.")
+
+    if args.batch_size % num_devices != 0:
+        raise ValueError(
+            f"Global batch size {args.batch_size} must be divisible by "
+            f"number of devices {num_devices}."
+        )
+
+    per_device_batch_size_for_init = args.batch_size // num_devices
+
+    rng = jax.random.key(args.seed)
+
+    # --- Initialize model ---
     rng, _rng = jax.random.split(rng)
     rngs = nnx.Rngs(_rng)
     genie = Genie(
@@ -133,11 +205,35 @@ def build_model(args: Args, rng: jax.Array) -> tuple[Genie, jax.Array]:
         decode=False,
         rngs=rngs,
     )
-    del genie.lam.decoder
-    return genie, rng
 
+    _, params, _ = nnx.split(genie, nnx.Param, ...)
+    param_counts = count_parameters_by_component(params)
+
+    if args.log and jax.process_index() == 0:
+        wandb_init_kwargs = {
+            "entity": args.entity,
+            "project": args.project,
+            "name": args.name,
+            "tags": args.tags,
+            "group": "debug",
+            "config": args,
+        }
+
+        if args.wandb_id:
+            wandb_init_kwargs.update(
+                {
+                    "id": args.wandb_id,
+                    "resume": "allow",
+                }
+            )
+        wandb.init(**wandb_init_kwargs)
+
+        wandb.config.update({"model_param_count": param_counts})
 
-def build_optimizer(genie: Genie, args: Args) -> tuple[nnx.Optimizer, optax.Schedule]:
+    print("Parameter counts:")
+    print(param_counts)
+
+    # --- Initialize optimizer ---
     lr_schedule = get_lr_schedule(
         args.lr_schedule,
         args.init_lr,
@@ -152,25 +248,18 @@ def build_optimizer(genie: Genie, args: Args) -> tuple[nnx.Optimizer, optax.Sche
         b1=0.9,
         b2=0.9,
         weight_decay=1e-4,
-        mu_dtype=args.param_dtype,  # moments in full precision
+        mu_dtype=args.dtype,
     )
     optimizer = nnx.Optimizer(genie, tx)
-    return optimizer, lr_schedule
-
+    del genie
 
-def build_mesh_and_sharding(
-    num_devices: int,
-) -> tuple[Mesh, NamedSharding, NamedSharding]:
+    # FIXME: switch to create_hybrid_device_mesh for runs spanning multiple nodes
     device_mesh_arr = create_device_mesh((num_devices,))
     mesh = Mesh(devices=device_mesh_arr, axis_names=("data",))
+
     replicated_sharding = NamedSharding(mesh, PartitionSpec())
     videos_sharding = NamedSharding(mesh, PartitionSpec("data", None, None, None, None))
-    return mesh, replicated_sharding, videos_sharding
-
 
-def shard_optimizer_states(
-    optimizer: nnx.Optimizer, replicated_sharding: NamedSharding
-) -> None:
     model_state = nnx.state(optimizer.model)
     model_sharded_state = jax.lax.with_sharding_constraint(
         model_state, replicated_sharding
@@ -182,31 +271,8 @@ def shard_optimizer_states(
     )
     nnx.update(optimizer, optimizer_sharded_state)
 
-
-def build_dataloader(args: Args) -> grain.DataLoaderIterator:
-    image_shape = (args.image_height, args.image_width, args.image_channels)
-    array_record_files = [
-        os.path.join(args.data_dir, x)
-        for x in os.listdir(args.data_dir)
-        if x.endswith(".array_record")
-    ]
-    grain_dataloader = get_dataloader(
-        array_record_files,
-        args.seq_len,
-        # NOTE: We deliberately pass the global batch size
-        # The dataloader shards the dataset across all processes
-        args.batch_size,
-        *image_shape,
-        num_workers=8,
-        prefetch_buffer_size=1,
-        seed=args.seed,
-    )
-    initial_state = grain_dataloader._create_initial_state()
-    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)
-    return grain_iterator
-
-
-def build_checkpoint_manager(args: Args) -> ocp.CheckpointManager:
+    # --- Initialize checkpoint manager ---
+    step = 0
     handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
     handler_registry.add(
         "model_state", ocp.args.PyTreeSave, ocp.handlers.PyTreeCheckpointHandler
@@ -224,6 +290,7 @@ def build_checkpoint_manager(args: Args) -> ocp.CheckpointManager:
         grain.checkpoint.CheckpointRestore,
         cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
     )
+
     checkpoint_options = ocp.CheckpointManagerOptions(
         save_interval_steps=args.log_checkpoint_interval,
         max_to_keep=3,
@@ -231,31 +298,40 @@ def build_checkpoint_manager(args: Args) -> ocp.CheckpointManager:
         step_format_fixed_length=6,
         cleanup_tmp_directories=True,
     )
+
     checkpoint_manager = ocp.CheckpointManager(
         args.ckpt_dir,
         options=checkpoint_options,
         handler_registry=handler_registry,
     )
-    return checkpoint_manager
-
-
-def restore_or_initialize_components(
-    args: Args,
-    checkpoint_manager: ocp.CheckpointManager,
-    optimizer: nnx.Optimizer,
-    grain_iterator: grain.DataLoaderIterator,
-    rng: jax.Array,
-    replicated_sharding: NamedSharding,
-    restore_step: Optional[int] = None,
-) -> tuple[int, nnx.Optimizer, grain.DataLoaderIterator, jax.Array]:
-    step = 0
-    if restore_step is None:
-        restore_step = checkpoint_manager.latest_step()
+
+    # --- Create DataLoaderIterator from dataloader ---
+    image_shape = (args.image_height, args.image_width, args.image_channels)
+    array_record_files = [
+        os.path.join(args.data_dir, x)
+        for x in os.listdir(args.data_dir)
+        if x.endswith(".array_record")
+    ]
+    grain_dataloader = get_dataloader(
+        array_record_files,
+        args.seq_len,
+        # NOTE: We deliberately pass the global batch size
+        # The dataloader shards the dataset across all processes
+        args.batch_size,
+        *image_shape,
+        num_workers=8,
+        prefetch_buffer_size=1,
+        seed=args.seed,
+    )
+    initial_state = grain_dataloader._create_initial_state()
+    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)
+
+    # --- Restore checkpoint ---
     if args.restore_ckpt:
         abstract_optimizer = nnx.eval_shape(lambda: optimizer)
         abstract_optimizer_state = nnx.state(abstract_optimizer)
         restored = checkpoint_manager.restore(
-            restore_step,
+            checkpoint_manager.latest_step(),
             args=ocp.args.Composite(
                 model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
                 dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),  # type: ignore
@@ -264,170 +340,29 @@ def restore_or_initialize_components(
         restored_optimizer_state = restored["model_state"]
         nnx.update(optimizer, restored_optimizer_state)
         grain_iterator = restored["dataloader_state"]
-        step = restore_step or 0
+        step = checkpoint_manager.latest_step() or 0
         print(f"Restored dataloader and model state from step {step}")
     else:
         # Restore from pre-trained tokenizer (and LAM)
-        rng, _rng = jax.random.split(rng)
-        optimizer = restore_genie_components(optimizer, replicated_sharding, _rng, args)
+        optimizer = restore_genie_components(optimizer, replicated_sharding, rng, args)
         # NOTE: We have to remove the (unused) tokenizer vq dropout due flax.nnx lazily initializing modules.
         # Specifically, the first dynamics model checkpoint will contain the vq dropout module,
         # but the first full restore will fail due to nnx not initializing the module when
         # dropout is set to 0.0.
         del optimizer.model.tokenizer.vq.drop
-    return step, optimizer, grain_iterator, rng
-
-
-def main(args: Args) -> None:
-    jax.distributed.initialize()
-    num_devices = jax.device_count()
-    if num_devices == 0:
-        raise ValueError("No JAX devices found.")
-    print(f"Running on {num_devices} devices.")
-
-    if args.batch_size % num_devices != 0:
-        raise ValueError(
-            f"Global batch size {args.batch_size} must be divisible by "
-            f"number of devices {num_devices}."
-        )
-
-    rng = jax.random.key(args.seed)
-
-    # --- Initialize model ---
-    genie, rng = build_model(args, rng)
-    _, params, _ = nnx.split(genie, nnx.Param, ...)
-    param_counts = count_parameters_by_component(params)
-
-    if args.log and jax.process_index() == 0:
-        wandb_init_kwargs = {
-            "entity": args.entity,
-            "project": args.project,
-            "name": args.name,
-            "tags": args.tags,
-            "group": "debug",
-            "config": args,
-        }
-
-        if args.wandb_id:
-            wandb_init_kwargs.update(
-                {
-                    "id": args.wandb_id,
-                    "resume": "allow",
-                }
-            )
-        wandb.init(**wandb_init_kwargs)
-
-        wandb.config.update({"model_param_count": param_counts})
-
-    print("Parameter counts:")
-    print(param_counts)
-
-    # --- Initialize optimizer ---
-    optimizer, lr_schedule = build_optimizer(genie, args)
-    del genie
-
-    # FIXME: switch to create_hybrid_device_mesh for runs spanning multiple nodes
-    mesh, replicated_sharding, videos_sharding = build_mesh_and_sharding(num_devices)
-
-    shard_optimizer_states(optimizer, replicated_sharding)
-
-    # --- Initialize checkpoint manager ---
-    checkpoint_manager = build_checkpoint_manager(args)
-
-    # --- Create DataLoaderIterator from dataloader ---
-    grain_iterator = build_dataloader(args)
-
-    # --- Restore checkpoint ---
-    step, optimizer, grain_iterator, rng = restore_or_initialize_components(
-        args, checkpoint_manager, optimizer, grain_iterator, rng, replicated_sharding
-    )
-
-    # --- Define loss and train step (close over args) ---
-    def dynamics_loss_fn(
-        model: Genie, inputs: dict
-    ) -> tuple[jax.Array, tuple[jax.Array, dict]]:
-        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"]
-        outputs["token_logits"] = outputs["token_logits"].astype(jnp.float32)
-        ce_loss = optax.softmax_cross_entropy_with_integer_labels(
-            outputs["token_logits"], outputs["video_tokens"]
-        )
-        ce_loss = (mask * ce_loss).sum() / mask.sum()
-        acc = outputs["token_logits"].argmax(-1) == outputs["video_tokens"]
-        acc = (mask * acc).sum() / mask.sum()
-        select_probs = jax.nn.softmax(outputs["token_logits"])
-        gt_val = 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_val, recon)).mean()
-        ssim = jnp.asarray(pix.ssim(gt_val, 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,
-            masked_token_accuracy=acc,
-            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)
-
-    @nnx.jit(donate_argnums=0)
-    def train_step(
-        optimizer: nnx.Optimizer, inputs: dict
-    ) -> tuple[jax.Array, jax.Array, dict]:
-        def loss_fn(model: Genie) -> tuple[jax.Array, tuple[jax.Array, dict]]:
-            return dynamics_loss_fn(model, inputs)
-
-        (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(
-            optimizer.model
-        )
-        optimizer.update(grads)
-        if args.log_gradients:
-            metrics["gradients_std/"] = jax.tree.map(
-                lambda x: x.std(), grads["params"]["dynamics"]
-            )
-        return loss, recon, metrics
 
     # --- TRAIN LOOP ---
     dataloader = (
         jax.make_array_from_process_local_data(videos_sharding, elem)
         for elem in grain_iterator
     )
-    if jax.process_index() == 0:
-        first_videos = next(dataloader)
-        sample_inputs = dict(videos=first_videos, mask_rng=rng)
-        compiled = train_step.lower(optimizer, sample_inputs).compile()
-        print_compiled_memory_stats(compiled.memory_analysis())
-        print_compiled_cost_analysis(compiled.cost_analysis())
-        # Do not skip the first batch during training
-        dataloader = itertools.chain([first_videos], dataloader)
     print(f"Starting training from step {step}...")
-    first_step = step
     while step < args.num_steps:
         for videos in dataloader:
             # --- Train step ---
             rng, _rng_mask = jax.random.split(rng, 2)
             inputs = dict(videos=videos, mask_rng=_rng_mask)
-            loss, recon, metrics = train_step(optimizer, inputs)
-            if step == first_step:
-                print_mem_stats("After params initialized")
+            loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
             metrics["lr"] = lr_schedule(step)
             print(f"Step {step}, loss: {loss}")
             step += 1
@@ -475,8 +410,3 @@ def main(args: Args) -> None:
                 break
 
     checkpoint_manager.close()
-
-
-if __name__ == "__main__":
-    args = tyro.cli(Args)
-    main(args)
diff --git a/train_lam.py b/train_lam.py
index e1b5aaa..388f355 100644
--- a/train_lam.py
+++ b/train_lam.py
@@ -1,10 +1,6 @@
-import os
-
-os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.98")
-
 from dataclasses import dataclass, field
-import itertools
-from typing import cast, Optional
+import os
+from typing import cast
 
 import einops
 from jax.sharding import Mesh, PartitionSpec, NamedSharding
@@ -22,13 +18,8 @@ import flax.nnx as nnx
 
 from models.lam import LatentActionModel
 from utils.dataloader import get_dataloader
-from utils.train_utils import (
-    get_lr_schedule,
-    count_parameters_by_component,
-    print_mem_stats,
-    print_compiled_memory_stats,
-    print_compiled_cost_analysis,
-)
+from utils.lr_utils import get_lr_schedule
+from utils.parameter_utils import count_parameters_by_component
 
 
 @dataclass
@@ -82,33 +73,146 @@ class Args:
     use_flash_attention: bool = True
 
 
-def build_model(args: Args, rng: jax.Array) -> tuple[LatentActionModel, jax.Array]:
+args = tyro.cli(Args)
+
+
+def lam_loss_fn(
+    model: LatentActionModel, inputs: dict
+) -> tuple[jax.Array, tuple[jax.Array, jax.Array, dict]]:
+    # --- Compute loss ---
+    gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
+    inputs["videos"] = gt.astype(args.dtype)
+    model.train()
+    outputs = model(inputs, training=True)
+    outputs["recon"] = outputs["recon"].astype(jnp.float32)
+    gt_future_frames = gt[:, 1:]
+    mse = jnp.square(gt_future_frames - outputs["recon"]).mean()
+    q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()
+    commitment_loss = jnp.square(
+        outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
+    ).mean()
+    loss = mse + q_loss + args.vq_beta * commitment_loss
+
+    # --- Compute validation metrics ---
+    gt = gt_future_frames.clip(0, 1).reshape(-1, *gt_future_frames.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()
+    count_fn = jax.vmap(lambda i: (outputs["indices"] == i).sum())
+    index_counts = count_fn(jnp.arange(args.num_latents))
+    metrics = dict(
+        loss=loss,
+        mse=mse,
+        q_loss=q_loss,
+        commitment_loss=commitment_loss,
+        psnr=psnr,
+        ssim=ssim,
+        codebook_usage=(index_counts != 0).mean(),
+    )
+    return loss, (outputs["recon"], index_counts, metrics)
+
+
+@nnx.jit
+def train_step(
+    lam: LatentActionModel,
+    optimizer: nnx.Optimizer,
+    inputs: dict,
+    action_last_active: jax.Array,
+    rng: jax.Array,
+) -> tuple[jax.Array, jax.Array, jax.Array, dict]:
+    def loss_fn(
+        model: LatentActionModel,
+    ) -> tuple[jax.Array, tuple[jax.Array, jax.Array, dict]]:
+        return lam_loss_fn(model, inputs)
+
+    # --- Update model ---
+    (loss, (recon, idx_counts, metrics)), grads = nnx.value_and_grad(
+        loss_fn, has_aux=True
+    )(lam)
+    optimizer.update(grads)
+
+    # --- Reset inactive latent actions ---
+    codebook = lam.vq.codebook
+    num_codes = len(codebook)
+    active_codes = idx_counts != 0.0
+    action_last_active = jnp.where(active_codes, 0, action_last_active + 1)
+    p_code = active_codes / active_codes.sum()
+    reset_idxs = jax.random.choice(rng, num_codes, shape=(num_codes,), p=p_code)
+    do_reset = action_last_active >= args.vq_reset_thresh
+    new_codebook = jnp.where(
+        jnp.expand_dims(do_reset, -1), codebook[reset_idxs], codebook.value
+    )
+    lam.vq.codebook.value = new_codebook
+    action_last_active = jnp.where(do_reset, 0, action_last_active)
+    return loss, recon, action_last_active, metrics
+
+
+if __name__ == "__main__":
+    jax.distributed.initialize()
+    num_devices = jax.device_count()
+    if num_devices == 0:
+        raise ValueError("No JAX devices found.")
+    print(f"Running on {num_devices} devices.")
+
+    if args.batch_size % num_devices != 0:
+        raise ValueError(
+            f"Global batch size {args.batch_size} must be divisible by "
+            f"number of devices {num_devices}."
+        )
+
+    per_device_batch_size_for_init = args.batch_size // num_devices
+
+    rng = jax.random.key(args.seed)
+
+    # --- Initialize model ---
     rng, _rng = jax.random.split(rng)
     rngs = nnx.Rngs(_rng)
-    return (
-        LatentActionModel(
-            in_dim=args.image_channels,
-            model_dim=args.model_dim,
-            ffn_dim=args.ffn_dim,
-            latent_dim=args.latent_dim,
-            num_latents=args.num_latents,
-            patch_size=args.patch_size,
-            num_blocks=args.num_blocks,
-            num_heads=args.num_heads,
-            dropout=args.dropout,
-            codebook_dropout=args.codebook_dropout,
-            param_dtype=args.param_dtype,
-            dtype=args.dtype,
-            use_flash_attention=args.use_flash_attention,
-            rngs=rngs,
-        ),
-        rng,
+    lam = LatentActionModel(
+        in_dim=args.image_channels,
+        model_dim=args.model_dim,
+        ffn_dim=args.ffn_dim,
+        latent_dim=args.latent_dim,
+        num_latents=args.num_latents,
+        patch_size=args.patch_size,
+        num_blocks=args.num_blocks,
+        num_heads=args.num_heads,
+        dropout=args.dropout,
+        codebook_dropout=args.codebook_dropout,
+        param_dtype=args.param_dtype,
+        dtype=args.dtype,
+        use_flash_attention=args.use_flash_attention,
+        rngs=rngs,
     )
 
+    # Count parameters
+    _, params, _ = nnx.split(lam, nnx.Param, ...)
+    param_counts = count_parameters_by_component(params)
+
+    if args.log and jax.process_index() == 0:
+        wandb_init_kwargs = {
+            "entity": args.entity,
+            "project": args.project,
+            "name": args.name,
+            "tags": args.tags,
+            "group": "debug",
+            "config": args,
+        }
 
-def build_optimizer(
-    model: LatentActionModel, args: Args
-) -> tuple[nnx.Optimizer, optax.Schedule]:
+        if args.wandb_id:
+            wandb_init_kwargs.update(
+                {
+                    "id": args.wandb_id,
+                    "resume": "allow",
+                }
+            )
+        wandb.init(**wandb_init_kwargs)
+
+        wandb.config.update({"model_param_count": param_counts})
+
+    print("Parameter counts:")
+    print(param_counts)
+
+    # --- Initialize optimizer ---
     lr_schedule = get_lr_schedule(
         args.lr_schedule,
         args.init_lr,
@@ -123,25 +227,17 @@ def build_optimizer(
         b1=0.9,
         b2=0.9,
         weight_decay=1e-4,
-        mu_dtype=args.param_dtype,  # moments in full precision
+        mu_dtype=args.dtype,
     )
-    optimizer = nnx.Optimizer(model, tx)
-    return optimizer, lr_schedule
-
+    optimizer = nnx.Optimizer(lam, tx)
 
-def build_mesh_and_sharding(
-    num_devices: int,
-) -> tuple[Mesh, NamedSharding, NamedSharding]:
+    # FIXME: switch to create_hybrid_device_mesh for runs spanning multiple nodes
     device_mesh_arr = create_device_mesh((num_devices,))
     mesh = Mesh(devices=device_mesh_arr, axis_names=("data",))
+
     replicated_sharding = NamedSharding(mesh, PartitionSpec())
     videos_sharding = NamedSharding(mesh, PartitionSpec("data", None, None, None, None))
-    return mesh, replicated_sharding, videos_sharding
-
 
-def shard_optimizer_states(
-    optimizer: nnx.Optimizer, replicated_sharding: NamedSharding
-) -> None:
     model_state = nnx.state(optimizer.model)
     model_sharded_state = jax.lax.with_sharding_constraint(
         model_state, replicated_sharding
@@ -153,31 +249,8 @@ def shard_optimizer_states(
     )
     nnx.update(optimizer, optimizer_sharded_state)
 
-
-def build_dataloader(args: Args) -> grain.DataLoaderIterator:
-    image_shape = (args.image_height, args.image_width, args.image_channels)
-    array_record_files = [
-        os.path.join(args.data_dir, x)
-        for x in os.listdir(args.data_dir)
-        if x.endswith(".array_record")
-    ]
-    grain_dataloader = get_dataloader(
-        array_record_files,
-        args.seq_len,
-        # NOTE: We deliberately pass the global batch size
-        # The dataloader shards the dataset across all processes
-        args.batch_size,
-        *image_shape,
-        num_workers=8,
-        prefetch_buffer_size=1,
-        seed=args.seed,
-    )
-    initial_state = grain_dataloader._create_initial_state()
-    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)
-    return grain_iterator
-
-
-def build_checkpoint_manager(args: Args) -> ocp.CheckpointManager:
+    # --- Initialize checkpoint manager ---
+    step = 0
     handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
     handler_registry.add(
         "model_state", ocp.args.PyTreeSave, ocp.handlers.PyTreeCheckpointHandler
@@ -195,6 +268,7 @@ def build_checkpoint_manager(args: Args) -> ocp.CheckpointManager:
         grain.checkpoint.CheckpointRestore,
         cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
     )
+
     checkpoint_options = ocp.CheckpointManagerOptions(
         save_interval_steps=args.log_checkpoint_interval,
         max_to_keep=3,
@@ -202,29 +276,40 @@ def build_checkpoint_manager(args: Args) -> ocp.CheckpointManager:
         step_format_fixed_length=6,
         cleanup_tmp_directories=True,
     )
+
     checkpoint_manager = ocp.CheckpointManager(
         args.ckpt_dir,
         options=checkpoint_options,
         handler_registry=handler_registry,
     )
-    return checkpoint_manager
 
+    # --- Create DataLoaderIterator from dataloader ---
+    image_shape = (args.image_height, args.image_width, args.image_channels)
+    array_record_files = [
+        os.path.join(args.data_dir, x)
+        for x in os.listdir(args.data_dir)
+        if x.endswith(".array_record")
+    ]
+    grain_dataloader = get_dataloader(
+        array_record_files,
+        args.seq_len,
+        # NOTE: We deliberately pass the global batch size
+        # The dataloader shards the dataset across all processes
+        args.batch_size,
+        *image_shape,
+        num_workers=8,
+        prefetch_buffer_size=1,
+        seed=args.seed,
+    )
+    initial_state = grain_dataloader._create_initial_state()
+    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)
 
-def restore_checkpoint_if_needed(
-    args: Args,
-    checkpoint_manager: ocp.CheckpointManager,
-    optimizer: nnx.Optimizer,
-    grain_iterator: grain.DataLoaderIterator,
-    restore_step: Optional[int] = None,
-) -> tuple[int, nnx.Optimizer, grain.DataLoaderIterator]:
-    step = 0
-    if restore_step is None:
-        restore_step = checkpoint_manager.latest_step()
+    # --- Restore checkpoint ---
     if args.restore_ckpt:
         abstract_optimizer = nnx.eval_shape(lambda: optimizer)
         abstract_optimizer_state = nnx.state(abstract_optimizer)
         restored = checkpoint_manager.restore(
-            restore_step,
+            checkpoint_manager.latest_step(),
             args=ocp.args.Composite(
                 model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
                 dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),  # type: ignore
@@ -233,170 +318,16 @@ def restore_checkpoint_if_needed(
         restored_optimizer_state = restored["model_state"]
         nnx.update(optimizer, restored_optimizer_state)
         grain_iterator = restored["dataloader_state"]
-        step = restore_step or 0
+        step = checkpoint_manager.latest_step() or 0
         print(f"Restored dataloader and model state from step {step}")
-    return step, optimizer, grain_iterator
-
-
-def enable_sowing(lam: LatentActionModel) -> None:
-    for model in [lam.encoder, lam.decoder]:
-        setattr(model, "sow_logits", True)
-        for blk in getattr(model, "blocks", []):
-            setattr(blk, "sow_weights", True)
-            setattr(blk, "sow_activations", True)
-
-
-def main(args: Args) -> None:
-    jax.distributed.initialize()
-    num_devices = jax.device_count()
-    if num_devices == 0:
-        raise ValueError("No JAX devices found.")
-    print(f"Running on {num_devices} devices.")
-
-    if args.batch_size % num_devices != 0:
-        raise ValueError(
-            f"Global batch size {args.batch_size} must be divisible by "
-            f"number of devices {num_devices}."
-        )
-
-    rng = jax.random.key(args.seed)
-
-    # --- Initialize model ---
-    lam, rng = build_model(args, rng)
-
-    # Count parameters
-    _, params, _ = nnx.split(lam, nnx.Param, ...)
-    param_counts = count_parameters_by_component(params)
-
-    if args.log and jax.process_index() == 0:
-        wandb_init_kwargs = {
-            "entity": args.entity,
-            "project": args.project,
-            "name": args.name,
-            "tags": args.tags,
-            "group": "debug",
-            "config": args,
-        }
-
-        if args.wandb_id:
-            wandb_init_kwargs.update(
-                {
-                    "id": args.wandb_id,
-                    "resume": "allow",
-                }
-            )
-        wandb.init(**wandb_init_kwargs)
-
-        wandb.config.update({"model_param_count": param_counts})
-
-    print("Parameter counts:")
-    print(param_counts)
-
-    # --- Initialize optimizer ---
-    optimizer, lr_schedule = build_optimizer(lam, args)
-    del lam
-
-    # FIXME: switch to create_hybrid_device_mesh for runs spanning multiple nodes
-    mesh, replicated_sharding, videos_sharding = build_mesh_and_sharding(num_devices)
-
-    shard_optimizer_states(optimizer, replicated_sharding)
-
-    # --- Initialize checkpoint manager ---
-    checkpoint_manager = build_checkpoint_manager(args)
-
-    # --- Create DataLoaderIterator from dataloader ---
-    grain_iterator = build_dataloader(args)
-
-    # --- Restore checkpoint ---
-    step, optimizer, grain_iterator = restore_checkpoint_if_needed(
-        args, checkpoint_manager, optimizer, grain_iterator
-    )
-
-    # --- Define loss and train step (close over args) ---
-    def lam_loss_fn(
-        model: LatentActionModel, inputs: dict
-    ) -> tuple[jax.Array, tuple[jax.Array, jax.Array, dict]]:
-        gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
-        inputs["videos"] = gt.astype(args.dtype)
-        model.train()
-        outputs = model(inputs, training=True)
-        outputs["recon"] = outputs["recon"].astype(jnp.float32)
-        gt_future_frames = gt[:, 1:]
-        mse = jnp.square(gt_future_frames - outputs["recon"]).mean()
-        q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()
-        commitment_loss = jnp.square(
-            outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
-        ).mean()
-        loss = mse + q_loss + args.vq_beta * commitment_loss
-
-        gt_val = gt_future_frames.clip(0, 1).reshape(-1, *gt_future_frames.shape[2:])
-        recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
-        psnr = jnp.asarray(pix.psnr(gt_val, recon)).mean()
-        ssim = jnp.asarray(pix.ssim(gt_val, recon)).mean()
-        count_fn = jax.vmap(lambda i: (outputs["indices"] == i).sum())
-        index_counts = count_fn(jnp.arange(args.num_latents))
-        metrics = dict(
-            loss=loss,
-            mse=mse,
-            q_loss=q_loss,
-            commitment_loss=commitment_loss,
-            psnr=psnr,
-            ssim=ssim,
-            codebook_usage=(index_counts != 0).mean(),
-        )
-        return loss, (outputs["recon"], index_counts, metrics)
-
-    @nnx.jit(donate_argnums=0)
-    def train_step(
-        optimizer: nnx.Optimizer,
-        inputs: dict,
-        action_last_active: jax.Array,
-        rng: jax.Array,
-    ) -> tuple[jax.Array, jax.Array, jax.Array, dict]:
-        def loss_fn(
-            model: LatentActionModel,
-        ) -> tuple[jax.Array, tuple[jax.Array, jax.Array, dict]]:
-            return lam_loss_fn(model, inputs)
-
-        # --- Update model ---
-        (loss, (recon, idx_counts, metrics)), grads = nnx.value_and_grad(
-            loss_fn, has_aux=True
-        )(optimizer.model)
-        optimizer.update(grads)
-
-        # --- Reset inactive latent actions ---
-        codebook = optimizer.model.vq.codebook
-        num_codes = len(codebook)
-        active_codes = idx_counts != 0.0
-        action_last_active = jnp.where(active_codes, 0, action_last_active + 1)
-        p_code = active_codes / active_codes.sum()
-        reset_idxs = jax.random.choice(rng, num_codes, shape=(num_codes,), p=p_code)
-        do_reset = action_last_active >= args.vq_reset_thresh
-        new_codebook = jnp.where(
-            jnp.expand_dims(do_reset, -1), codebook[reset_idxs], codebook.value
-        )
-        optimizer.model.vq.codebook.value = new_codebook
-        action_last_active = jnp.where(do_reset, 0, action_last_active)
-        return loss, recon, action_last_active, metrics
 
     # --- TRAIN LOOP ---
     dataloader = (
         jax.make_array_from_process_local_data(videos_sharding, elem)
         for elem in grain_iterator
     )
-    action_last_active = jnp.zeros(args.num_latents, dtype=jnp.int32)
-    if jax.process_index() == 0:
-        first_videos = next(dataloader)
-        sample_inputs = dict(videos=first_videos)
-        compiled = train_step.lower(
-            optimizer, sample_inputs, action_last_active, rng
-        ).compile()
-        print_compiled_memory_stats(compiled.memory_analysis())
-        print_compiled_cost_analysis(compiled.cost_analysis())
-        # Do not skip the first batch during training
-        dataloader = itertools.chain([first_videos], dataloader)
     print(f"Starting training from step {step}...")
-    first_step = step
+    action_last_active = jnp.zeros(args.num_latents, dtype=jnp.int32)
     while step < args.num_steps:
         for videos in dataloader:
             # --- Train step ---
@@ -405,10 +336,8 @@ def main(args: Args) -> None:
             inputs = dict(videos=videos, rng=_rng)
             rng, _rng = jax.random.split(rng)
             loss, recon, action_last_active, metrics = train_step(
-                optimizer, inputs, action_last_active, _rng
+                lam, optimizer, inputs, action_last_active, _rng
             )
-            if step == first_step:
-                print_mem_stats("After params initialized")
             metrics["lr"] = lr_schedule(step)
             print(f"Step {step}, loss: {loss}")
             step += 1
@@ -459,8 +388,3 @@ def main(args: Args) -> None:
                 break
 
     checkpoint_manager.close()
-
-
-if __name__ == "__main__":
-    args = tyro.cli(Args)
-    main(args)
diff --git a/train_tokenizer.py b/train_tokenizer.py
index 4b327aa..b096a7e 100644
--- a/train_tokenizer.py
+++ b/train_tokenizer.py
@@ -1,12 +1,8 @@
-import os
-
-os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.98")
-
 from dataclasses import dataclass, field
-from typing import cast, Optional
+import os
+from typing import cast
 
 import einops
-import itertools
 from jax.sharding import Mesh, PartitionSpec, NamedSharding
 from jax.experimental.mesh_utils import create_device_mesh
 import optax
@@ -22,13 +18,8 @@ import flax.nnx as nnx
 
 from models.tokenizer import TokenizerVQVAE
 from utils.dataloader import get_dataloader
-from utils.train_utils import (
-    get_lr_schedule,
-    count_parameters_by_component,
-    print_mem_stats,
-    print_compiled_memory_stats,
-    print_compiled_cost_analysis,
-)
+from utils.lr_utils import get_lr_schedule
+from utils.parameter_utils import count_parameters_by_component
 
 
 @dataclass
@@ -78,87 +69,90 @@ class Args:
     log_checkpoint_interval: int = 10000
     log_checkpoint_keep_period: int = 20000
     log_gradients: bool = False
+    val_data_dir: str = ""
+    val_interval: int = 20_000
+    val_steps: int = 50
     wandb_id: str = ""
     use_flash_attention: bool = True
 
 
-def build_model(args: Args, rng: jax.Array) -> tuple[TokenizerVQVAE, jax.Array]:
-    rng, _rng = jax.random.split(rng)
-    rngs = nnx.Rngs(_rng)
-    return (
-        TokenizerVQVAE(
-            in_dim=args.image_channels,
-            model_dim=args.model_dim,
-            ffn_dim=args.ffn_dim,
-            latent_dim=args.latent_dim,
-            num_latents=args.num_latents,
-            patch_size=args.patch_size,
-            num_blocks=args.num_blocks,
-            num_heads=args.num_heads,
-            dropout=args.dropout,
-            codebook_dropout=args.codebook_dropout,
-            param_dtype=args.param_dtype,
-            dtype=args.dtype,
-            use_flash_attention=args.use_flash_attention,
-            rngs=rngs,
-        ),
-        rng,
+args = tyro.cli(Args)
+
+
+def tokenizer_loss_fn(
+    model: TokenizerVQVAE, inputs: dict
+) -> tuple[jax.Array, tuple[jax.Array, dict]]:
+    # --- Compute loss ---
+    # FIXME (f.srambical): Can we even do native int8 training without casting the video at all?
+    # FIXME (f.srambical): If the tokenizer is the reason for the dynamics model being memory-bound,
+    # should we at least train the tokenizer natively in int8?
+    gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
+    inputs["videos"] = gt.astype(args.dtype)
+    outputs = model(inputs, training=True)
+    outputs["recon"] = outputs["recon"].astype(jnp.float32)
+    mse = jnp.square(gt - outputs["recon"]).mean()
+    q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()
+    commitment_loss = jnp.square(
+        outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
+    ).mean()
+    loss = mse + q_loss + args.vq_beta * commitment_loss
+
+    # --- Compute validation metrics ---
+    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 = jnp.unique_counts(
+        jnp.ravel(outputs["indices"]), size=args.num_latents, fill_value=0
     )
-
-
-def build_optimizer(
-    model: TokenizerVQVAE, args: Args
-) -> tuple[nnx.Optimizer, optax.Schedule]:
-    lr_schedule = get_lr_schedule(
-        args.lr_schedule,
-        args.init_lr,
-        args.max_lr,
-        args.decay_end,
-        args.num_steps,
-        args.warmup_steps,
-        args.wsd_decay_steps,
+    codebook_usage = (index_counts != 0).mean()
+    metrics = dict(
+        loss=loss,
+        mse=mse,
+        q_loss=q_loss,
+        commitment_loss=commitment_loss,
+        psnr=psnr,
+        ssim=ssim,
+        codebook_usage=codebook_usage,
     )
-    tx = optax.adamw(
-        learning_rate=lr_schedule,
-        b1=0.9,
-        b2=0.9,
-        weight_decay=1e-4,
-        mu_dtype=args.param_dtype,  # moments in full precision
-    )
-    optimizer = nnx.Optimizer(model, tx)
-    return optimizer, lr_schedule
-
+    return loss, (outputs["recon"], metrics)
 
-def build_mesh_and_sharding(
-    num_devices: int,
-) -> tuple[Mesh, NamedSharding, NamedSharding]:
-    device_mesh_arr = create_device_mesh((num_devices,))
-    mesh = Mesh(devices=device_mesh_arr, axis_names=("data",))
-    replicated_sharding = NamedSharding(mesh, PartitionSpec())
-    videos_sharding = NamedSharding(mesh, PartitionSpec("data", None, None, None, None))
-    return mesh, replicated_sharding, videos_sharding
 
+@nnx.jit
+def train_step(
+    tokenizer: TokenizerVQVAE, 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()
+        return tokenizer_loss_fn(model, inputs)
 
-def shard_optimizer_states(
-    optimizer: nnx.Optimizer, replicated_sharding: NamedSharding
-) -> None:
-    model_state = nnx.state(optimizer.model)
-    model_sharded_state = jax.lax.with_sharding_constraint(
-        model_state, replicated_sharding
+    (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(
+        tokenizer
     )
-    nnx.update(optimizer.model, model_sharded_state)
-    optimizer_state = nnx.state(optimizer, nnx.optimizer.OptState)
-    optimizer_sharded_state = jax.lax.with_sharding_constraint(
-        optimizer_state, replicated_sharding
-    )
-    nnx.update(optimizer, optimizer_sharded_state)
+    optimizer.update(grads)
+    if args.log_gradients:
+        metrics["encoder_gradients_std/"] = jax.tree.map(
+            lambda x: x.std(), grads["params"]["encoder"]
+        )
+        metrics["vq_gradients_std/"] = jax.tree.map(
+            lambda x: x.std(), grads["params"]["vq"]
+        )
+        metrics["decoder_gradients_std/"] = jax.tree.map(
+            lambda x: x.std(), grads["params"]["decoder"]
+        )
+    return loss, recon, metrics
 
+@nnx.jit
+def val_step(tokenizer: TokenizerVQVAE, inputs: dict) -> tuple[jax.Array, jax.Array, dict]:
+    tokenizer.eval()
+    (loss, (recon, metrics)) = tokenizer_loss_fn(tokenizer, inputs)
+    return loss, recon, metrics
 
-def build_dataloader(args: Args) -> grain.DataLoaderIterator:
+def create_dataloader_iterator(data_dir):
     image_shape = (args.image_height, args.image_width, args.image_channels)
     array_record_files = [
-        os.path.join(args.data_dir, x)
-        for x in os.listdir(args.data_dir)
+        os.path.join(data_dir, x)
+        for x in os.listdir(data_dir)
         if x.endswith(".array_record")
     ]
     grain_dataloader = get_dataloader(
@@ -174,71 +168,38 @@ def build_dataloader(args: Args) -> grain.DataLoaderIterator:
     )
     initial_state = grain_dataloader._create_initial_state()
     grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)
-    return grain_iterator
 
-
-def build_checkpoint_manager(args: Args) -> ocp.CheckpointManager:
-    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
-    handler_registry.add(
-        "model_state", ocp.args.PyTreeSave, ocp.handlers.PyTreeCheckpointHandler
-    )
-    handler_registry.add(
-        "model_state", ocp.args.PyTreeRestore, ocp.handlers.PyTreeCheckpointHandler
-    )
-    handler_registry.add(
-        "dataloader_state",
-        grain.checkpoint.CheckpointSave,
-        cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
-    )
-    handler_registry.add(
-        "dataloader_state",
-        grain.checkpoint.CheckpointRestore,
-        cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
-    )
-    checkpoint_options = ocp.CheckpointManagerOptions(
-        save_interval_steps=args.log_checkpoint_interval,
-        max_to_keep=3,
-        keep_period=args.log_checkpoint_keep_period,
-        step_format_fixed_length=6,
-        cleanup_tmp_directories=True,
-    )
-    checkpoint_manager = ocp.CheckpointManager(
-        args.ckpt_dir,
-        options=checkpoint_options,
-        handler_registry=handler_registry,
-    )
-    return checkpoint_manager
+    return grain_iterator
 
 
-def restore_checkpoint_if_needed(
-    args: Args,
-    checkpoint_manager: ocp.CheckpointManager,
-    optimizer: nnx.Optimizer,
-    grain_iterator: grain.DataLoaderIterator,
-    restore_step: Optional[int] = None,
-) -> tuple[int, nnx.Optimizer, grain.DataLoaderIterator]:
+# TODO mihir find appropriate name for this fn
+def calculate_validation_metrics(val_dataloader):
     step = 0
-    if restore_step is None:
-        restore_step = checkpoint_manager.latest_step()
-    if args.restore_ckpt:
-        abstract_optimizer = nnx.eval_shape(lambda: optimizer)
-        abstract_optimizer_state = nnx.state(abstract_optimizer)
-        restored = checkpoint_manager.restore(
-            restore_step,
-            args=ocp.args.Composite(
-                model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
-                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),  # type: ignore
-            ),
-        )
-        restored_optimizer_state = restored["model_state"]
-        nnx.update(optimizer, restored_optimizer_state)
-        grain_iterator = restored["dataloader_state"]
-        step = restore_step or 0
-        print(f"Restored dataloader and model state from step {step}")
-    return step, optimizer, grain_iterator
+    loss_per_step = []
+    metrics_per_step = []
+    for videos in val_dataloader:
+        inputs = dict(videos=videos)
+        loss, recon, metrics = val_step(tokenizer, inputs)
+        loss_per_step.append(loss)
+        metrics_per_step.append(metrics)
+        print(f"Step {step}, loss: {loss}")
+        step += 1
+        if step > args.val_steps:
+            break
+
+    if step < args.val_steps:
+        print(f"Warning: Your validation dataset is too small to make val_steps many steps. Made {step} steps, expected {args.val_steps}")
+
+    val_loss = np.mean(loss_per_step)
+    val_metrics = {
+        f"val_{key}": np.mean([m[key] for m in metrics_per_step])
+        for key in metrics_per_step[0].keys()
+    }
+    return val_loss, recon, val_metrics
+
 
 
-def main(args: Args) -> None:
+if __name__ == "__main__":
     jax.distributed.initialize()
     num_devices = jax.device_count()
     if num_devices == 0:
@@ -251,10 +212,29 @@ def main(args: Args) -> None:
             f"number of devices {num_devices}."
         )
 
+    per_device_batch_size_for_init = args.batch_size // num_devices
+
     rng = jax.random.key(args.seed)
 
     # --- Initialize model ---
-    tokenizer, rng = build_model(args, rng)
+    rng, _rng = jax.random.split(rng)
+    rngs = nnx.Rngs(_rng)
+    tokenizer = TokenizerVQVAE(
+        in_dim=args.image_channels,
+        model_dim=args.model_dim,
+        ffn_dim=args.ffn_dim,
+        latent_dim=args.latent_dim,
+        num_latents=args.num_latents,
+        patch_size=args.patch_size,
+        num_blocks=args.num_blocks,
+        num_heads=args.num_heads,
+        dropout=args.dropout,
+        codebook_dropout=args.codebook_dropout,
+        param_dtype=args.param_dtype,
+        dtype=args.dtype,
+        use_flash_attention=args.use_flash_attention,
+        rngs=rngs,
+    )
 
     _, params, _ = nnx.split(tokenizer, nnx.Param, ...)
     param_counts = count_parameters_by_component(params)
@@ -284,120 +264,149 @@ def main(args: Args) -> None:
     print(param_counts)
 
     # --- Initialize optimizer ---
-    optimizer, lr_schedule = build_optimizer(tokenizer, args)
-    del tokenizer
+    lr_schedule = get_lr_schedule(
+        args.lr_schedule,
+        args.init_lr,
+        args.max_lr,
+        args.decay_end,
+        args.num_steps,
+        args.warmup_steps,
+        args.wsd_decay_steps,
+    )
+    tx = optax.adamw(
+        learning_rate=lr_schedule,
+        b1=0.9,
+        b2=0.9,
+        weight_decay=1e-4,
+        mu_dtype=args.dtype,
+    )
+    optimizer = nnx.Optimizer(tokenizer, tx)
 
     # FIXME: switch to create_hybrid_device_mesh for runs spanning multiple nodes
-    mesh, replicated_sharding, videos_sharding = build_mesh_and_sharding(num_devices)
+    device_mesh_arr = create_device_mesh((num_devices,))
+    mesh = Mesh(devices=device_mesh_arr, axis_names=("data",))
 
-    shard_optimizer_states(optimizer, replicated_sharding)
+    replicated_sharding = NamedSharding(mesh, PartitionSpec())
+    videos_sharding = NamedSharding(mesh, PartitionSpec("data", None, None, None, None))
 
-    # --- Initialize checkpoint manager ---
-    checkpoint_manager = build_checkpoint_manager(args)
+    model_state = nnx.state(optimizer.model)
+    model_sharded_state = jax.lax.with_sharding_constraint(
+        model_state, replicated_sharding
+    )
+    nnx.update(optimizer.model, model_sharded_state)
+    optimizer_state = nnx.state(optimizer, nnx.optimizer.OptState)
+    optimizer_sharded_state = jax.lax.with_sharding_constraint(
+        optimizer_state, replicated_sharding
+    )
+    nnx.update(optimizer, optimizer_sharded_state)
 
-    # --- Create DataLoaderIterator from dataloader ---
-    grain_iterator = build_dataloader(args)
+    # --- Initialize checkpoint manager ---
+    step = 0
+    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
+    handler_registry.add(
+        "model_state", ocp.args.PyTreeSave, ocp.handlers.PyTreeCheckpointHandler
+    )
+    handler_registry.add(
+        "model_state", ocp.args.PyTreeRestore, ocp.handlers.PyTreeCheckpointHandler
+    )
+    handler_registry.add(
+        "train_dataloader_state",
+        grain.checkpoint.CheckpointSave,
+        cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
+    )
+    handler_registry.add(
+        "train_dataloader_state",
+        grain.checkpoint.CheckpointRestore,
+        cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
+    )
+    handler_registry.add(
+        "val_dataloader_state",
+        grain.checkpoint.CheckpointSave,
+        cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
+    )
+    handler_registry.add(
+        "val_dataloader_state",
+        grain.checkpoint.CheckpointRestore,
+        cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
+    )
 
-    # --- Restore checkpoint ---
-    step, optimizer, grain_iterator = restore_checkpoint_if_needed(
-        args, checkpoint_manager, optimizer, grain_iterator
+    checkpoint_options = ocp.CheckpointManagerOptions(
+        save_interval_steps=args.log_checkpoint_interval,
+        max_to_keep=3,
+        keep_period=args.log_checkpoint_keep_period,
+        step_format_fixed_length=6,
+        cleanup_tmp_directories=True,
     )
 
-    # --- Define loss and train step (close over args) ---
-    def tokenizer_loss_fn(
-        model: TokenizerVQVAE, inputs: dict
-    ) -> tuple[jax.Array, tuple[jax.Array, dict]]:
-        gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
-        inputs["videos"] = gt.astype(args.dtype)
-        model.train()
-        outputs = model(inputs, training=True)
-        outputs["recon"] = outputs["recon"].astype(jnp.float32)
-        mse = jnp.square(gt - outputs["recon"]).mean()
-        q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()
-        commitment_loss = jnp.square(
-            outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
-        ).mean()
-        loss = mse + q_loss + args.vq_beta * commitment_loss
-
-        gt_clipped = 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_clipped, recon)).mean()
-        ssim = jnp.asarray(pix.ssim(gt_clipped, recon)).mean()
-        _, index_counts = jnp.unique_counts(
-            jnp.ravel(outputs["indices"]), size=args.num_latents, fill_value=0
-        )
-        codebook_usage = (index_counts != 0).mean()
-        metrics = dict(
-            loss=loss,
-            mse=mse,
-            q_loss=q_loss,
-            commitment_loss=commitment_loss,
-            psnr=psnr,
-            ssim=ssim,
-            codebook_usage=codebook_usage,
-        )
-        return loss, (outputs["recon"], metrics)
+    checkpoint_manager = ocp.CheckpointManager(
+        args.ckpt_dir,
+        options=checkpoint_options,
+        handler_registry=handler_registry,
+    )
 
-    @nnx.jit(donate_argnums=0)
-    def train_step(
-        optimizer: nnx.Optimizer, inputs: dict
-    ) -> tuple[jax.Array, jax.Array, dict]:
-        def loss_fn(model: TokenizerVQVAE) -> tuple[jax.Array, tuple[jax.Array, dict]]:
-            return tokenizer_loss_fn(model, inputs)
+    # --- Create DataLoaderIterator from dataloader ---
+    train_iterator = create_dataloader_iterator(args.data_dir)
+    val_iterator = create_dataloader_iterator(args.val_data_dir)
 
-        (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(
-            optimizer.model
+    # --- Restore checkpoint ---
+    if args.restore_ckpt:
+        abstract_optimizer = nnx.eval_shape(lambda: optimizer)
+        abstract_optimizer_state = nnx.state(abstract_optimizer)
+        restored = checkpoint_manager.restore(
+            checkpoint_manager.latest_step(),
+            args=ocp.args.Composite(
+                model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
+                train_dataloader_state=grain.checkpoint.CheckpointRestore(train_iterator),  # type: ignore
+                val_dataloader_state=grain.checkpoint.CheckpointRestore(val_iterator),  # type: ignore
+            ),
         )
-        optimizer.update(grads)
-        if args.log_gradients:
-            metrics["encoder_gradients_std/"] = jax.tree.map(
-                lambda x: x.std(), grads["params"]["encoder"]
-            )
-            metrics["vq_gradients_std/"] = jax.tree.map(
-                lambda x: x.std(), grads["params"]["vq"]
-            )
-            metrics["decoder_gradients_std/"] = jax.tree.map(
-                lambda x: x.std(), grads["params"]["decoder"]
-            )
-        return loss, recon, metrics
+        restored_optimizer_state = restored["model_state"]
+        nnx.update(optimizer, restored_optimizer_state)
+        train_iterator = restored["train_dataloader_state"]
+        val_iterator = restored["val_dataloader_state"]
+        step = checkpoint_manager.latest_step() or 0
+        print(f"Restored dataloader and model state from step {step}")
 
+    
     # --- TRAIN LOOP ---
-    dataloader = (
+    dataloader_train = (
         jax.make_array_from_process_local_data(videos_sharding, elem)
-        for elem in grain_iterator
+        for elem in train_iterator
+    )
+    dataloader_val = (
+        jax.make_array_from_process_local_data(videos_sharding, elem)
+        for elem in val_iterator
     )
-    if jax.process_index() == 0:
-        first_videos = next(dataloader)
-        sample_inputs = dict(videos=first_videos)
-        compiled = train_step.lower(optimizer, sample_inputs).compile()
-        print_compiled_memory_stats(compiled.memory_analysis())
-        print_compiled_cost_analysis(compiled.cost_analysis())
-        # Do not skip the first batch during training
-        dataloader = itertools.chain([first_videos], dataloader)
     print(f"Starting training from step {step}...")
-    first_step = step
     while step < args.num_steps:
-        for videos in dataloader:
+        for videos in dataloader_train:
             # --- Train step ---
             inputs = dict(videos=videos)
-            loss, recon, metrics = train_step(optimizer, inputs)
-            if step == first_step:
-                print_mem_stats("After params initialized")
+            loss, recon, metrics = train_step(tokenizer, optimizer, inputs)
             metrics["lr"] = lr_schedule(step)
             print(f"Step {step}, loss: {loss}")
             step += 1
 
+            # --- Validation loss ---
+            if args.val_data_dir and step % args.val_interval == 0:
+                val_loss, val_recon, val_metrics = calculate_validation_metrics(dataloader_val)
+
             # --- Logging ---
             if args.log:
                 if step % args.log_interval == 0 and jax.process_index() == 0:
-                    wandb.log(
-                        {
-                            "loss": loss,
-                            "step": step,
-                            **metrics,
+                    log_dict = {
+                        "loss": loss,
+                        "step": step,
+                        **metrics
                         }
-                    )
+                    if args.val_data_dir and step % args.val_interval == 0:
+                        log_dict.update({
+                            "val_loss": val_loss,
+                            **val_metrics
+                        })
+                    wandb.log(log_dict)
                 if step % args.log_image_interval == 0:
+                    # TODO mihir: add validation recons here
                     gt_seq = inputs["videos"][0].astype(jnp.float32) / 255.0
                     recon_seq = recon[0].clip(0, 1)
                     comparison_seq = jnp.concatenate((gt_seq, recon_seq), axis=1)
@@ -423,8 +432,11 @@ def main(args: Args) -> None:
                     step,
                     args=ocp.args.Composite(
                         model_state=ocp.args.PyTreeSave(optimizer_state),  # type: ignore
-                        dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
-                            grain_iterator  # type: ignore
+                        train_dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
+                            train_iterator  # type: ignore
+                        ),
+                        val_dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
+                            val_iterator  # type: ignore
                         ),
                     ),
                 )
@@ -433,8 +445,3 @@ def main(args: Args) -> None:
                 break
 
     checkpoint_manager.close()
-
-
-if __name__ == "__main__":
-    args = tyro.cli(Args)
-    main(args)
diff --git a/utils/lr_utils.py b/utils/lr_utils.py
new file mode 100644
index 0000000..b0b7355
--- /dev/null
+++ b/utils/lr_utils.py
@@ -0,0 +1,43 @@
+import optax
+
+
+def get_lr_schedule(
+    lr_schedule: str,
+    init_lr: float,
+    max_lr: float,
+    decay_end: float,
+    total_steps: int,
+    warmup_steps: int,
+    wsd_decay_steps: int,
+) -> optax.Schedule:
+    supported_schedules = ["wsd", "cos"]
+    if lr_schedule == "cos":
+        assert (
+            warmup_steps <= total_steps
+        ), "Warmup steps can't be greater than total steps."
+        return optax.warmup_cosine_decay_schedule(
+            init_value=init_lr,
+            peak_value=max_lr,
+            warmup_steps=warmup_steps,
+            decay_steps=total_steps,  # Note: decay_steps includes the warmup steps, so we need to pass total value
+            end_value=decay_end,
+        )
+    elif lr_schedule == "wsd":
+        assert (
+            warmup_steps + wsd_decay_steps <= total_steps
+        ), "Warmup and decay period is longer than total steps."
+        schedules = [
+            optax.linear_schedule(
+                init_value=init_lr, end_value=max_lr, transition_steps=warmup_steps
+            ),
+            optax.constant_schedule(value=max_lr),
+            optax.linear_schedule(
+                init_value=max_lr, end_value=decay_end, transition_steps=wsd_decay_steps
+            ),
+        ]
+        boundaries = [warmup_steps, total_steps - wsd_decay_steps]
+        return optax.join_schedules(schedules, boundaries)
+    else:
+        raise ValueError(
+            f"Learning rate schedule not supported. Please use one of {supported_schedules}"
+        )
diff --git a/utils/nn.py b/utils/nn.py
index 2c2696c..53cb68a 100644
--- a/utils/nn.py
+++ b/utils/nn.py
@@ -11,7 +11,6 @@ class SpatioTemporalPositionalEncoding(nnx.Module):
     """
     Applies separate sinusoidal positional encodings to the temporal and spatial dimensions.
     """
-
     def __init__(self, d_model: int, max_len: int = 5000):
         self.d_model = d_model
         self.max_len = max_len
@@ -60,8 +59,6 @@ class STBlock(nnx.Module):
         dtype: jnp.dtype,
         use_flash_attention: bool,
         rngs: nnx.Rngs,
-        sow_weights: bool,
-        sow_activations: bool,
     ):
         self.dim = dim
         self.ffn_dim = ffn_dim
@@ -70,13 +67,11 @@ class STBlock(nnx.Module):
         self.param_dtype = param_dtype
         self.dtype = dtype
         self.use_flash_attention = use_flash_attention
-        self.sow_weights = sow_weights
-        self.sow_activations = sow_activations
 
         self.spatial_norm = nnx.LayerNorm(
             num_features=self.dim,
             param_dtype=self.param_dtype,
-            dtype=self.param_dtype,  # layer norm in full precision
+            dtype=self.dtype,
             rngs=rngs,
         )
         self.spatial_attention = nnx.MultiHeadAttention(
@@ -96,7 +91,7 @@ class STBlock(nnx.Module):
         self.temporal_norm = nnx.LayerNorm(
             num_features=self.dim,
             param_dtype=self.param_dtype,
-            dtype=self.param_dtype,  # layer norm in full precision
+            dtype=self.dtype,
             rngs=rngs,
         )
         self.temporal_attention = nnx.MultiHeadAttention(
@@ -116,7 +111,7 @@ class STBlock(nnx.Module):
         self.ffn_norm = nnx.LayerNorm(
             num_features=self.dim,
             param_dtype=self.param_dtype,
-            dtype=self.param_dtype,  # layer norm in full precision
+            dtype=self.dtype,
             rngs=rngs,
         )
         self.ffn_dense1 = nnx.Linear(
@@ -138,13 +133,13 @@ class STBlock(nnx.Module):
     def __call__(self, x_BTNM: jax.Array) -> jax.Array:
         # --- Spatial attention ---
         z_BTNM = self.spatial_norm(x_BTNM)
-        z_BTNM = self.spatial_attention(z_BTNM, sow_weights=self.sow_weights)
+        z_BTNM = self.spatial_attention(z_BTNM)
         x_BTNM = x_BTNM + z_BTNM
 
         # --- Temporal attention ---
         x_BNTM = x_BTNM.swapaxes(1, 2)
         z_BNTM = self.temporal_norm(x_BNTM)
-        z_BNTM = self.temporal_attention(z_BNTM, sow_weights=self.sow_weights)
+        z_BNTM = self.temporal_attention(z_BNTM)
         x_BNTM = x_BNTM + z_BNTM
         x_BTNM = x_BNTM.swapaxes(1, 2)
 
@@ -154,8 +149,7 @@ class STBlock(nnx.Module):
         z_BTND = jax.nn.gelu(z_BTND)
         z_BTNM = self.ffn_dense2(z_BTND)
         x_BTNM = x_BTNM + z_BTNM
-        if self.sow_activations:
-            self.sow(nnx.Intermediate, "activations", x_BTNM)
+
         return x_BTNM
 
 
@@ -168,9 +162,8 @@ class STTransformer(nnx.Module):
         I: number of input features
         M: model dimension
         D: FFN dimension
-        V: vocabulary size
+        O: number of output features
     """
-
     def __init__(
         self,
         input_dim: int,
@@ -184,9 +177,6 @@ class STTransformer(nnx.Module):
         dtype: jnp.dtype,
         use_flash_attention: bool,
         rngs: nnx.Rngs,
-        sow_weights: bool = False,
-        sow_activations: bool = False,
-        sow_logits: bool = False,
         max_len: int = 5000,
     ):
         self.input_dim = input_dim
@@ -199,14 +189,11 @@ class STTransformer(nnx.Module):
         self.param_dtype = param_dtype
         self.dtype = dtype
         self.use_flash_attention = use_flash_attention
-        self.sow_logits = sow_logits
-        self.sow_weights = sow_weights
-        self.sow_activations = sow_activations
 
         self.input_norm1 = nnx.LayerNorm(
             num_features=self.input_dim,
             param_dtype=self.param_dtype,
-            dtype=self.param_dtype,  # layer norm in full precision
+            dtype=self.dtype,
             rngs=rngs,
         )
         self.input_dense = nnx.Linear(
@@ -219,7 +206,7 @@ class STTransformer(nnx.Module):
         self.input_norm2 = nnx.LayerNorm(
             num_features=self.model_dim,
             param_dtype=self.param_dtype,
-            dtype=self.param_dtype,  # layer norm in full precision
+            dtype=self.dtype,
             rngs=rngs,
         )
 
@@ -237,8 +224,6 @@ class STTransformer(nnx.Module):
                     dtype=self.dtype,
                     use_flash_attention=self.use_flash_attention,
                     rngs=rngs,
-                    sow_weights=self.sow_weights,
-                    sow_activations=self.sow_activations,
                 )
             )
 
@@ -258,11 +243,8 @@ class STTransformer(nnx.Module):
         for block in self.blocks:
             x_BTNM = block(x_BTNM)
 
-        x_BTNV = self.output_dense(x_BTNM)
-        if self.sow_logits:
-            self.sow(nnx.Intermediate, "logits", x_BTNV)
-        return x_BTNV
-
+        x_BTNO = self.output_dense(x_BTNM)
+        return x_BTNO
 
 class TransformerBlock(nnx.Module):
     def __init__(
@@ -276,8 +258,6 @@ class TransformerBlock(nnx.Module):
         use_flash_attention: bool,
         decode: bool,
         rngs: nnx.Rngs,
-        sow_weights: bool,
-        sow_activations: bool,
     ):
         self.model_dim = model_dim
         self.ffn_dim = ffn_dim
@@ -287,25 +267,23 @@ class TransformerBlock(nnx.Module):
         self.dtype = dtype
         self.use_flash_attention = use_flash_attention
         self.decode = decode
-        self.sow_weights = sow_weights
-        self.sow_activations = sow_activations
 
         self.temporal_norm = nnx.LayerNorm(
             num_features=self.model_dim,
             param_dtype=self.param_dtype,
-            dtype=self.param_dtype,  # layer norm in full precision
+            dtype=self.dtype,
             rngs=rngs,
         )
         self.spatial_norm = nnx.LayerNorm(
             num_features=self.model_dim,
             param_dtype=self.param_dtype,
-            dtype=self.param_dtype,  # layer norm in full precision
+            dtype=self.dtype,
             rngs=rngs,
         )
         self.ffn_norm = nnx.LayerNorm(
             num_features=self.model_dim,
             param_dtype=self.param_dtype,
-            dtype=self.param_dtype,  # layer norm in full precision
+            dtype=self.dtype,
             rngs=rngs,
         )
         self.temporal_attention = nnx.MultiHeadAttention(
@@ -350,20 +328,18 @@ class TransformerBlock(nnx.Module):
         )
 
     @nnx.remat
-    def __call__(
-        self, x_BTNM: jax.Array, pos_index: Tuple[jax.Array, jax.Array] | None = None
-    ) -> jax.Array:
+    def __call__(self, x_BTNM: jax.Array, pos_index: Tuple[jax.Array, jax.Array] | None = None) -> jax.Array:
         # --- Spatial attention ---
         B, T, N, M = x_BTNM.shape
         z_FNM = einops.rearrange(x_BTNM, "b t n m -> (b t) n m")
         z_FNM = self.spatial_norm(z_FNM)
-        z_FNM = self.spatial_attention(z_FNM, sow_weights=self.sow_weights)
+        z_FNM = self.spatial_attention(z_FNM)
         z_BTNM = einops.rearrange(z_FNM, "(b t) n m -> b t n m", t=T)
         x_BTNM = x_BTNM + z_BTNM
         # --- Temporal attention ---
         z_PTM = einops.rearrange(x_BTNM, "b t n m -> (b n) t m")
         z_PTM = self.temporal_norm(z_PTM)
-        z_PTM = self.temporal_attention(z_PTM, sow_weights=self.sow_weights)
+        z_PTM = self.temporal_attention(z_PTM)
         z_BTNM = einops.rearrange(z_PTM, "(b n) t m -> b t n m", n=N)
         x_BTNM = x_BTNM + z_BTNM
         # --- Feedforward ---
@@ -372,12 +348,9 @@ class TransformerBlock(nnx.Module):
         z_BTND = jax.nn.gelu(z_BTND)
         z_BTNM = self.ffn_dense2(z_BTND)
         x_BTNM = x_BTNM + z_BTNM
-        if self.sow_activations:
-            self.sow(nnx.Intermediate, "activations", x_BTNM)
 
         return x_BTNM
 
-
 class Transformer(nnx.Module):
     """
     Dimension keys:
@@ -387,11 +360,10 @@ class Transformer(nnx.Module):
         I: number of input features
         M: model dimension
         D: FFN dimension
-        V: vocabulary size
+        O: number of output features
         F: number of frames in batch
         P: number of patch positions in batch
     """
-
     def __init__(
         self,
         input_dim: int,
@@ -406,9 +378,6 @@ class Transformer(nnx.Module):
         use_flash_attention: bool,
         decode: bool,
         rngs: nnx.Rngs,
-        sow_logits: bool = False,
-        sow_weights: bool = False,
-        sow_activations: bool = False,
         max_len: int = 5000,
     ):
         self.input_dim = input_dim
@@ -421,14 +390,11 @@ class Transformer(nnx.Module):
         self.param_dtype = param_dtype
         self.dtype = dtype
         self.use_flash_attention = use_flash_attention
-        self.sow_logits = sow_logits
-        self.sow_weights = sow_weights
-        self.sow_activations = sow_activations
 
         self.input_norm1 = nnx.LayerNorm(
             num_features=self.input_dim,
             param_dtype=self.param_dtype,
-            dtype=self.param_dtype,  # layer norm in full precision
+            dtype=self.dtype,
             rngs=rngs,
         )
         self.input_dense = nnx.Linear(
@@ -441,7 +407,7 @@ class Transformer(nnx.Module):
         self.input_norm2 = nnx.LayerNorm(
             num_features=self.model_dim,
             param_dtype=self.param_dtype,
-            dtype=self.param_dtype,  # layer norm in full precision
+            dtype=self.dtype,
             rngs=rngs,
         )
 
@@ -459,8 +425,6 @@ class Transformer(nnx.Module):
                     dtype=self.dtype,
                     use_flash_attention=self.use_flash_attention,
                     decode=decode,
-                    sow_weights=self.sow_weights,
-                    sow_activations=self.sow_activations,
                     rngs=rngs,
                 )
             )
@@ -472,9 +436,7 @@ class Transformer(nnx.Module):
             rngs=rngs,
         )
 
-    def __call__(
-        self, x_BTNI: jax.Array, pos_index: Tuple[jax.Array, jax.Array] | None = None
-    ) -> jax.Array:
+    def __call__(self, x_BTNI: jax.Array, pos_index: Tuple[jax.Array, jax.Array] | None = None) -> jax.Array:
         x_BTNI = self.input_norm1(x_BTNI)
         x_BTNM = self.input_dense(x_BTNI)
         x_BTNM = self.input_norm2(x_BTNM)
@@ -483,11 +445,8 @@ class Transformer(nnx.Module):
             x_BTNM = block(x_BTNM, pos_index)
 
         x_BTNV = self.output_dense(x_BTNM)
-        if self.sow_logits:
-            self.sow(nnx.Intermediate, "logits", x_BTNV)
         return x_BTNV
 
-
 def normalize(x: jax.Array) -> jax.Array:
     return x / (jnp.linalg.norm(x, ord=2, axis=-1, keepdims=True) + 1e-8)
 
@@ -499,19 +458,12 @@ class VectorQuantizer(nnx.Module):
         K: number of latents
         L: latent dimension
     """
-
     def __init__(
-        self,
-        latent_dim: int,
-        num_latents: int,
-        dropout: float,
-        dtype: jnp.dtype,
-        rngs: nnx.Rngs,
+        self, latent_dim: int, num_latents: int, dropout: float, rngs: nnx.Rngs
     ):
         self.latent_dim = latent_dim
         self.num_latents = num_latents
         self.dropout = dropout
-        self.dtype = dtype
 
         self.codebook = nnx.Param(
             normalize(
@@ -526,18 +478,15 @@ class VectorQuantizer(nnx.Module):
         self, x_DL: jax.Array, training: bool
     ) -> Tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
         # --- Compute distances ---
-        x_DL = x_DL.astype(self.dtype)
-        codebook = self.codebook.value.astype(self.dtype)
-
         x_DL = normalize(x_DL)
-        normalized_codebook_KL = normalize(codebook)
+        normalized_codebook_KL = normalize(self.codebook.value)
         distance_DK = -jnp.matmul(x_DL, normalized_codebook_KL.T)
         if training:
             distance_DK = self.drop(distance_DK)
 
         # --- Get indices and embeddings ---
         indices_D = jnp.argmin(distance_DK, axis=-1)
-        z_DL = codebook[indices_D]
+        z_DL = self.codebook[indices_D]
 
         # --- Straight through estimator ---
         z_q_DL = x_DL + jax.lax.stop_gradient(z_DL - x_DL)
@@ -559,9 +508,7 @@ def _create_flash_attention_fn(use_flash_attention: bool, is_causal: bool) -> Ca
     function due to strict shape checking.
     """
 
-    def attention_fn(
-        query_BTHD, key_BSHD, value_BSHD, bias=None, mask_B111=None, **kwargs
-    ):
+    def attention_fn(query_BTHD, key_BSHD, value_BSHD, bias=None, mask_B111=None, **kwargs):
         implementation = "cudnn" if use_flash_attention else None
 
         def _merge_batch_dims(x):
@@ -590,14 +537,7 @@ def _create_flash_attention_fn(use_flash_attention: bool, is_causal: bool) -> Ca
 
         mask_11TS = attention_mask[jnp.newaxis, jnp.newaxis, :, :]
 
-        bias_4d = (
-            jnp.pad(
-                _merge_batch_dims(bias),
-                ((0, 0), (0, 0), (0, pad_size_Q), (0, pad_size_K)),
-            )
-            if bias is not None
-            else None
-        )
+        bias_4d = jnp.pad(_merge_batch_dims(bias), ((0, 0), (0, 0), (0, pad_size_Q), (0, pad_size_K))) if bias is not None else None
 
         # NOTE: jax.nn.dot_product_attention does not support dropout
         output_4d = jax.nn.dot_product_attention(
diff --git a/utils/parameter_utils.py b/utils/parameter_utils.py
new file mode 100644
index 0000000..741d231
--- /dev/null
+++ b/utils/parameter_utils.py
@@ -0,0 +1,40 @@
+from jax.tree_util import tree_map, tree_reduce
+
+
+def _count_leaf(x):
+    """Count parameters in a single leaf node."""
+    if hasattr(x, "size"):
+        return x.size
+    return 0
+
+
+def _count_component(component_params):
+    """Count total parameters in a component."""
+    return tree_reduce(
+        lambda x, y: x + y, tree_map(_count_leaf, component_params), initializer=0
+    )
+
+
+def count_parameters_by_component(params):
+    """Count parameters for each component of the model.
+
+    Args:
+        params: Model parameters from nnx.split(model, nnx.Param, ...)
+
+    Returns:
+        Dictionary with parameter counts for each component
+    """
+    component_names = list(params.keys())
+    print(f"Counting all components: {component_names}")
+
+    counts = {}
+    total_params = 0
+
+    for name in component_names:
+        component_params = params[name]
+        count = _count_component(component_params)
+        counts[name] = count
+        total_params += count
+
+    counts["total"] = total_params
+    return counts
diff --git a/utils/train_utils.py b/utils/train_utils.py
deleted file mode 100644
index 4b633e8..0000000
--- a/utils/train_utils.py
+++ /dev/null
@@ -1,119 +0,0 @@
-import jax
-import optax
-import operator
-
-
-def get_lr_schedule(
-    lr_schedule: str,
-    init_lr: float,
-    max_lr: float,
-    decay_end: float,
-    total_steps: int,
-    warmup_steps: int,
-    wsd_decay_steps: int,
-) -> optax.Schedule:
-    supported_schedules = ["wsd", "cos"]
-    if lr_schedule == "cos":
-        assert (
-            warmup_steps <= total_steps
-        ), "Warmup steps can't be greater than total steps."
-        return optax.warmup_cosine_decay_schedule(
-            init_value=init_lr,
-            peak_value=max_lr,
-            warmup_steps=warmup_steps,
-            decay_steps=total_steps,  # Note: decay_steps includes the warmup steps, so we need to pass total value
-            end_value=decay_end,
-        )
-    elif lr_schedule == "wsd":
-        assert (
-            warmup_steps + wsd_decay_steps <= total_steps
-        ), "Warmup and decay period is longer than total steps."
-        schedules = [
-            optax.linear_schedule(
-                init_value=init_lr, end_value=max_lr, transition_steps=warmup_steps
-            ),
-            optax.constant_schedule(value=max_lr),
-            optax.linear_schedule(
-                init_value=max_lr, end_value=decay_end, transition_steps=wsd_decay_steps
-            ),
-        ]
-        boundaries = [warmup_steps, total_steps - wsd_decay_steps]
-        return optax.join_schedules(schedules, boundaries)
-    else:
-        raise ValueError(
-            f"Learning rate schedule not supported. Please use one of {supported_schedules}"
-        )
-
-
-def _count_component(component_params):
-    """Count total parameters in a component."""
-    params_sizes = jax.tree.map(jax.numpy.size, component_params)
-    total_parameters = jax.tree.reduce(operator.add, params_sizes)
-    return total_parameters
-
-
-def count_parameters_by_component(params):
-    """Count parameters for each component of the model.
-
-    Args:
-        params: Model parameters from nnx.split(model, nnx.Param, ...)
-
-    Returns:
-        Dictionary with parameter counts for each component
-    """
-    component_names = list(params.keys())
-    print(f"Counting all components: {component_names}")
-
-    counts = {}
-    total_params = 0
-
-    for name in component_names:
-        component_params = params[name]
-        count = _count_component(component_params)
-        counts[name] = count
-        total_params += count
-
-    counts["total"] = total_params
-    return counts
-
-
-def bytes_to_gb(num_bytes):
-    return num_bytes / (1024**3)
-
-
-def print_compiled_memory_stats(compiled_stats):
-    """from: https://github.com/AI-Hypercomputer/maxtext/blob/b18829fbaa48aec7ac350a03e62248e24c6a76b2/MaxText/max_utils.py#L739"""
-    output_gb = bytes_to_gb(compiled_stats.output_size_in_bytes)
-    temp_gb = bytes_to_gb(compiled_stats.temp_size_in_bytes)
-    argument_gb = bytes_to_gb(compiled_stats.argument_size_in_bytes)
-    alias_gb = bytes_to_gb(compiled_stats.alias_size_in_bytes)
-    host_temp_gb = bytes_to_gb(compiled_stats.host_temp_size_in_bytes)
-    total_gb = output_gb + temp_gb + argument_gb - alias_gb
-    print(
-        f"Total memory size: {total_gb:.1f} GB, Output size: {output_gb:.1f} GB, Temp size: {temp_gb:.1f} GB, "
-        f"Argument size: {argument_gb:.1f} GB, Host temp size: {host_temp_gb:.1f} GB."
-    )
-
-
-def print_compiled_cost_analysis(cost_stats):
-    flops = float(cost_stats.get("flops", 0.0))
-    bytes_accessed = float(cost_stats.get("bytes accessed", 0.0))
-    gb = bytes_to_gb(bytes_accessed) if bytes_accessed else 0.0
-    intensity = (flops / bytes_accessed) if bytes_accessed else float("nan")
-    print(
-        f"FLOPs: {flops:.3e}, Bytes: {bytes_accessed:.3e} ({gb:.1f} GB), "
-        f"Intensity: {intensity:.1f} FLOPs/byte"
-    )
-
-
-def print_mem_stats(label: str):
-    """from: https://github.com/AI-Hypercomputer/maxtext/blob/7898576359bacde81be25cb3038e348aac1f943b/MaxText/max_utils.py#L713"""
-    print(f"
Memstats: {label}:")
-    try:
-        for d in jax.local_devices():
-            stats = d.memory_stats()
-            used = round(stats["bytes_in_use"] / 2**30, 2)
-            limit = round(stats["bytes_limit"] / 2**30, 2)
-            print(f"	Using (GB) {used} / {limit} ({used/limit:%}) on {d}")
-    except (RuntimeError, KeyError, TypeError) as ex:
-        print(f"	Memstats unavailable, error: {ex}")