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

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:1
#SBATCH --time=3-00:00:00
#SBATCH --cpus-per-task=8
#SBATCH --partition=gpu_p
#SBATCH --reservation=haicu_stefan
#SBATCH --qos=gpu_long
#SBATCH --output=/lustre/groups/haicu/workspace/alfred.nguyen/jafar_worksapce/logs/coinrun/dynamics/%x_%j.log
#SBATCH --error=/lustre/groups/haicu/workspace/alfred.nguyen/jafar_worksapce/logs/coinrun/dynamics/%x_%j.log
#SBATCH --job-name=train_dynamics_coinrun_og_reproduction

# Log the sbatch script
cat $0

source .venv/bin/activate

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

tags="coinrun_og dynanmics 10m_dataset helmholtz_reproduction dyn_repro"

npy_records_dir="/lustre/groups/haicu/workspace/alfred.nguyen/jafar_worksapce/data/coinrun_episodes/coinrun_episodes"

tokenizer_ckpt_dir="/lustre/groups/haicu/workspace/alfred.nguyen/jafar_worksapce/checkpoints/coinrun/tokenizer/train_tokenizer_coinrun_og_reproduction_28246778/tokenizer_1756303195_110000"
lam_ckpt_dir="/lustre/groups/haicu/workspace/alfred.nguyen/jafar_worksapce/checkpoints/coinrun/lam/train_lam_coinrun_og_reproduction_28246647/lam_1756303037_200000"
CHECKPOINT_DIR="/lustre/groups/haicu/workspace/alfred.nguyen/jafar_worksapce/checkpoints/coinrun/dynamics/${job_name}/${slurm_job_id}"
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

srun python jasmine/train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --tokenizer_checkpoint="${tokenizer_ckpt_dir}" \
    --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
[tab] slurm/dev/franz/berlin/coinrun/sample/maskgit/sample_mila_submission_case_study_vanilla.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --job-name=coinrun_sample_maskgit_mila_submission_case_study_vanilla

# Activate virtual environment
source .venv/bin/activate

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10M_npy_arr_rec/array_record/test"
CHECKPOINT_PATH="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/dynamics/dynamics_case_study_dataset_10M_30031"

current_branch=$(git rev-parse --abbrev-ref HEAD)
if [ "$current_branch" != "main" ]; then
    echo "This script must be run from the main branch. Current branch is $current_branch. Exiting."
    exit 1
fi

echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --seed=1 \
    --maskgit_steps=1 \
    --tokenizer_ffn_dim=512 \
    --tokenizer_num_blocks=8 \
    --dyna_ffn_dim=512 \
    --dyna_num_blocks=12 \
    --output_dir=gifs/dynamics_case_study_dataset_10M_vanilla \
    --checkpoint $CHECKPOINT_PATH \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=32 \
    --patch_size=4 \
    --start_frame=4 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=maskgit
[tab] slurm/dev/franz/berlin/coinrun/sample/maskgit/sample_mila_submission.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --job-name=coinrun_sample_maskgit_mila_submission

# Activate virtual environment
source .venv/bin/activate

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_500m_seed_w_increment/val"
CHECKPOINT_PATH="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/dynamics/dynamics_coinrun_mila_submission/29973"


echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --checkpoint $CHECKPOINT_PATH \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=4 \
    --patch_size=16 \
    --start_frame=4 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=maskgit
[tab] slurm/dev/franz/berlin/coinrun/sample/maskgit/sample_mila_submission_case_study_action_prepend.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/dynamics_sample/%x_%j.log
#SBATCH --job-name=coinrun_sample_maskgit_mila_submission_case_study_action_prepend

# Activate virtual environment
source .venv/bin/activate

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10M_npy_arr_rec/array_record/test"
CHECKPOINT_PATH="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/dynamics/dynamics_case_study_dataset_10M_action_prepend_30032"

current_branch=$(git rev-parse --abbrev-ref HEAD)
if [ "$current_branch" != "prepend-action-maskgit" ]; then
    echo "This script must be run from the prepend-action-maskgit branch. Current branch is $current_branch. Exiting."
    exit 1
fi

echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --seed=1 \
    --maskgit_steps=1 \
    --tokenizer_ffn_dim=512 \
    --tokenizer_num_blocks=8 \
    --dyna_ffn_dim=512 \
    --dyna_num_blocks=12 \
    --output_dir=gifs/dynamics_case_study_dataset_10M_action_prepend \
    --checkpoint $CHECKPOINT_PATH \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=32 \
    --patch_size=4 \
    --start_frame=4 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=maskgit
[tab] slurm/dev/franz/berlin/atari/sample/spawn_sampler.sh (shellscript)
#!/usr/bin/env bash

set -euo pipefail

REPO_ROOT="/home/franz.srambical/jafar"

DATA_ROOT="${DATA_ROOT:-/fast/project/HFMI_SynergyUnit/jafar_ws/data/atari}"

SAMPLE_SCRIPT="${SAMPLE_SCRIPT:-$REPO_ROOT/slurm/dev/franz/berlin/atari/sample/sample_atari.sh}"

if [ "$#" -gt 0 ]; then
    ENV_LIST=("$@")
else
    if [ ! -d "$DATA_ROOT" ]; then
        echo "DATA_ROOT does not exist: $DATA_ROOT" >&2
        exit 1
    fi
    mapfile -t ENV_LIST < <(find "$DATA_ROOT" -mindepth 1 -maxdepth 1 -type d -printf '%f
' | sort)
fi

echo "Submitting Sample jobs for environments: ${ENV_LIST[*]}"

for env in "${ENV_LIST[@]}"; do
    env_dir="$DATA_ROOT/$env"
    if [ ! -d "$env_dir/val" ]; then
        echo "Skipping $env (missing val under $env_dir)"
        continue
    fi

    job_name="sample_atari_${env}_maskgit"

    job_id=$(sbatch --parsable \
        --job-name="$job_name" \
        --chdir="$REPO_ROOT" \
        --export=ALL,ENV_NAME="$env" \
        "$SAMPLE_SCRIPT")

    echo "Submitted $job_name (ENV_NAME=$env) as job $job_id"
done
[tab] slurm/dev/franz/berlin/atari/sample/sample_atari.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/atari/sample/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/atari/sample/%x_%j.log
#SBATCH --job-name=atari_sample_maskgit

# Activate virtual environment
source .venv/bin/activate

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/atari/${ENV_NAME}/val/"

# TODO: demon_attack is still training
# TODO: add pong and chopper_command after fixing their tokenizer and finishing dynamics model training
declare -A TOKENIZER_CKPT_MAP=(
    [alien]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/alien/dynamics/dynamics_atari_alien_dev_fleuret_hparams_patch_size_4/30523"
    [amidar]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/amidar/dynamics/dynamics_atari_amidar_dev_fleuret_hparams_patch_size_4/30524"
    [assault]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/assault/dynamics/dynamics_atari_assault_dev_fleuret_hparams_patch_size_4/30525"
    [asterix]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/asterix/dynamics/dynamics_atari_asterix_dev_fleuret_hparams_patch_size_4/30526"
    [bank_heist]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/bank_heist/dynamics/dynamics_atari_bank_heist_dev_fleuret_hparams_patch_size_4/30527"
    [battle_zone]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/battle_zone/dynamics/dynamics_atari_battle_zone_dev_fleuret_hparams_patch_size_4/30528"
    [boxing]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/boxing/dynamics/dynamics_atari_boxing_dev_fleuret_hparams_patch_size_4/30529"
    [breakout]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/breakout/dynamics/dynamics_atari_breakout_dev_fleuret_hparams_patch_size_4/30530"
    [crazy_climber]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/crazy_climber/dynamics/dynamics_atari_crazy_climber_dev_fleuret_hparams_patch_size_4/30531"
    [demon_attack]=""
    [pong]=""
    [chopper_command]=""
)

# Allow explicit override via env var; otherwise use map
CHECKPOINT_PATH="${SAMPLE_CHECKPOINT:-${CHECKPOINT_PATH:-${TOKENIZER_CKPT_MAP[$ENV_NAME]}}}"

if [ -z "${ENV_NAME:-}" ]; then
    echo "ENV_NAME is not set; cannot locate data/checkpoints" >&2
    exit 1
fi

if [ -z "${CHECKPOINT_PATH}" ]; then
    echo "No checkpoint configured for ENV_NAME=$ENV_NAME. Set SAMPLE_CHECKPOINT or populate TOKENIZER_CKPT_MAP." >&2
    exit 1
fi

echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --checkpoint $CHECKPOINT_PATH \
    --num_patch_latents=512 \
    --latent_patch_dim=512 \
    --tokenizer_dim=256 \
    --tokenizer_ffn_dim=256 \
    --tokenizer_num_heads=4 \
    --patch_size=4 \
    --latent_action_dim=512 \
    --image_height=84 \
    --image_width=84 \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=4 \
    --start_frame=4 \
    --dyna_type=maskgit