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

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=48:00: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/tokenizer/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/tokenizer/%x_%j.log
#SBATCH --job-name=train_tokenizer_default
#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_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/tokenizer/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

export PYTHONUNBUFFERED=1

srun python train_tokenizer.py \
    --save_ckpt \
    --image_height=64 \
    --image_width=64 \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=120 \
    --init_lr=0 \
    --patch_size=16 \
    --log \
    --name=coinrun-tokenizer-default-$slurm_job_id \
    --tags tokenizer coinrun default \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \
[tab] slurm/jobs/mihir/horeka/coinrun/default_runs/train_lam_default.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=48:00: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/lam/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/lam/%x_%j.log
#SBATCH --job-name=train_lam_default
#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_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/lam/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

export PYTHONUNBUFFERED=1

srun python train_lam.py \
    --save_ckpt \
    --image_height=64 \
    --image_width=64 \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=120 \
    --log \
    --name=coinrun-lam-default-$slurm_job_id \
    --tags lam coinrun default \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val
[tab] slurm/jobs/mihir/horeka/coinrun/default_runs/train_lam_default copy.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=48:00: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/lam/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/lam/%x_%j.log
#SBATCH --job-name=train_lam_default
#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_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/lam/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

export PYTHONUNBUFFERED=1

srun python train_lam.py \
    --save_ckpt \
    --image_height=64 \
    --image_width=64 \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=120 \
    --log \
    --name=coinrun-lam-default-$slurm_job_id \
    --tags lam coinrun default \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val
[tab] slurm/jobs/mihir/horeka/coinrun/default_runs/train_lam_as_15.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=48:00: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/lam/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/lam/%x_%j.log
#SBATCH --job-name=train_lam_default
#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_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/lam/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

export PYTHONUNBUFFERED=1

srun python train_lam.py \
    --save_ckpt \
    --image_height=64 \
    --image_width=64 \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=120 \
    --log \
    --name=coinrun-lam-default-$slurm_job_id \
    --tags lam coinrun default \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val
[tab] slurm/jobs/mihir/horeka/coinrun/train_dyn_single_gpu_gt_actions.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=00:20:00
#SBATCH --partition=dev_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_lam_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_test
array_records_dir_train=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_test/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_test/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/3498707
tokenizer_checkpoint=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer/interactive/3498707

env | grep SLURM

export PYTHONUNBUFFERED=1

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=50 \
    --log_checkpoint_interval=2 \
    --dyna_type=maskgit \
    --log \
    --name=coinrun-dyn-dev-gt-actions-$slurm_job_id \
    --tags dyn coinrun dev \
    --entity instant-uv \
    --project jafar \
    --warmup_steps 0 \
    --wsd_decay_steps 0 \
    --num_steps 10 \
    --data_dir $array_records_dir_train \
    --tokenizer_checkpoint $tokenizer_checkpoint \
    --use_gt_actions \
    --val_data_dir $array_records_dir_val \
    --val_interval 2 \
    --val_steps 5
[content] input_pipeline/generate_coinrun_dataset.py (python)

    if len(chunks) > 0:
        print(f"Warning: Dropping {len(chunks)} chunks for consistent number of chunks per file.",
        "Consider changing the chunk_size and chunks_per_file parameters to prevent data-loss.")
[content] input_pipeline/generate_coinrun_dataset.py (python)
def generate_episodes(num_episodes, split):
    episode_idx = 0
    episode_metadata = []
    obs_chunks = []
    act_chunks = []
    file_idx = 0
    output_dir_split = os.path.join(args.output_dir, split)
    while episode_idx < num_episodes:
        seed = np.random.randint(0, 10000)
        env = ProcgenGym3Env(num=1, env_name="coinrun", start_level=seed)

        observations_seq = []
        actions_seq = []
        episode_obs_chunks = []
        episode_act_chunks = []
[content] input_pipeline/generate_coinrun_dataset.py (python)
        # --- Run episode ---
        step_t = 0
        for step_t in range(args.max_episode_length):
            action = types_np.sample(env.ac_space, bshape=(env.num,))
            env.act(action)
            _, obs, first = env.observe()
            observations_seq.append(obs["rgb"])
            actions_seq.append(action)
            if len(observations_seq) == args.chunk_size:
                episode_obs_chunks.append(observations_seq)
                episode_act_chunks.append(actions_seq)
                observations_seq = []
                actions_seq = []
            if first:
                break
[content] input_pipeline/generate_coinrun_dataset.py (python)
        # --- Save episode ---
        if step_t + 1 >= args.min_episode_length:
            if observations_seq:
                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_obs_chunks.append(observations_seq)
                episode_act_chunks.append(actions_seq)

            obs_chunks_data = [
                np.concatenate(seq, axis=0).astype(np.uint8)
                for seq in episode_obs_chunks
            ]
            act_chunks_data = [
                np.concatenate(act, axis=0) for act in episode_act_chunks
            ]
            obs_chunks.extend(obs_chunks_data)
            act_chunks.extend(act_chunks_data)
[content] input_pipeline/generate_coinrun_dataset.py (python)
            ep_metadata, obs_chunks, file_idx, act_chunks = save_chunks(
                obs_chunks, file_idx, args.chunks_per_file, output_dir_split, act_chunks
            )
            episode_metadata.extend(ep_metadata)
[content] input_pipeline/generate_coinrun_dataset.py (python)
            print(f"Episode {episode_idx} completed, length: {step_t + 1}.")
            episode_idx += 1
        else:
            print(f"Episode too short ({step_t + 1}), resampling...")

    if len(obs_chunks) > 0:
        print(
            f"Warning: Dropping {len(obs_chunks)} chunks for consistent number of chunks per file.",
            "Consider changing the chunk_size and chunks_per_file parameters to prevent data-loss.",
        )

    print(f"Done generating {split} split")
    return episode_metadata


def get_action_space():
    env = ProcgenGym3Env(num=1, env_name="coinrun", start_level=0)
    return env.ac_space.eltype.n


def main():
    # Set random seed and create dataset directories
    np.random.seed(args.seed)
    # --- 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_actions": get_action_space(),
        "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["avg_seq_len"] for ep in train_episode_metadata]
        ),
        "avg_episode_len_val": np.mean(
            [ep["avg_seq_len"] for ep in val_episode_metadata]
        ),
        "avg_episode_len_test": np.mean(
            [ep["avg_seq_len"] 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(os.path.join(args.output_dir, "metadata.json"), "w") as f:
        json.dump(metadata, f)
[content] train_dynamics.py (python)
    val_data_dir: str = ""
    val_interval: int = 20_000
    val_steps: int = 50
    eval_full_frame: bool = False
    val_maskgit_steps: int = 25
    val_temperature: float = 1
    val_sample_argmax: bool = False
[content] train_dynamics.py (python)
        if val_iterator:
            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
[content] train_dynamics.py (python)
    return step, optimizer, train_iterator, val_iterator, rng


def _calculate_step_metrics(
    outputs: dict[str, jax.Array],
    gt: jax.Array,
    num_latent_actions: int,
    num_patch_latents: int,
) -> tuple[jax.Array, dict]:
    mask = outputs["mask"]
    outputs["token_logits"] = outputs["token_logits"].astype(jnp.float32)
    ce_loss = optax.softmax_cross_entropy_with_integer_labels(
        outputs["token_logits"], outputs["video_tokens"]
    )
    ce_loss = (mask * ce_loss).sum() / mask.sum()
    acc = outputs["token_logits"].argmax(-1) == outputs["video_tokens"]
    acc = (mask * acc).sum() / mask.sum()
    select_probs = jax.nn.softmax(outputs["token_logits"])
    gt_val = gt.clip(0, 1).reshape(-1, *gt.shape[2:])
    recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
    psnr = jnp.asarray(pix.psnr(gt_val, recon)).mean()
    ssim = jnp.asarray(pix.ssim(gt_val, recon)).mean()
    _, index_counts_tokenizer = jnp.unique_counts(
        jnp.ravel(outputs["video_tokens"]),
        size=num_patch_latents,
        fill_value=0,
    )
    codebook_usage_tokenizer = (index_counts_tokenizer != 0).mean()
    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_tokenizer=codebook_usage_tokenizer,
    )
    if "lam_indices" in outputs.keys():
        _, index_counts_lam = jnp.unique_counts(
            jnp.ravel(outputs["lam_indices"]),
            size=num_latent_actions,
            fill_value=0,
        )
        codebook_usage_lam = (index_counts_lam != 0).mean()
        metrics["codebook_usage_lam"] = codebook_usage_lam
    return ce_loss, metrics
[content] train_dynamics.py (python)
    step, optimizer, train_iterator, val_iterator, rng = (
        restore_or_initialize_components(
            args,
            checkpoint_manager,
            optimizer,
            train_iterator,
            rng,
            replicated_sharding,
            val_iterator,
        )
[content] train_dynamics.py (python)
    dataloader_train = (
        {
            "videos": jax.make_array_from_process_local_data(
                videos_sharding, local_data=elem["videos"]
            ),
            "actions": (
                jax.make_array_from_process_local_data(
                    actions_sharding, elem["actions"]
                )
                if args.use_gt_actions
                else None
            ),
        }
        for elem in train_iterator
    )
    dataloader_val = None
    if val_iterator:
        dataloader_val = (
            {
                "videos": jax.make_array_from_process_local_data(
                    videos_sharding, elem["videos"]
                ),
                "actions": (
                    jax.make_array_from_process_local_data(
                        actions_sharding, elem["actions"]
                    )
                    if args.use_gt_actions
                    else None
                ),
            }
            for elem in val_iterator
        )
[content] train_lam.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
            )
        restored = checkpoint_manager.restore(restore_step, args=restore_args)
[content] train_lam.py (python)
    if val_iterator:
        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)
    @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 = []
        batch = None
        recon = None
        for batch in val_dataloader:
            loss, recon, metrics = val_step(lam, batch)
            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([float(m[key]) for m in metrics_per_step])
            for key in metrics_per_step[0].keys()
        }
        val_metrics["val_loss"] = val_loss
        return val_metrics, batch, recon
[content] train_lam.py (python)
    dataloader_train = (
        {
            "videos": jax.make_array_from_process_local_data(
                videos_sharding, elem["videos"]
            ),
        }
        for elem in train_iterator
    )
    dataloader_val = None
    if val_iterator:
        dataloader_val = (
            {
                "videos": jax.make_array_from_process_local_data(
                    videos_sharding, elem["videos"]
                ),
            }
            for elem in val_iterator
        )
[content] train_lam.py (python)
            # --- Validation loss ---
            val_results = {}
            if dataloader_val and step % args.val_interval == 0:
                print("Calculating validation metrics...")
                val_metrics, val_gt_batch, val_recon = calculate_validation_metrics(
                    dataloader_val, optimizer.model
                )
                print(f"Step {step}, validation loss: {val_metrics['val_loss']}")
                val_results = {
                    "metrics": val_metrics,
                    "gt_batch": val_gt_batch,
                    "recon": val_recon,
                }
[content] train_lam.py (python)
                    log_dict = {"loss": loss, "step": step, **metrics}
                    if val_results:
                        log_dict.update(val_results["metrics"])
                    wandb.log(log_dict)
[content] train_lam.py (python)
                    )
                    if val_results:
                        val_results["gt_seq_val"] = (
                            val_results["gt_batch"]["videos"][0, 1:].astype(jnp.float32)
                            / 255.0
                        )
                        val_results["recon_seq_val"] = val_results["recon"][0].clip(
                            0, 1
                        )
                        val_results["val_comparison_seq"] = jnp.concatenate(
                            (val_results["gt_seq_val"], val_results["recon_seq_val"]),
                            axis=1,
                        )
                        val_results["val_comparison_seq"] = einops.rearrange(
                            val_results["val_comparison_seq"] * 255,
                            "t h w c -> h (t w) c",
                        )
[content] train_lam.py (python)
                        )
                        if val_results:
                            log_images.update(
                                {
                                    "val_image": wandb.Image(
                                        np.asarray(val_results["gt_seq_val"][0])
                                    ),
                                    "val_recon": wandb.Image(
                                        np.asarray(val_results["recon_seq_val"][0])
                                    ),
                                    "val_true_vs_recon": wandb.Image(
                                        np.asarray(
                                            val_results["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_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
            )
        restored = checkpoint_manager.restore(restore_step, args=restore_args)
[content] train_tokenizer.py (python)
    if val_iterator:
        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_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, training=False)
        return loss, recon, metrics

    def calculate_validation_metrics(val_dataloader, tokenizer):
        step = 0
        loss_per_step = []
        metrics_per_step = []
        batch = None
        recon = None
        for batch in val_dataloader:
            loss, recon, metrics = val_step(tokenizer, batch)
            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([float(m[key]) for m in metrics_per_step])
            for key in metrics_per_step[0].keys()
        }
        val_metrics["val_loss"] = val_loss
        return val_metrics, batch, recon
[content] train_tokenizer.py (python)
    dataloader_train = (
        {
            "videos": jax.make_array_from_process_local_data(
                videos_sharding, elem["videos"]
            ),
        }
        for elem in train_iterator
    )
    dataloader_val = None
    if val_iterator:
        dataloader_val = (
            {
                "videos": jax.make_array_from_process_local_data(
                    videos_sharding, elem["videos"]
                ),
            }
            for elem in val_iterator
        )
[content] train_tokenizer.py (python)
            # --- Validation loss ---
            val_results = {}
            if dataloader_val and step % args.val_interval == 0:
                print("Calculating validation metrics...")
                val_metrics, val_gt_batch, val_recon = calculate_validation_metrics(
                    dataloader_val, optimizer.model
                )
                print(f"Step {step}, validation loss: {val_metrics['val_loss']}")
                val_results = {
                    "metrics": val_metrics,
                    "gt_batch": val_gt_batch,
                    "recon": val_recon,
                }
[content] train_tokenizer.py (python)
                    )
                    if args.val_data_dir and step % args.val_interval == 0:
                        val_results["gt_seq_val"] = (
                            val_results["gt_batch"]["videos"][0].astype(jnp.float32)
                            / 255.0
                        )
                        val_results["recon_seq_val"] = val_results["recon"][0].clip(
                            0, 1
                        )
                        val_results["val_comparison_seq"] = jnp.concatenate(
                            (val_results["gt_seq_val"], val_results["recon_seq_val"]),
                            axis=1,
                        )
                        val_results["val_comparison_seq"] = einops.rearrange(
                            val_results["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(val_results["gt_seq_val"][0])
                                    ),
                                    val_recon=wandb.Image(
                                        np.asarray(val_results["recon_seq_val"][0])
                                    ),
                                    val_true_vs_recon=wandb.Image(
                                        np.asarray(
                                            val_results["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)
                    log_dict = {"loss": loss, "step": step, **metrics}
                    if val_results:
                        log_dict.update(val_results["metrics"])
                    wandb.log(log_dict)
[content] train_dynamics.py (python)
                val_metrics, val_gt_batch, val_recon, val_recon_full_frame = (
                    calculate_validation_metrics(
                        dataloader_val, optimizer.model, _rng_mask_val
                    )
[tab] slurm/jobs/mihir/horeka/coinrun/train_dyn_single_gpu_gt_actions.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=00:20:00
#SBATCH --partition=dev_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_lam_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_test
array_records_dir_train=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_test/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_test/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/3498707
tokenizer_checkpoint=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer/interactive/3498707

env | grep SLURM

export PYTHONUNBUFFERED=1

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=50 \
    --log_checkpoint_interval=2 \
    --dyna_type=maskgit \
    --log \
    --name=coinrun-dyn-dev-gt-actions-$slurm_job_id \
    --tags dyn coinrun dev \
    --entity instant-uv \
    --project jafar \
    --warmup_steps 0 \
    --wsd_decay_steps 0 \
    --num_steps 10 \
    --data_dir $array_records_dir_train \
    --tokenizer_checkpoint $tokenizer_checkpoint \
    --val_data_dir $array_records_dir_val \
    --val_interval 2 \
    --val_steps 5
[tab] slurm/jobs/mihir/horeka/coinrun/train_tokenizer_single_gpu.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=00:20:00
#SBATCH --partition=dev_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/tokenizer/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/coinrun/tokenizer/%x_%j.log
#SBATCH --job-name=train_tokenizer_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_test
array_records_dir_train=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_test/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_test/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/tokenizer/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

export PYTHONUNBUFFERED=1

srun python train_tokenizer.py \
    --save_ckpt \
    --image_height=64 \
    --image_width=64 \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=120 \
    --init_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=50 \
    --log_checkpoint_interval=2 \
    --log \
    --name=coinrun-tokenizer-dataset-test-$slurm_job_id \
    --tags tokenizer coinrun dev \
    --entity instant-uv \
    --project jafar \
    --warmup_steps 0 \
    --wsd_decay_steps 0 \
    --num_steps 10 \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \
    --val_interval 3 \
    --val_steps 5