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"])
        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
[tab] experiments/dynamics_grain_tok_restore.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

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

export XLA_FLAGS=--xla_gpu_autotune_level=0
export PYTHONUNBUFFERED=1
srun python train_dynamics.py \
    --save_ckpt \
    --dynamics_type causal \
    --log_checkpoint_interval 5 \
    --batch_size 12 \
    --tokenizer_checkpoint $tokenizer_ckpt_dir \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] jasmine.py (python)
            masked_token_idxs = jnp.where(mask, 0, token_idxs)  # Use 0 as mask token index
            
            # Use vanilla forward pass with proper input formatting
            dynamics_inputs = {
                "video_tokens": masked_token_idxs,
                "latent_actions": action_tokens,
            }
            final_logits, _ = self.dynamics(dynamics_inputs, training=False)
            
            # Apply temperature scaling
[content] jasmine.py (python)
            vid_embed = jnp.where(mask_expanded, mask_token, vid_embed)

            # --- Predict transition ---
            act_embed = self.dynamics.action_up(action_tokens)
            vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
[content] jasmine.py (python)
            
            # Use vanilla forward pass with proper input formatting
            dynamics_inputs = {
                "video_tokens": token_idxs,
                "latent_actions": action_tokens,
            }
            
            # We need to manually apply the mask logic since we're not in training mode
            # First, get the embedded tokens
            vid_embed = self.dynamics.patch_embed(token_idxs)
            # Apply mask tokens by replacing masked embeddings
[content] jasmine.py (python)
            
            # Create a dummy mask_rng since we're not in training mode
            # The dynamics model will return mask=None when training=False
            dynamics_inputs = {
                "video_tokens": token_idxs,
                "latent_actions": action_tokens,
                "mask_rng": jnp.array([0]),  # Dummy value, won't be used in inference
            }
            
            # Call the vanilla forward pass
            final_logits, _ = self.dynamics(dynamics_inputs, training=False)
            
            # Apply temperature scaling
[content] jasmine.py (python)
            mask_token = self.dynamics.mask_token.value  # (1, 1, 1, D,)
            mask_expanded = mask[..., None]  # (B, S, N, 1)
            vid_embed = jnp.where(mask_expanded, mask_token, vid_embed)

            # --- Predict transition ---
            act_embed = self.dynamics.action_up(action_tokens)
            vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
[content] jasmine.py (python)
            
            # Apply mask tokens manually since we need them for MaskGIT sampling
            # First, get the embedded tokens using the vanilla patch_embed
            vid_embed = self.dynamics.patch_embed(token_idxs)
            # Apply mask tokens by replacing masked embeddings (same logic as vanilla forward pass)
[content] jasmine.py (python)
            
            # Use vanilla forward pass with provided mask for sampling
            dynamics_inputs = {
                "video_tokens": token_idxs,
                "latent_actions": action_tokens,
            }
            final_logits, _ = self.dynamics(dynamics_inputs, training=False, mask=mask)
            
            # Apply temperature scaling
[content] utils/nn.py (python)
        
        # For autoregressive decoding, slice to last frame before temporal attention
        if self.decode:
            z = z[:, -1:, :, :]  # Keep only the last frame for temporal attention
[content] utils/nn.py (python)
        
        # For autoregressive decoding, slice to last frame before temporal attention
        if self.decode:
            z = z[:, -1:, :, :]  # Keep only the last frame for temporal attention
            x = x[:, -1:, :, :]  # Slice x to match z
[content] models/dynamics_causal.py (python)
            # For autoregressive decoding, we only need to pass the token from the previous step
            # current positionernally uses a KV cache to remember previous tokens.
            vid_embed = vid_embed[:, -1, :, :]
[content] models/dynamics_maskgit.py (python)
        self, batch: Dict[str, jax.Array], training: bool = True
    ) -> tuple[jax.Array, jax.Array | None]:
        # --- Mask videos ---
        vid_embed = self.patch_embed(batch["video_tokens"])
        if training:
            batch_size = vid_embed.shape[0]
            _rng_prob, *_rngs_mask = jax.random.split(batch["mask_rng"], batch_size + 1)
            mask_prob = jax.random.uniform(
                _rng_prob, shape=(batch_size,), minval=self.mask_limit
            )
            per_sample_shape = vid_embed.shape[1:-1]
            mask = jax.vmap(
                lambda rng, prob: jax.random.bernoulli(rng, prob, per_sample_shape),
                in_axes=(0, 0),
            )(jnp.asarray(_rngs_mask), mask_prob)
            mask = mask.at[:, 0].set(False)
            vid_embed = jnp.where(
                jnp.expand_dims(mask, -1), self.mask_token.value, vid_embed
            )
        else:
            mask = None

        # --- Predict transition ---
[content] sample.py (python)
    @nnx.jit
    def _autoreg_sample(rng, video_batch, action_batch):
        vid = video_batch[:, : args.start_frame + 1]
        rng, _rng = jax.random.split(rng)
        batch = dict(videos=vid, latent_actions=action_batch, rng=_rng)
        generated_vid = _sampling_fn(jasmine, batch)
        return generated_vid

    # --- Get video + latent actions ---
    array_record_files = [
        os.path.join(args.data_dir, x)
        for x in os.listdir(args.data_dir)
        if x.endswith(".array_record")
    ]
    dataloader = get_dataloader(
        array_record_files,
        args.seq_len,
        args.batch_size,
        args.image_height,
        args.image_width,
        args.image_channels,
        # We don't use workers in order to avoid grain shutdown issues (https://github.com/google/grain/issues/398)
        num_workers=0,
        prefetch_buffer_size=1,
        seed=args.seed,
    )
    dataloader = iter(dataloader)
    video_batch = next(dataloader)
    video_batch = video_batch.astype(args.dtype) / 255.0
[tab] models/dynamics.py (python)
from typing import Dict

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

from utils.nn import STTransformer


class DynamicsMaskGIT(nnx.Module):
    """MaskGIT 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,
        mask_limit: float,
        param_dtype: jnp.dtype,
        dtype: jnp.dtype,
        use_flash_attention: 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.mask_limit = mask_limit
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention

        self.dynamics = 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,
            rngs=rngs,
        )
        self.patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=rngs)
        self.mask_token = nnx.Param(
            nnx.initializers.lecun_uniform()(rngs.params(), (1, 1, 1, self.model_dim))
        )
        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]:
        # --- Mask videos ---
        vid_embed = self.patch_embed(batch["video_tokens"])
        if training:
            batch_size = vid_embed.shape[0]
            _rng_prob, *_rngs_mask = jax.random.split(batch["mask_rng"], batch_size + 1)
            mask_prob = jax.random.uniform(
                _rng_prob, shape=(batch_size,), minval=self.mask_limit
            )
            per_sample_shape = vid_embed.shape[1:-1]
            mask = jax.vmap(
                lambda rng, prob: jax.random.bernoulli(rng, prob, per_sample_shape),
                in_axes=(0, 0),
            )(jnp.asarray(_rngs_mask), mask_prob)
            mask = mask.at[:, 0].set(False)
            vid_embed = jnp.where(
                jnp.expand_dims(mask, -1), self.mask_token.value, vid_embed
            )
        else:
            mask = None

        # --- Predict transition ---
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        logits = self.dynamics(vid_embed)
        return logits, mask
[tab] experiments/sample.sh (shellscript)
source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/causal_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] experiments/dynamics_grain_tok_restore.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/maskgit_dynamics_openai_grain_tok_lam_restore"
tokenizer_ckpt_dir="$PWD/checkpoints/tokenizer_openai_grain_checkpointing"

export XLA_FLAGS=--xla_gpu_autotune_level=0
export PYTHONUNBUFFERED=1
srun python train_dynamics.py \
    --save_ckpt \
    --dynamics_type maskgit \
    --log_checkpoint_interval 5 \
    --batch_size 12 \
    --tokenizer_checkpoint $tokenizer_ckpt_dir \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] experiments/dynamics_grain_full_restore.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

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 train_dynamics.py \
    --batch_size 48 \
    --restore_ckpt \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] experiments/lam_grain_checkpointing.sh (shellscript)
#!/usr/bin/env bash
source .venv/bin/activate

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

mkdir -p $ckpt_dir

export XLA_FLAGS=--xla_gpu_autotune_level=0
PYTHONUNBUFFERED=1 python train_lam.py \
    --batch_size 48 \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] models/dynamics.py (python)

        param_dtype: jnp.dtype,
        dtype: jnp.dtype,
        use_flash_attention: bool,
        rngs: nnx.Rngs,
    ):
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.num_latents = num_latents
[content] models/dynamics.py (python)

        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.transformer = STTransformer(
[content] models/dynamics.py (python)

            self.model_dim,
            self.ffn_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
[content] models/dynamics.py (python)

            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,
        )
[content] models/dynamics.py (python)
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"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        logits = self.transformer(vid_embed)
        return logits, None
[content] utils/nn.py (python)

            num_heads=self.num_heads,
            in_features=self.model_dim,
            qkv_features=self.model_dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
[content] utils/nn.py (python)
    @nnx.remat
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = self.spatial_pos_enc(x)
        z = self.spatial_norm(z)
        z = self.spatial_attention(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = self.temporal_pos_enc(x)
        z = self.temporal_norm(z)
        z = self.temporal_attention(z)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = self.ffn_norm(x)
        z = self.ffn_dense1(z)
        z = jax.nn.gelu(z)
        z = self.ffn_dense2(z)
        x = x + z

        return x
[content] utils/nn.py (python)

        self.temporal_attention = nnx.MultiHeadAttention(
            num_heads=self.num_heads,
            in_features=self.model_dim,
            qkv_features=self.model_dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(
                self.use_flash_attention, is_causal=False
            ),
            rngs=rngs,
            decode=False,
        )
[content] utils/nn.py (python)

        self.input_dim = input_dim
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.out_dim = out_dim
        self.num_blocks = num_blocks
        self.num_heads = num_heads
[content] utils/nn.py (python)

                TransformerBlock(
                    model_dim=self.model_dim,
                    ffn_dim=self.ffn_dim,
                    num_heads=self.num_heads,
                    dropout=self.dropout,
                    param_dtype=self.param_dtype,
[content] utils/nn.py (python)

                    use_flash_attention=self.use_flash_attention,
                    rngs=rngs,
                )
            )
        self.output_dense = nnx.Linear(
            in_features=self.model_dim,
[content] utils/nn.py (python)


        self.input_norm1 = nnx.LayerNorm(
            num_features=self.input_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )
        self.input_dense = nnx.Linear(
            in_features=self.input_dim,
            out_features=self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )
        self.input_norm2 = nnx.LayerNorm(
            num_features=self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )