p-doom/crowd-code-dataset-0.2 huggingface dataset:
[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_coinrun

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

tokenizer_checkpoint=/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/coinrun_train_tokenizer_mae/32632


srun python jasmine/train_dynamics_diffusion.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --diffusion_denoise_steps=128 \
    --name=train_dyn_diffusion_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 \
    --tokenizer_checkpoint=$tokenizer_checkpoint \
[tab] slurm/jobs/mihir/berlin/diffusion/coinrun/ablations/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_coinrun

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

tokenizer_checkpoint=/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/coinrun_train_tokenizer_mae/32632


srun python jasmine/train_dynamics_diffusion.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --diffusion_denoise_steps=128 \
    --name=train_dyn_diffusion_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 \
    --tokenizer_checkpoint=$tokenizer_checkpoint \
[tab] slurm/jobs/mihir/berlin/diffusion/coinrun/ablations/dyn_dit_coinrun_lookup_embed.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_coinrun

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

tokenizer_checkpoint=/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/coinrun_train_tokenizer_mae/32632


srun python jasmine/train_dynamics_diffusion.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --diffusion_denoise_steps=128 \
    --name=train_dyn_diffusion_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 \
    --tokenizer_checkpoint=$tokenizer_checkpoint \
[tab] slurm/jobs/mihir/berlin/diffusion/coinrun/ablations/dyn_dit_coinrun_lookup_embed copy.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_coinrun

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

tokenizer_checkpoint=/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/coinrun_train_tokenizer_mae/32632

branch_name=$(git rev-parse --abbrev-ref HEAD)
if [[ "$branch_name" != "feat/diffusion-lookup-embed" ]]; then
  echo "Error: This script must be run from the 'feat/diffusion-lookup-embed' branch. Current branch: $branch_name"
  exit 1
fi

srun python jasmine/train_dynamics_diffusion.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --diffusion_denoise_steps=128 \
    --name=train_dyn_diffusion_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 \
    --tokenizer_checkpoint=$tokenizer_checkpoint \
[tab] slurm/jobs/mihir/berlin/diffusion/coinrun/ablations/dyn_dit_coinrun_prepend_conditioning.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_coinrun

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

tokenizer_checkpoint=/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/coinrun_train_tokenizer_mae/32632

branch_name=$(git rev-parse --abbrev-ref HEAD)
if [[ "$branch_name" != "feat/diffusion-lookup-embed" ]]; then
  echo "Error: This script must be run from the 'feat/diffusion-lookup-embed' branch. Current branch: $branch_name"
  exit 1
fi

srun python jasmine/train_dynamics_diffusion.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --diffusion_denoise_steps=128 \
    --name=train_dyn_diffusion_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 \
    --tokenizer_checkpoint=$tokenizer_checkpoint \
[tab] slurm/jobs/mihir/berlin/diffusion/coinrun/ablations/dyn_dit_coinrun_prepend_conditioning.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_coinrun

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

tokenizer_checkpoint=/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/coinrun_train_tokenizer_mae/32632

branch_name=$(git rev-parse --abbrev-ref HEAD)
if [[ "$branch_name" != "feat/diffusion-prepend-conditioning" ]]; then
  echo "Error: This script must be run from the 'feat/diffusion-prepend-conditioning' branch. Current branch: $branch_name"
  exit 1
fi

srun python jasmine/train_dynamics_diffusion.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --diffusion_denoise_steps=128 \
    --name=train_dyn_diffusion_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 \
    --tokenizer_checkpoint=$tokenizer_checkpoint \
[content] jasmine/genie.py (python)
            t = denoise_step / diffusion_steps

            # corrupt the context frames like in Dreamer 4 section 3.2
            t_vector = jnp.ones((B, seq_len)) * t_corrupt_context
            t_vector = t_vector.at[:, frame_i].set(t)

            noise_context_BSNL = jax.random.normal(
                _rng_noise_context, (B, seq_len, N, L)
            )
            currupted_latents_BSNL = (
                latents_BSNL * t_corrupt_context
                + (1 - t_corrupt_context) * noise_context_BSNL
            )
            frame_mask = jnp.arange(seq_len) < frame_i
            frame_mask_1S11 = frame_mask.reshape(1, seq_len, 1, 1)
            corrupted_tok_latents_BSNL = jnp.where(
                frame_mask_1S11, currupted_latents_BSNL, latents_BSNL
            )

            # --- Predict transition ---
            pred_latents_BSNL = dynamics_diffusion.diffusion_transformer(
                corrupted_tok_latents_BSNL, t_vector, act_embed_BSM
[tab] slurm/jobs/mihir/berlin/diffusion/coinrun/ablations/dyn_dit_coinrun_lookup_embed.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_coinrun

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

tokenizer_checkpoint=/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/coinrun_train_tokenizer_mae/32632

branch_name=$(git rev-parse --abbrev-ref HEAD)
if [[ "$branch_name" != "feat/diffusion-lookup-embed" ]]; then
  echo "Error: This script must be run from the 'feat/diffusion-lookup-embed' branch. Current branch: $branch_name"
  exit 1
fi

srun python jasmine/train_dynamics_diffusion.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --diffusion_denoise_steps=128 \
    --name=train_dyn_diffusion_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 \
    --tokenizer_checkpoint=$tokenizer_checkpoint \
[content] jasmine/train_dynamics_diffusion.py (python)
    gt_val = gt.clip(0, 1).reshape(-1, *gt.shape[2:])
    recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
    psnr = jnp.asarray(pix.psnr(gt_val, recon)).mean()
    ssim = jnp.asarray(pix.ssim(gt_val, recon)).mean()
    metrics = dict(
        psnr=psnr,
        ssim=ssim,
    )
[content] jasmine/train_tokenizer_vqvae.py (python)

        gt_clipped = gt.clip(0, 1).reshape(-1, *gt.shape[2:])
        recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
        psnr = jnp.asarray(pix.psnr(gt_clipped, recon)).mean()
        ssim = jnp.asarray(pix.ssim(gt_clipped, recon)).mean()
[content] jasmine/models/genie.py (python)
  tokenizer_outputs = self.tokenizer.vq_encode(videos_BTHWC, training=False)
            token_indices_BTN = tokenizer_outputs["indices"]
            outputs = dict(
                video_tokens=jax.lax.stop_gradient(token_indices_BTN),
                latent_actions=(
                    action_embeddings_BTm11L
                    if self.use_gt_actions
                    else latent_actions_BTm11L
                ),
            )
            outputs["mask_rng"] = batch["rng"]
            dyna_logits_BTNV, dyna_mask = self.dynamics(outputs)
            outputs["token_logits"] = dyna_logits_BTNV
            outputs["mask"] = dyna_mask
            mle_indices_BTN = jnp.argmax(outputs["token_logits"], axis=-1)
            outputs["recon"] = self.tokenizer.decode(mle_indices_BTN, (H, W))
[content] jasmine/models/genie.py (python)

                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,
                num_blocks=self.tokenizer_num_blocks,
                num_heads=self.tokenizer_num_heads,
                dropout=0.0,
                max_mask_ratio=0.0,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=rngs,
            )