p-doom/crowd-code-dataset-0.2 huggingface dataset:
[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/spaceinvaders \
    --output_path=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_atari/array_records_test \
    --env_name=atari
[content] input_pipeline/pngs_to_array_records.py (python)

    chunks = []
    file_idx = 0
    results = []
    for bucket_idx in range(0, len(pool_args), num_processes):
        args_batch = pool_args[bucket_idx : bucket_idx + num_processes]
        with mp.Pool(processes=num_processes) as pool:
            for episode_chunks in pool.starmap(preprocess_pngs, args_batch):
                chunks.extend(episode_chunks)
        results_batch, chunks, file_idx = save_chunks(chunks, file_idx, args.chunks_per_file, args.output_path) 
        results.extend(results_batch)
[content] input_pipeline/pngs_to_array_records.py (python)

        "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,