p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] models/dynamics.py (python)

            vid_embed_BTNM = self.patch_embed(video_tokens_BTN)
            act_embed_BTm11M = self.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
            )
[content] models/dynamics.py (python)

            vid_embed_BTNM = self.patch_embed(video_tokens_BTN)
            act_embed_BTm11M = self.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 = self.transformer(vid_embed_BTNp1M)
            step_logits_BV = step_logits_BTNp1V[:, -1, step_n, :]
            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
            )
[content] models/dynamics.py (python)

    def _eval_full_frame_pred():
        vid_embed_BTNM = self.patch_embed(video_tokens_BTN)
        act_embed_BTm11M = self.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 = self.transformer(vid_embed_BTNp1M)
        step_logits_BV = step_logits_BTNp1V[:, -1, step_n, :]
        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
        )
[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 \
    --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/train_dynamics.py (python)
    def val_step(genie: Genie, inputs: dict) -> tuple[jax.Array, jax.Array, dict, jax.Array, jax.Array, dict]:
        """Evaluate model and compute metrics"""
        genie.eval()
        (loss, (recon, metrics)) = dynamics_loss_fn(genie, inputs, training=False, pred_full_frame=False)
        (loss_full_frame, (recon_full_frame, metrics_full_frame)) = dynamics_loss_fn(genie, inputs, training=False, pred_full_frame=True)
        return loss, recon, metrics, loss_full_frame, recon_full_frame, metrics_full_frame


    def calculate_validation_metrics(val_dataloader, genie, rng):
        step = 0
        loss_per_step = []
        metrics_per_step = []
        loss_full_frame_per_step = []
        metrics_full_frame_per_step = []
        for videos in val_dataloader:
            rng, _rng_mask = jax.random.split(rng, 2)
            inputs = dict(videos=videos, mask_rng=_rng_mask)
            loss, recon, metrics, loss_full_frame, recon_full_frame, metrics_full_frame = val_step(genie, inputs)
            loss_per_step.append(loss)
            metrics_per_step.append(metrics)
            loss_full_frame_per_step.append(loss_full_frame)
            metrics_full_frame_per_step.append(metrics_full_frame)
            step += 1
            if step > args.val_steps:
                break

        if step < args.val_steps:
            print(f"Warning: Your validation dataset is too small to make val_steps many steps. Made {step} steps, expected {args.val_steps}")

        val_metrics = {
            f"val_{key}": np.mean([float(m[key]) for m in metrics_per_step])
            for key in metrics_per_step[0].keys()
        }
        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()
        }
        val_losses = {
            "val_loss": np.mean(loss_per_step),
            "val_loss_full_frame": np.mean(loss_full_frame_per_step)
[content] /hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jasmine/train_dynamics.py (python)
                params = nnx.state(optimizer.model)
                genie = build_model(args, rng, params)
                val_metrics, val_gt_batch, val_recon, val_full_frame = calculate_validation_metrics(dataloader_val, genie, _rng_mask_val)
[tab] slurm/jobs/mihir/horeka/preprocessing/preprocess_minecraft_multi_chunk.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=48:00:00
#SBATCH --partition=cpuonly
#SBATCH --cpus-per-task=16
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/preprocess/minecraft/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/preprocess/minecraft/%x_%j.log
#SBATCH --job-name=preprocess_minecraft_multi_chunk

source .venv/bin/activate

python input_pipeline/video_to_array_records.py \
        --input_path /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/tmp/vid \
        --output_path /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/tmp/multi_chunk \
        --env_name minecraft \
        --chunk_size 160 \
        --chunks_per_file 100