【強化学習#10】DQN

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
np.random.seed(1)

import numpy as np

import matplotlib.pyplot as plt

import seaborn as sns

np.random.seed(1)

from torch import nn
from torch import optim
import torch
import numpy as np
torch.manual_seed(1)

from torch import nn

from torch import optim

import torch

import numpy as np

torch.manual_seed(1)

class Environment:
    
    def __init__(self, size=3, lucky=[]):
        
        self.size = size
        self.lucky = lucky
        self.goal = (size-1, size-1)
        self.states = [(x, y) for x in range(size) for y in range(size)]
            
    def next_state(self, s, a):
        
        s_next = (s[0] + a[0], s[1] + a[1])
        
        if s == self.goal:
            return s
        
        if s_next not in self.states:
            return s
        
        if s_next in self.lucky:
            if np.random.random() < 0.8:
                return self.goal
            else:
                return s_next
        
        return s_next
    
    def reward(self, s, s_next):
        
        if s == self.goal:
            return -1
        
        if s_next == self.goal:
            return 0
        
        return -1

class Environment:

    def __init__(self, size=3, lucky=[]):

        self.size = size

        self.lucky = lucky

        self.goal = (size-1, size-1)

        self.states = [(x, y) for x in range(size) for y in range(size)]

    def next_state(self, s, a):

        s_next = (s[0] + a[0], s[1] + a[1])

        if s == self.goal:

            return s

        if s_next not in self.states:

            return s

        if s_next in self.lucky:

            if np.random.random() < 0.8:

                return self.goal

            else:

                return s_next

        return s_next

    def reward(self, s, s_next):

        if s == self.goal:

            return -1

        if s_next == self.goal:

            return 0

        return -1

class Agent():
    
    def __init__(self, environment):
        
        self.actions = [(-1, 0), (0, -1), (1, 0), (0, 1)]
        self.environment = environment
                
    def action(self, s, a, prob=False):

        s_next = self.environment.next_state(s, a)
        r = self.environment.reward(s, s_next)

        return r, s_next

class Agent():

    def __init__(self, environment):

        self.actions = [(-1, 0), (0, -1), (1, 0), (0, 1)]

        self.environment = environment

    def action(self, s, a, prob=False):

​

        s_next = self.environment.next_state(s, a)

        r = self.environment.reward(s, s_next)

​

        return r, s_next

class NN:
    def __init__(self, agent):
        self.model = self.model()
        self.criterion = nn.MSELoss()
        self.actions = agent.actions
        
    def model(self):
        model = nn.Sequential()
        model.add_module('fc1', nn.Linear(4,  16))
        model.add_module('relu1', nn.ReLU())
        model.add_module('fc2', nn.Linear(16, 8))
        model.add_module('relu1', nn.ReLU())
        model.add_module('fc3', nn.Linear(8, 1))
        self.optimizer = optim.Adam(model.parameters())
        return model
    
    def train_model(self, sa, labels, num_train=1000):
        for _ in range(num_train):
            qvalue = self.model(torch.tensor(sa).float())
            loss = self.criterion(qvalue, torch.tensor(labels).float())
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
    
    def q_max(self, state):
        sa = []
        for action in self.actions:
            sa.append(state+action)
        q = self.model(torch.tensor([np.array(sa)]).float()).detach()
        a_max = np.argmax(q)
        return self.actions[a_max], q[0,a_max,0]

class NN:

    def __init__(self, agent):

        self.model = self.model()

        self.criterion = nn.MSELoss()

        self.actions = agent.actions

    def model(self):

        model = nn.Sequential()

        model.add_module('fc1', nn.Linear(4,  16))

        model.add_module('relu1', nn.ReLU())

        model.add_module('fc2', nn.Linear(16, 8))

        model.add_module('relu1', nn.ReLU())

        model.add_module('fc3', nn.Linear(8, 1))

        self.optimizer = optim.Adam(model.parameters())

        return model

    def train_model(self, sa, labels, num_train=1000):

        for _ in range(num_train):

            qvalue = self.model(torch.tensor(sa).float())

            loss = self.criterion(qvalue, torch.tensor(labels).float())

            self.optimizer.zero_grad()

            loss.backward()

            self.optimizer.step()

    def q_max(self, state):

        sa = []

        for action in self.actions:

            sa.append(state+action)

        q = self.model(torch.tensor([np.array(sa)]).float()).detach()

        a_max = np.argmax(q)

        return self.actions[a_max], q[0,a_max,0]

def get_episode(agent, nn_model, epsilon=0.1):
    
    s = agent.environment.states[np.random.randint(agent.environment.size**2-1)]
    
    episode = []
    while True:
        
        if np.random.random() < epsilon:
            a = agent.actions[np.random.randint(2,4)]
        else:
            a, _ = nn_model.q_max(s)
            
        r, s_next = agent.action(s, a)
        
        episode.append((s, a, r, s_next))        
        if s_next == agent.environment.goal:
            break
        s = s_next
        
    return episode

def get_episode(agent, nn_model, epsilon=0.1):

    s = agent.environment.states[np.random.randint(agent.environment.size**2-1)]

    episode = []

    while True:

        if np.random.random() < epsilon:

            a = agent.actions[np.random.randint(2,4)]

        else:

            a, _ = nn_model.q_max(s)

        r, s_next = agent.action(s, a)

        episode.append((s, a, r, s_next))        

        if s_next == agent.environment.goal:

            break

        s = s_next

    return episode

def train(agent, nn_model, epsilon=0.1, num=100, num_train=1000):
    for c in range(num):
        print(f'num : {c+1} ')
        
        examples = []
        for _ in range(100):
            episode = get_episode(agent, nn_model, epsilon)
            examples += episode
        np.random.shuffle(examples)
        
        sa = []
        labels = []
        for s, a, r, s_next in examples:
            sa.append(s+a)
            _, q_next = nn_model.q_max(s_next)
            labels.append([r + q_next.detach()])
            
        nn_model.train_model(sa, labels, num_train)
    
    show_values(agent, model1)
    show_policy(agent, model1)

def train(agent, nn_model, epsilon=0.1, num=100, num_train=1000):

    for c in range(num):

        print(f'num : {c+1} ')

        examples = []

        for _ in range(100):

            episode = get_episode(agent, nn_model, epsilon)

            examples += episode

        np.random.shuffle(examples)

        sa = []

        labels = []

        for s, a, r, s_next in examples:

            sa.append(s+a)

            _, q_next = nn_model.q_max(s_next)

            labels.append([r + q_next.detach()])

        nn_model.train_model(sa, labels, num_train)

    show_values(agent, model1)

    show_policy(agent, model1)

def show_maze(environment):
    size = environment.size
    fig = plt.figure(figsize=(3,3))

    plt.plot([-0.5, -0.5], [-0.5, size-0.5], color='k')
    plt.plot([-0.5, size-0.5], [size-0.5, size-0.5], color='k')
    plt.plot([size-0.5, -0.5], [-0.5, -0.5], color='k')
    plt.plot([size-0.5, size-0.5], [size-0.5, -0.5], color='k')
    
    for i in range(size):
        for j in range(size):
            plt.text(i, j, "{}".format(i+size*j), size=20, ha="center", va="center")
            if (i,j) in environment.lucky:
                x = np.array([i-0.5,i-0.5,i+0.5,i+0.5])
                y = np.array([j-0.5,j+0.5,j+0.5,j-0.5])
                plt.fill(x,y, color="lightgreen")

    plt.axis("off")

def show_maze(environment):

    size = environment.size

    fig = plt.figure(figsize=(3,3))

​

    plt.plot([-0.5, -0.5], [-0.5, size-0.5], color='k')

    plt.plot([-0.5, size-0.5], [size-0.5, size-0.5], color='k')

    plt.plot([size-0.5, -0.5], [-0.5, -0.5], color='k')

    plt.plot([size-0.5, size-0.5], [size-0.5, -0.5], color='k')

    for i in range(size):

        for j in range(size):

            plt.text(i, j, "{}".format(i+size*j), size=20, ha="center", va="center")

            if (i,j) in environment.lucky:

                x = np.array([i-0.5,i-0.5,i+0.5,i+0.5])

                y = np.array([j-0.5,j+0.5,j+0.5,j-0.5])

                plt.fill(x,y, color="lightgreen")

​

    plt.axis("off")

def show_values(agent, nn_model):

    fig = plt.figure(figsize=(3,3))
    result = np.zeros([agent.environment.size, agent.environment.size])
    for (x, y) in agent.environment.states:
        a_max, q_max =  nn_model.q_max((x, y))
        result[y][x]  = q_max
        
    sns.heatmap(result, square=True, cbar=False, annot=True, fmt='3.2f', cmap='autumn_r').invert_yaxis()
    plt.axis("off")

def show_values(agent, nn_model):

​

    fig = plt.figure(figsize=(3,3))

    result = np.zeros([agent.environment.size, agent.environment.size])

    for (x, y) in agent.environment.states:

        a_max, q_max =  nn_model.q_max((x, y))

        result[y][x]  = q_max

    sns.heatmap(result, square=True, cbar=False, annot=True, fmt='3.2f', cmap='autumn_r').invert_yaxis()

    plt.axis("off")

 def show_policy(agent, nn_model):
    size = agent.environment.size
    fig = plt.figure(figsize=(3,3))

    plt.plot([-0.5, -0.5], [-0.5, size-0.5], color='k')
    plt.plot([-0.5, size-0.5], [size-0.5, size-0.5], color='k')
    plt.plot([size-0.5, -0.5], [-0.5, -0.5], color='k')
    plt.plot([size-0.5, size-0.5], [size-0.5, -0.5], color='k')

    for i in range(size):
        for j in range(size):
            if (i,j) in agent.environment.lucky:
                x = np.array([i-0.5,i-0.5,i+0.5,i+0.5])
                y = np.array([j-0.5,j+0.5,j+0.5,j-0.5])
                plt.fill(x,y, color="lightgreen")

    rotation = {(-1, 0): 180, (0, 1): 90, (1, 0): 0, (0, -1): 270}
    for s in agent.environment.states:
        if s == agent.environment.goal:
            direction=None
        else:
            a_max, q_max =  nn_model.q_max(s)
            direction = rotation[a_max]
        
        if direction != None:
            bbox_props = dict(boxstyle='rarrow')
            plt.text(s[0], s[1], '     ', bbox=bbox_props, size=8,
                     ha='center', va='center', rotation=direction)
                        
    plt.axis("off")

 def show_policy(agent, nn_model):

    size = agent.environment.size

    fig = plt.figure(figsize=(3,3))

​

    plt.plot([-0.5, -0.5], [-0.5, size-0.5], color='k')

    plt.plot([-0.5, size-0.5], [size-0.5, size-0.5], color='k')

    plt.plot([size-0.5, -0.5], [-0.5, -0.5], color='k')

    plt.plot([size-0.5, size-0.5], [size-0.5, -0.5], color='k')

​

    for i in range(size):

        for j in range(size):

            if (i,j) in agent.environment.lucky:

                x = np.array([i-0.5,i-0.5,i+0.5,i+0.5])

                y = np.array([j-0.5,j+0.5,j+0.5,j-0.5])

                plt.fill(x,y, color="lightgreen")

​

    rotation = {(-1, 0): 180, (0, 1): 90, (1, 0): 0, (0, -1): 270}

    for s in agent.environment.states:

        if s == agent.environment.goal:

            direction=None

        else:

            a_max, q_max =  nn_model.q_max(s)

            direction = rotation[a_max]

        if direction != None:

            bbox_props = dict(boxstyle='rarrow')

            plt.text(s[0], s[1], '     ', bbox=bbox_props, size=8,

                     ha='center', va='center', rotation=direction)

    plt.axis("off")

env1 = Environment(size=4, lucky=[(1,2), (2,3)])
agent1 = Agent(env1)
show_maze(env1)

env1 = Environment(size=4, lucky=[(1,2), (2,3)])

agent1 = Agent(env1)

show_maze(env1)

train(agent1, num=10000)

train(agent1, num=10000)