[content] input_pipeline/generate_atari_dataset.py (python)
    env_id: str = "BreakoutNoFrameskip-v4"
    num_episodes_train: int = 10000
    num_episodes_val: int = 500
    num_episodes_test: int = 500
    output_dir: str = "data/atari_episodes"
    min_episode_length: int = 1000
    max_episode_length: int = 1000
    chunk_size: int = 160
    chunks_per_file: int = 100
    seed: int = 0


args = tyro.cli(Args)
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."
    )


def make_env(env_id: str, seed: int):
    env = gym.make(env_id)
    env = gym.wrappers.RecordEpisodeStatistics(env)

    env = NoopResetEnv(env, noop_max=30)
    env = MaxAndSkipEnv(env, skip=4)
    env = EpisodicLifeEnv(env)
    if "FIRE" in env.unwrapped.get_action_meanings():
        env = FireResetEnv(env)
    env = ClipRewardEnv(env)

    env.action_space.seed(seed)
    return env


# --- Generate episodes ---
def generate_episodes(num_episodes, split):
    episode_idx = 0
    episode_metadata = []
    obs_chunks = []
    act_chunks = []
    file_idx = 0
    output_dir_split = os.path.join(args.output_dir, split)
    os.makedirs(output_dir_split, exist_ok=True)

    rng = np.random.default_rng(args.seed)
    while episode_idx < num_episodes:
        seed = int(rng.integers(0, 10_000))
        env = make_env(args.env_id, seed)

        observations_seq = []
        actions_seq = []
        episode_obs_chunks = []
        episode_act_chunks = []

        # --- Run episode ---
        step_t = 0
        obs, _ = env.reset(seed=seed)
        for step_t in range(args.max_episode_length):
            action = env.action_space.sample()
            obs, reward, terminated, truncated, info = env.step(action)
            # store as (1, H, W, C) to match concatenation logic
            observations_seq.append(obs[None, ...].astype(np.uint8))
            actions_seq.append(np.array([action], dtype=np.int64))

            if len(observations_seq) == args.chunk_size:
                episode_obs_chunks.append(observations_seq)
                episode_act_chunks.append(actions_seq)
                observations_seq = []
                actions_seq = []

            if terminated or truncated:
                break

        env.close()

        # --- Save episode ---
        if step_t + 1 >= args.min_episode_length:
            if observations_seq:
                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_obs_chunks.append(observations_seq)
                episode_act_chunks.append(actions_seq)

            obs_chunks_data = [
                np.concatenate(seq, axis=0).astype(np.uint8) for seq in episode_obs_chunks
            ]
            act_chunks_data = [np.concatenate(act, axis=0) for act in episode_act_chunks]
            obs_chunks.extend(obs_chunks_data)
            act_chunks.extend(act_chunks_data)

            ep_metadata, obs_chunks, file_idx, act_chunks = save_chunks(
                obs_chunks, file_idx, args.chunks_per_file, output_dir_split, act_chunks
            )
            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...")

    if len(obs_chunks) > 0:
        print(
            f"Warning: Dropping {len(obs_chunks)} chunks for consistent number of chunks per file.",
            "Consider changing the chunk_size and chunks_per_file parameters to prevent data-loss.",
        )

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


def get_action_space():
    env = make_env(args.env_id, args.seed)
    n = env.action_space.n if hasattr(env.action_space, "n") else None
    env.close()
    return int(n) if n is not None else 0


def main():
    # Set random seed and create dataset directories
    np.random.seed(args.seed)

    # --- 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": args.env_id,
        "num_actions": get_action_space(),
        "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": float(np.mean([ep["avg_seq_len"] for ep in train_episode_metadata]))
        if len(train_episode_metadata) > 0
        else 0.0,
        "avg_episode_len_val": float(np.mean([ep["avg_seq_len"] for ep in val_episode_metadata]))
        if len(val_episode_metadata) > 0
        else 0.0,
        "avg_episode_len_test": float(np.mean([ep["avg_seq_len"] for ep in test_episode_metadata]))
        if len(test_episode_metadata) > 0
        else 0.0,
        "episode_metadata_train": train_episode_metadata,
        "episode_metadata_val": val_episode_metadata,
        "episode_metadata_test": test_episode_metadata,
    }
    os.makedirs(args.output_dir, exist_ok=True)
    with open(os.path.join(args.output_dir, "metadata.json"), "w") as f:
        json.dump(metadata, f)

    print("Done generating dataset.")


if __name__ == "__main__":
    main()