p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] slurm/jobs/mihir/berlin/diffusion/train_dyn_dit.sh (shellscript)

    --save_ckpt \
    --image_height=64 \
    --image_width=64 \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=110 \
    --patch_size=4 \
    --init_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=50 \
    --log_checkpoint_interval=50 \
    --dyna_type=diffusion \
[content] slurm/jobs/mihir/berlin/diffusion/train_dyn_dit.sh (shellscript)

    --diffusion_denoise_steps=1000 \
    --z_loss_weight=1.0 \
    --name=train_dyn_dit \
    --tags dyn diffusion \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --tokenizer_checkpoint $tokenizer_checkpoint \
    --val_data_dir $array_records_dir_val \
    --log_interval 1 \
    --val_interval 50 \
    --eval_full_frame \
    --val_steps 5
[content] slurm/jobs/mihir/berlin/diffusion/train_dyn_dit.sh (shellscript)
array_records_dir_train=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_500m_gt_actions_split/train
array_records_dir_val=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_500m_gt_actions_split/val

tokenizer_checkpoint=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer/interactive/3583033
[content] jasmine/train_dynamics.py (python)

        masked_token_top1_accuracy=masked_token_top_1_acc,
        masked_token_top2_accuracy=masked_token_top_2_acc,
        masked_token_top5_accuracy=masked_token_top_5_acc,
        masked_token_top16_accuracy=masked_token_top_16_acc,
        select_logit=outputs["token_logits"].max(-1).mean(),
        select_p=select_probs.max(-1).mean(),
        entropy=jax.scipy.special.entr(select_probs).sum(-1).mean(),
        z_loss=z_loss_metric,
[content] slurm/jobs/mihir/berlin/diffusion/train_dyn_dit.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=00:40:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/berlin/diffusion/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/berlin/diffusion/%x_%j.log
#SBATCH --job-name=train_dyn_dit
#SBATCH --requeue
#SBATCH --signal=b:usr1@300 # 5 min before timeout
[content] jasmine/genie.py (python)
eq_len - T, N)
        pad = jnp.zeros(pad_shape, dtype=token_idxs_BTN.dtype)
        token_idxs_BSN = jnp.concatenate([token_idxs_BTN, pad], axis=1)
        logits_BSNV = jnp.zeros((*token_idxs_BSN.shape, self.num_patch_latents))
        dynamics_state = nnx.state(self.dynamics)
[content] jasmine/genie.py (python)
assert isinstance(self.dynamics, DynamicsDiffusion)
        # --- Encode videos and actions ---
        videos_BTHWC = batch["videos"]
        tokenizer_out = self.tokenizer.vq_encode(videos_BTHWC, training=False)
        token_idxs_BTN = tokenizer_out["indices"]
        B, T, N = token_idxs_BTN.shape