[content] jasmine/train_dynamics.py (python)
    diffusion_use_ramp_weight: bool = True
    z_loss_weight: float = 0.0
    param_dtype = jnp.float32
    dtype = jnp.bfloat16
    use_flash_attention: bool = True
    use_gt_actions: bool = False
    # Logging
    log: bool = True
    entity: str = ""
    project: str = ""
    name: str = "train_dynamics"
    tags: list[str] = field(default_factory=lambda: ["dynamics"])
    log_interval: int = 50
    log_image_interval: int = 1000
    ckpt_dir: str = ""
    log_checkpoint_interval: int = 5000
    log_checkpoint_keep_period: int = 20_000
    log_gradients: bool = False
    val_data_dir: str = ""
    val_interval: int = 20_000
    val_steps: int = 50
    eval_full_frame: bool = True
    val_maskgit_steps: int = 25
    val_temperature: float = 1
    val_sample_argmax: bool = False
    wandb_id: str = ""


def build_model(args: Args, rng: jax.Array) -> tuple[Genie, jax.Array]:
    rng, _rng = jax.random.split(rng)
    rngs = nnx.Rngs(_rng)
    genie = Genie(
        # Tokenizer
        in_dim=args.image_channels,
        tokenizer_dim=args.tokenizer_dim,
        tokenizer_ffn_dim=args.tokenizer_ffn_dim,
        latent_patch_dim=args.latent_patch_dim,
        num_patch_latents=args.num_patch_latents,
        patch_size=args.patch_size,
        tokenizer_num_blocks=args.tokenizer_num_blocks,
        tokenizer_num_heads=args.tokenizer_num_heads,
        # LAM
        lam_dim=args.lam_dim,
        lam_ffn_dim=args.lam_ffn_dim,
        latent_action_dim=args.latent_action_dim,
        num_actions=args.num_actions,
        lam_patch_size=args.lam_patch_size,
        lam_num_blocks=args.lam_num_blocks,
        lam_num_heads=args.lam_num_heads,
        lam_co_train=not args.lam_checkpoint,
        use_gt_actions=args.use_gt_actions,
        # Dynamics
        dyna_type=args.dyna_type,
        dyna_dim=args.dyna_dim,
        dyna_ffn_dim=args.dyna_ffn_dim,
        dyna_num_blocks=args.dyna_num_blocks,
        dyna_num_heads=args.dyna_num_heads,
        dropout=args.dropout,
        mask_limit=args.mask_limit,
        diffusion_denoise_steps=args.diffusion_denoise_steps,
        param_dtype=args.param_dtype,
        dtype=args.dtype,
        use_flash_attention=args.use_flash_attention,
        decode=False,
        rngs=rngs,
    )
    if args.use_gt_actions:
        assert (
            not args.lam_checkpoint
        ), "Cannot use LAM when using ground-truth actions."
    else:
        assert genie.lam is not None
        del genie.lam.decoder
    return genie, rng


def build_optimizer(genie: Genie, args: Args) -> nnx.ModelAndOptimizer:
    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.param_dtype,  # moments in full precision
    )
    optimizer = nnx.ModelAndOptimizer(genie, tx)
    return optimizer


def build_mesh_and_sharding(
    num_devices: int,
) -> tuple[Mesh, NamedSharding, 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))
    actions_sharding = NamedSharding(mesh, PartitionSpec("data", None))
    return mesh, replicated_sharding, videos_sharding, actions_sharding


def shard_optimizer_states(
    optimizer: nnx.ModelAndOptimizer, replicated_sharding: NamedSharding
) -> None:
    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)


def build_dataloader(args: Args, data_dir: str) -> grain.DataLoaderIterator:
    image_shape = (args.image_height, args.image_width, args.image_channels)
    array_record_files = [
        os.path.join(data_dir, x)
        for x in os.listdir(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) -> Optional[ocp.CheckpointManager]:
    if args.restore_ckpt or args.save_ckpt:
        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),
        )
        if args.val_data_dir:
            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
                ),
            )
        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
    else:
        return None


def restore_or_initialize_components(
    args: Args,
    checkpoint_manager: Optional[ocp.CheckpointManager],
    optimizer: nnx.ModelAndOptimizer,
    train_iterator: grain.DataLoaderIterator,
    rng: jax.Array,
    replicated_sharding: NamedSharding,
    val_iterator: Optional[grain.DataLoaderIterator],
    restore_step: Optional[int] = None,
) -> tuple[
    int,
    nnx.ModelAndOptimizer,
    grain.DataLoaderIterator,
    grain.DataLoaderIterator,
    jax.Array,
]:
    step = 0
    if checkpoint_manager and restore_step is None:
        restore_step = checkpoint_manager.latest_step()
    if args.restore_ckpt:
        assert checkpoint_manager is not None
        abstract_optimizer = nnx.eval_shape(lambda: optimizer)
        abstract_optimizer_state = nnx.state(abstract_optimizer)
        if val_iterator:
            restore_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
            )
        else:
            restore_args = ocp.args.Composite(
                model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
                train_dataloader_state=grain.checkpoint.CheckpointRestore(train_iterator),  # type: ignore
            )
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(), args=restore_args
        )
        restored_optimizer_state = restored["model_state"]
        nnx.update(optimizer, restored_optimizer_state)
        train_iterator = restored["train_dataloader_state"]
        if val_iterator:
            val_iterator = restored["val_dataloader_state"]
        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)
    return step, optimizer, train_iterator, val_iterator, rng


def _calculate_top_k_accuracy(
    token_logits_BTNV: jax.Array,
    video_tokens_BTN: jax.Array,
    mask_BTN: jax.Array,
    k: int,
) -> jax.Array:
    _, topk_indices_BTNK = jax.lax.top_k(token_logits_BTNV, k)
    topk_correct = jnp.any(
        topk_indices_BTNK == video_tokens_BTN[..., jnp.newaxis], axis=-1
    )
    topk_acc = (mask_BTN * topk_correct).sum() / mask_BTN.sum()
    return topk_acc


def _calculate_step_metrics(
    outputs: dict[str, jax.Array],
    gt: jax.Array,
    num_actions: int,
    num_patch_latents: int,
) -> tuple[jax.Array, dict]:

    if args.dyna_type == "diffusion":
        if "x_pred" in outputs.keys():
            mse_BTHWC = (outputs["x_pred"] - outputs["x_gt"]) ** 2
            mse_BT = jnp.mean(mse_BTHWC, axis=(2, 3, 4))
            mse = jnp.mean(mse_BT)
            if args.diffusion_use_ramp_weight:
                ramp_weight = 0.9 * outputs["signal_level"] + 0.1
            else:
                ramp_weight = 1.0
            loss = jnp.mean(mse_BT * ramp_weight)
        else:
            mse = 0.0  # TODO mihir: fix this
            loss = 0.0
    else:
        mask_BTN = 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_BTN * ce_loss).sum() / mask_BTN.sum()
        z_val = jax.nn.logsumexp(outputs["token_logits"], axis=-1)
        z_loss_metric = (mask_BTN * (z_val**2)).sum() / mask_BTN.sum()

        masked_token_top_1_acc = _calculate_top_k_accuracy(
            outputs["token_logits"], outputs["video_tokens"], mask_BTN, 1
        )
        masked_token_top_2_acc = _calculate_top_k_accuracy(
            outputs["token_logits"], outputs["video_tokens"], mask_BTN, 2
        )
        masked_token_top_5_acc = _calculate_top_k_accuracy(
            outputs["token_logits"], outputs["video_tokens"], mask_BTN, 5
        )
        masked_token_top_16_acc = _calculate_top_k_accuracy(
            outputs["token_logits"], outputs["video_tokens"], mask_BTN, 16
        )

        select_probs = jax.nn.softmax(outputs["token_logits"])
        _, index_counts_tokenizer = jnp.unique_counts(
            jnp.ravel(outputs["video_tokens"]),
            size=num_patch_latents,
            fill_value=0,
        )
        codebook_usage_tokenizer = (index_counts_tokenizer != 0).mean()
    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()
    metrics = dict(
        psnr=psnr,
        ssim=ssim,
    )
    if args.dyna_type == "diffusion":
        metrics["mse"] = mse