p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] models/dynamics_causal.py (python)
from typing import Dict

import jax
import jax.numpy as jnp
import flax.nnx as nnx

from utils.nn import STTransformer


class DynamicsCausal(nnx.Module):
    """Causal dynamics model"""

    def __init__(
        self,
        model_dim: int,
        ffn_dim: int,
        num_latents: int,
        latent_action_dim: int,
        num_blocks: int,
        num_heads: int,
        dropout: float,
        param_dtype: jnp.dtype,
        dtype: jnp.dtype,
        use_flash_attention: bool,
        decode: bool,
        rngs: nnx.Rngs,
    ):
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.num_latents = num_latents
        self.latent_action_dim = latent_action_dim
        self.num_blocks = num_blocks
        self.num_heads = num_heads
        self.dropout = dropout
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention
        self.decode = decode

        self.transformer = STTransformer(
            self.model_dim,
            self.model_dim,
            self.ffn_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
            spatial_causal=True,
            decode=self.decode,
            rngs=rngs,
        )
        self.patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=rngs)
        self.action_up = nnx.Linear(
            self.latent_action_dim,
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

    def __call__(
        self, batch: Dict[str, jax.Array], training: bool = True
    ) -> tuple[jax.Array, jax.Array | None]:
        vid_embed = self.patch_embed(batch["video_tokens"])
        # breakpoint()
        act_embed = self.action_up(batch["latent_actions"])
        breakpoint()
        # FIXME: this is the culprit!
        # vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        vid_embed += act_embed
        vid_embed_padded = jnp.pad(vid_embed, ((0, 0), (0, 0), (1, 0), (0, 0)))
        # breakpoint()

        logits = self.transformer(vid_embed_padded)[:, :, :-1]

        mask = jnp.ones(vid_embed.shape[:-1])
        return logits, mask
[content] models/dynamics_causal.py (python)
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        vid_embed_padded = jnp.pad(vid_embed, ((0, 0), (0, 0), (1, 0), (0, 0)))
[content] README.md (markdown)
<h1 align="center">Jafar: A JAX-based Genie Implementation 🧞</h1>

<p align="center">
    <a href= "https://github.com/FLAIROx/jafar/blob/main/LICENSE">
        <img src="https://img.shields.io/badge/license-Apache2.0-blue.svg" /></a>
    <a href= "https://github.com/psf/black">
        <img src="https://img.shields.io/badge/code%20style-black-000000.svg" /></a>
</p>

Jafar is a JAX-based implementation of the DeepMind paper "[Genie: Generative Interactive Environments](https://arxiv.org/abs/2402.15391)" (Bruce et al., 2024).

Jafar supports training of all Genie components and can complete the CoinRun reproducibility experiment (Appendix F) on a single L40S GPU in under a week.

<h2 name="start" id="start">Setup 🧗 </h2>

Jafar was built with `python 3.10` and `jax 0.4.30`. To install requirements, run:

```bash
pip install -r requirements.txt
pre-commit install
```

Before training the models, generate the CoinRun dataset by running:

```bash
python generate_dataset.py --num_episodes 10000
```

Note: this is a large dataset (around 100GB) and may take a while to generate.

For performant distributed training, we additionally preprocess the dataset into `TFRecord`s:

```bash
python preprocess_dataset.py
```

<h2 name="train" id="train">Quick Start 🚀 </h2>

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.

To train the video tokenizer (similar for the LAM), run:

```bash
python train_tokenizer.py --ckpt_dir <path>
```

Once the tokenizer and LAM are trained, the dynamics model can be trained with:

```bash
python train_dynamics.py --tokenizer_checkpoint <path> --lam_checkpoint <path>
```

Logging with `wandb` is supported. To enable logging, set the `WANDB_API_KEY` environment variable or run:

```bash
wandb login
```

Training can then be logged by setting the `--log` flag:

```bash
python train_tokenizer.py --log --entity <wandb-entity> --project <wandb-project>
```

<h2 name="cite" id="cite">Citing Jafar 📜 </h2>

Jafar was built by [Matthew Jackson](https://matthewtjackson.com) and [Timon Willi](https://www.timonwilli.com).

If you use Jafar in your work, please cite us and the original Genie paper as follows:
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
            # If mask is 4D (batch, heads, seq_len, seq_len), pad the last two dimensions
            if mask_bool.ndim == 4:
                expanded_mask = jnp.pad(
                    mask_bool, ((0, 0), (0, 0), (0, pad_size), (0, pad_size)), 
                    constant_values=False
                )
            else:
                # For other shapes, assume it's (seq_len, seq_len) and pad accordingly
                expanded_mask = jnp.pad(
                    mask_bool, ((0, pad_size), (0, pad_size)), constant_values=False
                )
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
        if mask is not None:
            # Handle the mask properly for decode mode
            if mask.ndim == 4:  # (batch, heads, seq_len, seq_len)
                # For decode mode, mask should already be the right shape
                # We need to pad it to match the padded sequence length
                mask_4d = jnp.pad(
                    mask, ((0, 0), (0, 0), (0, pad_size), (0, pad_size)), 
                    constant_values=False
                )
            else:
                # For other cases, use the original logic
                mask_4d = _fuse_masks(mask, attention_mask)
                mask_4d = mask_4d[jnp.newaxis, jnp.newaxis, :, :]  # (1, 1, seq_len, seq_len)
        else:
            mask_4d = attention_mask[jnp.newaxis, jnp.newaxis, :, :]  # (1, 1, seq_len, seq_len)
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
        mask_4d = (
            _fuse_masks(mask, attention_mask) if mask is not None else attention_mask
        )
        mask_4d = mask_4d[jnp.newaxis, jnp.newaxis, :, :]  # (1, 1, seq_len, seq_len)
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
        def _fuse_masks(mask: jax.Array, attention_mask: jax.Array) -> jax.Array:
            mask_bool = mask.astype(jnp.bool_)
            expanded_mask = jnp.pad(
                mask_bool, ((0, pad_size), (0, pad_size)), constant_values=False
            )
            return jnp.logical_and(attention_mask, expanded_mask)
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
        # Handle mask from nnx MultiHeadAttention
        if mask is not None:
            # Mask from nnx is always 4D: (batch, heads, seq_len, seq_len)
            # Pad it to match the padded sequence length
            mask_4d = jnp.pad(
                mask, ((0, 0), (0, 0), (0, pad_size), (0, pad_size)), 
                constant_values=False
            )
        else:
            # No mask provided, use the attention mask for padded positions
            mask_4d = attention_mask[jnp.newaxis, jnp.newaxis, :, :]  # (1, 1, seq_len, seq_len)
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
        # for temporal attention (using kv cache)
        # FIRST PASS: qkv.shape (1, 921, 1, 8, 64)
        # SECOND PASS: qkv.shape 
        if query.shape == (1, 921, 1, 512):
            breakpoint()
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
            # mask.shape (1, 921, 1, 1, 1) - (batch, kv_length, heads, q_length, k_length)
            if mask.ndim == 5:
                # For decode mode: pad the kv_length dimension (dim 1) and k_length dimension (dim 4)
                mask_4d = jnp.pad(
                    mask, ((0, 0), (0, pad_size), (0, 0), (0, 0), (0, pad_size)), 
                    constant_values=False
                )
            else:
                # Fallback for other mask shapes
                mask_4d = jnp.pad(
                    mask, ((0, 0), (0, 0), (0, pad_size), (0, pad_size)), 
                    constant_values=False
                )
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
            print(f"DEBUG: query.shape {query.shape}")
            print(f"DEBUG: key.shape {key.shape}")
            print(f"DEBUG: value.shape {value.shape}")
            # mask.shape (1, 921, 1, 1, 1) - (batch, kv_length, heads, q_length, k_length)
            if mask.ndim == 5:
                # For decode mode: pad the kv_length dimension (dim 1) and k_length dimension (dim 4)
                mask_4d = jnp.pad(
                    mask, ((0, 0), (0, pad_size), (0, 0), (0, 0), (0, pad_size)), 
                    constant_values=False
                )
            else:
                # Fallback for other mask shapes
                mask_4d = jnp.pad(
                    mask, ((0, 0), (0, 0), (0, pad_size), (0, pad_size)), 
                    constant_values=False
                )
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
        # Reshape from (1, 921, 1, 512) to (1, 921, 512) to fix the extra dimension issue
        original_shape = x.shape
        x_reshaped = x.reshape(original_shape[0], original_shape[1], -1)
        z = self.temporal_pos_enc(x_reshaped)
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
        if bias is not None:
            # Handle bias tensor - it might need to be broadcasted to match the padded dimensions
            bias_4d = _rearrange(bias)
            # If bias has fewer dimensions than expected, pad it
            if bias_4d.shape[-2] != target_seq_len or bias_4d.shape[-1] != target_seq_len:
                bias_4d = _pad(bias_4d)
        else:
            bias_4d = None
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
            # Handle bias tensor - ensure it has the correct shape for attention
            bias_4d = _rearrange(bias)
            # The bias should have shape (batch, target_seq, source_seq) for attention
            # If it doesn't match the padded dimensions, we need to handle it properly
            if len(bias_4d.shape) == 4:  # (batch, target_seq, source_seq, heads)
                if bias_4d.shape[-3] != target_seq_len or bias_4d.shape[-2] != target_seq_len:
                    bias_4d = _pad(bias_4d)
            else:
                # If bias has unexpected shape, skip it for now
                bias_4d = None
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
            # For attention bias, we expect shape (batch, target_seq, source_seq, heads)
            # If the bias doesn't match the expected dimensions, we need to handle it
            if bias_4d.shape[-3] != target_seq_len or bias_4d.shape[-2] != target_seq_len:
                # If bias dimensions don't match, we might need to broadcast or skip
                # For now, let's skip the bias if dimensions don't match
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
        # implementation = "cudnn" if use_flash_attention else None
        implementation = None

        def _rearrange(x):
            return einops.rearrange(x, "... l h d -> (...) l h d")

        def _pad(x):
            return jnp.pad(x, ((0, 0), (0, pad_size), (0, 0), (0, 0)))

        original_shape = query.shape
        original_seq_len = query.shape[-3]

        # Pad to nearest multiple of 4
        target_seq_len = ((original_seq_len + 3) // 4) * 4
        pad_size = target_seq_len - original_seq_len

        # query_4d = _pad(_rearrange(query))
        # key_4d = _pad(_rearrange(key))
        # value_4d = _pad(_rearrange(value))
        query_4d = _rearrange(query)
        key_4d = _rearrange(key)
        value_4d = _rearrange(value)

        attention_mask = jnp.ones((target_seq_len, target_seq_len), dtype=jnp.bool_)
        attention_mask = attention_mask.at[original_seq_len:, :].set(False)
        attention_mask = attention_mask.at[:, original_seq_len:].set(False)

        # Handle causal mask for cached decoder self-attention (from nnx.MultiHeadAttention)
        if mask is not None:
            # mask.shape (1, 921, 1, 1, 1)
            # mask_4d = _pad(_rearrange(mask))
            mask_4d = _rearrange(mask)
            mask_4d = mask_4d.astype(jnp.bool)
        else:
            # FIXME (f.srambical): Investigate whether/why this is needed
            mask_4d = attention_mask[jnp.newaxis, jnp.newaxis, :, :]  # (1, 1, seq_len, seq_len)
        if query.shape == (1, 921, 1, 8, 64):
            breakpoint()

        # bias_4d = _pad(_rearrange(bias)) if bias is not None else None
        bias_4d = _rearrange(bias) if bias is not None else None
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)

        # attention_mask = jnp.ones((target_seq_len, target_seq_len), dtype=jnp.bool_)
        # attention_mask = attention_mask.at[original_seq_len:, :].set(False)
        # attention_mask = attention_mask.at[:, original_seq_len:].set(False)
[content] /fast/home/franz.srambical/jafar/utils/nn.py (python)
        # query_4d.shape (921, 4, 8, 64)

        attention_mask = jnp.ones((target_seq_len, target_seq_len), dtype=jnp.bool_)
        attention_mask = attention_mask.at[original_seq_len:, :].set(False)
        attention_mask = attention_mask.at[:, original_seq_len:].set(False)

        # Handle causal mask for cached decoder self-attention (from nnx.MultiHeadAttention)
        if mask is not None:
            # mask.shape (1, 921, 1, 1, 1)
            mask_4d = _pad(_rearrange(mask))
            mask_4d = mask_4d.astype(jnp.bool)
            # mask_4d.shape (921, 4, 1, 1)
[tab] experiments/sample.sh (shellscript)
source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/dynamics_openai_grain_tok_lam_restore"

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python sample.py \
    --dynamics_type "causal" \
    --batch_size 1 \
    --seq_len 2 \
    --start_frame 1 \
    --checkpoint $ckpt_dir \
    --data_dir $data_dir
[tab] slurm/jobs/mihir/horeka/overfit_batch_tiny/sample.sh (shellscript)
#!/usr/bin/env bash

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

# Activate virtual environment
source .venv/bin/activate

# Set workspace and checkpoint directory (update slurm_job_id as needed)
ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared'
# Replace the following with the actual job id/checkpoint you want to sample from
slurm_job_id=3301029

# job_name=train_dynamics_minecraft_overfit_sample_tiny
CHECKPOINT_DIR=$ws_dir/checkpoints/${slurm_job_id}

# Example: If you want to use a specific checkpoint, set it here
# CHECKPOINT_PATH=$ws_dir/checkpoints/3299272/dynamics-tiny-overfit-big-lr-3299272_50000/
# Or use the latest in the directory
# CHECKPOINT_PATH=$(ls -d $CHECKPOINT_DIR/*/ | sort | tail -n 1)
CHECKPOINT_PATH=$CHECKPOINT_DIR/genie_1751067601_200000/
# CHECKPOINT_PATH=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/0000/genie_1751301068_2000/
# CHECKPOINT_PATH=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/../checkpoints/3307618/genie_1751322003_15500/
# CHECKPOINT_PATH=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3307619/genie_1751322003_200000/
CHECKPOINT_PATH=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3309699/genie_1751384516_200000/


echo "Sampling from checkpoint: $CHECKPOINT_PATH"

python sample.py \
    --checkpoint "$CHECKPOINT_PATH" \
    --tokenizer_dim=384 \
    --latent_patch_dim=32 \
    --num_patch_latents=1024 \
    --patch_size=4 \
    --tokenizer_num_blocks=8 \
    --tokenizer_num_heads=8 \
    --lam_dim=384 \
    --latent_action_dim=32 \
    --lam_patch_size=16 \
    --lam_num_blocks=8 \
    --lam_num_heads=8 \
    --dyna_dim=128 \
    --dyna_num_blocks=2 \
    --dyna_num_heads=4 \
    --maskgit_steps=1000 \
    --num_latent_actions=6 \
    --seq_len=16 \
    --start_frame=0

# python sample.py \
    # --checkpoint "$CHECKPOINT_PATH" \
    # --data_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/coinrun_episodes
[tab] /fast/home/franz.srambical/jafar/models/dynamics_causal.py (python)
from typing import Dict

import jax
import jax.numpy as jnp
import flax.nnx as nnx

from utils.nn import STTransformer


class DynamicsCausal(nnx.Module):
    """Causal dynamics model"""

    def __init__(
        self,
        model_dim: int,
        ffn_dim: int,
        num_latents: int,
        latent_action_dim: int,
        num_blocks: int,
        num_heads: int,
        dropout: float,
        param_dtype: jnp.dtype,
        dtype: jnp.dtype,
        use_flash_attention: bool,
        decode: bool,
        rngs: nnx.Rngs,
    ):
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.num_latents = num_latents
        self.latent_action_dim = latent_action_dim
        self.num_blocks = num_blocks
        self.num_heads = num_heads
        self.dropout = dropout
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention
        self.decode = decode

        self.transformer = STTransformer(
            self.model_dim,
            self.model_dim,
            self.ffn_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
            spatial_causal=True,
            decode=self.decode,
            rngs=rngs,
        )
        self.patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=rngs)
        self.action_up = nnx.Linear(
            self.latent_action_dim,
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

    def __call__(
        self, batch: Dict[str, jax.Array], training: bool = True
    ) -> tuple[jax.Array, jax.Array | None]:
        vid_embed = self.patch_embed(batch["video_tokens"])
        act_embed = self.action_up(batch["latent_actions"])
        # FIXME: this is the culprit!
        # vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        vid_embed += act_embed
        vid_embed_padded = jnp.pad(vid_embed, ((0, 0), (0, 0), (1, 0), (0, 0)))

        # vid_embed_padded.shape (1, 1, 921, 512)
        logits = self.transformer(vid_embed_padded)[:, :, :-1]

        mask = jnp.ones(vid_embed.shape[:-1])
        return logits, mask