예제 소스

moutainCar-dqn.py

import gymnasium as gym
import sys
import numpy as np
import tensorflow as tf

from tensorflow.keras import models, layers, optimizers
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.initializers import RandomUniform
from tensorflow.keras.optimizers import Adam

from collections import deque
import random
import matplotlib.pyplot as plt
import cv2

class DQN(tf.keras.Model):
    def __init__(self, action_size):
        super(DQN, self).__init__()
        self.fc1 = Dense(24, activation='relu')
        self.fc2 = Dense(24, activation='relu')
        self.fc_out = Dense(action_size, kernel_initializer=RandomUniform(-1e-3, 1e-3))

    def call(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        q = self.fc_out(x)
        return q

class DQNAgent:
    def __init__(self, state_size, action_size):
        self.state_size = state_size
        self.action_size = action_size
        self.discount_factor = 0.99
        self.learning_rate = 0.001
        self.epsilon = 1.0
        self.epsilon_decay = 0.999
        self.epsilon_min = 0.01
        self.batch_size = 64
        self.train_start = 1000
        self.memory = deque(maxlen=2000)
        self.model = DQN(action_size)
        self.target_model = DQN(action_size)
        self.optimizer = Adam(learning_rate=self.learning_rate)
		
        self.update_target_model()

    def update_target_model(self):
        self.target_model.set_weights(self.model.get_weights())
    
    def remember(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))

    def choose_action(self, state):
        return random.randrange(self.action_size) if (np.random.rand() <= self.epsilon) else np.argmax(self.model.predict(state))

    def train_model(self):
        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay

        mini_batch = random.sample(self.memory, self.batch_size)

        states = np.array([sample[0][0] for sample in mini_batch])
        actions = np.array([sample[1] for sample in mini_batch])
        rewards = np.array([sample[2] for sample in mini_batch])
        next_states = np.array([sample[3][0] for sample in mini_batch])
        dones = np.array([sample[4] for sample in mini_batch])

        with tf.GradientTape() as tape:
            predicts = self.model(states)
            one_hot_action = tf.one_hot(actions, self.action_size)
            predicts = tf.reduce_sum(one_hot_action * predicts, axis=1)

            target_predicts = self.target_model(next_states)
            target_predicts = tf.stop_gradient(target_predicts)

            max_q = np.amax(target_predicts, axis=-1)
            targets = rewards + (1 - dones) * self.discount_factor * max_q
            loss = tf.reduce_mean(tf.square(targets - predicts))

        trainable_vars = self.model.trainable_variables
        grads = tape.gradient(loss, trainable_vars)
        grads_and_vars = [(g, v) for g, v in zip(grads, trainable_vars) if g is not None]
        self.optimizer.apply_gradients(grads_and_vars)



if __name__ == "__main__":
    # Create the environment
    env = gym.make('MountainCar-v0', render_mode="rgb_array")
    state_size = env.observation_space.shape[0]
    action_size = env.action_space.n

    agent = DQNAgent(state_size, action_size)

    scores, episodes = [], []
    score_avg = 0

    num_episode = 300
    for e in range(num_episode):
        done = False
        score = 0

        state, _ = env.reset()
        state = state.reshape(1, -1)

        while not done:
            img = env.render()
            # if img is not None:
            #     # 이미지를 시각화합니다.
            #     plt.imshow(img)
            #     plt.axis('off')  # 축을 숨깁니다.
            #     plt.show(block=False)  # Non-blocking
            #     # plt.pause(0.000001)
            # if plt.waitforbuttonpress(0.0001):  # 키 입력 대기 (0.1초 동안)
            #     key = plt.gcf().canvas.get_tk_widget().tk.call("focus_get")
            #     if key == 'q':  # 'q' 키를 누르면 종료
            #         break
            if img is not None:
                cv2.imshow('moutainCar', img)
                cv2.waitKeyEx(delay=1)

            action = agent.choose_action(state)

            next_state, reward, terminated, truncated, info = env.step(action)
            done = terminated or truncated
            next_state = next_state.reshape(1, -1)

            score += reward
            reward = 0.1 if not done or score == 500 else -1

            agent.remember(state, action, reward, next_state, done)

            if len(agent.memory) >= agent.train_start:
                agent.train_model()

            state = next_state

            if done:
                agent.update_target_model()
                cv2.destroyAllWindows()

                score_avg = 0.9 * score_avg + 0.1 * score if score_avg != 0 else score
                print('episode: {:3d} | score avg {:3.2f} | memory length: {:4d} | epsilon: {:.4f}'.format(e, score_avg, len(agent.memory), agent.epsilon))

                scores.append(score_avg)
                episodes.append(e)
                plt.plot(episodes, scores, 'b')
                plt.xlabel('episode')
                plt.ylabel('average score')
                plt.savefig('MountainCar_graph.png')

                if score_avg > 98:
                    agent.model.save_weights('./save_model_mountaincar/model', save_format='tf')
                    env.close()
                    plt.close()  # 모든 플롯을 닫습니다.
                    sys.exit()