p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] jasmine/utils/nn.py (python)
        self.time_step_embedder = nnx.Embed(
            num_embeddings=self.diffusion_steps,
            features=self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )
[content] jasmine/models/dynamics.py (python)
        self.time_step_embedder = nnx.Embed(
            num_embeddings=self.diffusion_steps,
            features=self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )
[content] 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/models/genie.py (python)

        action_tokens_BSm11L = jnp.reshape(action_tokens_EL, (B, seq_len - 1, 1, -1))
        act_embed_BSm11L = self.dynamics.action_up(action_tokens_BSm11L)
        act_embed_BS1L = jnp.pad(act_embed_BSm11L, ((0, 0), (1, 0), (0, 0), (0, 0)))
[content] jasmine/models/genie.py (python)
        t_corrupt_context = 1 - diffusion_corrupt_context_factor
        t_corrupt_context = jnp.argmin(
            jnp.abs(jnp.arange(diffusion_steps) / diffusion_steps - t_corrupt_context)
        )
        t_corrupt_context = t_corrupt_context / diffusion_steps

        @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)
        def denoise_step_fn(
            carry: tuple[jax.Array, jax.Array, jax.Array], denoise_step: jax.Array
        ) -> tuple[jax.Array, jax.Array, jax.Array]:
            latents_BSNL, frame_i, rng = carry
            rng, _rng_noise_context = jax.random.split(rng)

            # We need to reconstruct the submodule inside scan body to prevent trace context mismatches
            dynamics_diffusion = DynamicsDiffusion(
                model_dim=self.dyna_dim,
                ffn_dim=self.dyna_ffn_dim,
                latent_patch_dim=self.latent_patch_dim,
                latent_action_dim=self.latent_action_dim,
                num_blocks=self.dyna_num_blocks,
                num_heads=self.dyna_num_heads,
                denoise_steps=self.diffusion_denoise_steps,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                use_flash_attention=self.use_flash_attention,
                rngs=nnx.Rngs(0),
                decode=self.decode,
            )
            nnx.update(dynamics_diffusion, dynamics_state)

            # corrupt the context frames like in Dreamer 4 section 3.2
            denoise_step_BS = jnp.ones((B, seq_len)) * diffusion_steps - 1
            denoise_step_BS = denoise_step_BS.at[:, frame_i].set(denoise_step)

            noise_context_BSNL = jax.random.normal(
                _rng_noise_context, (B, seq_len, N, L)
            )
            corrupted_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, corrupted_latents_BSNL, latents_BSNL
            )

            # --- Predict transition ---
            action_tokens_BSm11L = jnp.reshape(
                action_tokens_EL, (B, seq_len - 1, 1, -1)
            )
[content] README.md (markdown)
Genie has three components: a [video tokenizer](models/tokenizer.py), a [latent action model](models/lam.py), and a [dynamics model](models/dynamics.py). Each of these components are trained separately, however, the dynamics model requires a pre-trained video tokenizer (and latent action model).
[content] README.md (markdown)
├── data
│   ├── jasmine_data
│   ├── pyproject.toml
│   ├── uv.lock
│   ├── _vizdoom
│   └── _vizdoom.ini
├── gifs
│   └── diffusion
├── jasmine
│   ├── baselines
│   ├── genie_bak.py
│   ├── genie_continous.py
│   ├── genie_discrete.py
│   ├── genie_sliding.py
│   ├── __init__.py
│   ├── models
│   ├── __pycache__
│   ├── sample_sliding.py
│   ├── test_dit.py
│   ├── test_masking.py
│   └── utils
├── LICENSE
├── pyproject.toml
├── README.md
├── slurm
│   ├── common
│   ├── dev
│   ├── jobs
│   ├── README.md
│   ├── templates
│   └── utils
├── tests
│   ├── data
│   └── test_dataloader.py
[content] README.md (markdown)
├── baselines
│   ├── diffusion
│   │   ├── sample_diffusion.py
│   │   ├── train_dynamics_diffusion.py
│   │   └── train_tokenizer_mae.py
│   ├── maskgit
│   │   ├── sample_maskgit.py
│   │   ├── train_dynamics_causal.py
│   │   ├── train_dynamics_maskgit.py
│   │   └── train_tokenizer_vqvae.py
│   └── train_lam.py
[content] README.md (markdown)
├── diffusion
│   │   ├── sample_diffusion.py
│   │   ├── train_dynamics_diffusion.py
│   │   └── train_tokenizer_mae.py
│   ├── maskgit
│   │   ├── sample_maskgit.py
│   │   ├── train_dynamics_causal.py
│   │   ├── train_dynamics_maskgit.py
│   │   └── train_tokenizer_vqvae.py
[tab] slurm/jobs/mihir/berlin/diffusion/breakout/train_tokenizer_mae.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/tokenizer/mae/%x_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/train/tokenizer/mae/%x_%j.log
#SBATCH --job-name=breakout_train_tokenizer_mae

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/tokenizer/mae/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

srun python jasmine/train_tokenizer.py \
    --save_ckpt \
    --image_height=10 \
    --image_width=10 \
    --tokenizer_type=mae \
    --patch_size=4 \
    --log_checkpoint_keep_period=5000 \
    --ckpt_dir $CHECKPOINT_DIR \
    --name=train_tokenizer_mae_breakout_$slurm_job_id \
    --tags tokenizer mae breakout \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \