p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] models/dynamics.py (python)
video_tokens_BTN, jnp.zeros((
                    video_tokens_BTN.shape[0], 
                    video_tokens_BTN.shape[1],
                    video_tokens_BTN.shape[2], 
                    self.num_latents))
[content] models/dynamics.py (python)
(_, final_logits_BTNV), _ = jax.lax.scan(
   229|                _pred_full_frame,
   230|                (video_tokens_BTN, jnp.zeros((
   231|                    video_tokens_BTN.shape[0],
   232|                    video_tokens_BTN.shape[1],
   233|                    video_tokens_BTN.shape[2],
   234|                    self.num_latents))),
   235|                jnp.arange(video_tokens_BTN.shape[1])
   236|            )
[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

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=250 \
    --log_checkpoint_interval=250 \
    --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 50 \
    --data_dir $array_records_dir_train \
    --tokenizer_checkpoint $tokenizer_checkpoint \
    --val_data_dir $array_records_dir_val \
    --val_interval 10 \
    --val_steps 5
[content] /hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jasmine/models/dynamics.py (python)
                video_tokens_BTN, final_logits_BTNV, patch_embed, action_up, transformer = carry
                vid_embed_BTNM = patch_embed(video_tokens_BTN)
                act_embed_BTm11M = action_up(latent_actions_BTm11L)
                padded_act_embed_BT1M = jnp.pad(
                    act_embed_BTm11M, ((0, 0), (1, 0), (0, 0), (0, 0))
                )
                vid_embed_BTNp1M = jnp.concatenate(
                    [padded_act_embed_BT1M, vid_embed_BTNM], axis=2
                )
                step_logits_BTNp1V = transformer(vid_embed_BTNp1M)
                step_logits_BV = step_logits_BTNp1V[:, -1, step_n, :]
                final_logits_BTNV = final_logits_BTNV.at[:, -1, step_n].set(step_logits_BV)
                sampled_token_idxs_B = jnp.argmax(step_logits_BV, axis=-1)
                video_tokens_BTN = video_tokens_BTN.at[:, -1, step_n].set(
                    sampled_token_idxs_B
                )
                return (video_tokens_BTN, final_logits_BTNV, patch_embed, action_up, transformer), None
            carry = (video_tokens_BTN, jnp.zeros((
                    video_tokens_BTN.shape[0], 
                    video_tokens_BTN.shape[1],
                    video_tokens_BTN.shape[2], 
                    self.num_latents)),
                    self.patch_embed,
                    self.action_up,
                    self.transformer
                    )
            (_, final_logits_BTNV, _, _, _), _ = jax.lax.scan(
                 _pred_full_frame, carry, jnp.arange(video_tokens_BTN.shape[1])
                 )
[content] /hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jasmine/models/dynamics.py (python)

        self.transformer = Transformer(
            self.model_dim,
            self.model_dim,
            self.ffn_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
            decode=decode,
            rngs=rngs,
        )
        self.patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=rngs)
        self.action_up = nnx.Linear(
            self.latent_action_dim,
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

    def __call__(
        self, batch: Dict[str, jax.Array], training: bool = True, pred_full_frame: bool = False,
    ) -> tuple[jax.Array, jax.Array | None]:
        assert not (training and pred_full_frame), "Cannot evaluate full frame prediction during training."
        video_tokens_BTN = batch["video_tokens"]
        latent_actions_BTm11L = batch["latent_actions"]
        if pred_full_frame:
            # --- Extract submodule states ---
            patch_embed_state = nnx.state(self.patch_embed)
            action_up_state = nnx.state(self.action_up)
            transformer_state = nnx.state(self.transformer)
            model_dim = self.model_dim
            latent_action_dim = self.latent_action_dim
            num_latents = self.num_latents
            ffn_dim = self.ffn_dim
            num_blocks = self.num_blocks
            num_heads = self.num_heads
            dropout = self.dropout
            param_dtype = self.param_dtype
            dtype = self.dtype
            use_flash_attention = self.use_flash_attention
            decode = True  # For full frame prediction

            def _pred_full_frame(carry, step_n):
                rngs = jax.random.key(0)
[content] /hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jasmine/models/dynamics.py (python)
                #nnx.update(patch_embed, patch_embed_state)
                action_up = nnx.Linear(
                    latent_action_dim, model_dim, param_dtype=param_dtype, dtype=dtype, rngs=rngs
                )
                # nnx.update(action_up, action_up_state)
                transformer = Transformer(
                    model_dim, model_dim, ffn_dim, num_latents, num_blocks, num_heads,
                    dropout, param_dtype, dtype, use_flash_attention=use_flash_attention,
                    decode=decode, rngs=rngs
                )
                # nnx.update(transformer, transformer_state)
[content] /hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jasmine/models/dynamics.py (python)
                video_tokens_BTN, final_logits_BTNV = carry
                # Reconstruct submodules inside scan body
                patch_embed = nnx.Embed(num_latents, model_dim, rngs=nnx.Rngs(0))
                nnx.update(patch_embed, patch_embed_state)
                action_up = nnx.Linear(
                    latent_action_dim, model_dim, param_dtype=param_dtype, dtype=dtype, rngs=nnx.Rngs(0)
                )
                nnx.update(action_up, action_up_state)
                transformer = Transformer(
                    model_dim, model_dim, ffn_dim, num_latents, num_blocks, num_heads,
                    dropout, param_dtype, dtype, use_flash_attention=use_flash_attention,
                    decode=decode, rngs=nnx.Rngs(0)
                )
                nnx.update(transformer, transformer_state)
[content] train_dynamics.py (python)

        val_metrics_full_frame = {
            f"val_full_frame_{key}": np.mean([float(m[key]) for m in metrics_full_frame_per_step])
            for key in metrics_full_frame_per_step[0].keys()
        }
[content] train_dynamics.py (python)
metrics_full_frame = {
            f"val_full_frame_{key}": np.mean([float(m[key]) for m in metrics_full_frame_per_step])
            for key in metrics_full_frame_per_step[0].keys()
        }
        val_
[content] train_dynamics.py (python)

                                    val_full_frame=wandb.Image(np.asarray(full_frame_seq_val[args.seq_len - 1])),
                                    val_true_vs_full_frame=wandb.Image(
                                        np.asarray(val_full_frame_comparison_seq.astype(np.uint8))
                                    )
[tab] scripts_cremers/train_tokenizer_overfit_batch_coinrun.sbatch (shellscript)
#!/bin/bash
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=2
#SBATCH --time=5:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:2,VRAM:12G
#SBATCH --mem=50G
#SBATCH --output=logs/logs_training/%x_%j.log
#SBATCH --error=logs/logs_training/%x_%j.log
#SBATCH --job-name=train_tokenizer_coinrun_overfit_batch

# Log the sbatch script
cat $0

tf_records_dir="/usr/stud/mahajanm/Projects/jafar/data_tfrecords"
ws_dir='/storage/user/mahajanm/Projects/world-modeling'
job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=$ws_dir/checkpoints/$job_name_$slurm_job_id
echo $CHECKPOINT_DIR
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

SLURM_STEP_NODELIST=$SLURM_NODELIST srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=6 \
    --image_width=64 \
    --image_height=64 \
    --min_lr=1e-4 \
    --max_lr=1e-4 \
    --log_image_interval=3 \
    --log \
    --entity instant-uv \
    --project jafar \
    --data_dir $tf_records_dir
[content] /hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jasmine/models/dynamics.py (python)

            def _pred_full_frame(carry, step_n):
                video_tokens_BTN, final_logits_BTNV = carry
                # We need to reconstruct submodules inside scan body to prevent trace context mismatches
                patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=nnx.Rngs(0))
                nnx.update(patch_embed, patch_embed_state)
                action_up = nnx.Linear(
                    self.latent_action_dim, self.model_dim, param_dtype=self.param_dtype, dtype=self.dtype, rngs=nnx.Rngs(0)
                )
                nnx.update(action_up, action_up_state)
                transformer = Transformer(
                    self.model_dim, self.model_dim, self.ffn_dim, self.num_latents, self.num_blocks, self.num_heads,
                    self.dropout, self.param_dtype, self.dtype, use_flash_attention=self.use_flash_attention,
                    decode=self.decode, rngs=nnx.Rngs(0)
                )
                nnx.update(transformer, transformer_state)

                vid_embed_BTNM = patch_embed(video_tokens_BTN)
                act_embed_BTm11M = action_up(latent_actions_BTm11L)
                padded_act_embed_BT1M = jnp.pad(
                    act_embed_BTm11M, ((0, 0), (1, 0), (0, 0), (0, 0))
                )
                vid_embed_BTNp1M = jnp.concatenate(
                    [padded_act_embed_BT1M, vid_embed_BTNM], axis=2
                )
                step_logits_BTNp1V = transformer(vid_embed_BTNp1M)
                step_logits_BV = step_logits_BTNp1V[:, -1, step_n, :]
                final_logits_BTNV = final_logits_BTNV.at[:, -1, step_n].set(step_logits_BV)
                sampled_token_idxs_B = jnp.argmax(step_logits_BV, axis=-1)
                video_tokens_BTN = video_tokens_BTN.at[:, -1, step_n].set(
                    sampled_token_idxs_B
                )
                return (video_tokens_BTN, final_logits_BTNV), None

            (_, final_logits_BTNV), _ = jax.lax.scan(
                _pred_full_frame,
                (video_tokens_BTN, jnp.zeros((
                    video_tokens_BTN.shape[0],
                    video_tokens_BTN.shape[1],
                    video_tokens_BTN.shape[2],
                    self.num_latents))),
                jnp.arange(video_tokens_BTN.shape[2])