p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/jobs/franz/berlin/coinrun/coinrun_lam/lam_coinrun_nan_investigation_z_loss.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1 
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/lam/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/lam/%x_%j.log
#SBATCH --job-name=train_lam_coinrun_nan_investigation_z_loss

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

tags="coinrun reproduction 10m_dataset lam lam_repro z_loss"

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m"

CHECKPOINT_DIR="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/lam/${job_name}_${slurm_job_id}"
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

srun python train_lam_full_prec.py \
    --save_ckpt \
    --wandb_id $SLURM_JOB_ID \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=36 \
    --image_height=64 \
    --image_width=64 \
    --max_lr=3e-5 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name="${job_name}_${slurm_job_id}" \
    --tags $tags \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir
[content] train_dynamics.py (python)
    # Logging
    log: bool = False
    entity: str = ""
    project: str = ""
    name: str = "train_dynamics"
    tags: list[str] = field(default_factory=lambda: ["dynamics"])
    log_interval: int = 5
    log_image_interval: int = 250
    ckpt_dir: str = ""
    log_checkpoint_interval: int = 25000
    log_checkpoint_keep_period: int = 20000
    log_gradients: bool = False
    wandb_id: str = ""


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)
[content] train_dynamics.py (python)
                    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)
                    comparison_seq = einops.rearrange(
                        comparison_seq * 255, "t h w c -> h (t w) c"
                    )
                    if jax.process_index() == 0:
                        log_images = dict(
                            image=wandb.Image(np.asarray(gt_seq[args.seq_len - 1])),
                            recon=wandb.Image(np.asarray(recon_seq[args.seq_len - 1])),
                            true_vs_recon=wandb.Image(
                                np.asarray(comparison_seq.astype(np.uint8))
                            ),
                        )
                        wandb.log(log_images)
[content] train_dynamics.py (python)
    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
    )
[content] train_dynamics.py (python)
# Auxiliary z-loss to stabilize logits: E[(logsumexp(logits))^2]
    log_z = jax.scipy.special.logsumexp(outputs["token_logits"], axis=-1)
    z_loss = (mask * jnp.square(log_z)).sum() / mask.sum()
    loss = ce_loss + args.z_loss_coeff * z_loss
[content] train_dynamics.py (python)
    # 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
    # )
[content] genie.py (python)
        # --- Random index path (bypass tokenizer and LAM) ---
        # Use optionally provided shape overrides; otherwise fall back to simple defaults
        BTN_shape = batch.get("BTN_shape")
        if BTN_shape is not None:
            B, T, N = map(int, BTN_shape)
        else:
            B, T, N = 1, 4, 64

        rng = batch["mask_rng"]
        rng, rng_tok, rng_act, rng_mask = jax.random.split(rng, 4)

        # Create random token indices in [0, num_patch_latents)
        token_indices_BTN = jax.random.randint(
            rng_tok,
            shape=(B, T, N),
            minval=0,
            maxval=self.num_latents,
            dtype=jnp.int32,
        )

        # Create random latent actions with shape (B, T-1, 1, latent_action_dim)
        latent_actions_BTm11L = jax.random.normal(
            rng_act,
            shape=(B, max(T - 1, 1), 1, self.latent_action_dim),
            dtype=self.dtype,
[content] genie.py (python)
        mle_indices_BTN = jnp.argmax(outputs["token_logits"], axis=-1)
        H, W = batch["videos"].shape[2:4]
        outputs["recon"] = self.tokenizer.decode(mle_indices_BTN, (H, W))
        outputs["lam_indices"] = action_indices_E
[content] genie.py (python)
        videos_BTHWC = batch["videos"]
        tokenizer_outputs = self.tokenizer.vq_encode(videos_BTHWC, training=False)
        token_indices_BTN = tokenizer_outputs["indices"]
        lam_outputs = self.lam.vq_encode(videos_BTHWC, training=False)
        z_q_BTm11L = lam_outputs["z_q"]
        action_indices_E = lam_outputs["indices"]
        latent_actions_BTm11L = jax.lax.cond(
            self.lam_co_train,
            lambda: z_q_BTm11L,
            lambda: jax.lax.stop_gradient(z_q_BTm11L),
[content] genie.py (python)
        # --- Random index path (bypass tokenizer and LAM) ---
        # Use optionally provided shape overrides; otherwise fall back to simple defaults
        BTN_shape = batch.get("BTN_shape")
        if BTN_shape is not None:
            B, T, N = map(int, BTN_shape)
        else:
            B, T, N = 1, 4, 64

        rng = batch["mask_rng"]
        rng, rng_tok, rng_act, rng_mask = jax.random.split(rng, 4)

        # Create random token indices in [0, num_patch_latents)
        token_indices_BTN = jax.random.randint(
            rng_tok,
            shape=(B, T, N),
            minval=0,
            maxval=self.num_patch_latents,
            dtype=jnp.int32,
        )

        # Create random latent actions with shape (B, T-1, 1, latent_action_dim)
        latent_actions_BTm11L = jax.random.normal(
            rng_act,
            shape=(B, max(T - 1, 1), 1, self.latent_action_dim),
            dtype=self.dtype,