[content] generate_dataset.py (python)

def generate_episodes(num_episodes, split):
    episode_idx = 0
    episode_metadata = []
    while episode_idx < 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,)))
            _ , 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 = os.path.join(args.output_dir, split, 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...")

    print(f"Done generating {split} split")
    return episode_metadata


def main():
    # Set random seed and create dataset directories
    np.random.seed(args.seed)
    output_dir = Path(args.output_dir)
    (output_dir / "train").mkdir(parents=True, exist_ok=True)
    (output_dir / "val").mkdir(parents=True, exist_ok=True)
    (output_dir / "test").mkdir(parents=True, exist_ok=True)

    # --- Generate episodes ---
    train_episode_metadata = generate_episodes(args.num_episodes_train, "train")
    val_episode_metadata = generate_episodes(args.num_episodes_val, "val")
    test_episode_metadata = generate_episodes(args.num_episodes_test, "test")

    # --- Save metadata ---
    metadata = {
        "env": "coinrun",
        "num_episodes_train": args.num_episodes_train,
        "num_episodes_val": args.num_episodes_val,
        "num_episodes_test": args.num_episodes_test,
        "avg_episode_len_train": np.mean([ep["length"] for ep in train_episode_metadata]),
        "avg_episode_len_val": np.mean([ep["length"] for ep in val_episode_metadata]),
        "avg_episode_len_test": np.mean([ep["length"] for ep in test_episode_metadata]),
        "episode_metadata_train": train_episode_metadata,
        "episode_metadata_val": val_episode_metadata,
        "episode_metadata_test": test_episode_metadata,

    }
    with open(output_dir / "metadata.json", "w") as f:
        json.dump(metadata, f)

    print(f"Done generating dataset.")

if __name__ == "__main__":
    main()
assert args.max_episode_length >= args.min_episode_length, "Maximum episode length must be greater than or equal to minimum episode length."

if args.min_episode_length < args.chunk_size:
    print("Warning: Minimum episode length is smaller than chunk size. Note that episodes shorter than the chunk size will be discarded.")

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

def _save_chunks(chunks, file_idx):
    ep_metadata = []
    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

        ep_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")
    return ep_metadata, chunks, file_idx

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

    # --- Run episode ---
    for step_t 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 len(observations_seq) == args.chunk_size:
            episode_chunks.append(observations_seq)
            observations_seq = []
        if first:
            break

    # --- Save episode ---
    if step_t + 1  >= args.min_episode_length:
        if len(observations_seq) < args.chunk_size:
            print(
                f"Warning: Inconsistent chunk_sizes. Episode has {len(observations_seq)} frames, "
                f"which is smaller than the requested chunk_size: {args.chunk_size}. "
                "This might lead to performance degradation during training."
            )
        episode_chunks.append(observations_seq)
        chunks_data = [np.concatenate(seq, axis=0).astype(np.uint8) for seq in episode_chunks]
        chunks.extend(chunks_data)

        ep_metadata, chunks, file_idx = _save_chunks(chunks, file_idx)
        episode_metadata.extend(ep_metadata)

        print(f"Episode {episode_idx} completed, length: {step_t + 1}.")
        episode_idx += 1
    else:
        print(f"Episode too short ({step_t + 1}), resampling...")
ep_metadata, chunks, file_idx = _save_chunks(chunks, file_idx)
episode_metadata.extend(ep_metadata)
# --- Save metadata ---
metadata = {
    "env": "coinrun",
    "num_episodes": args.num_episodes,
    "avg_episode_len": np.mean([ep["avg_seq_len"] 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 {args.num_episodes} episodes, saved over {len(episode_metadata)} files")