Atari_Breakout

import gym
from gym.wrappers.monitoring.video_recorder import VideoRecorder
import math
import random
import os, sys
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import pickle
import gc
from collections import namedtuple
from itertools import count
from PIL import Image
from tqdm import tqdm, trange

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision.transforms as T

seed = 1234
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.enabled = False
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True

env = gym.make('Breakout-v0')
# env = gym.make('BipedalWalker-v3')

is_ipython = 'inline' in matplotlib.get_backend()
plt.ion()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

num_stack = 4
batch_size = 32
num_episode = 10000
Replay_capacity = 80000

Learning_start = 50000

EPS_START = 1.
EPS_END = 0.1
EPS_DECAY = 1000000
EPS_MEAN = 100

TARGET_EPS = 0.05
TARGET_UPDATE = 10000

GAMMA = 0.99

lr = 0.00025

Transition = namedtuple('Transition',
                        ('current_state', 'action', 'next_state', 'reward'))


class ReplayMemory(object):
    def __init__(self, capacity):
        self.capacity = capacity
        self.memory = []
        self.position = 0

    def push(self, *args):
        """Saves a transition."""
        if len(self.memory) < self.capacity:
            self.memory.append(None)
        self.memory[self.position] = Transition(*args)
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)


def get_screen():
    screen = env.render(mode='rgb_array').transpose((2, 0, 1))
    screen = screen[:, 32:210 - 15, 8:160 - 8]
    screen = torch.from_numpy(screen.astype(np.float))
    screen = torch.mean(screen, dim=0, keepdim=True)

    resize = T.Compose([T.ToPILImage(),
                        T.Resize((84, 84), interpolation=Image.BILINEAR),
                        T.ToTensor()])
    return resize(screen)


def stack_screen():
    global state_list, screen
    state_list.append(screen)
    if len(state_list) > num_stack:
        state_list = state_list[1:]
    state_tensor = torch.zeros(num_stack, 84, 84)
    for idx, sc in enumerate(state_list):
        state_tensor[idx] = sc
    return state_list, state_tensor.unsqueeze(dim=0)


class DQN(nn.Module):
    def __init__(self, outputs):
        super(DQN, self).__init__()
        self.conv1 = nn.Conv2d(4, 32, kernel_size=8, stride=4)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
        self.relu2 = nn.ReLU()
        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)
        self.relu3 = nn.ReLU()
        self.ln1 = nn.Linear(7 * 7 * 64, 512)
        self.relu4 = nn.ReLU()
        self.ln2 = nn.Linear(512, outputs)

    def forward(self, input):
        _ = self.conv1(input)
        _ = self.relu1(_)
        _ = self.conv2(_)
        _ = self.relu2(_)
        _ = self.conv3(_)
        _ = self.relu3(_)
        _ = nn.Flatten()(_)
        _ = self.ln1(_)
        _ = self.relu4(_)
        output = self.ln2(_)
        return output


def select_action(state, learn_frames, policy='Policy'):
    sample = random.random()
    eps_threshold = np.clip(EPS_START - (EPS_START - EPS_END) * (learn_frames / EPS_DECAY), 0.1, 1)
    if policy == 'Target':
        eps_threshold = TARGET_EPS

    if sample > eps_threshold:
        with torch.no_grad():
            # t.max(1) will return largest column value of each row.
            # second column on max result is index of where max element was
            # found, so we pick action with the larger expected reward.
            return policy_net(state.to(device)).max(1)[1].view(1, 1)
    else:
        return torch.tensor([[random.randrange(n_actions)]], device=device, dtype=torch.long)


def plot_scores():
    plt.figure(2)
    plt.clf()
    scores_t = torch.tensor(episode_scores, dtype=torch.float)
    plt.title('Training & Test')
    plt.xlabel('Episode')
    plt.ylabel('Score')
    plt.plot(scores_t.numpy())
    # Take 100 episode averages and plot them too
    if len(scores_t) >= EPS_MEAN:
        means = scores_t.unfold(0, EPS_MEAN, 1).mean(1).view(-1)
        means = torch.cat((torch.zeros(EPS_MEAN - 1), means))
        plt.plot(means.numpy())

    if len(test_scores) > 0:
        plt.plot(test_ep, test_scores, 'r*')

    plt.pause(0.001)  # pause a bit so that plots are updated
    # if is_ipython:
    #     display.clear_output(wait=True)
    #     display.display(plt.gcf())


memory = ReplayMemory(Replay_capacity)

n_actions = env.action_space.n

policy_net = DQN(n_actions).to(device)
target_net = DQN(n_actions).to(device)
target_net.load_state_dict(policy_net.state_dict())
target_net.eval()
optimizer = optim.RMSprop(policy_net.parameters(), lr=lr, eps=0.001, alpha=0.95)
crit = nn.MSELoss()

learn_frames = 0
episode_scores = []
test_scores = []
test_ep = []
num_frames = 0

print(f'~{Learning_start} case: Making initial replay memory by random policy')
print(
    f'{Learning_start}~{Learning_start + EPS_DECAY} frames: Learning by e-greedy policy with decreased linearly from {EPS_START} to {EPS_END}')
print(f'Target policy is e-greedy with e = {TARGET_EPS}')


for i_episode in trange(num_episode + 220):
    video_recorder = None
    update = False

    env.reset()

    state_list = []
    state_tensor = []
    screen = get_screen()
    state_list, state_tensor = stack_screen()

    score = 0

    if i_episode % 100 == 0:
        gc.collect()

    while True:
        if len(state_list) < num_stack:
            # env.step(torch.tensor([[random.randrange(n_actions)]], device=device, dtype=torch.long))
            env.step(0)
            screen = get_screen()
            state_list, state_tensor = stack_screen()
            continue
        num_frames = num_frames + 1

        current_state = state_tensor

        if num_frames < Learning_start:
            action = select_action(state_tensor, 0)
        else:
            learn_frames = learn_frames + 1
            if learn_frames % TARGET_UPDATE == 0 and learn_frames != 0:
                update = True
            action = select_action(state_tensor, learn_frames)

        _, reward, done, _ = env.step(action.item())

        if reward != 0:
            print(reward)

        score = score + reward
        screen = get_screen()
        env.render()
        state_list, state_tensor = stack_screen()
        reward = torch.tensor([reward])

        if not done:
            next_state = state_tensor
        else:
            next_state = None

        memory.push(current_state, action, next_state, reward)

        if num_frames == Learning_start:
            initset_done_ep = i_episode
            print(
                f'Learning starts at {Learning_start} case ({i_episode}th iteration(ep.), but not included in a progressive graph).')
            # with open('C:\\Users\\YKW\\PycharmProjects\\Reinforcement\\Breakout_v0\\fixed_init_states.pickle',
            #           'wb') as f:
            #     pickle.dump(memory.memory, f)

        if num_frames == Learning_start + EPS_DECAY:
            print(f'Epsilon decreasing stops with 0.1 at {Learning_start + EPS_DECAY}th, {i_episode}th episode')

        if num_frames < Learning_start:
            if done:
                break
            else:
                continue

        if learn_frames % 4 == 0:
            transitions = memory.sample(batch_size)
            batch = Transition(*zip(*transitions))

            non_final_mask = torch.tensor(tuple(map(lambda s: s is not None,
                                                    batch.next_state)), device=device, dtype=torch.bool)
            non_final_next_states = torch.cat([s for s in batch.next_state if s is not None]).to(device)

            state_batch = torch.cat(batch.current_state).to(device)
            action_batch = torch.cat(batch.action).to(device)
            reward_batch = torch.cat(batch.reward).to(device)

            state_action_values = policy_net(state_batch).gather(1, action_batch)

            next_state_values = torch.zeros(batch_size).to(device)
            next_state_values[non_final_mask] = target_net(non_final_next_states).max(1)[0].detach()

            expected_state_action_values = (next_state_values * GAMMA) + reward_batch

            loss = crit(state_action_values, expected_state_action_values.unsqueeze(1))

            optimizer.zero_grad()
            loss.backward()
            for param in policy_net.parameters():
                param.grad.data.clamp_(-1, 1)
            optimizer.step()

        if done:
            episode_scores.append(score)
            plot_scores()
            break

    # if i_episode % TARGET_UPDATE == 0 and num_frames >= Learning_start:
    if update:
        target_net.load_state_dict(policy_net.state_dict())

        video_path = f'C:\\Users\\YKW\\PycharmProjects\\Reinforcement\\Breakout_v0\\record\\{str(int(i_episode - initset_done_ep))}.mp4'
        video_recorder = VideoRecorder(env, path=video_path, enabled=True)

        with torch.no_grad():

            env.reset()

            test_score = 0
            state_list = []
            state_tensor = []
            screen = get_screen()
            state_list, state_tensor = stack_screen()
            while True:
                if len(state_list) < num_stack:
                    env.step(0)
                    screen = get_screen()
                    state_list, state_tensor = stack_screen()
                    continue

                current_state = state_tensor

                action = select_action(state_tensor, learn_frames, policy='Target')
                _, reward, done, _ = env.step(action.item())
                video_recorder.capture_frame()
                test_score = test_score + reward
                screen = get_screen()
                env.render()
                state_list, state_tensor = stack_screen()

                if done:
                    test_scores.append(test_score)
                    test_ep.append(i_episode - initset_done_ep)
                    plot_scores()
                    plt.savefig('C:\\Users\\YKW\\PycharmProjects\\Reinforcement\\Breakout_v0\\progressive.png')
                    break

        video_recorder.close()
        video_recorder.enabled = False
        os.rename(f'C:\\Users\\YKW\\PycharmProjects\\Reinforcement\\Breakout_v0\\record\\{str(int(i_episode - initset_done_ep))}.mp4',
                  f'C:\\Users\\YKW\\PycharmProjects\\Reinforcement\\Breakout_v0\\record\\{str(int(test_score))}_{str(int(i_episode - initset_done_ep))}.mp4')