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

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=10:00:00
#SBATCH --partition=large
#SBATCH --cpus-per-task=16
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/preprocess/coinrun/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/preprocess/coinrun/%x_%j.log
#SBATCH --job-name=preprocess_coinrun_chunked

source .venv/bin/activate

python input_pipeline/generate/generate_dataset.py \
    --min_episode_length 1000 \
    --output_dir /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_test \
    --num_episodes 35
[tab] input_pipeline/generate/generate_dataset.py (python)
"""
Generates a dataset of random-action CoinRun episodes.
Episodes are saved individually as memory-mapped files for efficient loading.
"""

from dataclasses import dataclass

from gym3 import types_np
import numpy as np
from procgen import ProcgenGym3Env
import tyro
import json
from input_pipeline.utils import save_chunks



@dataclass
class Args:
    num_episodes: int = 10000
    output_dir: str = "data/coinrun_episodes"
    min_episode_length: int = 1000
    max_episode_length: int = 1000
    chunk_size: int = 160
    chunks_per_file: int = 100


args = tyro.cli(Args)
assert args.max_episode_length >= args.min_episode_length, "Maximum episode length must be greater than or equal to minimum episode length."

if args.min_episode_length < args.chunk_size:
    print("Warning: Minimum episode length is smaller than chunk size. Note that episodes shorter than the chunk size will be discarded.")

# --- Generate episodes ---
episode_idx = 0
episode_metadata = []
chunks = []
file_idx = 0
while episode_idx < args.num_episodes:
    seed = np.random.randint(0, 10000)
    env = ProcgenGym3Env(num=1, env_name="coinrun", start_level=seed)
    
    observations_seq = []
    episode_chunks = []

    # --- Run episode ---
    for step_t in range(args.max_episode_length):
        env.act(types_np.sample(env.ac_space, bshape=(env.num,)))
        rew, obs, first = env.observe()
        observations_seq.append(obs["rgb"])
        if len(observations_seq) == args.chunk_size:
            episode_chunks.append(observations_seq)
            observations_seq = []
        if first:
            break

    # --- Save episode ---
    if step_t + 1  >= args.min_episode_length:
        if len(observations_seq) < args.chunk_size:
            print(
                f"Warning: Inconsistent chunk_sizes. Episode has {len(observations_seq)} frames, "
                f"which is smaller than the requested chunk_size: {args.chunk_size}. "
                "This might lead to performance degradation during training."
            )
        episode_chunks.append(observations_seq)
        chunks_data = [np.concatenate(seq, axis=0).astype(np.uint8) for seq in episode_chunks]
        chunks.extend(chunks_data)

        ep_metadata, chunks, file_idx = save_chunks(chunks, file_idx, args.chunks_per_file, args.output_dir)
        episode_metadata.extend(ep_metadata)

        print(f"Episode {episode_idx} completed, length: {step_t + 1}.")
        episode_idx += 1
    else:
        print(f"Episode too short ({step_t + 1}), resampling...")
ep_metadata, chunks, file_idx = _save_chunks(chunks, file_idx)
episode_metadata.extend(ep_metadata)
# --- Save metadata ---
metadata = {
    "env": "coinrun",
    "num_episodes": args.num_episodes,
    "avg_episode_len": np.mean([ep["avg_seq_len"] for ep in episode_metadata]),
    "episode_metadata": episode_metadata,
}
with open(output_dir / "metadata.json", "w") as f:
    json.dump(metadata, f)

print(f"Dataset generated with {args.num_episodes} episodes, saved over {len(episode_metadata)} files")
[content] input_pipeline/generate/generate_dataset.py (python)
# getting the name of the directory
# where the this file is present.
current = os.path.dirname(os.path.realpath(__file__))

# Getting the parent directory name
# where the current directory is present.
parent = os.path.dirname(current)

# adding the parent directory to 
# the sys.path.
sys.path.append(parent)
[tab] input_pipeline/generate_coinrun_dataset.py (python)
"""
Generates a dataset of random-action CoinRun episodes.
Episodes are saved individually as memory-mapped files for efficient loading.
"""

from dataclasses import dataclass

from gym3 import types_np
import numpy as np
from procgen import ProcgenGym3Env
import tyro
import json
from utils import save_chunks


@dataclass
class Args:
    num_episodes: int = 10000
    output_dir: str = "data/coinrun_episodes"
    min_episode_length: int = 1000
    max_episode_length: int = 1000
    chunk_size: int = 160
    chunks_per_file: int = 100


args = tyro.cli(Args)
assert args.max_episode_length >= args.min_episode_length, "Maximum episode length must be greater than or equal to minimum episode length."

if args.min_episode_length < args.chunk_size:
    print("Warning: Minimum episode length is smaller than chunk size. Note that episodes shorter than the chunk size will be discarded.")

# --- Generate episodes ---
episode_idx = 0
episode_metadata = []
chunks = []
file_idx = 0
while episode_idx < args.num_episodes:
    seed = np.random.randint(0, 10000)
    env = ProcgenGym3Env(num=1, env_name="coinrun", start_level=seed)
    
    observations_seq = []
    episode_chunks = []

    # --- Run episode ---
    for step_t in range(args.max_episode_length):
        env.act(types_np.sample(env.ac_space, bshape=(env.num,)))
        rew, obs, first = env.observe()
        observations_seq.append(obs["rgb"])
        if len(observations_seq) == args.chunk_size:
            episode_chunks.append(observations_seq)
            observations_seq = []
        if first:
            break

    # --- Save episode ---
    if step_t + 1  >= args.min_episode_length:
        if len(observations_seq) < args.chunk_size:
            print(
                f"Warning: Inconsistent chunk_sizes. Episode has {len(observations_seq)} frames, "
                f"which is smaller than the requested chunk_size: {args.chunk_size}. "
                "This might lead to performance degradation during training."
            )
        episode_chunks.append(observations_seq)
        chunks_data = [np.concatenate(seq, axis=0).astype(np.uint8) for seq in episode_chunks]
        chunks.extend(chunks_data)

        ep_metadata, chunks, file_idx = save_chunks(chunks, file_idx, args.chunks_per_file, args.output_dir)
        episode_metadata.extend(ep_metadata)

        print(f"Episode {episode_idx} completed, length: {step_t + 1}.")
        episode_idx += 1
    else:
        print(f"Episode too short ({step_t + 1}), resampling...")
ep_metadata, chunks, file_idx = save_chunks(chunks, file_idx)
episode_metadata.extend(ep_metadata)
# --- Save metadata ---
metadata = {
    "env": "coinrun",
    "num_episodes": args.num_episodes,
    "avg_episode_len": np.mean([ep["avg_seq_len"] for ep in episode_metadata]),
    "episode_metadata": episode_metadata,
}
with open(args.output_dir / "metadata.json", "w") as f:
    json.dump(metadata, f)

print(f"Dataset generated with {args.num_episodes} episodes, saved over {len(episode_metadata)} files")
[content] input_pipeline/generate_coinrun_dataset.py (python)
def main():
    # Set random seed and create dataset directories
    np.random.seed(args.seed)
    output_dir = Path(args.output_dir)
    (output_dir / "train").mkdir(parents=True, exist_ok=True)
    (output_dir / "val").mkdir(parents=True, exist_ok=True)
    (output_dir / "test").mkdir(parents=True, exist_ok=True)

    # --- Generate episodes ---
    train_episode_metadata = generate_episodes(args.num_episodes_train, "train")
    val_episode_metadata = generate_episodes(args.num_episodes_val, "val")
    test_episode_metadata = generate_episodes(args.num_episodes_test, "test")

    # --- Save metadata ---
    metadata = {
        "env": "coinrun",
        "num_episodes_train": args.num_episodes_train,
        "num_episodes_val": args.num_episodes_val,
        "num_episodes_test": args.num_episodes_test,
        "avg_episode_len_train": np.mean([ep["length"] for ep in train_episode_metadata]),
        "avg_episode_len_val": np.mean([ep["length"] for ep in val_episode_metadata]),
        "avg_episode_len_test": np.mean([ep["length"] for ep in test_episode_metadata]),
        "episode_metadata_train": train_episode_metadata,
        "episode_metadata_val": val_episode_metadata,
        "episode_metadata_test": test_episode_metadata,

    }
    with open(output_dir / "metadata.json", "w") as f:
        json.dump(metadata, f)

    print(f"Done generating dataset.")

if __name__ == "__main__":
    main()
[content] train_tokenizer.py (python)
        if args.val_data_dir:
            restore_args = ocp.args.Composite(
                    model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
                    train_dataloader_state=grain.checkpoint.CheckpointRestore(train_iterator),  # type: ignore
                    val_dataloader_state=grain.checkpoint.CheckpointRestore(val_iterator),  # type: ignore
                )
        else: 
            restore_args = ocp.args.Composite(
                    model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
                    train_dataloader_state=grain.checkpoint.CheckpointRestore(train_iterator),  # type: ignore
                )
[content] train_tokenizer.py (python)
    if args.val_data_dir:
        handler_registry.add(
            "val_dataloader_state",
            grain.checkpoint.CheckpointSave,
            cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
        )
        handler_registry.add(
            "val_dataloader_state",
            grain.checkpoint.CheckpointRestore,
            cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
        )
[content] train_tokenizer.py (python)
            # --- Validation loss ---
            if args.val_data_dir and step % args.val_interval == 0:
                print(f"Calculating validation metrics...")
                val_loss, val_metrics, val_gt_batch, val_recon = calculate_validation_metrics(dataloader_val, tokenizer)
                print(f"Step {step}, validation loss: {val_loss}")
[content] train_tokenizer.py (python)
                    log_dict = {
                        "loss": loss,
                        "step": step,
                        **metrics
                        }
                    if args.val_data_dir and step % args.val_interval == 0:
                        log_dict.update({
                            "val_loss": val_loss,
                            **val_metrics
                        })
                    wandb.log(log_dict)
[content] train_tokenizer.py (python)
                    if args.val_data_dir and step % args.val_interval == 0:
                        gt_seq_val = val_gt_batch["videos"][0].astype(jnp.float32) / 255.0
                        recon_seq_val = val_recon[0].clip(0, 1)
                        val_comparison_seq = jnp.concatenate((gt_seq, recon_seq), axis=1)
                        val_comparison_seq = einops.rearrange(
                            val_comparison_seq * 255, "t h w c -> h (t w) c"
                        )
[content] train_tokenizer.py (python)
                        if args.val_data_dir and step % args.val_interval == 0:
                            log_images.update(
                                dict(
                                    val_image=wandb.Image(np.asarray(gt_seq_val[0])),
                                    val_recon=wandb.Image(np.asarray(recon_seq_val[0])),
                                    val_true_vs_recon=wandb.Image(
                                        np.asarray(val_comparison_seq.astype(np.uint8))
                                    )
                                )
                            )
[content] train_tokenizer.py (python)
                if args.val_data_dir:
                    ckpt_manager_args = ocp.args.Composite(
                        model_state=ocp.args.PyTreeSave(optimizer_state),  # type: ignore
                        train_dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
                            train_iterator  # type: ignore
                        ),
                        val_dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
                            val_iterator  # type: ignore
                        )
                    )
                else: 
                    ckpt_manager_args = ocp.args.Composite(
                        model_state=ocp.args.PyTreeSave(optimizer_state),  # type: ignore
                        train_dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
                            train_iterator  # type: ignore
                        )
                    )
                checkpoint_manager.save(
                    step,
                    args=ckpt_manager_args
                    )
[content] train_tokenizer.py (python)
@nnx.jit
def val_step(tokenizer: TokenizerVQVAE, inputs: dict) -> tuple[jax.Array, jax.Array, dict]:
    tokenizer.eval()
    (loss, (recon, metrics)) = tokenizer_loss_fn(tokenizer, inputs)
    return loss, recon, metrics

def calculate_validation_metrics(val_dataloader, tokenizer):
    step = 0
    loss_per_step = []
    metrics_per_step = []
    for videos in val_dataloader:
        inputs = dict(videos=videos)
        loss, recon, metrics = val_step(tokenizer, inputs)
        loss_per_step.append(loss)
        metrics_per_step.append(metrics)
        step += 1
        if step > args.val_steps:
            break

    if step < args.val_steps:
        print(f"Warning: Your validation dataset is too small to make val_steps many steps. Made {step} steps, expected {args.val_steps}")

    val_loss = np.mean(loss_per_step)
    val_metrics = {
        f"val_{key}": np.mean([m[key] for m in metrics_per_step])
        for key in metrics_per_step[0].keys()
    }
    return val_loss, val_metrics, inputs, recon
[content] train_lam.py (python)
    if args.val_data_dir:
        handler_registry.add(
            "val_dataloader_state",
            grain.checkpoint.CheckpointSave,
            cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
        )
        handler_registry.add(
            "val_dataloader_state",
            grain.checkpoint.CheckpointRestore,
            cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
        )
[content] train_lam.py (python)
+    if args.val_data_dir:
+        step, optimizer, train_iterator, val_iterator = restore_checkpoint_if_needed(
+            args, checkpoint_manager, optimizer, train_iterator, val_iterator
+        )
+    else:
+        step, optimizer, train_iterator, _ = restore_checkpoint_if_needed(
+            args, checkpoint_manager, optimizer, train_iterator
+        )
[content] train_lam.py (python)
            # --- Validation loss ---
            if args.val_data_dir and step % args.val_interval == 0:
                print(f"Calculating validation metrics...")
                val_loss, val_metrics, val_gt_batch, val_recon = calculate_validation_metrics(dataloader_val, lam)
                print(f"Step {step}, validation loss: {val_loss}")
[content] train_lam.py (python)
@nnx.jit
def val_step(lam: LatentActionModel, inputs: dict) -> tuple[jax.Array, jax.Array, dict]:
    lam.eval()
    (loss, (recon, _, metrics)) = lam_loss_fn(lam, inputs, training=False)
    return loss, recon, metrics

def calculate_validation_metrics(val_dataloader, lam):
    step = 0
    loss_per_step = []
    metrics_per_step = []
    for videos in val_dataloader:
        inputs = dict(videos=videos)
        loss, recon, metrics = val_step(lam, inputs)
        loss_per_step.append(loss)
        metrics_per_step.append(metrics)
        step += 1
        if step > args.val_steps:
            break

    if step < args.val_steps:
        print(f"Warning: Your validation dataset is too small to make val_steps many steps. Made {step} steps, expected {args.val_steps}")

    val_loss = np.mean(loss_per_step)
    val_metrics = {
        f"val_{key}": np.mean([m[key] for m in metrics_per_step])
        for key in metrics_per_step[0].keys()
    }
    return val_loss, val_metrics, inputs, recon
[content] train_lam.py (python)
                    log_dict = {
                        "loss": loss,
                        "step": step,
                        **metrics
                        }
                    if args.val_data_dir and step % args.val_interval == 0:
                        log_dict.update({
                            "val_loss": val_loss,
                            **val_metrics
                        })
                    wandb.log(log_dict)
[content] train_lam.py (python)
                    if args.val_data_dir and step % args.val_interval == 0:
                        gt_seq_val = val_gt_batch["videos"][0, 1:].astype(jnp.float32) / 255.0
                        recon_seq_val = val_recon[0].clip(0, 1)
                        val_comparison_seq = jnp.concatenate((gt_seq, recon_seq), axis=1)
                        val_comparison_seq = einops.rearrange(
                            val_comparison_seq * 255, "t h w c -> h (t w) c"
                        )
[content] train_lam.py (python)
                        if args.val_data_dir and step % args.val_interval == 0:
                            log_images.update(
                                dict(
                                    val_image=wandb.Image(np.asarray(gt_seq_val[0])),
                                    val_recon=wandb.Image(np.asarray(recon_seq_val[0])),
                                    val_true_vs_recon=wandb.Image(
                                        np.asarray(val_comparison_seq.astype(np.uint8))
                                    )
                                )
                            )
[content] train_lam.py (python)
               if args.val_data_dir:
                    ckpt_manager_args = ocp.args.Composite(
                        model_state=ocp.args.PyTreeSave(optimizer_state),  # type: ignore
                        train_dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
                            train_iterator  # type: ignore
                        ),
                        val_dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
                            val_iterator  # type: ignore
                        )
                    )
                else: 
                    ckpt_manager_args = ocp.args.Composite(
                        model_state=ocp.args.PyTreeSave(optimizer_state),  # type: ignore
                        train_dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
                            train_iterator  # type: ignore
                        )
                    )
                checkpoint_manager.save(
                    step,
                    args=ckpt_manager_args
                    )
[content] train_lam.py (python)
if args.val_data_dir and step % args.val_interval == 0:
                        gt_seq_val = val_gt_batch["videos"][0, 1:].astype(jnp.float32) / 255.0
                        recon_seq_val = val_recon[0].clip(0, 1)
                        val_comparison_seq = jnp.concatenate((gt_seq, recon_seq), axis=1)
                        val_comparison_seq = einops.rearrange(
                            val_comparison_seq * 255, "t h w c -> h (t w) c"
                        )
[content] train_lam.py (python)
if args.val_data_dir and step % args.val_interval == 0:
                            log_images.update(
                                dict(
                                    val_image=wandb.Image(np.asarray(gt_seq_val[0])),
                                    val_recon=wandb.Image(np.asarray(recon_seq_val[0])),
                                    val_true_vs_recon=wandb.Image(
                                        np.asarray(val_comparison_seq.astype(np.uint8))
                                    )
                                )
                            )
[content] train_lam.py (python)
if args.val_data_dir:
                        ckpt_manager_args = ocp.args.Composite(
                            model_state=ocp.args.PyTreeSave(optimizer_state),  # type: ignore
                            train_dataloader_state=grain.
[content] train_lam.py (python)
.CheckpointSave(  # type: ignore
                                train_iterator  # type: ignore
                            ),
                            val_dataloader_state=grain.checkpoint.CheckpointS
[content] train_dynamics.py (python)
    if args.restore_ckpt:
        abstract_optimizer = nnx.eval_shape(lambda: optimizer)
        abstract_optimizer_state = nnx.state(abstract_optimizer)
        if args.val_data_dir:
            restore_args = ocp.args.Composite(
                    model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
                    train_dataloader_state=grain.checkpoint.CheckpointRestore(train_iterator),  # type: ignore
                    val_dataloader_state=grain.checkpoint.CheckpointRestore(val_iterator),  # type: ignore
                )
        else: 
            restore_args = ocp.args.Composite(
                    model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
                    train_dataloader_state=grain.checkpoint.CheckpointRestore(train_iterator),  # type: ignore
                )
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=restore_args
        )
        restored_optimizer_state = restored["model_state"]
        nnx.update(optimizer, restored_optimizer_state)
        train_iterator = restored["train_dataloader_state"]
        if args.val_data_dir:
            val_iterator = restored["val_dataloader_state"]
        step = checkpoint_manager.latest_step() or 0
        print(f"Restored dataloader and model state from step {step}")
[content] train_dynamics.py (python)
           # --- Validation loss ---
            if args.val_data_dir and step % args.val_interval == 0:
                rng, _rng_mask_val = jax.random.split(rng, 2)
                print(f"Calculating validation metrics...")
                val_loss, val_metrics, val_gt_batch, val_recon = calculate_validation_metrics(dataloader_val, optimizer.model, _rng_mask_val)
                print(f"Step {step}, validation loss: {val_loss}")
[content] train_dynamics.py (python)
                    log_dict = {
                        "loss": loss,
                        "step": step,
                        **metrics
                        }
                    if args.val_data_dir and step % args.val_interval == 0:
                        log_dict.update({
                            "val_loss": val_loss,
                            **val_metrics
                        })
                    wandb.log(log_dict)
[content] train_dynamics.py (python)
                    if args.val_data_dir and step % args.val_interval == 0:
                        gt_seq_val = val_gt_batch["videos"][0].astype(jnp.float32) / 255.0
                        recon_seq_val = val_recon[0].clip(0, 1)
                        val_comparison_seq = jnp.concatenate((gt_seq, recon_seq), axis=1)
                        val_comparison_seq = einops.rearrange(
                            val_comparison_seq * 255, "t h w c -> h (t w) c"
                        )
                    # NOTE: Process-dependent control flow deliberately happens
                    # after indexing operation since it must not contain code
                    # sections that lead to cross-accelerator communication.
[content] train_dynamics.py (python)
                        if args.val_data_dir and step % args.val_interval == 0:
                            log_images.update(
                                dict(
                                    val_image=wandb.Image(np.asarray(gt_seq_val[0])),
                                    val_recon=wandb.Image(np.asarray(recon_seq_val[0])),
                                    val_true_vs_recon=wandb.Image(
                                        np.asarray(val_comparison_seq.astype(np.uint8))
                                    )
                                )
                            )
[content] train_dynamics.py (python)
                if args.val_data_dir:
                    ckpt_manager_args = ocp.args.Composite(
                        model_state=ocp.args.PyTreeSave(optimizer_state),  # type: ignore
                        train_dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
                            train_iterator  # type: ignore
                        ),
                        val_dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
                            val_iterator  # type: ignore
                        )
                    )
                else: 
                    ckpt_manager_args = ocp.args.Composite(
                        model_state=ocp.args.PyTreeSave(optimizer_state),  # type: ignore
                        train_dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
                            train_iterator  # type: ignore
                        )
                    )
                checkpoint_manager.save(
                    step,
                    args=ckpt_manager_args
                    )
[tab] slurm/jobs/mihir/horeka/coinrun/train_dyn_single_gpu.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=00:20:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/dyn/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/dyn/%x_%j.log
#SBATCH --job-name=train_dyn_1e-4
#SBATCH --requeue
#SBATCH --signal=b:usr1@300 # 5 min before timeout

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

# array_records_dir_train=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_10m_gt_actions_split/train
array_records_dir_train=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_10m_chunked
array_records_dir_val=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_10m_gt_actions_split/val

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dyn/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

lam_checkpoint=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/lam/interactive/3468835
tokenizer_checkpoint=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer/interactive/3468835

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --image_height=64 \
    --image_width=64 \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=110 \
    --init_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name=coinrun-dyn-dev-$slurm_job_id \
    --tags dyn coinrun dev \
    --entity instant-uv \
    --project jafar \
    --warmup_steps 0 \
    --wsd_decay_steps 0 \
    --num_steps 1000 \
    --data_dir $array_records_dir_train \
    --tokenizer_checkpoint $tokenizer_checkpoint
[content] slurm/jobs/mihir/horeka/coinrun/train_dyn_single_gpu.sh (shellscript)
lam_checkpoint=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/lam/interactive/3468835
tokenizer_checkpoint=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer/interactive/3468835
[content] train_dynamics.py (python)
@nnx.jit
def val_step(genie: Genie, inputs: dict) -> tuple[jax.Array, jax.Array, dict]:
    """Evaluate model and compute metrics"""
    genie.eval()
    (loss, (recon, metrics)) = dynamics_loss_fn(genie, inputs)
    return loss, recon, metrics


def calculate_validation_metrics(val_dataloader, genie, rng):
    step = 0
    loss_per_step = []
    metrics_per_step = []
    for videos in val_dataloader:
        rng, _rng_mask = jax.random.split(rng, 2)
        inputs = dict(videos=videos, mask_rng=_rng_mask)
        loss, recon, metrics = val_step(genie, inputs)
        loss_per_step.append(loss)
        metrics_per_step.append(metrics)
        step += 1
        if step > args.val_steps:
            break

    if step < args.val_steps:
        print(f"Warning: Your validation dataset is too small to make val_steps many steps. Made {step} steps, expected {args.val_steps}")

    val_loss = np.mean(loss_per_step)
    val_metrics = {
        f"val_{key}": np.mean([m[key] for m in metrics_per_step])
        for key in metrics_per_step[0].keys()
    }
    return val_loss, val_metrics, inputs, recon