[tab] slurm/dev/franz/berlin/atari/dynamics_fleuret_hparams_patch_size_4/dynamics.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/dynamics/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/atari/dynamics/%x_%j.log
#SBATCH --job-name=dynamics_atari_alien_dev_fleuret_hparams_patch_size_4
#SBATCH --requeue
#SBATCH --signal=b:usr1@300 # 5 min before timeout

# --- signal trap to requeue job before timeout ---
requeue_job() {
    echo "[$(date)] caught sigusr1 (timeout warning), requeueing slurm job $SLURM_JOB_ID..."
    # optional: trigger checkpoint saving here
    # e.g., touch $checkpoint_dir/requeue_trigger
    scontrol requeue $SLURM_JOB_ID
    exit 0
}

trap requeue_job sigusr1

# set checkpoint flag based on restart count
restart_count=$(scontrol show job $SLURM_JOB_ID | grep -o 'Restarts=[0-9]*' | cut -d'=' -f2)

if [ $restart_count -eq 0 ]; then
    restore_ckpt_flag="--no-restore-ckpt"
else
    restore_ckpt_flag="--restore-ckpt"
fi



# Log the sbatch script
cat $0

source .venv/bin/activate

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

# Require ENV_NAME to be provided via sbatch --export=ENV_NAME=...
if [ -z "${ENV_NAME:-}" ]; then
    echo "ENV_NAME must be provided (e.g., sbatch --export=ALL,ENV_NAME=alien $0)" >&2
    exit 2
fi

tags="atari ${ENV_NAME} dynamics dev fleuret_hparams patch_size_4"

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

# Map environments to tokenizer checkpoint directories (default if no override provided)
declare -A TOKENIZER_CKPT_MAP=(
    [alien]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/alien/tokenizer/tokenizer_atari_alien_dev_lr_3e-6_fleuret_hparams_patch_size_4_30280"
    [amidar]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/amidar/tokenizer/tokenizer_atari_amidar_dev_lr_3e-6_fleuret_hparams_patch_size_4_30281"
    [assault]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/assault/tokenizer/tokenizer_atari_assault_dev_lr_3e-6_fleuret_hparams_patch_size_4_30282"
    [asterix]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/asterix/tokenizer/tokenizer_atari_asterix_dev_lr_3e-6_fleuret_hparams_patch_size_4_30283"
    [bank_heist]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/bank_heist/tokenizer/tokenizer_atari_bank_heist_dev_lr_3e-6_fleuret_hparams_patch_size_4_30284"
    [battle_zone]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/battle_zone/tokenizer/tokenizer_atari_battle_zone_dev_lr_3e-6_fleuret_hparams_patch_size_4_30285"
    [boxing]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/boxing/tokenizer/tokenizer_atari_boxing_dev_lr_3e-6_fleuret_hparams_patch_size_4_30286"
    [breakout]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/breakout/tokenizer/tokenizer_atari_breakout_dev_lr_3e-6_fleuret_hparams_patch_size_4_30287"
    [crazy_climber]="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/crazy_climber/tokenizer/tokenizer_atari_crazy_climber_dev_lr_3e-6_fleuret_hparams_patch_size_4_30289"
)

# Allow explicit override via TOKENIZER_CKPT_DIR, else look up by ENV_NAME
if [ -n "${TOKENIZER_CKPT_DIR:-}" ]; then
    tokenizer_ckpt_dir="${TOKENIZER_CKPT_DIR}"
else
    if [ -n "${TOKENIZER_CKPT_MAP[$ENV_NAME]:-}" ]; then
        tokenizer_ckpt_dir="${TOKENIZER_CKPT_MAP[$ENV_NAME]}"
    else
        echo "No tokenizer checkpoint mapping found for ENV_NAME=$ENV_NAME. Provide TOKENIZER_CKPT_DIR or add an entry to TOKENIZER_CKPT_MAP." >&2
        exit 2
    fi
fi

CHECKPOINT_DIR="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/atari/${ENV_NAME}/dynamics/${job_name}/${slurm_job_id}"
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

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

srun python jasmine/train_dynamics.py \
    --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 \
    --save_ckpt \
    $restore_ckpt_flag \
    --wandb_id $SLURM_JOB_ID \
    --ckpt_dir $CHECKPOINT_DIR \
    --name="${job_name}_${slurm_job_id}" \
    --tags ${tags} \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint="${tokenizer_ckpt_dir}" \
    --val_data_dir="${array_records_dir}/val" \
    --data_dir="${array_records_dir}/train" &

child_pid=$!

wait $child_pid