p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/jobs/mihir/berlin/diffusion/breakout/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 \
    --dyna_type=diffusion \
    --diffusion_denoise_steps=128 \
    --name=train_dyn_dit_breakout_$slurm_job_id \
    --tags dyn diffusion breakout \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \
[tab] slurm/jobs/mihir/berlin/diffusion/coinrun/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=36 \
    --dyna_type=diffusion \
    --diffusion_denoise_steps=128 \
    --name=train_dyn_dit_coinrun_$slurm_job_id \
    --tags dyn diffusion coinrun \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \
[content] jasmine/genie.py (python)

            self.tokenizer = TokenizerMAE(
                in_dim=self.in_dim,
                model_dim=self.tokenizer_dim,
                ffn_dim=self.tokenizer_ffn_dim,
                latent_dim=self.latent_patch_dim,
                num_latents=self.num_patch_latents,
                patch_size=self.patch_size,
            )
[content] jasmine/genie.py (python)

                num_blocks=self.tokenizer_num_blocks,
                num_heads=self.tokenizer_num_heads,
                dropout=0.0,
                codebook_dropout=0.0,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=rngs,
[content] jasmine/models/dynamics.py (python)
t_full = t[:, :, None, None, None]  # [B, T, 1, 1, 1]
        x_1 = videos
        x_0 = jax.random.normal(noise_rng, (B, T, H, W, C))
        x_t = (1 - (1 - 1e-5) * t_full) * x_0 + t_full * x_1
[tab] slurm/jobs/mihir/berlin/diffusion/breakout/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] jasmine/genie.py (python)
 in_dim=args.image_channels,
        model_dim=args.tokenizer_dim,
        ffn_dim=args.tokenizer_ffn_dim,
        latent_dim=args.latent_patch_dim,
        num_latents=args.num_patch_latents,
        patch_size=args.patch_size,
        num_blocks=args.tokenizer_num_blocks,
        num_heads=args.tokenizer_num_heads,
        dropout=args.dropout,
        codebook_dropout=args.dropout,
        param_dtype=args.param_dtype,
        dtype=args.dtype,
        use_flash_attention=args.use_flash_attention,
        rngs=rngs_tokenizer,
    )
[content] jasmine/genie.py (python)
in_dim=args.image_channels,
            model_dim=args.tokenizer_dim,
            ffn_dim=args.tokenizer_ffn_dim,
            latent_dim=args.latent_patch_dim,
            num_latents=args.num_patch_latents,
            patch_size=args.patch_size,
            num_blocks=args.tokenizer_num_blocks,
            num_heads=args.tokenizer_num_heads,
[content] jasmine/genie.py (python)

            dropout=args.dropout,
            codebook_dropout=args.dropout,
            param_dtype=args.param_dtype,
            dtype=args.dtype,
            use_flash_attention=args.use_flash_attention,
            rngs=rngs_tokenizer,