p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/jobs/mihir/horeka/yolo-runs/tester.sh (shellscript)

# 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=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords_chunked
ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'

job_name="debug"
slurm_job_id="debug-mihir"

CHECKPOINT_DIR=$ws_dir/checkpoints/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

XLA_FLAGS=--xla_gpu_autotune_level=0 srun python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-debug-run-$slurm_job_id \
    --tags dynamics debug \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir
[tab] models/dynamics.py (python)
from typing import Dict, Any

import jax
import jax.numpy as jnp
import flax.linen as nn
import einops

from utils.nn import STTransformer, CausalTransformer


class DynamicsMaskGIT(nn.Module):
    """MaskGIT dynamics model"""

    model_dim: int
    num_latents: int
    num_blocks: int
    num_heads: int
    dropout: float
    mask_limit: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype

    def setup(self):
        self.dynamics = STTransformer(
            self.model_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
        )
        self.patch_embed = nn.Embed(self.num_latents, self.model_dim)
        self.mask_token = self.param(
            "mask_token",
            nn.initializers.lecun_uniform(),
            (1, 1, 1, self.model_dim),
        )
        self.action_up = nn.Dense(
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        # --- Mask videos ---
        vid_embed = self.patch_embed(batch["video_tokens"])
        if training:
            rng1, rng2 = jax.random.split(batch["mask_rng"])
            mask_prob = jax.random.uniform(rng1, minval=self.mask_limit)
            mask = jax.random.bernoulli(rng2, mask_prob, vid_embed.shape[:-1])
            mask = mask.at[:, 0].set(False)
            vid_embed = jnp.where(jnp.expand_dims(mask, -1), self.mask_token, vid_embed)
        else:
            mask = None

        # --- Predict transition ---
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        logits = self.dynamics(vid_embed)
        return dict(token_logits=logits, mask=mask)


class DynamicsAutoregressive(nn.Module):
    """Autoregressive (causal) dynamics model"""

    model_dim: int
    num_latents: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype

    def setup(self):
        self.dynamics = CausalTransformer(
            self.model_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
        )
        self.patch_embed = nn.Embed(self.num_latents, self.model_dim)
        self.action_up = nn.Dense(
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        vid_embed = self.patch_embed(batch["video_tokens"])
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        logits = self.dynamics(vid_embed)
        return dict(token_logits=logits, mask=None)
[content] genie.py (python)
        # self.dynamics = DynamicsMaskGIT(
            # model_dim=self.dyna_dim,
            # num_latents=self.num_patch_latents,
            # 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,
        # )
[content] genie.py (python)

        # self.dynamics = DynamicsMaskGIT(
            # model_dim=self.dyna_dim,
            # num_latents=self.num_patch_latents,
            # 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,
        # )
[content] models/dynamics.py (python)
        self.dynamics = STTransformer(
            self.model_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
        )
[tab] /hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/models/tokenizer.py (python)
from typing import Dict, Any, Tuple

import flax.linen as nn
import jax.numpy as jnp

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
    param_dtype: jnp.dtype
    dtype: jnp.dtype

    def setup(self):
        self.encoder = STTransformer(
            self.model_dim,
            self.latent_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
        )
        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,
            self.param_dtype,
            self.dtype,
        )

    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 = recon.astype(jnp.float32)
        recon = nn.sigmoid(recon)
        recon = recon.astype(self.dtype)
        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, int]):
        z = self.vq.codebook[indices]
        recon = self.decoder(z)
        recon = recon.astype(jnp.float32)
        recon = nn.sigmoid(recon)
        recon = recon.astype(self.dtype)
        return unpatchify(recon, self.patch_size, *video_hw)
[tab] slurm/jobs/mihir/horeka/yolo-runs/train_dynamics_yolorun.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24: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/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/%x_%j.log
#SBATCH --job-name=train_dyn_yolorun

# 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=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords_chunked
# tf_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/open_ai_minecraft_tfrecord

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

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

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --grad_clip_threshold=10 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-grad-norm-test-run-$slurm_job_id \
    --tags dynamics debug \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir
[content] slurm/jobs/mihir/horeka/yolo-runs/train_dynamics_yolorun.sbatch (shellscript)
srun python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-debug-run-$slurm_job_id \
    --tags dynamics debug \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir
[content] slurm/jobs/mihir/horeka/yolo-runs/train_dynamics_yolorun.sbatch (shellscript)
CHECKPOINT_DIR=$ws_dir/checkpoints/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4
[tab] diff.log (log)
71,72c71,72
<     name: str = "train_dynamics_maskgit"
<     tags: list[str] = field(default_factory=lambda: ["dynamics", "maskgit"])
---
>     name: str = "train_dynamics_causal"
>     tags: list[str] = field(default_factory=lambda: ["dynamics", "causal"])
94d93
<     mask = outputs["mask"]
95a95,97
>     logits = outputs["token_logits"][:, :, :-1]
>     targets = outputs["video_tokens"][:, :, 1:]
>     mask = outputs["mask"][:, :, 1:] 
97c99
<         outputs["token_logits"], outputs["video_tokens"]
---
>         logits, targets
100c102
<     acc = outputs["token_logits"].argmax(-1) == outputs["video_tokens"]
---
>     acc = logits.argmax(-1) == targets
102c104
<     select_probs = jax.nn.softmax(outputs["token_logits"])
---
>     select_probs = jax.nn.softmax(logits)
118c120
<         select_logit=outputs["token_logits"].max(-1).mean(),
---
>         select_logit=logits.max(-1).mean(),
[tab] models/dynamics.py (python)
from typing import Dict, Any

import jax
import jax.numpy as jnp
import flax.linen as nn

from utils.nn import STTransformer


class DynamicsMaskGIT(nn.Module):
    """MaskGIT dynamics model"""

    model_dim: int
    num_latents: int
    num_blocks: int
    num_heads: int
    dropout: float
    mask_limit: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype

    def setup(self):
        self.dynamics = STTransformer(
            self.model_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            spacial_bert=True,
        )
        self.patch_embed = nn.Embed(self.num_latents, self.model_dim)
        self.mask_token = self.param(
            "mask_token",
            nn.initializers.lecun_uniform(),
            (1, 1, 1, self.model_dim),
        )
        self.action_up = nn.Dense(
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        # --- Mask videos ---
        vid_embed = self.patch_embed(batch["video_tokens"])
        if training:
            rng1, rng2 = jax.random.split(batch["mask_rng"])
            mask_prob = jax.random.uniform(rng1, minval=self.mask_limit)
            mask = jax.random.bernoulli(rng2, mask_prob, vid_embed.shape[:-1])
            mask = mask.at[:, 0].set(False)
            vid_embed = jnp.where(jnp.expand_dims(mask, -1), self.mask_token, vid_embed)
        else:
            mask = None

        # --- Predict transition ---
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        logits = self.dynamics(vid_embed)
        return dict(token_logits=logits, mask=mask)


class DynamicsAutoregressive(nn.Module):
    """Autoregressive (causal) dynamics model"""

    model_dim: int
    num_latents: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype

    def setup(self):
        self.dynamics = STTransformer(
            self.model_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            spacial_bert=False,
        )
        self.patch_embed = nn.Embed(self.num_latents, self.model_dim)
        self.action_up = nn.Dense(
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        vid_embed = self.patch_embed(batch["video_tokens"])
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        logits = self.dynamics(vid_embed)
        mask = jnp.ones(vid_embed.shape[:-1])
        return dict(token_logits=logits, mask=mask)
[content] genie.py (python)
        """
        Autoregressively samples up to `seq_len` future frames using the causal transformer backend.

        - Input frames are tokenized once.
        - Future frames are generated one at a time, each conditioned on all previous frames.
        - All frames are detokenized in a single pass at the end.

        Args:
            batch: Dict with at least "videos" (B, T, H, W, C)
            seq_len: total number of frames to generate (including context)
            temperature: sampling temperature
            sample_argmax: if True, use argmax instead of sampling

        Returns:
            Generated video frames (B, seq_len, H, W, C)
        """
        # --- Encode context frames ---
        tokenizer_out = self.tokenizer.vq_encode(batch["videos"], training=False)
        token_idxs = tokenizer_out["indices"]  # (B, T, N)
        B, T, N = token_idxs.shape

        # --- Prepare initial token sequence ---
        # Pad with zeros for future frames
        pad_shape = (B, seq_len - T, N)
        token_idxs_full = jnp.concatenate(
            [token_idxs, jnp.zeros(pad_shape, dtype=token_idxs.dtype)], axis=1
        )  # (B, seq_len, N)

        # --- Prepare latent actions ---
        # If you have action tokens, use them; otherwise, use zeros
        if "latent_actions" in batch:
            action_tokens = self.lam.vq.get_codes(batch["latent_actions"])
        else:
            action_tokens = jnp.zeros((B, seq_len, self.num_latent_actions), dtype=jnp.int32)

        # --- Autoregressive generation loop ---
        rng = batch.get("rng", None)
        for t in range(T, seq_len):
            # Feed all tokens up to t (i.e., frames 0..t-1) to the dynamics model
            dyna_inputs = {
                "video_tokens": token_idxs_full[:, :t, :],  # (B, t, N)
                "latent_actions": action_tokens[:, :t, ...],  # (B, t, ...)
                "mask_rng": batch.get("mask_rng", None),
            }
            dyna_outputs = self.dynamics(dyna_inputs, training=False)
            # dyna_outputs["token_logits"]: (B, t, N, vocab_size)
            # We want the logits for the last time step (frame t-1 predicting t)
            next_token_logits = dyna_outputs["token_logits"][:, -1, :, :]  # (B, N, vocab_size)

            # Sample or argmax for each patch
            if sample_argmax:
                next_token = jnp.argmax(next_token_logits, axis=-1)  # (B, N)
            else:
                if rng is not None:
                    rng, step_rng = jax.random.split(rng)
                else:
                    step_rng = jax.random.PRNGKey(0)
                next_token = jax.random.categorical(
                    step_rng, next_token_logits / temperature, axis=-1
                )  # (B, N)

            # Insert the generated tokens into the sequence
            token_idxs_full = token_idxs_full.at[:, t, :].set(next_token)

        # --- Decode all tokens at once at the end ---
        final_frames = self.tokenizer.decode(
            token_idxs_full, video_hw=batch["videos"].shape[2:4]
        )
        return final_frames