p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] input_pipeline/generate_atari_dataset.py (python)

"""
Generates a dataset of random-action Atari episodes.
Episodes are saved individually as memory-mapped files for efficient loading.
Replicates the behavior of generate_coinrun_dataset.py but for Atari.
"""

from dataclasses import dataclass

import gymnasium as gym
import numpy as np
import tyro
import json
import os
from cleanrl_utils.atari_wrappers import (
    ClipRewardEnv,
    EpisodicLifeEnv,
    FireResetEnv,
    MaxAndSkipEnv,
    NoopResetEnv,
)
from utils import save_chunks  # type: ignore


"""
Old dataset-only generator removed in favor of integrated Rainbow + capture mode.
"""
[tab] slurm/jobs/alfred/berlin/coinrun/coinrun_bigrun/generate_data.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/ali/coinrun/tokenizer/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/ali/coinrun/tokenizer/%x_%j.log
#SBATCH --job-name=train_tokenizer_coinrun_reproduction
#SBATCH --requeue
#SBATCH --signal=b:usr1@300 # 5 min before timeout


python generate_dataset_arr_records.py --num_episodes 10000 --output_dir /fast/project/HFMI_SynergyUnit/jafar_ws/data/coinrun/array_records_10k
[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(chunks, file_idx, chunks_per_file, output_dir):
    os.makedirs(output_dir, exist_ok=True)

    metadata = []
    while len(chunks) >= chunks_per_file:
        chunk_batch = chunks[:chunks_per_file]
        chunks = 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 chunk in chunk_batch:
            seq_len = chunk.shape[0]
            seq_lens.append(seq_len)
            chunk_record = {
                "raw_video": chunk.tobytes(),
                "sequence_length": seq_len,
            }
            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, chunks, file_idx