예제 소스
Atari Pong: 화면 추적 상태 벡터로 DQN 학습하기 · 파일 12개
Pong-dqn.py
import os
import ale_py
import gymnasium as gym
import sys
gym.register_envs(ale_py)
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
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from PyQt5.QtWidgets import QApplication, QLabel, QVBoxLayout, QWidget
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtCore import Qt
class AtariQtView(QWidget):
"""단일 QLabel에 프레임만 갱신 (OpenCV imshow 대비 깜박임 완화)."""
def __init__(self, title="Pong"):
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 int(np.argmax(self.model.predict(state, verbose=0)))
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('Pong-v4', 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_Pong", "model.weights.h5")
if os.path.exists(weights_path):
flat_dim = int(np.prod(env.observation_space.shape))
dummy = np.zeros((1, flat_dim), dtype=np.float32)
agent.model(dummy)
agent.target_model(dummy)
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 = AtariQtView("Pong DQN")
viewer.show()
graph_fig, graph_ax = plt.subplots(figsize=(6, 4))
scores, episodes = [], []
score_avg = 0
max_score = -np.inf
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:
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.batch_size:
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(os.path.dirname(weights_path) or ".", 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("Pong_graph.png")
plt.close(graph_fig)
env.close()
viewer.close()
app.quit()