p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] genie.py (python)
DynamicsMaskGIT(
                model_dim=self.dyna_dim,
                ffn_dim=self.dyna_ffn_dim,
                num_latents=self.num_patch_latents,
                latent_action_dim=self.latent_action_dim,
                num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
                dropout=self.dropout,
                mask_limit=self.mask_limit,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=rngs,
            )
[tab] slurm/jobs/mihir/horeka/coinrun/train_dyn_single_gpu.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=00:20:00
#SBATCH --partition=dev_accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/lam/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/lam/%x_%j.log
#SBATCH --job-name=train_lam_1e-4
#SBATCH --requeue
#SBATCH --signal=b:usr1@300 # 5 min before timeout

# Log the sbatch script
cat $0

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

# array_records_dir_train=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_test
array_records_dir_train=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_10m_split/train
# array_records_dir_train=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_10m_chunked
array_records_dir_val=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_10m_split/val

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dyn/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

lam_checkpoint=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/lam/interactive/3483568
tokenizer_checkpoint=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer/interactive/3483568

env | grep SLURM

export PYTHONUNBUFFERED=1

srun python train_dynamics.py \
    --save_ckpt \
    --image_height=64 \
    --image_width=64 \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=110 \
    --init_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=50 \
    --log_checkpoint_interval=50 \
    --dyna_type=maskgit \
    --log \
    --name=coinrun-dyn-dev-$slurm_job_id \
    --tags dyn coinrun dev \
    --entity instant-uv \
    --project jafar \
    --warmup_steps 0 \
    --wsd_decay_steps 0 \
    --num_steps 10 \
    --data_dir $array_records_dir_train \
    --tokenizer_checkpoint $tokenizer_checkpoint \
    --val_data_dir $array_records_dir_val \
    --val_interval 2 \
    --val_steps 5
[content] genie.py (python)

            # We need to reconstruct the submodule inside scan body to prevent trace context mismatches
            dynamics = DynamicsMaskGIT(
                model_dim=self.dyna_dim,
                ffn_dim=self.dyna_ffn_dim,
                num_latents=self.num_patch_latents,
                latent_action_dim=self.latent_action_dim,
                num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
                dropout=self.dropout,
                mask_limit=self.mask_limit,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=nnx.Rngs(0),
            )
            nnx.update(dynamics, dynamics_state)
[content] genie.py (python)
                model_dim=self.dyna_dim,
                ffn_dim=self.dyna_ffn_dim,
                num_latents=self.num_patch_latents,
                latent_action_dim=self.latent_action_dim,
                num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
                dropout=self.dropout,
                mask_limit=self.mask_limit,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=nnx.Rngs(0),
            )
[content] genie.py (python)
DynamicsCausal(
                model_dim=self.dyna_dim,
                ffn_dim=self.dyna_ffn_dim,
                num_latents=self.num_patch_latents,
                latent_action_dim=self.latent_action_dim,
                num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                decode=decode,
                rngs=rngs,
            )
[content] train_dynamics.py (python)

        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,
        )
[terminal_command] TERMINAL ()
salloc --time=01:00:00 --partition=dev_accelerated --nodes=1 --gres=gpu:1  --cpus-per-task=8salloc: Pending job allocation 3498224salloc: job 3498224 queued and waiting for resourcessalloc: job 3498224 has been allocated resourcessalloc: Granted job allocation 3498224salloc: Waiting for resource configurationsalloc: Nodes hkn0402 are ready for job[tum_cte0515@hkn0402 jasmine]$ salloc: Job 3498224 has exceeded its time limit and its allocation has been revoked.                                   srun: Job step aborted: Waiting up to 32 seconds for job step to finish.slurmstepd: error: *** STEP 3498224.interactive ON hkn0402 CANCELLED AT 2025-09-16T14:59:42 DUE TO TIME LIMIT ***srun: error: hkn0402: task 0: Killed
[content] train_dynamics.py (python)

            lam_indices = model.vq_encode(inputs, training=False)
            inputs["latent_actions"] = lam_indices
            inputs["videos"] = inputs["videos"][:, :-1]
            inputs["rng"] = inputs["mask_rng"]
            frames, video_tokens, logits = model.sample(
                inputs, args.seq_len  # TODO (mihir) add other args
            )
            outputs = {
                "recon": frames,
                "token_logits": logits,
                "video_tokens": video_tokens,
                "mask": jnp.zeros_like(video_tokens).at[:, -1].set(True),
                "lam_indices": lam_indices,
            }
[content] train_dynamics.py (python)

                                dict(
                                    val_image=wandb.Image(
                                        np.asarray(
                                            val_results["gt_seq_val"][args.seq_len - 1]
                                        )
                                    ),
                                    val_recon=wandb.Image(
                                        np.asarray(
                                            val_results["recon_seq_val"][
                                                args.seq_len - 1
                                            ]
                                        )
                                    ),
                                    val_true_vs_recon=wandb.Image(
                                        np.asarray(
                                            val_results["val_comparison_seq"].astype(
                                                np.uint8
                                            )
                                        )
                                    ),
                                    val_full_frame=wandb.Image(
                                        np.asarray(
                                            val_results["full_frame_seq_val"][
                                                args.seq_len - 1
                                            ]
                                        )
                                    ),
                                    val_true_vs_full_frame=wandb.Image(
                                        np.asarray(
                                            val_results[
                                                "val_full_frame_comparison_seq"
                                            ].astype(np.uint8)
                                        )
                                    ),
                                )
[tab] slurm/jobs/mihir/horeka/coinrun/sample_maskgit.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=01:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/dynamics/sampling/maskgit/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/dynamics/sampling/maskgit/%x_%j.log
#SBATCH --job-name=coinrun_sample_maskgit

# Unload modules that may interfere
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4

# Activate virtual environment
source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/array_records
CHECKPOINT_PATH=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/coinrun/maskgit/train_dynamics_maskgit_1_node/3415713


echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python sample.py \
    --checkpoint "$CHECKPOINT_PATH" \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=4 \
    --start_frame=1 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=maskgit