예제 소스

cartpol-dqn.py

import gymnasium as gym
import sys
import os

import random
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from collections import deque

import tensorflow as tf
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 PyQt5.QtWidgets import QApplication, QLabel, QVBoxLayout, QWidget
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtCore import Qt


class CartPoleQtView(QWidget):
    def __init__(self, title="CartPole"):
        super().__init__()
        self.setWindowTitle(title)
        self._label = QLabel(self)
        self._label.setAlignment(Qt.AlignCenter)
        self._label.setScaledContents(False)
        layout = QVBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)
        layout.addWidget(self._label)

    def show_frame(self, rgb):
        if rgb is None:
            return
        rgb = np.ascontiguousarray(np.asarray(rgb, dtype=np.uint8))
        if rgb.ndim != 3 or rgb.shape[2] != 3:
            return
        h, w, _ = rgb.shape
        qimg = QImage(rgb.data, w, h, 3 * w, QImage.Format_RGB888).copy()
        self._label.setPixmap(QPixmap.fromImage(qimg))
        app = QApplication.instance()
        if app is not None:
            app.processEvents()


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__":
    env = gym.make('CartPole-v1', render_mode="rgb_array")
    state_size = int(env.observation_space.shape[0])
    action_size = int(env.action_space.n)

    agent = DQNAgent(state_size, action_size)

    weights_path = os.path.join("save_model", "model.weights.h5")
    _warm = np.zeros((1, state_size), dtype=np.float32)
    agent.model(_warm)
    agent.target_model(_warm)
    if os.path.exists(weights_path):
        agent.model.load_weights(weights_path)
        agent.update_target_model()
        print(f"Loaded weights from {weights_path}")

    app = QApplication.instance() or QApplication(sys.argv)
    viewer = CartPoleQtView("CartPole DQN")
    viewer.show()
    graph_fig, graph_ax = plt.subplots(figsize=(6, 4))

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

    num_episode = 300
    max_score = -np.inf
    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:
                viewer.show_frame(img)

            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()

                score_avg = 0.9 * score_avg + 0.1 * score if score_avg != 0 else score
                if score_avg > max_score:
                    max_score = score_avg
                    os.makedirs("save_model", exist_ok=True)
                    agent.model.save_weights(weights_path)
                print('episode: {:3d} | score avg {:3.2f} | max score {:3.2f} | memory length: {:4d} | epsilon: {:.4f}'.format(e, score_avg, max_score, len(agent.memory), agent.epsilon))

                scores.append(score_avg)
                episodes.append(e)
                graph_ax.clear()
                graph_ax.plot(episodes, scores, "b")
                graph_ax.set_xlabel("episode")
                graph_ax.set_ylabel("average score")
                graph_fig.tight_layout()
                graph_fig.savefig("cartpole_graph.png")

    env.close()
    plt.close(graph_fig)
    viewer.close()
    app.quit()