p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] input_pipeline/generate_atari_dataset.py (python)
                            (
                                ep_metadata,
                                file_idx_by_split[current_split],
                                obs_chunks,
                                act_chunks,
                            ) = save_chunks(
                                file_idx_by_split[current_split],
                                args.chunks_per_file,
                                split_dir,
                                obs_chunks,
                                act_chunks,
                            )
[tab] input_pipeline/utils.py (python)
import os
import pickle
import numpy as np
from array_record.python.array_record_module import ArrayRecordWriter


def save_chunks(file_idx, chunks_per_file, output_dir, obs_chunks, act_chunks=None):
    os.makedirs(output_dir, exist_ok=True)

    metadata = []
    while len(obs_chunks) >= chunks_per_file:
        chunk_batch = obs_chunks[:chunks_per_file]
        obs_chunks = obs_chunks[chunks_per_file:]
        act_chunk_batch = None
        if act_chunks:
            act_chunk_batch = act_chunks[:chunks_per_file]
            act_chunks = act_chunks[chunks_per_file:]
        episode_path = os.path.join(output_dir, f"data_{file_idx:04d}.array_record")
        writer = ArrayRecordWriter(str(episode_path), "group_size:1")
        seq_lens = []
        for idx, chunk in enumerate(chunk_batch):
            seq_len = chunk.shape[0]
            seq_lens.append(seq_len)
            chunk_record = {
                "raw_video": chunk.tobytes(),
                "sequence_length": seq_len,
            }
            if act_chunk_batch:
                assert len(chunk) == len(
                    act_chunk_batch[idx]
                ), f"Observation data length and action sequence length do not match: {len(chunk)} != {len(act_chunk_batch[idx])}"
                chunk_record["actions"] = act_chunk_batch[idx]
            writer.write(pickle.dumps(chunk_record))
        writer.close()
        file_idx += 1
        metadata.append(
            {
                "path": episode_path,
                "num_chunks": len(chunk_batch),
                "avg_seq_len": np.mean(seq_lens),
            }
        )
        print(f"Created {episode_path} with {len(chunk_batch)} video chunks")

    return metadata, file_idx, obs_chunks, act_chunks
[content] input_pipeline/generate_atari_dataset.py (python)
                        # Append the final post-action observation once to capture all observations
                        final_obs_to_append = next_obs.copy()
                        # If an episode finished, next_obs for that index is already the reset obs.
                        # Swap in the true terminal obs from infos["final_observation"] for any done envs.
                        done_mask = np.logical_or(terminations, truncations)
                        for idx, done in enumerate(done_mask):
                            if done:
                                final_obs_to_append[idx] = infos["final_observation"][idx]
                        observations_seq.append(final_obs_to_append.astype(np.uint8))

                        # Ensure actions and observations have the same length for saving
                        if len(actions_seq) + 1 == len(observations_seq):
                            actions_seq.append(actions_seq[-1])
[content] genie.py (python)
        lam_outputs = self.lam.vq_encode(videos_BTHWC, training=False)
        z_q_BTm11L = lam_outputs["z_q"]
        action_indices_E = lam_outputs["indices"]
        latent_actions_BTm11L = jax.lax.cond(
            self.lam_co_train,
            lambda: z_q_BTm11L,
            lambda: jax.lax.stop_gradient(z_q_BTm11L),
        )
[tab] slurm/jobs/franz/berlin/coinrun/coinrun_lam/lam_coinrun_nan_investigation_100k_to_107k.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/lam/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/lam/%x_%j.log
#SBATCH --job-name=lam_coinrun_nan_investigation_100k_to_107k


# 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=/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

tags="lam coinrun 38M nan_investigation"

CHECKPOINT_PATH="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/lam/train_lam_coinrun_reproduction_20067/100000_ckpt"

env | grep SLURM

srun python train_lam.py \
    --save_ckpt \
    --restore_ckpt \
    --wandb_id $SLURM_JOB_ID \
    --ckpt_dir $CHECKPOINT_PATH \
    --batch_size=36 \
    --image_height=64 \
    --image_width=64 \
    --max_lr=3e-5 \
    --decay_end=3e-5 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --num_steps=107885 \
    --log \
    --name="${job_name}_${slurm_job_id}" \
    --tags $tags \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir
[tab] slurm/jobs/franz/berlin/coinrun/coinrun_tokenizer/coinrun_tokenizer_repoduction.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/tokenizer/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/tokenizer/%x_%j.log
#SBATCH --job-name=train_tokenizer_coinrun_reproduction_codebook_normal_init

# Log the sbatch script
cat $0

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

tags="coinrun reproduction 10m_dataset tokenizer codebook_normal_init"

array_records_dir="/fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10m"

CHECKPOINT_DIR="/fast/project/HFMI_SynergyUnit/jafar_ws/checkpoints/coinrun/tokenizer/${job_name}_${slurm_job_id}"
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

srun python train_tokenizer.py \
    --wandb_id $SLURM_JOB_ID \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=48 \
    --image_height=64 \
    --image_width=64 \
    --init_lr=0 \
    --max_lr=3e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name="${job_name}_${slurm_job_id}" \
    --tags $tags \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir &

child_pid=$!

wait $child_pid
[content] slurm/jobs/franz/berlin/atari/atari_data_gen.sh (shellscript)

#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/tokenizer/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/coinrun/tokenizer/%x_%j.log
#SBATCH --job-name=train_tokenizer_coinrun_reproduction_codebook_normal_init