main.py

import gymnasium as gym
import utils
from ddqn import DDQNAgent
import numpy as np
import os
import torch
import warnings
from argparse import ArgumentParser
import pandas as pd
from preprocess import AtariEnv
from ale_py import ALEInterface, LoggerMode
from config import environments

warnings.simplefilter("ignore")
ALEInterface.setLoggerMode(LoggerMode.Error)


def run_ddqn(args):

    def make_env():
        return AtariEnv(
            args.env,
            shape=(84, 84),
            repeat=4,
            clip_rewards=True,
        ).make()

    envs = gym.vector.AsyncVectorEnv([make_env for _ in range(args.n_envs)])
    save_prefix = args.env.split("/")[-1]

    print(f"\nEnvironment: {save_prefix}")
    print(f"Obs.Space: {envs.single_observation_space.shape}")
    print(f"Act.Space: {envs.single_action_space.n}")

    agent = DDQNAgent(
        args.env,
        envs.single_observation_space.shape,
        envs.single_action_space.n,
        mem_size=30000,
        batch_size=64,
        eps_dec=5e-7,
        replace_target_count=1000,
    )

    if args.continue_training:
        if os.path.exists(f"weights/{save_prefix}_dqn.pt"):
            agent.load_checkpoints()

    best_score = -np.inf
    avg_score = np.nan
    score = np.zeros(args.n_envs)
    history, metrics = [], []

    fixed_states = torch.tensor(utils.collect_fixed_states(envs)).to(agent.q1.device)

    states, _ = envs.reset()
    for i in range(args.n_steps):
        actions = [agent.choose_action(state) for state in states]

        next_states, rewards, term, trunc, _ = envs.step(actions)

        for j in range(args.n_envs):
            agent.store_transition(
                states[j],
                actions[j],
                rewards[j],
                next_states[j],
                term[j] or trunc[j],
            )

            score[j] += rewards[j]
            if term[j] or trunc[j]:
                history.append(score[j])
                score[j] = 0

            agent.learn()
        states = next_states

        if len(history) > 0:
            avg_score = np.mean(history[-100:])

        if avg_score > best_score:
            best_score = avg_score
            agent.save_checkpoint()

        with torch.no_grad():
            avg_q_value = (
                torch.minimum(agent.q1(fixed_states), agent.q2(fixed_states))
                .mean()
                .cpu()
                .numpy()
            )

        metrics.append(
            {
                "episode": i + 1,
                "average_score": avg_score,
                "best_score": best_score,
                "average_q_value": avg_q_value,
            }
        )

        ep_str = f"[Epoch {i + 1:05}/{args.n_steps}]"
        g_str = f"  Games = {len(history)}"
        avg_str = f"  Avg.Score = {avg_score:.2f}"
        q_str = f"  Avg.Q = {avg_q_value:.2e}"
        eps_str = f"  Eps. = {agent.epsilon:.2f}"
        print(ep_str + g_str + avg_str + q_str + eps_str, end="\r")

    torch.save(agent.q1.state_dict(), f"weights/{save_prefix}_q1_final.pt")
    torch.save(agent.q2.state_dict(), f"weights/{save_prefix}_q2_final.pt")

    save_results(args.env, metrics, agent)


def save_results(env_name, metrics, agent):
    save_prefix = env_name.split("/")[-1]
    df = pd.DataFrame(metrics)
    df.to_csv(f"csv/{save_prefix}_metrics.csv", index=False)
    utils.plot_metrics(save_prefix, df)
    save_best_version(env_name, agent)


def save_best_version(env_name, agent, seeds=100):
    # actually, its usally better to use the final weights...
    # agent.load_checkpoint()

    best_total_reward = float("-inf")
    best_frames = None

    env = AtariEnv(
        env_name,
        shape=(84, 84),
        repeat=4,
        clip_rewards=False,
    ).make()

    save_prefix = env_name.split("/")[-1]

    for s in range(seeds):
        state, _ = env.reset(seed=s)

        frames = []
        total_reward = 0

        term, trunc = False, False
        while not term and not trunc:
            frames.append(env.render())

            action = agent.choose_action(state)
            next_state, reward, term, trunc, _ = env.step(action)

            total_reward += reward
            state = next_state

        if total_reward > best_total_reward:
            best_total_reward = total_reward
            best_frames = frames

    save_prefix = env_name.split("/")[-1]
    utils.save_animation(best_frames, f"environments/{save_prefix}.gif")


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument(
        "-e", "--env", default=None, help="Environment name from Gymnasium"
    )
    parser.add_argument(
        "--n_steps",
        default=100000,
        type=int,
        help="Number of learning steps to run during training",
    )
    parser.add_argument(
        "--n_envs",
        default=32,
        type=int,
        help="Number of parallel environments",
    )
    parser.add_argument(
        "--continue_training",
        default=True,
        type=bool,
        help="Continue training from saved weights.",
    )
    args = parser.parse_args()

    for fname in ["metrics", "environments", "weights", "csv"]:
        if not os.path.exists(fname):
            os.makedirs(fname)

    if args.env:
        run_ddqn(args)
    else:
        for env_name in environments:
            args.env = env_name
            run_ddqn(args)