예제 소스
OpenAI Gym과 강화학습(Reinforcement Learning) · 파일 12개
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()