p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] jasmine/train_dynamics.py (python)
    max_noise_level: float = 0.7
    noise_buckets: int = 10
    dropout: float = 0.0
    mask_limit: float = 0.5
    param_dtype = jnp.float32
    dtype = jnp.bfloat16
    use_flash_attention: bool = True
    use_gt_actions: bool = False
    # Logging
    log: bool = False
    entity: str = ""
    project: str = ""
    name: str = "train_dynamics"
    tags: list[str] = field(default_factory=lambda: ["dynamics"])
    log_interval: int = 5
    log_image_interval: int = 250
    ckpt_dir: str = ""
    log_checkpoint_interval: int = 25000
    log_checkpoint_keep_period: int = 20000
    log_gradients: bool = False
    val_data_dir: str = ""
    val_interval: int = 20_000
    val_steps: int = 50
    eval_full_frame: bool = False
    val_maskgit_steps: int = 25
    val_temperature: float = 1
    val_sample_argmax: bool = False
    wandb_id: str = ""


def build_model(args: Args, rng: jax.Array) -> tuple[Genie, jax.Array]:
    rng, _rng = jax.random.split(rng)
    rngs = nnx.Rngs(_rng)
    genie = Genie(
        # Tokenizer
        in_dim=args.image_channels,
        tokenizer_dim=args.tokenizer_dim,
        tokenizer_ffn_dim=args.tokenizer_ffn_dim,
        latent_patch_dim=args.latent_patch_dim,
        num_patch_latents=args.num_patch_latents,
        patch_size=args.patch_size,
        tokenizer_num_blocks=args.tokenizer_num_blocks,
        tokenizer_num_heads=args.tokenizer_num_heads,
        # LAM
        lam_dim=args.lam_dim,
        lam_ffn_dim=args.lam_ffn_dim,
        latent_action_dim=args.latent_action_dim,
        num_actions=args.num_actions,
        lam_patch_size=args.lam_patch_size,
        lam_num_blocks=args.lam_num_blocks,
        lam_num_heads=args.lam_num_heads,
        lam_co_train=not args.lam_checkpoint,
        use_gt_actions=args.use_gt_actions,
        # Dynamics
        dyna_type=args.dyna_type,
        dyna_dim=args.dyna_dim,
        dyna_ffn_dim=args.dyna_ffn_dim,
        dyna_num_blocks=args.dyna_num_blocks,
        dyna_num_heads=args.dyna_num_heads,
        dropout=args.dropout,
        max_noise_level=args.max_noise_level,
        noise_buckets=args.noise_buckets,
[content] jasmine/train_dynamics.py (python)
                batch=inputs,
                seq_len=args.seq_len,
                noise_level=0.0,
                temperature=args.val_temperature,
                sample_argmax=args.val_sample_argmax,
                maskgit_steps=args.val_maskgit_steps,
[content] slurm/jobs/franz/berlin/coinrun/sample.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=01:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/dynamics/sampling/maskgit/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/dynamics/sampling/maskgit/%x_%j.log
#SBATCH --job-name=coinrun_sample_maskgit

# Unload modules that may interfere
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4

# Activate virtual environment
source .venv/bin/activate

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_500m_seed_w_increment/val
CHECKPOINT_PATH=/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/dynamics/dynamics_coinrun_exposure_bias_maskgit_no_noise/29660


echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --checkpoint "$CHECKPOINT_PATH" \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=4 \
    --start_frame=1 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=maskgit
[tab] slurm/jobs/franz/berlin/coinrun/sample_dynamics_fully_trained_tokenizer.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=01:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/dynamics/sampling/maskgit/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/dynamics/sampling/maskgit/%x_%j.log
#SBATCH --job-name=coinrun_sample_maskgit

# Activate virtual environment
source .venv/bin/activate

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_500m_seed_w_increment/val"
CHECKPOINT_PATH="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/dynamics/dynamics_coinrun_exposure_bias_maskgit_no_noise/29660"


echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --checkpoint $CHECKPOINT_PATH \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=4 \
    --start_frame=1 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=maskgit
[tab] slurm/jobs/franz/berlin/coinrun/sample_dynamics_from_fully_trained_tokenizer.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=01:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/dynamics/sampling/maskgit/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/dynamics/sampling/maskgit/%x_%j.log
#SBATCH --job-name=coinrun_sample_maskgit

# Activate virtual environment
source .venv/bin/activate

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_500m_seed_w_increment/val"
CHECKPOINT_PATH="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/dynamics/dynamics_coinrun_exposure_bias_maskgit_no_noise/29660"


echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --checkpoint $CHECKPOINT_PATH \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=4 \
    --start_frame=1 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=maskgit
[content] slurm/jobs/franz/berlin/coinrun/sample_dynamics_from_140k_tokenizer.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=01:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/dynamics/sampling/maskgit/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/dynamics/sampling/maskgit/%x_%j.log
#SBATCH --job-name=coinrun_sample_maskgit

# Activate virtual environment
source .venv/bin/activate

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_500m_seed_w_increment/val"
CHECKPOINT_PATH="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/dynamics/dynamics_coinrun_exposure_bias_maskgit_no_noise/29660"


echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/sample.py \
    --checkpoint $CHECKPOINT_PATH \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=4 \
    --start_frame=1 \
    --image_height=64 \
    --image_width=64 \
    --dyna_type=maskgit