[content] train_lam.py (python)
# ==============================
# Notebook-friendly helper API
# ==============================

@dataclass
class LamConfig:
    seed: int = 0
    seq_len: int = 16
    image_channels: int = 3
    image_height: int = 90
    image_width: int = 160
    batch_size: int = 36
    data_dir: str = ""
    ckpt_dir: str = ""

    # Optimization
    vq_beta: float = 0.25
    init_lr: float = 0.0
    max_lr: float = 3e-5
    decay_end: float = 0.0
    wsd_decay_steps: int = 10000
    warmup_steps: int = 5000
    lr_schedule: str = "wsd"
    vq_reset_thresh: int = 50

    # Model
    model_dim: int = 512
    ffn_dim: int = 2048
    latent_dim: int = 32
    num_latents: int = 6
    patch_size: int = 16
    num_blocks: int = 4
    num_heads: int = 8
    dropout: float = 0.0
    codebook_dropout: float = 0.0
    param_dtype: jnp.dtype = jnp.float32
    dtype: jnp.dtype = jnp.bfloat16
    use_flash_attention: bool = True


class LamSession:
    def __init__(self, config: LamConfig, build_dataloader: bool = True, build_checkpoint_manager: bool = True) -> None:
        self.config = config
        self.rng = jax.random.key(config.seed)

        # Model
        self.rng, _rng = jax.random.split(self.rng)
        rngs = nnx.Rngs(_rng)
        self.model = LatentActionModel(
            in_dim=config.image_channels,
            model_dim=config.model_dim,
            ffn_dim=config.ffn_dim,
            latent_dim=config.latent_dim,
            num_latents=config.num_latents,
            patch_size=config.patch_size,
            num_blocks=config.num_blocks,
            num_heads=config.num_heads,
            dropout=config.dropout,
            codebook_dropout=config.codebook_dropout,
            param_dtype=config.param_dtype,
            dtype=config.dtype,
            use_flash_attention=config.use_flash_attention,
            rngs=rngs,
        )

        # Optimizer and schedule
        # Schedule type conforms to optax schedule; annotate loosely to satisfy linter
        self.lr_schedule = get_lr_schedule(
            config.lr_schedule,
            config.init_lr,
            config.max_lr,
            config.decay_end,
            1_000_000,
            config.warmup_steps,
            config.wsd_decay_steps,
        )
        tx = optax.adamw(
            learning_rate=self.lr_schedule,
            b1=0.9,
            b2=0.9,
            weight_decay=1e-4,
            mu_dtype=config.param_dtype,
        )
        self.optimizer = nnx.Optimizer(self.model, tx)

        # Data
        self.loader_iterator: Optional[grain.DataLoaderIterator] = None
        if build_dataloader and config.data_dir:
            self.loader_iterator = self._build_dataloader()

        # Checkpointing
        self.ckpt_mgr: Optional[ocp.CheckpointManager] = None
        if build_checkpoint_manager and config.ckpt_dir:
            self.ckpt_mgr = self._build_checkpoint_manager(config.ckpt_dir)

        self.step = 0
        self.action_last_active = jnp.zeros(config.num_latents, dtype=jnp.int32)

    def _build_dataloader(self):
        image_shape = (self.config.image_height, self.config.image_width, self.config.image_channels)
        array_record_files = [
            os.path.join(self.config.data_dir, x)
            for x in os.listdir(self.config.data_dir)
            if x.endswith(".array_record")
        ]
        dataloader = get_dataloader(
            array_record_files,
            self.config.seq_len,
            self.config.batch_size,
            *image_shape,
            num_workers=8,
            prefetch_buffer_size=1,
            seed=self.config.seed,
        )
        initial_state = dataloader._create_initial_state()
        return grain.DataLoaderIterator(dataloader, initial_state)

    def _build_checkpoint_manager(self, ckpt_dir: str):
        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),
        )
        options = ocp.CheckpointManagerOptions(step_format_fixed_length=6)
        return ocp.CheckpointManager(directory=ckpt_dir, options=options, handler_registry=handler_registry)

    def restore_latest(self) -> int:
        if self.ckpt_mgr is None:
            raise ValueError("Checkpoint manager not initialized; set ckpt_dir in config.")
        if self.loader_iterator is None:
            raise ValueError("Dataloader not initialized; set data_dir in config.")
        abstract_optimizer = nnx.eval_shape(lambda: self.optimizer)
        abstract_optimizer_state = nnx.state(abstract_optimizer)
        restored = self.ckpt_mgr.restore(
            self.ckpt_mgr.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
                dataloader_state=grain.checkpoint.CheckpointRestore(self.loader_iterator),  # type: ignore
            ),
        )
        nnx.update(self.optimizer, restored["model_state"])  # type: ignore
        self.loader_iterator = restored["dataloader_state"]
        self.step = self.ckpt_mgr.latest_step() or 0
        return self.step

    def save(self) -> None:
        if self.ckpt_mgr is None:
            raise ValueError("Checkpoint manager not initialized; set ckpt_dir in config.")
        if self.loader_iterator is None:
            raise ValueError("Dataloader not initialized; set data_dir in config.")
        optimizer_state = nnx.state(self.optimizer)
        self.ckpt_mgr.save(
            self.step,
            args=ocp.args.Composite(
                model_state=ocp.args.PyTreeSave(optimizer_state),  # type: ignore
                dataloader_state=grain.checkpoint.CheckpointSave(self.loader_iterator),  # type: ignore
            ),
        )

    def next_batch(self) -> Dict[str, jax.Array]:
        if self.loader_iterator is None:
            raise ValueError("Dataloader not initialized; set data_dir in config.")
        videos = next(self.loader_iterator)
        gt = jnp.asarray(videos, dtype=jnp.float32) / 255.0
        videos = gt.astype(self.config.dtype)
        self.rng, _rng = jax.random.split(self.rng)
        return {"videos": videos, "rng": _rng}

    def enable_sowing(self) -> None:
        for model in [self.model.encoder, self.model.decoder]:
            model.sow_logits = True
            for blk in getattr(model, "blocks", []):
                blk.sow_weights = True
                blk.sow_activations = True

    def forward(self, batch: Dict[str, jax.Array], training: bool = True) -> Dict[str, jax.Array]:
        return self.model(batch, training=training)

    def _loss(self, model: LatentActionModel, inputs: Dict[str, jax.Array]):
        gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
        inputs = dict(inputs)
        inputs["videos"] = gt.astype(self.config.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 + self.config.vq_beta * commitment_loss

        gt_img = gt_future_frames.clip(0, 1).reshape(-1, *gt_future_frames.shape[2:])
        recon_img = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
        psnr = jnp.asarray(pix.psnr(gt_img, recon_img)).mean()
        ssim = jnp.asarray(pix.ssim(gt_img, recon_img)).mean()
        count_fn = jax.vmap(lambda i: (outputs["indices"] == i).sum())
        index_counts = count_fn(jnp.arange(self.config.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)

    def train_step(self, batch: Dict[str, jax.Array]):
        def loss_fn(model: LatentActionModel):
            return self._loss(model, batch)

        (loss, (recon, idx_counts, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(self.model)
        self.optimizer.update(grads)

        codebook = self.model.vq.codebook
        num_codes = len(codebook)
        active_codes = idx_counts != 0.0
        self.action_last_active = jnp.where(active_codes, 0, self.action_last_active + 1)
        denom = jnp.maximum(active_codes.sum(), 1)
        p_code = active_codes / denom
        self.rng, choice_rng = jax.random.split(self.rng)
        reset_idxs = jax.random.choice(choice_rng, num_codes, shape=(num_codes,), p=p_code)
        do_reset = self.action_last_active >= self.config.vq_reset_thresh
        new_codebook = jnp.where(jnp.expand_dims(do_reset, -1), codebook[reset_idxs], codebook.value)
        self.model.vq.codebook.value = new_codebook
        self.action_last_active = jnp.where(do_reset, 0, self.action_last_active)

        metrics = dict(metrics)
        metrics["lr"] = self.lr_schedule(self.step)
        self.step += 1
        return loss, recon, metrics