p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/jobs/mihir/berlin/diffusion/dyn_dit_breakout_dev.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --mem=100GB
#SBATCH --output=/fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/diffusion/%x_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/diffusion/%x_%j.log
#SBATCH --job-name=train_dyn_dit

cat $0

source .venv/bin/activate

array_records_dir_train=/fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m/train
array_records_dir_val=/fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m/val

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/fast/project/HFMI_ProFOUND/mihir/checkpoints/dyn/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

export PYTHONUNBUFFERED=1

srun python jasmine/train_dynamics.py \
    --save_ckpt \
    --image_height=10 \
    --image_width=10 \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=128 \
    --patch_size=4 \
    --log_interval 1 \
    --val_steps 1 \
    --val_interval 10 \
    --dyna_type=diffusion \
    --diffusion_denoise_steps=32 \
    --name=train_dyn_dit_breakout_dev \
    --tags dyn diffusion breakout \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \
[content] /fast/home/mihir.mahajan/Projects/jasmine/jasmine/models/dynamics.py (python)

        self.action_up = nnx.Linear(
            self.latent_action_dim,
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )
[content] jasmine/genie.py (python)

            act_embed_BTm11M = self.dynamics.action_up(latent_actions_BTm11L)
            padded_act_embed_BTM = jnp.pad(
                act_embed_BTm11M, ((0, 0), (1, 0), (0, 0), (0, 0))
            ).reshape(B, T, self.model_dim)
[content] jasmine/genie.py (python)
        if self.use_gt_actions:
            assert self.action_embed is not None
            latent_actions_BT1L = self.action_embed(batch["actions"]).reshape(
                *batch["actions"].shape[:2], 1, self.latent_action_dim
            )
            latent_actions_BTm11L = latent_actions_BT1L[:, :-1]
            action_tokens_EL = latent_actions_BTm11L.reshape(-1, self.latent_action_dim)
        else:
            assert self.lam is not None
            latent_actions_E = batch["latent_actions"]
            action_tokens_EL = self.lam.vq.get_codes(latent_actions_E)
[tab] slurm/jobs/mihir/berlin/diffusion/dyn_dit_breakout.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --mem=100GB
#SBATCH --output=/fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/diffusion/%x_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/diffusion/%x_%j.log
#SBATCH --job-name=train_dyn_dit

cat $0

source .venv/bin/activate

array_records_dir_train=/fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m/train
array_records_dir_val=/fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m/val

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/fast/project/HFMI_ProFOUND/mihir/checkpoints/dyn/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

PYTHONUNBUFFERED=1

srun python jasmine/train_dynamics.py \
    --save_ckpt \
    --image_height=10 \
    --image_width=10 \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=128 \
    --patch_size=4 \
    --val_interval 1000 \
    --dyna_type=diffusion \
    --diffusion_denoise_steps=64 \
    --name=train_dyn_dit_breakout \
    --tags dyn diffusion breakout \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \
[terminal_command] TERMINAL ()
                                                                                                                                                                      9                                                                                           8                                                                                           9                                                                                           7                                                                                                                                                                           [mihir.mahajan@hai-login2.haicore.berlin:~/Projects/jasmine] $ python Python 3.12.9 | packaged by conda-forge | (main, Mar  4 2025, 22:48:41) [GCC 13.3.0] on linuxType "help", "copyright", "credits" or "license" for more information.>>> 16*464>>> [mihir.mahajan@hai-login2.haicore.berlin:~/Projects/jasmine] $ queue[mihir.mahajan@hai-login2.haicore.berlin:~/Projects/jasmine] $ queue[mihir.mahajan@hai-login2.haicore.berlin:~/Projects/jasmine] $ scancel 32381[mihir.mahajan@hai-login2.haicore.berlin:~/Projects/jasmine] $ queue
[tab] slurm/jobs/mihir/berlin/diffusion/dyn_dit_coinrun.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --mem=100GB
#SBATCH --output=/fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/diffusion/%x_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/diffusion/%x_%j.log
#SBATCH --job-name=train_dyn_dit

cat $0

source .venv/bin/activate

array_records_dir_train=/fast/project/HFMI_ProFOUND/mihir/data/data_coinrun/coinrun_episodes_50m/train
array_records_dir_val=/fast/project/HFMI_ProFOUND/mihir/data/data_coinrun/coinrun_episodes_50m/val

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/fast/project/HFMI_ProFOUND/mihir/checkpoints/dyn/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

srun python jasmine/train_dynamics.py \
    --save_ckpt \
    --image_height=64 \
    --image_width=64 \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=32 \
    --val_interval 1000 \
    --dyna_type=diffusion \
    --diffusion_denoise_steps=64 \
    --name=train_dyn_dit_coinrun \
    --tags dyn diffusion coinrun \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \