dqn.py

import time
import gym
import torch
import numpy as np
import tensorflow as tf # BUG: should be import before torchvision
import torchvision
import matplotlib as mpl

mpl.use('Agg')
import matplotlib.pyplot as plt
import torch.nn.functional as F
import torch.nn as nn
import torch.optim as optim
import argparse
from gym.wrappers import Monitor
import datetime
import os

from model import Model
from environment import Environment
from utils import Memory, EpsilonScheduler, make_log_dir, save_gif
import pdb

##### Arguments #####
parser = argparse.ArgumentParser(
    description='An implementation of the 2015 DeepMind DQN Paper')

parser.add_argument('--weights',
                    type=str,
                    help='weights file for pretrained weights')
parser.add_argument('--nosave',
                    default=False,
                    action='store_true',
                    help='do not save a record of the run')
parser.add_argument('--game',
                    type=str,
                    default='breakout')

args = parser.parse_args()

##### Hyper-Parameters #####
BATCH_SIZE = 32  # size of minibatch
MEM_SIZE = int(1e6)  # size of replay memory
TARGET_UPDATE_EVERY = 10000  # in units of minibatch updates
GAMMA = 0.99  # discount factor
UPDATE_FREQ = 4  # perform minibatch update once every UPDATE_FREQ
INIT_MEMORY_SIZE = 200000  # initial size of memory before minibatch updates begin
#scheduler = EpsilonScheduler(schedule=[(0, 1),(INIT_MEMORY_SIZE,1),(1e6, 0.1)])
scheduler = EpsilonScheduler(schedule=[(0, 1), (INIT_MEMORY_SIZE, 1), (2e6, 0.1), (30e6, 0.01)])


STORAGE_DEVICES = [
    'cuda:0'
]  # list of devices to use for episode storage (need about 10GB for 1 million memories)
DEVICE = 'cuda:0'  # list of devices for computation
EPISODES = int(1e5)  # total training episodes
NUM_TEST = 20
TEST_EVERY = 1000  # (episodes)
PLOT_EVERY = 10  # (episodes)
SAVE_EVERY = 1000  # (episodes)
EXPERIMENT_DIR = "experiments"
#GAME = 'breakout'
GAME = args.game

if not args.nosave:
    root_dir, weight_dir, video_dir = make_log_dir(EXPERIMENT_DIR, GAME)
    with open(os.path.join(EXPERIMENT_DIR, "current.txt"), "w") as f:
        f.write(os.path.abspath(video_dir))

##### Gym Environment #####
env = Environment(game=GAME)
mem = Memory(MEM_SIZE, storage_devices=STORAGE_DEVICES, target_device=DEVICE)

##### Q Functions #####
q_func = Model(env.action_space.n).to(DEVICE)
if args.weights:
    q_func.load_state_dict(torch.load(args.weights))

target_q_func = Model(env.action_space.n).to(DEVICE)
target_q_func.load_state_dict(q_func.state_dict())

##### Optimizer #####
#optimizer = optim.RMSprop(q_func.parameters(), lr=2.5e-4, alpha=0.95, momentum=0.95, eps=1e-2)
optimizer = optim.Adam(
    q_func.parameters(), lr=0.00001,
    eps=1.5e-4)  # 0.00001 for breakout, 0.00025 is faster for pong

##### Loss Function #####
loss_func = nn.SmoothL1Loss()

##### Logger #####
# for tensorboard
current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
summary_writer = tf.summary.create_file_writer(os.path.join(root_dir,'logs'))

if not args.nosave:
    env = Monitor(env,
                  directory=video_dir,
                  video_callable=lambda count: count % 500 == 0,
                  force=True)

##### Testing #####

def test(active_target_count, save=False):
    print("[TESTING]")
    total_reward = 0
    unclipped_reward = 0

    for i in range(NUM_TEST):
        if i == 0 and save:
            frames = []

        env.reset(eval=True)  # performs random actions to start
        state, _, done, _ = env.step(env.action_space.sample())
        frame = 0

        while not done:
            if i == 0 and save:
                frames.append(state[0, 0])

            q_values = q_func(state.to(DEVICE))

            if np.random.random(
            ) > 0.01:  # small epsilon-greedy, sometimes 0.05
                action = torch.argmax(q_values, dim=1).item()
            else:
                action = env.action_space.sample()

            lives = env.ale.lives()
            next_state, reward, done, info = env.step(action)
            if env.ale.lives() != lives:  # lost life
                pass

            unclipped_reward += info['unclipped_reward']
            total_reward += reward
            state = next_state
            frame += 1

        if i == 0 and save:
            frames.append(state[0, 0])
            save_gif(
                frames, "{}.gif".format(
                    os.path.join(video_dir, str(scheduler.step_count()))))

    total_reward /= NUM_TEST
    unclipped_reward /= NUM_TEST

    with summary_writer.as_default():
        tf.summary.scalar('Total Test Reward', total_reward, step=scheduler.step_count())
        tf.summary.scalar('Total Unclipped Test Reward', unclipped_reward, step=scheduler.step_count())

    print(
        f"[TESTING] Total Reward: {total_reward}, Unclipped Reward: {unclipped_reward}"
    )

    return total_reward

##### Main Process #####
start_time = time.time()
active_target_count = 1 # NEW: for average
avg_reward = 0
avg_unclipped_reward = 0
avg_q = 0
num_parameter_updates = 0
for episode in range(EPISODES):
    avg_loss = 0
    total_reward = 0
    unclipped_reward = 0
    frame = 0

    env.reset()
    state, _, done, _ = env.step(env.action_space.sample())

    while not done:
        q_values = q_func(state.to(DEVICE))
        if np.random.random() > scheduler.epsilon():  # epsilon-random policy
            action = torch.argmax(q_values, dim=1)
        else:
            action = env.action_space.sample()

        avg_q = 0.9 * avg_q + 0.1 * q_values.mean().item(
        )  # record average q value

        lives = env.ale.lives()  # get lives before action
        next_state, reward, done, info = env.step(action)
        #reward = info['unclipped_reward'] # NOTE: Use unclipped reward 

        # hack to make learning faster (count loss of life as end of episode for memory purposes)
        mem.store(state[0, 0], action, reward, done
                  or (env.ale.lives() != lives))

        state = next_state
        total_reward += reward
        unclipped_reward += info['unclipped_reward']
        frame += 1
        scheduler.step(1)

        if mem.size() < INIT_MEMORY_SIZE:
            continue

        if scheduler.step_count() % UPDATE_FREQ == 0:
            states, next_states, actions, rewards, terminals = mem.sample(
                BATCH_SIZE)
            mask = (1 - terminals).float()
            y = rewards + mask * GAMMA * torch.max(
                target_q_func(next_states), dim=1).values.view(-1, 1).detach()
            x = q_func(states)[range(BATCH_SIZE), actions.squeeze()]
            loss = loss_func(x, y.squeeze())
            optimizer.zero_grad()
            loss.backward()

            for param in q_func.parameters():  # gradient clipping
                param.grad.data.clamp_(-1, 1)

            optimizer.step()
            avg_loss += loss.item()
            num_parameter_updates += 1

        if num_parameter_updates % TARGET_UPDATE_EVERY == 0:  # reset target to source
            target_q_func.load_state_dict(q_func.state_dict())

    avg_loss /= frame
    avg_reward = 0.9 * avg_reward + 0.1 * total_reward
    avg_unclipped_reward = 0.9 * avg_unclipped_reward + 0.1 * unclipped_reward

    print(
        f"[EPISODE {episode}] Loss: {avg_loss:4f}, " +
        f"Total Reward: {total_reward}, Total Unclipped Reward: {unclipped_reward}, " +
        f"Frames: {frame}, Epsilon: {scheduler.epsilon():4f}, " +
        f"Total Frames: {scheduler.step_count()}, Memory Size: {mem.size()}, " +
        f"Average Q: {avg_q:4f}, " +
        f"Elapsed Time: {int(time.time()-start_time)} sec"
    )

    if episode % PLOT_EVERY == 0:
        with summary_writer.as_default():
            tf.summary.scalar('Total Train Reward', avg_reward, step=scheduler.step_count())
            tf.summary.scalar('Total Unclipped Train Reward', avg_unclipped_reward, step=scheduler.step_count())
            tf.summary.scalar('Epsilon', scheduler.epsilon(), step=scheduler.step_count())
            tf.summary.scalar('Episode Length', frame, step=scheduler.step_count())
            tf.summary.scalar('Average Loss', avg_loss, step=scheduler.step_count())
            tf.summary.scalar('Average Q', avg_q, step=scheduler.step_count())
            tf.summary.scalar('Active Target Count', active_target_count, step=scheduler.step_count())
            tf.summary.scalar('Q', q_values.mean().item(), step=scheduler.step_count())

    if episode % TEST_EVERY == 0 and episode != 0:
        test_reward = test(active_target_count, save=not args.nosave)

    if episode % SAVE_EVERY == 0 and episode != 0 and not args.nosave:
        path = f"episode-{episode}.pt"
        weight_path = os.path.join(weight_dir, path)
        info_path = os.path.join(root_dir, "info.txt")

        torch.save(q_func.state_dict(), weight_path)

        with open(info_path, "a+") as f:
            f.write(",".join([
                str(x) for x in [
                    path,
                    scheduler.step_count(),
                    scheduler.epsilon(), episode, test_reward
                ]
            ]) + "\n")