p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/jobs/mihir/berlin/data_gen/doom/doom_chunked_1m_60x80.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=16
#SBATCH --output=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --error=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --job-name=generate_doom_1m_60x80

source .venv/bin/activate

python jasmine_data/ViZDoomPPO/load_model_generate_dataset.py \
    --output_dir /home/mihir.mahajan/data/data_doom/doom_episodes_1m_60x80 \
    --agent_path /home/mihir.mahajan/Projects/jasmine/data/jasmine_data/ViZDoomPPO/logs/models/deathmatch_simple_bak/best_model.zip \
    --num_episodes_train 1_000 \
    --target_width 80 \
    --target_height 60 \
[tab] slurm/jobs/mihir/berlin/data_gen/doom/doom_chunked_10m_60x80.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=16
#SBATCH --output=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --error=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --job-name=generate_doom_10m_60x80


source .venv/bin/activate

python jasmine_data/ViZDoomPPO/load_model_generate_dataset.py \
    --output_dir /home/mihir.mahajan/data/data_doom/doom_episodes_10m_60x80 \
    --agent_path /home/mihir.mahajan/Projects/jasmine/data/jasmine_data/ViZDoomPPO/logs/models/deathmatch_simple_bak/best_model.zip \
    --num_episodes_train 10_000 \
    --target_width 80 \
    --target_height 60 \
[tab] slurm/jobs/mihir/berlin/data_gen/doom/doom_chunked_50m_60x80.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=16
#SBATCH --output=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --error=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --job-name=generate_doom_50m_60x80

source .venv/bin/activate

python jasmine_data/ViZDoomPPO/load_model_generate_dataset.py \
    --output_dir /home/mihir.mahajan/data/data_doom/doom_episodes_50m_60x80 \
    --agent_path /home/mihir.mahajan/Projects/jasmine/data/jasmine_data/ViZDoomPPO/logs/models/deathmatch_simple_bak/best_model.zip \
    --num_episodes_train 50_000 \
    --target_width 80 \
    --target_height 60 \
[tab] slurm/jobs/mihir/berlin/data_gen/doom/doom_chunked_50m_120x160.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=16
#SBATCH --output=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --error=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --job-name=generate_doom_50m_120x160

source .venv/bin/activate

python jasmine_data/ViZDoomPPO/load_model_generate_dataset.py \
    --output_dir /home/mihir.mahajan/data/data_doom/doom_episodes_50m_120x160 \
    --agent_path /home/mihir.mahajan/Projects/jasmine/data/jasmine_data/ViZDoomPPO/logs/models/deathmatch_simple_bak/best_model.zip \
    --num_episodes_train 50_000 \
    --target_width 160 \
    --target_height 120 \
[tab] slurm/jobs/mihir/berlin/data_gen/breakout_chunked.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --gres=gpu:0
#SBATCH --cpus-per-task=16
#SBATCH --output=/home/mihir.mahajan/logs/generate/breakout/%x_%j.log
#SBATCH --error=/home/mihir.mahajan/logs/generate/breakout/%x_%j.log
#SBATCH --job-name=generate_breakout_chunked

source .venv/bin/activate

python jasmine_data/minatar/generate_minatar_breakout_dataset.py \
    --num_episodes_train 5_000_000 \
    --min_episode_length 16 \
    --max_episode_length 16 \
    --chunk_size 32 \
    --output_dir /home/mihir.mahajan/data/data_breakout/breakout_episodes_80m
[tab] slurm/jobs/mihir/berlin/data_gen/coinrun_chunked.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --gres=gpu:0
#SBATCH --cpus-per-task=16
#SBATCH --output=/home/mihir.mahajan/logs/generate/coinrun/%x_%j.log
#SBATCH --error=/home/mihir.mahajan/logs/generate/coinrun/%x_%j.log
#SBATCH --job-name=generate_coinrun_chunked

source .venv/bin/activate

python jasmine_data/coinrun/generate_coinrun_dataset.py \
    --min_episode_length 500 \
    --num_episodes_train 100_000 \
    --num_episodes_val 100 \
    --num_episodes_test 100 \
    --output_dir /home/mihir.mahajan/data/data_coinrun/coinrun_episodes_50m
[tab] slurm/jobs/mihir/berlin/data_gen/doom/doom_chunked_50m_120x160.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=16
#SBATCH --output=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --error=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --job-name=generate_doom_50m_120x160

source .venv/bin/activate

python jasmine_data/ViZDoomPPO/load_model_generate_dataset.py \
    --output_dir /fast/project/HFMI_ProFOUND/mihir/data/data_doom/doom_episodes_50m_120x160 \
    --agent_path /home/mihir.mahajan/Projects/jasmine/data/jasmine_data/ViZDoomPPO/logs/models/deathmatch_simple_bak/best_model.zip \
    --num_episodes_train 50_000 \
    --target_width 160 \
    --target_height 120 \
[tab] /home/mihir.mahajan/logs/generate/breakout/generate_breakout_chunked_32312.log (log)
Traceback (most recent call last):
  File "/fast/home/mihir.mahajan/Projects/jasmine/data/jasmine_data/minatar/generate_minatar_breakout_dataset.py", line 12, in <module>
    from minatar import Environment
ModuleNotFoundError: No module named 'minatar'
[tab] slurm/jobs/mihir/berlin/data_gen/doom/doom_chunked_50m_60x80.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=16
#SBATCH --output=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --error=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --job-name=generate_doom_50m_60x80

source .venv/bin/activate

python jasmine_data/ViZDoomPPO/load_model_generate_dataset.py \
    --output_dir /fast/project/HFMI_ProFOUND/mihir/data/data_doom/doom_episodes_50m_60x80 \
    --agent_path /home/mihir.mahajan/Projects/jasmine/data/jasmine_data/ViZDoomPPO/logs/models/deathmatch_simple_bak/best_model.zip \
    --num_episodes_train 50_000 \
    --target_width 80 \
    --target_height 60 \
[tab] slurm/jobs/mihir/berlin/data_gen/doom/doom_chunked_10m_60x80.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=16
#SBATCH --output=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --error=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --job-name=generate_doom_10m_60x80


source .venv/bin/activate

python jasmine_data/ViZDoomPPO/load_model_generate_dataset.py \
    --output_dir /fast/project/HFMI_ProFOUND/mihir/data/data_doom/doom_episodes_10m_60x80 \
    --agent_path /home/mihir.mahajan/Projects/jasmine/data/jasmine_data/ViZDoomPPO/logs/models/deathmatch_simple_bak/best_model.zip \
    --num_episodes_train 10_000 \
    --target_width 80 \
    --target_height 60 \
[tab] slurm/jobs/mihir/berlin/data_gen/doom/doom_chunked_1m_60x80.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=16
#SBATCH --output=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --error=/home/mihir.mahajan/logs/generate/doom/%x_%j.log
#SBATCH --job-name=generate_doom_1m_60x80

source .venv/bin/activate

python jasmine_data/ViZDoomPPO/load_model_generate_dataset.py \
    --output_dir /fast/project/HFMI_ProFOUND/mihir/data/data_doom/doom_episodes_1m_60x80 \
    --agent_path /home/mihir.mahajan/Projects/jasmine/data/jasmine_data/ViZDoomPPO/logs/models/deathmatch_simple_bak/best_model.zip \
    --num_episodes_train 1_000 \
    --target_width 80 \
    --target_height 60 \
[tab] slurm/jobs/mihir/berlin/data_gen/breakout_chunked.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --gres=gpu:0
#SBATCH --cpus-per-task=16
#SBATCH --output=/home/mihir.mahajan/logs/generate/breakout/%x_%j.log
#SBATCH --error=/home/mihir.mahajan/logs/generate/breakout/%x_%j.log
#SBATCH --job-name=generate_breakout_chunked

source .venv/bin/activate

python jasmine_data/minatar/generate_minatar_breakout_dataset.py \
    --num_episodes_train 5_000_000 \
    --min_episode_length 16 \
    --max_episode_length 16 \
    --chunk_size 32 \
    --output_dir /fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m
[content] jasmine/models/dynamics.py (python)

    """Diffusion 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,
[content] jasmine/models/dynamics.py (python)
    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,
    ):
[content] jasmine/models/dynamics.py (python)

        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
[content] jasmine/models/dynamics.py (python)
    # Sample t.
    t = jax.random.randint(time_key, (images.shape[0],), minval=0, maxval=FLAGS.model['denoise_timesteps']).astype(jnp.float32)
    t /= FLAGS.model['denoise_timesteps']
    force_t_vec = jnp.ones(images.shape[0], dtype=jnp.float32) * force_t
    t = jnp.where(force_t_vec != -1, force_t_vec, t)         # If force_t is not -1, then use force_t.
    t_full = t[:, None, None, None] # [batch, 1, 1, 1]
[content] jasmine/utils/nn.py (python)
class PatchEmbed(nn.Module):
    """ 2D Image to Patch Embedding """
    patch_size: int
    hidden_size: int
    tc: TrainConfig
    bias: bool = True

    @nn.compact
    def __call__(self, x):
        B, H, W, C = x.shape
        patch_tuple = (self.patch_size, self.patch_size)
        num_patches = (H // self.patch_size)
        x = nn.Conv(self.hidden_size, patch_tuple, patch_tuple, use_bias=self.bias, padding="VALID",
                     kernel_init=self.tc.kern_init('patch'), bias_init=self.tc.kern_init('patch_bias', zero=True),
                     dtype=self.tc.dtype)(x) # (B, P, P, hidden_size)
        x = rearrange(x, 'b h w c -> b (h w) c', h=num_patches, w=num_patches)
        return x
[content] jasmine/utils/nn.py (python)
        self.blocks: List[TransformerBlock] = []
        for _ in range(self.num_blocks):
            self.blocks.append(
                TransformerBlock(
                    model_dim=self.model_dim,
                    ffn_dim=self.ffn_dim,
                    num_heads=self.num_heads,
                    dropout=self.dropout,
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                    use_flash_attention=self.use_flash_attention,
                    decode=decode,
                    sow_weights=self.sow_weights,
                    sow_activations=self.sow_activations,
                    rngs=rngs,
                )
            )
[content] jasmine/utils/nn.py (python)
class TimestepEmbedder(nn.Module):
    """
    Embeds scalar timesteps into vector representations.
    """
    hidden_size: int
    tc: TrainConfig
    frequency_embedding_size: int = 256

    @nn.compact
    def __call__(self, t):
        x = self.timestep_embedding(t)
        x = nn.Dense(self.hidden_size, kernel_init=nn.initializers.normal(0.02), 
                     bias_init=self.tc.kern_init('time_bias'), dtype=self.tc.dtype)(x)
        x = nn.silu(x)
        x = nn.Dense(self.hidden_size, kernel_init=nn.initializers.normal(0.02), 
                     bias_init=self.tc.kern_init('time_bias'))(x)
        return x
    
    # t is between [0, 1].
    def timestep_embedding(self, t, max_period=10000):
        """
        Create sinusoidal timestep embeddings.
        :param t: a 1-D Tensor of N indices, one per batch element.
                            These may be fractional.
        :param dim: the dimension of the output.
        :param max_period: controls the minimum frequency of the embeddings.
        :return: an (N, D) Tensor of positional embeddings.
        """
        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
        t = jax.lax.convert_element_type(t, jnp.float32)
        # t = t * max_period
        dim = self.frequency_embedding_size
        half = dim // 2
        freqs = jnp.exp( -math.log(max_period) * jnp.arange(start=0, stop=half, dtype=jnp.float32) / half)
        args = t[:, None] * freqs[None]
        embedding = jnp.concatenate([jnp.cos(args), jnp.sin(args)], axis=-1)
        embedding = embedding.astype(self.tc.dtype)
        return embedding
[content] jasmine/utils/nn.py (python)

        self.dense1 = nnx.Linear(self.hidden_size, kernel_init=nnx.initializers.normal(0.02), bias_init=nnx.initializers.normal(0.02))
        self.dense2 = nnx.Linear(self.hidden_size, kernel_init=nnx.initializers.normal(0.02), bias_init=nnx.initializers.normal(0.02))
[content] jasmine/utils/nn.py (python)
        c = nn.silu(c)
        c = nn.Dense(2 * self.hidden_size, kernel_init=self.tc.kern_init(zero=True), 
                     bias_init=self.tc.kern_init('bias', zero=True), dtype=self.tc.dtype)(c)
        shift, scale = jnp.split(c, 2, axis=-1)
        x = nn.LayerNorm(use_bias=False, use_scale=False, dtype=self.tc.dtype)(x)
        x = modulate(x, shift, scale)
        x = nn.Dense(self.patch_size * self.patch_size * self.out_channels, 
                     kernel_init=self.tc.kern_init('final', zero=True), 
                     bias_init=self.tc.kern_init('final_bias', zero=True), dtype=self.tc.dtype)(x)
[content] jasmine/utils/nn.py (python)
        x = jnp.reshape(x, (batch_size, num_patches_side, num_patches_side, 
                            self.patch_size, self.patch_size, self.out_channels))
        x = jnp.einsum('bhwpqc->bhpwqc', x)
        x = rearrange(x, 'B H P W Q C -> B (H P) (W Q) C', H=int(num_patches_side), W=int(num_patches_side))
        assert x.shape == (batch_size, input_size, input_size, self.out_channels)
[content] jasmine/utils/nn.py (python)
    def __call__(self, x, c):
        # Calculate adaLn modulation parameters.
        c = nn.silu(c)
        c = nn.Dense(6 * self.hidden_size, **self.tc.default_config())(c)
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = jnp.split(c, 6, axis=-1)
        
        # Attention Residual.
        x_norm = nn.LayerNorm(use_bias=False, use_scale=False, dtype=self.tc.dtype)(x)
        x_modulated = modulate(x_norm, shift_msa, scale_msa)
        channels_per_head = self.hidden_size // self.num_heads
        k = nn.Dense(self.hidden_size, **self.tc.default_config())(x_modulated)
        q = nn.Dense(self.hidden_size, **self.tc.default_config())(x_modulated)
        v = nn.Dense(self.hidden_size, **self.tc.default_config())(x_modulated)
        k = jnp.reshape(k, (k.shape[0], k.shape[1], self.num_heads, channels_per_head))
        q = jnp.reshape(q, (q.shape[0], q.shape[1], self.num_heads, channels_per_head))
        v = jnp.reshape(v, (v.shape[0], v.shape[1], self.num_heads, channels_per_head))
        q = q / q.shape[3] # (1/d) scaling.
        w = jnp.einsum('bqhc,bkhc->bhqk', q, k) # [B, HW, HW, num_heads]
        w = w.astype(jnp.float32)
        w = nn.softmax(w, axis=-1)
        y = jnp.einsum('bhqk,bkhc->bqhc', w, v) # [B, HW, num_heads, channels_per_head]
        y = jnp.reshape(y, x.shape) # [B, H, W, C] (C = heads * channels_per_head)
        attn_x = nn.Dense(self.hidden_size, **self.tc.default_config())(y)
        x = x + (gate_msa[:, None] * attn_x)

        # MLP Residual.
        x_norm2 = nn.LayerNorm(use_bias=False, use_scale=False, dtype=self.tc.dtype)(x)
        x_modulated2 = modulate(x_norm2, shift_mlp, scale_mlp)
        mlp_x = MlpBlock(mlp_dim=int(self.hidden_size * self.mlp_ratio), tc=self.tc, 
                         dropout_rate=self.dropout, train=self.train)(x_modulated2)
        x = x + (gate_mlp[:, None] * mlp_x)
        return x
[content] jasmine/utils/nn.py (python)
ize, kernel_init=nnx.initializers.normal(0.02), bias_init=nnx.initializers.normal(0.02))
        self.dense5 = nnx.Linear(self.hidden_size, kernel_init=nnx.initializers.normal(0.02), bias_init=nnx.initializers.normal(0.02))
        self.dense6 = nnx.Linear(self.hidden_size, kernel_init=nnx.initializers.normal(0.02), bias_init=nnx.initializers.normal(0.02))
        self.layer_norm = nnx.LayerNorm(use_bias=False, use_scale=False, dtype=self.dtype)
[content] jasmine/utils/nn.py (python)
self.dense1 = nnx.Linear(6 * self.hidden_size, kernel_init=nnx.initializers.normal(0.02), bias_init=nnx.initializers.normal(0.02))
        self.dense2 = nnx.Linear(self.hidden_size, kernel_init=nnx.initializers.normal(0.02), bias_init=nnx.initializers.normal(0.02))
        self.dense3 = nnx.Linear(self.hidden_size, kernel_init=nnx.initializers.normal(0.02), bias_init=nnx.initializers.normal(0.02))
        self.dense4 = nnx.Line
[content] jasmine/utils/nn.py (python)

        self.layer_norm2 = nnx.LayerNorm(use_bias=False, use_scale=False, dtype=self.dtype)
        self.layer_norm3 = nnx.LayerNorm(use_bias=False, use_scale=False, dtype=self.dtype)
        self.layer_norm4 = nnx.LayerNorm(use_bias=False, use_scale=False, dtype=self.dtype)
        self.layer_norm5 = nnx.LayerNorm(use_bias=False, use_scale=False, dtype=self.dtype)
        self.layer_norm6 = nnx.LayerNorm(use_bias=False, use_scale=False, dtype=self.dtype)
[content] jasmine/models/dynamics.py (python)

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

    batch = {
        "images": jnp.zeros((1, 10, 64, 64, 3)),
        "step_t": 0.5,
        "rng": jax.random.PRNGKey(0),
    }
    x_t, mask = model(batch)
    print(x_t.shape)
    print(mask.shape)
[content] jasmine/test_dit.py (python)
if __name__ == "__main__":
    model = DynamicsDiffusion(
        model_dim=128,
        ffn_dim=128,
        num_latents=128,
        latent_action_dim=128,
        num_blocks=128,
        num_heads=128,
    )
    batch = {
        "images": jnp.zeros((1, 10, 64, 64, 3)),
        "step_t": 0.5,
        "rng": jax.random.PRNGKey(0),
    }
    x_t, mask = model(batch)
    print(x_t.shape)
    print(mask.shape)