[tab] input_pipeline/generate_breakout_dataset_agent.py (python)
"""
Train a DQN agent on MinAtar-Breakout and collect episodes to build a dataset.
Episodes are chunked and stored using save_chunks.
"""

from dataclasses import dataclass
import numpy as np
import os
import json
import tyro
from typing import List, Tuple

from minatar import Environment          # pip install minatar
from gym import spaces
import gym
from stable_baselines3 import DQN        # pip install stable-baselines3[extra]

from utils import save_chunks            # reuse your existing utility


# ---- Wrap MinAtar in a Gym API ----
class MinAtarBreakout(gym.Env):
    def __init__(self):
        super().__init__()
        self.env = Environment("breakout", sticky_action_prob=0.1)
        self.action_space = spaces.Discrete(self.env.num_actions())
        self.observation_space = spaces.Box(
            low=0, high=1, shape=self.env.state().shape, dtype=np.float32
        )

    def reset(self):
        self.env.reset()
        return self.env.state()

    def step(self, action):
        reward, done = self.env.act(action)
        return self.env.state(), reward, done, {}


@dataclass
class Args:
    output_dir: str = "data/breakout_episodes"
    num_episodes: int = 2000
    min_episode_length: int = 100
    max_episode_length: int = 500
    chunk_size: int = 50
    chunks_per_file: int = 100
    train_timesteps: int = 200_000   # DQN training timesteps
    collection_interval: int = 10_000  # Collect episodes every N training steps
    seed: int = 0


args = tyro.cli(Args)


def collect_episodes(model: DQN, num_eps: int) -> Tuple[List[np.ndarray], List[np.ndarray], List[dict]]:
    obs_chunks, act_chunks, ep_metadata = [], [], []
    for ep in range(num_eps):
        env = MinAtarBreakout()
        obs = env.reset()
        done = False
        ep_obs, ep_act = [], []
        steps = 0
        while not done and steps < args.max_episode_length:
            action, _ = model.predict(obs, deterministic=False)
            ep_obs.append(obs.astype(np.uint8))
            ep_act.append(action)
            obs, reward, done, _ = env.step(action)
            steps += 1

        if steps >= args.min_episode_length:
            # Chunk episode
            episode_obs_chunks, episode_act_chunks = [], []
            for i in range(0, len(ep_obs), args.chunk_size):
                o_chunk = np.stack(ep_obs[i:i + args.chunk_size])
                a_chunk = np.array(ep_act[i:i + args.chunk_size])
                episode_obs_chunks.append(o_chunk)
                episode_act_chunks.append(a_chunk)

            obs_chunks.extend(episode_obs_chunks)
            act_chunks.extend(episode_act_chunks)

            avg_len = np.mean([len(ep_obs)])
            ep_metadata.append({"avg_seq_len": avg_len})
        else:
            print(f"Episode too short ({steps}), discarded.")

    return obs_chunks, act_chunks, ep_metadata


def main():
    np.random.seed(args.seed)
    os.makedirs(args.output_dir, exist_ok=True)

    # --- Create training environment and model ---
    env = MinAtarBreakout()
    model = DQN(
        "MlpPolicy",
        env,
        learning_starts=1_000,
        buffer_size=50_000,
        train_freq=4,
        verbose=1,
        seed=args.seed,
    )

    episode_metadata_all = []
    obs_chunks, act_chunks = [], []
    file_idx = 0

    # --- Train and periodically collect ---
    timesteps = 0
    while timesteps < args.train_timesteps:
        model.learn(total_timesteps=args.collection_interval, reset_num_timesteps=False)
        timesteps += args.collection_interval

        # Collect episodes from the current policy
        o_chunks, a_chunks, ep_meta = collect_episodes(model, num_eps=20)
        obs_chunks.extend(o_chunks)
        act_chunks.extend(a_chunks)
        episode_metadata_all.extend(ep_meta)

        # Save chunks to disk
        obs_chunks, act_chunks_meta, file_idx, act_chunks = save_chunks(
            obs_chunks, file_idx, args.chunks_per_file, args.output_dir, act_chunks
        )
        episode_metadata_all.extend(act_chunks_meta)

        print(f"Collected and saved dataset at {timesteps} timesteps.")

    # --- Save metadata ---
    metadata = {
        "env": "MinAtar-Breakout",
        "num_actions": env.action_space.n,
        "train_timesteps": args.train_timesteps,
        "num_episodes": args.num_episodes,
        "avg_episode_len": float(np.mean([ep["avg_seq_len"] for ep in episode_metadata_all])),
        "episode_metadata": episode_metadata_all,
    }

    with open(os.path.join(args.output_dir, "metadata.json"), "w") as f:
        json.dump(metadata, f)

    print("Dataset generation complete.")


if __name__ == "__main__":
    main()