p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/jobs/mihir/horeka/coinrun/sample_causal.sbatch (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/causal/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/dynamics/sampling/causal/%x_%j.log
#SBATCH --job-name=coinrun_sample_causal

# 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=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/array_records
CHECKPOINT_PATH=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/coinrun/causal/train_dynamics_causal_1_node/3415712


echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python 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=causal
[content] genie.py (python)

        self,
        batch: Dict[str, jax.Array],
        seq_len: int,
        steps: int = 25,
        temperature: float = 1,
        sample_argmax: bool = False,
    ) -> tuple[jax.Array, jax.Array]:
[content] train_dynamics.py (python)

        metrics = dict(
            cross_entropy_loss=ce_loss,
            masked_token_accuracy=acc,
            select_logit=outputs["token_logits"].max(-1).mean(),
            select_p=select_probs.max(-1).mean(),
            entropy=jax.scipy.special.entr(select_probs).sum(-1).mean(),
            psnr=psnr,
            ssim=ssim,
            codebook_usage_lam=codebook_usage_lam,
            codebook_usage_tokenizer=codebook_usage_tokenizer,
        )
[content] train_dynamics.py (python)
        _, index_counts_lam = jnp.unique_counts(
            jnp.ravel(outputs["lam_indices"]),
            size=args.num_latent_actions,
            fill_value=0,
        )
        _, index_counts_tokenizer = jnp.unique_counts(
            jnp.ravel(outputs["video_tokens"]),
            size=args.num_patch_latents,
            fill_value=0,
        )
        codebook_usage_lam = (index_counts_lam != 0).mean()
        codebook_usage_tokenizer = (index_counts_tokenizer != 0).mean()
[content] genie.py (python)

                # We need to reconstruct submodules inside scan body to prevent trace context mismatches
                patch_embed = nnx.Embed(
                    self.num_latents, self.model_dim, rngs=nnx.Rngs(0)
                )
                nnx.update(patch_embed, patch_embed_state)
                action_up = nnx.Linear(
                    self.latent_action_dim,
                    self.model_dim,
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                    rngs=nnx.Rngs(0),
                )
                nnx.update(action_up, action_up_state)
                transformer = Transformer(
                    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,
                    decode=self.decode,
                    rngs=nnx.Rngs(0),
                )
                nnx.update(transformer, transformer_state)