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/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=spaceinvaders
[content] input_pipeline/video_to_array_records.py (python)

    train_episode_metadata = save_split(pool_args_train, args.chunks_per_file, os.path.join(args.output_path, "train"))
    val_episode_metadata = save_split(pool_args_val, args.chunks_per_file, os.path.join(args.output_path, "val"))
    test_episode_metadata = save_split(pool_args_test, args.chunks_per_file, os.path.join(args.output_path, "test"))
[tab] slurm/jobs/mihir/horeka/preprocessing/preprocess_minecraft_multi_chunk.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/minecraft/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/preprocess/minecraft/%x_%j.log
#SBATCH --job-name=preprocess_minecraft_multi_chunk

source .venv/bin/activate

python input_pipeline/preprocess/video_to_array_records.py \
        --input_path /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft \
        --output_path /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/tmp/multi_chunk \
        --env_name minecraft \
        --chunk_size 160 \
        --chunks_per_file 100
[content] input_pipeline/video_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,
[content] models/dynamics.py (python)

        vid_embed_BTNM = self.patch_embed(video_tokens_BTN)
        act_embed_BTm11M = self.action_up(latent_actions_BTm11L)
        padded_act_embed_BT1M = jnp.pad(
            act_embed_BTm11M, ((0, 0), (1, 0), (0, 0), (0, 0))
        )
        vid_embed_BTNp1M = jnp.concatenate(
            [padded_act_embed_BT1M, vid_embed_BTNM], axis=2
        )