p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] generate_dataset.py (python)
import numpy as np
from procgen import ProcgenGym3Env
import tyro
import pickle
import json
from array_record.python.array_record_module import ArrayRecordWriter 



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


args = tyro.cli(Args)
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)

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

    # --- Run episode ---
    for _ 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 first:
            break

    # --- Save episode ---
    if len(observations_seq) >= args.min_episode_length:
        observations_data = np.concatenate(observations_seq, axis=0).astype(np.uint8)

        file_chunks = []
        for start_idx in range(0, observations_data.shape[0] - args.chunk_size + 1, args.chunk_size):
            chunk = observations_data[start_idx : start_idx + args.chunk_size]

            chunk_record = {
                "raw_video": chunk.tobytes(),
                "sequence_length": args.chunk_size,
            }

            file_chunks.append(chunk_record)
        # --- Save as ArrayRecord ---
        for file_idx in range(0, len(file_chunks), args.chunks_per_file):
            batch_chunks = file_chunks[file_idx : file_idx + args.chunks_per_file]
            episode_path = output_dir / f"episode_{episode_idx:04d}_part_{file_idx:04d}.array_record"  
            writer = ArrayRecordWriter(str(episode_path), "group_size:1")
            for chunk in batch_chunks:
                writer.write(pickle.dumps(chunk))
            writer.close()

            episode_metadata.append({"path": episode_path, "chunk_size": args.chunk_size, "num_chunks": len(batch_chunks)})
            print(f"Created {episode_path} with {len(batch_chunks)} video chunks")
        print(f"Episode {episode_idx} completed, length: {len(observations_seq)}. Saved across {len(file_chunks) // args.chunks_per_file} array_records files.")
        episode_idx += 1
    else:
        print(f"Episode too short ({len(observations_seq)}), resampling...")

# --- Save metadata ---
metadata = {
    "env": "coinrun",
    "num_episodes": args.num_episodes,
    "avg_episode_len": np.mean([ep["length"] 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 {len(episode_metadata)} valid episodes")
[tab] 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 pathlib import Path

from gym3 import types_np
import numpy as np
from procgen import ProcgenGym3Env
import tyro
import pickle
import json
from array_record.python.array_record_module import ArrayRecordWriter 



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


args = tyro.cli(Args)
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)

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

    # --- Run episode ---
    for _ 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 first:
            break

    # --- Save episode ---
    if len(observations_seq) >= args.min_episode_length:
        observations_data = np.concatenate(observations_seq, axis=0).astype(np.uint8)

        file_chunks = []
        for start_idx in range(0, observations_data.shape[0] - args.chunk_size + 1, args.chunk_size):
            chunk = observations_data[start_idx : start_idx + args.chunk_size]

            chunk_record = {
                "raw_video": chunk.tobytes(),
                "sequence_length": args.chunk_size,
            }

            file_chunks.append(chunk_record)
        # --- Save as ArrayRecord ---
        for file_idx in range(0, len(file_chunks), args.chunks_per_file):
            batch_chunks = file_chunks[file_idx : file_idx + args.chunks_per_file]
            episode_path = output_dir / f"episode_{episode_idx:04d}_part_{file_idx:04d}.array_record"  
            writer = ArrayRecordWriter(str(episode_path), "group_size:1")
            for chunk in batch_chunks:
                writer.write(pickle.dumps(chunk))
            writer.close()

            episode_metadata.append({"path": episode_path, "chunk_size": args.chunk_size, "num_chunks": len(batch_chunks)})
            print(f"Created {episode_path} with {len(batch_chunks)} video chunks")
        print(f"Episode {episode_idx} completed, length: {len(observations_seq)}. Saved across {len(file_chunks) // args.chunks_per_file} array_records files.")
        episode_idx += 1
    else:
        print(f"Episode too short ({len(observations_seq)}), resampling...")

# --- Save metadata ---
metadata = {
    "env": "coinrun",
    "num_episodes": args.num_episodes,
    "avg_episode_len": np.mean([ep["length"] 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 {len(episode_metadata)} valid episodes")
[content] input_pipeline/preprocess/pngs_to_array_records_bucket.py (python)

        while len(chunks) >= args.chunks_per_file:
            chunk_batch = chunks[:args.chunks_per_file]
            chunks = chunks[args.chunks_per_file:]
            episode_path = output_dir / f"coinrun_episodes_{file_idx:04d}.array_record"  
            # --- Save as ArrayRecord ---
            writer = ArrayRecordWriter(str(episode_path), "group_size:1")
            for chunk in chunk_batch:
                chunk_record = {
                    "raw_video": chunk.tobytes(),
                    "sequence_length": chunk.shape[0],
                }
                writer.write(pickle.dumps(chunk_record))
            writer.close()
            file_idx += 1
[tab] slurm/jobs/mihir/horeka/preprocessing/preprocess_atari.sbatch (shellscript)
#!/usr/bin/env bash

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

source .venv/bin/activate

python input_pipeline/preprocess/pngs_to_array_records.py \
    --input_path=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_atari/per-game/atari_v1/screens \
    --output_path=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_atari/array_records \
    --env_name=atari
[content] input_pipeline/preprocess/pngs_to_array_records.py (python)
      print(
            f"Warning: Inconsistent chunk_sizes. Episode has {current_episode_len} frames, "
            f"which is smaller than the requested chunk_size: {chunk_size}. "
            "This might lead to performance degradation during training."
        )
[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 generate_dataset_chunk.py --min_episode_length 1000 --output_dir /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_coinrun/coinrun_episodes_10m_chunked_2
[content] generate_dataset.py (python)

        if len(frames) < chunk_size:
            print(
                f"Warning: Inconsistent chunk_sizes. Episode has {len(frames)} frames, "
                f"which is smaller than the requested chunk_size: {chunk_size}. "
                "This might lead to performance degradation during training."
            )
[content] generate_dataset.py (python)

        while len(chunks) >= args.chunks_per_file:
            chunk_batch = chunks[:args.chunks_per_file]
            chunks = chunks[args.chunks_per_file:]
            episode_path = output_dir / f"coinrun_episodes_{file_idx:04d}.array_record"  
            # --- Save as ArrayRecord ---
            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

            episode_metadata.append({"path": str(episode_path), "avg_seq_len": np.mean(seq_lens), "num_chunks": len(chunk_batch)})
            print(f"Created {episode_path} with {len(chunk_batch)} video chunks")