p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/jobs/mihir/horeka/jafar_og_reproduction/og_coinrun_dynamics_reproduction.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=2-00:00:00
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --partition=accelerated
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/jafar_og_reproduction/dynamics/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/jafar_og_reproduction/dynamics/%x_%j.log
#SBATCH --job-name=train_dynamics_coinrun_og_reproduction

# 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_og dynanmics 10m_dataset repro_mihir"

npy_records_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_10m"

# TODO mihir: update the tokenizer and lam checkpoints
tokenizer_ckpt_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/mihir/jafar_og_reproduction/tokenizer/train_tokenizer_coinrun_og_reproduction/3466286/tokenizer_1757013407_280000"
lam_ckpt_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/mihir/jafar_og_reproduction/lam/train_lam_coinrun_og_reproduction/3466287/lam_1757013407_200000"

CHECKPOINT_DIR="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/mihir/jafar_og_reproduction/dynamics/${job_name}/${slurm_job_id}"
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

srun python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --tokenizer_checkpoint="${tokenizer_ckpt_dir}" \
    --patch_size=16 \
    --lam_checkpoint="${lam_ckpt_dir}" \
    --log_image_interval=1000 \
    --log \
    --name="${job_name}_${slurm_job_id}" \
    --tags ${tags} \
    --entity instant-uv \
    --project jafar \
    --data_dir $npy_records_dir \
    --wandb_id $slurm_job_id
[tab] slurm/jobs/mihir/horeka/jafar_og_reproduction/og_coinrun_dynamics_reproduction-h100.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=2-00:00:00
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --partition=accelerated-h100
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/jafar_og_reproduction/dynamics/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/jafar_og_reproduction/dynamics/%x_%j.log
#SBATCH --job-name=train_dynamics_coinrun_og_reproduction

# 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_og dynanmics 10m_dataset repro_mihir"

npy_records_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_10m"

# TODO mihir: update the tokenizer and lam checkpoints
tokenizer_ckpt_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/mihir/jafar_og_reproduction/tokenizer/train_tokenizer_coinrun_og_reproduction/3466286"
lam_ckpt_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/mihir/jafar_og_reproduction/lam/train_lam_coinrun_og_reproduction/3466287"

CHECKPOINT_DIR="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/mihir/jafar_og_reproduction/dynamics/${job_name}/${slurm_job_id}"
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

srun python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --tokenizer_checkpoint="${tokenizer_ckpt_dir}" \
    --patch_size=16 \
    --lam_checkpoint="${lam_ckpt_dir}" \
    --log_image_interval=1000 \
    --log \
    --name="${job_name}_${slurm_job_id}" \
    --tags ${tags} \
    --entity instant-uv \
    --project jafar \
    --data_dir $npy_records_dir \
    --wandb_id $slurm_job_id
[tab] slurm/jobs/mihir/horeka/jafar_og_reproduction/og_coinrun_dynamics_reproduction-h100 copy.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=2-00:00:00
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --partition=accelerated-h100
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/jafar_og_reproduction/dynamics/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/jafar_og_reproduction/dynamics/%x_%j.log
#SBATCH --job-name=train_dynamics_coinrun_og_reproduction

# 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_og dynanmics 10m_dataset repro_mihir"

npy_records_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_10m"

# TODO mihir: update the tokenizer and lam checkpoints
tokenizer_ckpt_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/mihir/jafar_og_reproduction/tokenizer/train_tokenizer_coinrun_og_reproduction/3466286"
lam_ckpt_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/mihir/jafar_og_reproduction/lam/train_lam_coinrun_og_reproduction/3466287"

CHECKPOINT_DIR="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/mihir/jafar_og_reproduction/dynamics/${job_name}/${slurm_job_id}"
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

srun python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --tokenizer_checkpoint="${tokenizer_ckpt_dir}" \
    --patch_size=16 \
    --lam_checkpoint="${lam_ckpt_dir}" \
    --log_image_interval=1000 \
    --log \
    --name="${job_name}_${slurm_job_id}" \
    --tags ${tags} \
    --entity instant-uv \
    --project jafar \
    --data_dir $npy_records_dir \
    --wandb_id $slurm_job_id
[tab] slurm/jobs/mihir/horeka/jafar_og_reproduction/og_coinrun_dynamics_reproduction-sample.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=2-00:00:00
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --partition=accelerated-h100
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/jafar_og_reproduction/dynamics/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/jafar_og_reproduction/dynamics/%x_%j.log
#SBATCH --job-name=train_dynamics_coinrun_og_reproduction

# 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_og dynanmics 10m_dataset repro_mihir"

npy_records_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_10m"

# TODO mihir: update the tokenizer and lam checkpoints
tokenizer_ckpt_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/mihir/jafar_og_reproduction/tokenizer/train_tokenizer_coinrun_og_reproduction/3466286"
lam_ckpt_dir="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/mihir/jafar_og_reproduction/lam/train_lam_coinrun_og_reproduction/3466287"

CHECKPOINT_DIR="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/mihir/jafar_og_reproduction/dynamics/${job_name}/${slurm_job_id}"
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

srun python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --tokenizer_checkpoint="${tokenizer_ckpt_dir}" \
    --patch_size=16 \
    --lam_checkpoint="${lam_ckpt_dir}" \
    --log_image_interval=1000 \
    --log \
    --name="${job_name}_${slurm_job_id}" \
    --tags ${tags} \
    --entity instant-uv \
    --project jafar \
    --data_dir $npy_records_dir \
    --wandb_id $slurm_job_id
[tab] models/tokenizer.py (python)
from typing import Dict, Any, Tuple

import flax.linen as nn

from utils.preprocess import patchify, unpatchify
from utils.nn import STTransformer, VectorQuantizer


class TokenizerVQVAE(nn.Module):
    """ST-ViVit VQ-VAE"""

    in_dim: int
    model_dim: int
    latent_dim: int
    num_latents: int
    patch_size: int
    num_blocks: int
    num_heads: int
    dropout: float
    codebook_dropout: float

    def setup(self):
        self.encoder = STTransformer(
            self.model_dim,
            self.latent_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
        )
        self.vq = VectorQuantizer(
            self.latent_dim,
            self.num_latents,
            self.codebook_dropout,
        )
        self.out_dim = self.in_dim * self.patch_size**2
        self.decoder = STTransformer(
            self.model_dim,
            self.out_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        H, W = batch["videos"].shape[2:4]
        outputs = self.vq_encode(batch["videos"], training)
        recon = self.decoder(outputs["z_q"])  # (B, T, H_down * W_down, C)
        recon = nn.sigmoid(recon)
        outputs["recon"] = unpatchify(recon, self.patch_size, H, W)
        return outputs

    def vq_encode(self, videos: Any, training: bool = True) -> Dict[str, Any]:
        # --- Preprocess + encode ---
        B, T = videos.shape[:2]
        x = patchify(videos, self.patch_size)
        N = x.shape[2]
        x = self.encoder(x)  # (B, T, N, E)

        # --- Vector quantize ---
        x = x.reshape(B * T * N, self.latent_dim)
        z_q, z, emb, indices = self.vq(x, training)
        z_q = z_q.reshape(B, T, N, self.latent_dim)
        indices = indices.reshape(B, T, N)
        return dict(z_q=z_q, z=z, emb=emb, indices=indices)

    def decode(self, indices: Any, video_hw: Tuple[int]):
        z = self.vq.codebook[indices]
        recon = self.decoder(z)
        recon = nn.sigmoid(recon)
        return unpatchify(recon, self.patch_size, *video_hw)