generate_data.py

import argparse
import pickle
import random
from pathlib import Path

import numpy as np
import torch
from stable_baselines3 import PPO
from tqdm import tqdm

from rths.env.wrappers import make_vec_env


GAME_CONFIG = {
    "pacman": {
        "env_id": "ALE/MsPacman-v5",
        "default_policy": "games/pacman/ppo_ms_pacman.zip",
        "default_output": "games/pacman/data/latent_transitions.pkl",
    },
    "amidar": {
        "env_id": "ALE/Amidar-v5",
        "default_policy": "games/amidar/ppo_amidar.zip",
        "default_output": "games/amidar/data/latent_transitions.pkl",
    },
    "qbert": {
        "env_id": "ALE/Qbert-v5",
        "default_policy": "games/qbert/ppo_qbert.zip",
        "default_output": "games/qbert/data/latent_transitions.pkl",
    },
}


def parse_args():
    parser = argparse.ArgumentParser(
        description="Generate latent transition dataset for a game (offline .pkl)."
    )
    parser.add_argument("--game", required=True, choices=sorted(GAME_CONFIG.keys()))
    parser.add_argument("--num-transitions", type=int, default=200000)
    parser.add_argument("--n-envs", type=int, default=8)
    parser.add_argument("--policy", default=None, help="Override policy .zip path")
    parser.add_argument("--output", default=None, help="Override output .pkl path")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--noop-prob", type=float, default=0.05)
    parser.add_argument("--random-prob", type=float, default=0.05)
    parser.add_argument("--noop-action-idx", type=int, default=0)
    parser.add_argument(
        "--random-actions",
        action="store_true",
        help="Use random actions as the base policy (still applies noop/random overrides)",
    )
    parser.add_argument("--device", default="auto")
    parser.add_argument("--log-every", type=int, default=10000)
    return parser.parse_args()


def _select_device(device_arg: str) -> torch.device:
    if device_arg == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(device_arg)


def _frame_from_obs(obs: np.ndarray) -> np.ndarray:
    # (np.ndarray, [84, 84]) (extract latest grayscale frame from stacked Atari obs)
    return obs[0, -1, :, :].astype(np.uint8, copy=True)


def main():
    args = parse_args()
    cfg = GAME_CONFIG[args.game]

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.noop_prob < 0.0 or args.random_prob < 0.0:
        raise ValueError("--noop-prob and --random-prob must be non-negative")
    if args.noop_prob + args.random_prob > 1.0:
        raise ValueError("--noop-prob + --random-prob must be <= 1.0")

    noop_prob_effective = 0.0 if args.game == "pacman" else args.noop_prob
    if args.game == "pacman" and args.noop_prob > 0.0:
        print(
            "Note: Pac-Man encoder data forces --noop-prob to 0 (NOOP is not sampled); "
            "policy NOOP steps are labeled with the last non-NOOP action per env."
        )

    env_id = cfg["env_id"]
    output_path = Path(args.output or cfg["default_output"]).resolve()
    output_path.parent.mkdir(parents=True, exist_ok=True)

    device = _select_device(args.device)
    print(f"Using device: {device}")
    print(f"Game: {args.game} | Env: {env_id}")
    print(f"Target transitions: {args.num_transitions}")
    print(f"Parallel envs: {args.n_envs}")
    print(
        f"Action mixing: noop={noop_prob_effective:.2%}, random={args.random_prob:.2%}, "
        f"base={1.0 - (noop_prob_effective + args.random_prob):.2%}"
    )

    env = make_vec_env(env_id, args.n_envs, args.seed)
    num_actions = env.action_space.n

    model = None
    if not args.random_actions:
        policy_path = Path(args.policy or cfg["default_policy"]).resolve()
        if not policy_path.is_file():
            raise FileNotFoundError(f"Policy file not found: {policy_path}")
        model = PPO.load(str(policy_path), env=env, device=device)
        print(f"Collecting with PPO policy: {policy_path}")
        print("PPO action sampling: stochastic (deterministic=False) for dataset diversity")
    else:
        print("Collecting with random actions")

    states = np.empty((args.num_transitions, 84, 84), dtype=np.uint8)  # (np.ndarray, [N,84,84]) (store state_t)
    next_states = np.empty((args.num_transitions, 84, 84), dtype=np.uint8)  # (np.ndarray, [N,84,84]) (store state_t+1)
    actions = np.empty((args.num_transitions,), dtype=np.int64)  # (np.ndarray, [N]) (store action_t)
    rewards = np.empty((args.num_transitions,), dtype=np.float32)  # (np.ndarray, [N]) (store reward_t)

    obs = env.reset()
    collected = 0
    env_steps = 0
    noop_count = 0
    random_count = 0
    policy_count = 0
    last_non_noop = (
        np.full(args.n_envs, args.noop_action_idx, dtype=np.int64)
        if args.game == "pacman"
        else None
    )

    with tqdm(total=args.num_transitions, desc="Collecting", unit="transitions", dynamic_ncols=True) as pbar:
        try:
            while collected < args.num_transitions:
                n_envs = obs.shape[0]
                if model is not None and not args.random_actions:
                    base_actions, _ = model.predict(obs, deterministic=False)
                    base_actions = np.asarray(base_actions, dtype=np.int64)  # (np.ndarray, [n_envs]) (policy actions)
                else:
                    base_actions = np.random.randint(0, num_actions, size=n_envs, dtype=np.int64)  # (np.ndarray, [n_envs]) (random base actions)

                roll = np.random.random(size=n_envs)  # (np.ndarray, [n_envs]) (mixture source sampling)
                noop_mask = roll < noop_prob_effective
                random_mask = (roll >= noop_prob_effective) & (
                    roll < (noop_prob_effective + args.random_prob)
                )
                base_mask = ~(noop_mask | random_mask)

                final_actions = base_actions.copy()  # (np.ndarray, [n_envs]) (actions sent to vec env)
                if np.any(random_mask):
                    final_actions[random_mask] = np.random.randint(
                        0, num_actions, size=int(np.sum(random_mask)), dtype=np.int64
                    )
                if np.any(noop_mask):
                    final_actions[noop_mask] = int(args.noop_action_idx)

                noop_count += int(np.sum(noop_mask))
                if model is not None and not args.random_actions:
                    policy_count += int(np.sum(base_mask))
                    random_count += int(np.sum(random_mask))
                else:
                    random_count += int(np.sum(base_mask) + np.sum(random_mask))

                next_obs, step_rewards, dones, _ = env.step(final_actions)
                env_steps += n_envs

                valid_idx = np.where(~dones.astype(bool))[0]
                remaining = args.num_transitions - collected
                take = min(len(valid_idx), remaining)
                if take > 0:
                    idx = valid_idx[:take]
                    states[collected : collected + take] = obs[idx, -1, :, :].astype(np.uint8, copy=False)
                    next_states[collected : collected + take] = next_obs[idx, -1, :, :].astype(np.uint8, copy=False)
                    if args.game == "pacman":
                        stored_actions = final_actions[idx].astype(np.int64, copy=True)
                        noop_i = int(args.noop_action_idx)
                        for row, e in enumerate(idx):
                            if stored_actions[row] == noop_i:
                                stored_actions[row] = last_non_noop[e]
                        actions[collected : collected + take] = stored_actions
                    else:
                        actions[collected : collected + take] = final_actions[idx]
                    rewards[collected : collected + take] = step_rewards[idx].astype(np.float32, copy=False)
                    collected += take
                    pbar.update(take)

                # Keep postfix counters live; exact modulo checks can be skipped with batched collection.
                if collected:
                    pbar.set_postfix(
                        policy=policy_count,
                        random=random_count,
                        noop=noop_count,
                        refresh=False,
                    )

                if args.game == "pacman" and last_non_noop is not None:
                    noop_i = int(args.noop_action_idx)
                    for e in range(n_envs):
                        if int(final_actions[e]) != noop_i:
                            last_non_noop[e] = int(final_actions[e])
                    done_mask = dones.astype(bool)
                    if done_mask.any():
                        last_non_noop[done_mask] = noop_i

                obs = next_obs
        finally:
            env.close()

    payload = {
        "game": args.game,
        "env_id": env_id,
        "seed": args.seed,
        "num_actions": int(num_actions),
        "num_transitions": int(collected),
        "states": states,
        "next_states": next_states,
        "actions": actions,
        "rewards": rewards,
        "noop_prob": float(noop_prob_effective),
        "random_prob": float(args.random_prob),
        "noop_action_idx": int(args.noop_action_idx),
    }

    with output_path.open("wb") as f:
        pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL)

    print(f"Saved dataset to {output_path}")
    print(f"Transitions: {collected} | Environment steps: {env_steps}")
    print(f"Action source counts: policy={policy_count}, random={random_count}, noop={noop_count}")


if __name__ == "__main__":
    main()