dqn_class.py

import gym
import collections
import random
import time
import os

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import numpy as np

import wandb
from common import utils
from common.utils import current_time

# Hyperparameters
learning_rate = 0.0005
gamma = 0.98
buffer_limit = 20000
batch_size = 64


class ReplayBuffer:
    def __init__(self):
        self.buffer = collections.deque(maxlen=buffer_limit)

    def put(self, transition):
        self.buffer.append(transition)

    def sample(self, n):
        mini_batch = random.sample(self.buffer, n)
        s_lst, a_lst, r_lst, s_prime_lst, done_mask_lst = [], [], [], [], []

        for transition in mini_batch:
            s, a, r, s_prime, done_mask = transition
            s_lst.append(s)
            a_lst.append([a])
            r_lst.append([r])
            s_prime_lst.append(s_prime)
            done_mask_lst.append([done_mask])

        s_batch = torch.tensor(s_lst, dtype=torch.float)
        a_batch = torch.tensor(a_lst)
        r_batch = torch.tensor(r_lst)
        s_prime_batch = torch.tensor(s_prime_lst, dtype=torch.float)
        done_batch = torch.tensor(done_mask_lst)

        return s_batch, a_batch, r_batch, s_prime_batch, done_batch

    def size(self):
        return len(self.buffer)


class Network(nn.Module):
    def __init__(self, obs_size, n_actions, hidden_size: int = 128):
        super(Network, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, n_actions)
        )

    def forward(self, x):
        return self.net(x)


class DQN(nn.Module):
    def __init__(self):
        super(DQN, self).__init__()
        self.net = Network(4, 2, 128)
        self.target_net = Network(4, 2, 128)
        self.target_net.load_state_dict(self.net.state_dict())

        self.optimizer = optim.Adam(self.net.parameters(), lr=learning_rate)

    def forward(self, x):
        return self.net(x)

    # a = argmax Q(a|s)
    def sample_action(self, obs, epsilon):
        out = self.forward(obs)
        coin = random.random()
        if coin < epsilon:
            return random.randint(0, 1)
        else:
            return out.argmax().item()

    def train_step(self, memory):
        total_loss = 0
        for i in range(30):
            s, a, r, s_prime, done_mask = memory.sample(batch_size)
            # TODO 변수명 정리
            # Q(a|s)
            q_out = self.net(s)
            q_a = q_out.gather(1, a)
            # q_a.shape : [batch_size, action_space]

            max_q_prime = self.target_net(s_prime).max(1)[0].unsqueeze(1)
            target = r + gamma * max_q_prime * done_mask

            # approximate q_a to target
            loss = F.smooth_l1_loss(q_a, target)
            total_loss += loss

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

        return total_loss

    def sync_target(self):
        self.target_net.load_state_dict(self.net.state_dict())


def main():
    # TODO wandb use -> hyperparameter
    # wandb.init(project="minimalrl")
    # wandb.run.name = "dqn_{}".format(current_time())
    print(current_time())
    main_time = time.time()

    # Make the model and env
    env = gym.make('CartPole-v1')

    dqn = DQN()
    # q = DQN()
    # q_target = DQN()
    # q_target.load_state_dict(q.state_dict())
    # optimizer = optim.Adam(q.parameters(), lr=learning_rate)
    # https://tutorials.pytorch.kr/beginner/saving_loading_models.html

    memory = ReplayBuffer()

    print_interval = 100
    train_steps = 0
    last_mean_reward = 0.0
    total_steps = 0
    total_reward = 0.0
    total_loss = 0.0
    best_mean_reward = -1000
    start_time = time.time()
    train_start_time = time.time()

    reward_log = np.array([])
    loss_log = np.array([])

    for n_epi in range(5000):  # TODO epoch -> hypermarameter
        # TODO epsilon decay || linear -> common
        # Linear annealing from 8% to 1%
        epsilon = max(0.01, 0.08 - 0.01 * (n_epi / 200))

        s = env.reset()
        done = False

        while not done:
            s_ = torch.from_numpy(s).float()
            a = dqn.sample_action(s_, epsilon)  # a = argmax Q(a|s)  # TODO DeterActor get_action() 만들기
            s_prime, r, done, info = env.step(a)

            r = r * ((2.4 - abs(s_prime[0])) / 2.4)  # center weighted reward
            r = float(r)

            done_mask = 0.0 if done else 1.0
            memory.put((s, a, r / 100.0, s_prime, done_mask))
            s = s_prime

            total_reward += r
            total_steps += 1

            if done:
                break

            # # Train per step
            # if memory.size() > 2000:
            #     total_loss += train(q, q_target, memory, optimizer)

        # Train per episode
        if memory.size() > 2000:
            total_loss += dqn.train_step(memory)
            train_steps += 1
            reward_log = np.append(reward_log, last_mean_reward)

        # Update q_target
        if n_epi % 20 == 0 and n_epi != 0:
            dqn.sync_target()

        if n_epi % print_interval == 0 and n_epi != 0:
            elapsed_time = time.time() - start_time
            mean_reward = total_reward / print_interval
            print(
                "Steps: {:6}".format(total_steps),
                "Episode: {:4}".format(n_epi),
                "Reward: {:5.1f}".format(mean_reward),
                "Loss: {:5.2f}".format(total_loss*100),
                "Epsilon: {:.2f}%".format(epsilon * 100),
                "Elapsed Time: {:5.2f}s".format(elapsed_time)
            )
            # TODO epsilon log
            # wandb.log(
            #     data={
            #         "Reward": mean_reward,
            #         "Loss": total_loss*100
            #     },
            #     step=train_steps
            # )

            # reward_log = np.append(reward_log, mean_reward)
            last_mean_reward = mean_reward
            loss_log = np.append(loss_log, float(total_loss)*100)


            # # save best model
            # if best_mean_reward < mean_reward:
            #     torch.save(q.state_dict(), "DQN-best.dat")
            #     best_mean_reward = mean_reward
            #     print("┌──────────────────────────────────┐")
            #     print("│ Best mean reward updated {:7.3f} |".format(best_mean_reward))
            #     print("└──────────────────────────────────┘")

            # # early stopping
            # if mean_reward > 400:
            #     break

            total_reward = 0.0
            total_loss = 0.0
            start_time = time.time()

    # TODO save_log -> common
    # Save log
    main_time = time.strftime('%y-%m-%d_%H%M', time.localtime(main_time))
    reward_path = os.path.join(".", "output_data", "dqn_reward_{}".format(main_time))
    loss_path = os.path.join(".", "output_data", "dqn_loss_{}".format(main_time))
    np.savetxt(reward_path, reward_log)
    np.savetxt(loss_path, loss_log)

    # print training time
    training_time = time.time() - train_start_time
    training_time = time.strftime('%H:%M:%S', time.gmtime(training_time))
    print("Training time : {}".format(training_time))

    # TODO render -> common
    # # Test render
    # q.load_state_dict(torch.load("DQN-best.dat"))
    # fps = 60
    #
    # for _ in range(10):
    #     s = env.reset()
    #     while True:
    #         start_time = time.time()
    #
    #         env.render()
    #         s_ = torch.from_numpy(s).float()
    #         a = q.sample_action(s_, 0)
    #         s, r, done, _ = env.step(a)
    #         if done:
    #             env.close()
    #             break
    #
    #         delay = 1/fps - (time.time() - start_time)
    #         if delay > 0:
    #             time.sleep(delay)

    # wandb.finish()


if __name__ == '__main__':
    main()


# TODO dqn, network, agent, hypermarameter 분리
# TODO device 추가