예제 소스
Atari Pong: 화면 추적 상태 벡터로 DQN 학습하기 · 파일 12개
Pong-vlm-dqn.py
#!/usr/bin/env python3
"""
Pong-v4: 화면 RGB → (Pong-vlm 과 동일한) 프레임차 기반 공·패들 추정 → **수치 상태 벡터**로 DQN 학습.
왜 VLM을 매 step 안 쓰나요?
----------------------------
강화학습은 수백만 step 이 필요합니다. 매 프레임 VLM(비전 LLM)을 호출하는 것은 시간·비용상 불가능에 가깝습니다.
대신 `Pong-vlm.py` 의 `InterceptAssist` 와 **같은 추적기**로 공 좌표·패들 y·motion 등을 만들고,
그것을 VLM `--vlm-input text` 가 설명하는 것과 **동일한 정보**를 담은 벡터로 정규화해 DQN에 넣습니다.
(나중에 VLM으로 생성한 텍스트 라벨을 붙여 지도학습하는 파이프라인의 **입력 특성**으로도 쓸 수 있습니다.)
선택 `--dump-text-log`: 매 step 한 줄씩 사람이 읽기 쉬운 상태 텍스트를 JSONL 로 쌓아, VLM 학습용 데이터 준비에 활용할 수 있습니다.
실행 예
-------
python Pong-vlm-dqn.py --episodes 400 --no-gui
python Pong-vlm-dqn.py --episodes 200 --dump-text-log pong_state.jsonl --no-gui
"""
from __future__ import annotations
import argparse
import importlib.util
import json
import os
import random
import sys
from typing import Any, Optional, Tuple
import ale_py
import gymnasium as gym
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt # noqa: E402
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.initializers import RandomUniform
from tensorflow.keras.layers import Dense
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import Adam
gym.register_envs(ale_py)
from PyQt5.QtCore import Qt # noqa: E402
from PyQt5.QtGui import QImage, QPixmap # noqa: E402
from PyQt5.QtWidgets import QApplication, QLabel, QVBoxLayout, QWidget # noqa: E402
_ROOT = os.path.dirname(os.path.abspath(__file__))
def _load_intercept_assist():
path = os.path.join(_ROOT, "Pong-vlm.py")
spec = importlib.util.spec_from_file_location("pong_vlm_module", path)
if spec is None or spec.loader is None:
raise RuntimeError(f"로드 실패: {path}")
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
return mod.InterceptAssist
InterceptAssist = _load_intercept_assist()
class AtariQtView(QWidget):
def __init__(self, title: str = "Pong VLM-DQN") -> None:
super().__init__()
self.setWindowTitle(title)
self._label = QLabel(self)
self._label.setAlignment(Qt.AlignCenter)
lay = QVBoxLayout(self)
lay.setContentsMargins(0, 0, 0, 0)
lay.addWidget(self._label)
def show_frame(self, rgb: np.ndarray) -> None:
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: int) -> None:
super().__init__()
self.fc1 = Dense(64, activation="relu")
self.fc2 = Dense(64, activation="relu")
self.fc_out = Dense(action_size, kernel_initializer=RandomUniform(-1e-3, 1e-3))
def call(self, x: tf.Tensor) -> tf.Tensor:
x = self.fc1(x)
x = self.fc2(x)
return self.fc_out(x)
class DQNAgent:
def __init__(self, state_size: int, action_size: int) -> None:
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.9995
self.epsilon_min = 0.05
self.batch_size = 64
self.train_start = 1000
from collections import deque
self.memory: Any = deque(maxlen=5000)
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) -> None:
self.target_model.set_weights(self.model.get_weights())
def remember(self, state: np.ndarray, action: int, reward: float, next_state: np.ndarray, done: bool) -> None:
self.memory.append((state, action, reward, next_state, done))
def choose_action(self, state: np.ndarray) -> int:
if np.random.rand() <= self.epsilon:
return random.randrange(self.action_size)
q = self.model.predict(state, verbose=0)
return int(np.argmax(q))
def train_model(self) -> None:
if self.epsilon > self.epsilon_min:
self.epsilon *= self.epsilon_decay
mini_batch = random.sample(self.memory, self.batch_size)
states = np.array([s[0][0] for s in mini_batch], dtype=np.float32)
actions = np.array([s[1] for s in mini_batch], dtype=np.int32)
rewards = np.array([s[2] for s in mini_batch], dtype=np.float32)
next_states = np.array([s[3][0] for s in mini_batch], dtype=np.float32)
dones = np.array([s[4] for s in mini_batch], dtype=np.float32)
with tf.GradientTape() as tape:
predicts = self.model(states)
one_hot = tf.one_hot(actions, self.action_size)
predicts = tf.reduce_sum(one_hot * predicts, axis=1)
target_q = self.target_model(next_states)
target_q = tf.stop_gradient(target_q)
max_q = tf.reduce_max(target_q, axis=-1)
targets = rewards + (1.0 - dones) * self.discount_factor * max_q
loss = tf.reduce_mean(tf.square(targets - predicts))
tv = self.model.trainable_variables
grads = tape.gradient(loss, tv)
pairs = [(g, v) for g, v in zip(grads, tv) if g is not None]
self.optimizer.apply_gradients(pairs)
def _norm_xy(
t: dict[str, Any],
h: float,
w: float,
prev_ball: Optional[Tuple[Optional[float], Optional[float]]],
) -> Tuple[np.ndarray, Tuple[Optional[float], Optional[float]]]:
"""VLM 텍스트 CSV 와 대응되는 10차원 벡터 + 다음 step 용 이전 공 좌표."""
paddle = t.get("paddle_py")
bcx = t.get("ball_cx")
bcy = t.get("ball_cy")
pred = t.get("ball_cy_pred")
ball_y = pred if pred is not None else bcy
mot = t.get("motion_sum") or 0.0
dy = t.get("delta_y")
ghost = float(t.get("ghost") or 0)
rej = t.get("reject")
px = float(bcx) / w if bcx is not None else 0.5
py = float(ball_y) / h if ball_y is not None else 0.5
pp = float(paddle) / h if paddle is not None else 0.5
pbx, pby = prev_ball if prev_ball is not None else (None, None)
vx = 0.0
vy = 0.0
if bcx is not None and bcy is not None and pbx is not None and pby is not None:
vx = float(np.clip((bcx - pbx) / max(w, 1.0), -1.0, 1.0))
vy = float(np.clip((bcy - pby) / max(h, 1.0), -1.0, 1.0))
dnorm = float(dy) / h if dy is not None else 0.0
mot_n = float(np.clip(mot / 8000.0, 0.0, 3.0))
gh_n = ghost / 25.0
vis = 1.0 if bcx is not None and ball_y is not None else 0.0
coast = 1.0 if rej == "coast" else 0.0
vec = np.array(
[px, py, vx, vy, pp, dnorm, mot_n, gh_n, vis, coast],
dtype=np.float32,
)
next_prev: Tuple[Optional[float], Optional[float]] = (
float(bcx) if bcx is not None else None,
float(bcy) if bcy is not None else None,
)
return vec, next_prev
def _telemetry_to_text(t: dict[str, Any], last_action: int, step: int) -> str:
by = t.get("ball_cy_pred") if t.get("ball_cy_pred") is not None else t.get("ball_cy")
bx = t.get("ball_cx")
parts = [
f"step={step}",
f"ball_xy=({bx},{by})",
f"paddle_y={t.get('paddle_py')}",
f"motion={t.get('motion_sum')}",
f"delta_y={t.get('delta_y')}",
f"reject={t.get('reject')}",
f"ghost={t.get('ghost')}",
f"last_action={last_action}",
]
return " ".join(str(p) for p in parts)
def main() -> None:
ap = argparse.ArgumentParser(description="Pong geometry state → DQN (same info as VLM text mode)")
ap.add_argument("--episodes", type=int, default=300)
ap.add_argument("--no-gui", action="store_true")
ap.add_argument("--weights", default=os.path.join("save_model_Pong_vlm_dqn", "model.weights.h5"))
ap.add_argument("--graph", default="Pong_vlm_dqn_graph.png")
ap.add_argument(
"--dump-text-log",
default="",
help="JSONL 경로. 각 줄에 text 상태 설명 + state 벡터 (VLM 지도학습용 원료)",
)
args = ap.parse_args()
env = gym.make("Pong-v4", render_mode="rgb_array")
action_size = int(env.action_space.n)
state_dim = 10
agent = DQNAgent(state_dim, action_size)
dummy = np.zeros((1, state_dim), dtype=np.float32)
agent.model(dummy)
agent.target_model(dummy)
if os.path.isfile(args.weights):
agent.model.load_weights(args.weights)
agent.update_target_model()
print(f"가중치 로드: {args.weights}", flush=True)
tracker = InterceptAssist(
deadzone=4.0,
min_motion=100.0,
min_ball_x_frac=0.22,
ghost_max=18,
hot_motion_frac=0.28,
)
app: Optional[QApplication] = None
viewer: Optional[AtariQtView] = None
if not args.no_gui:
app = QApplication.instance() or QApplication(sys.argv)
viewer = AtariQtView("Pong VLM-DQN")
viewer.show()
graph_fig, graph_ax = plt.subplots(figsize=(6, 4))
scores: list[float] = []
episodes_axis: list[int] = []
score_avg = 0.0
max_score = -np.inf
dump_f = None
if args.dump_text_log:
os.makedirs(os.path.dirname(args.dump_text_log) or ".", exist_ok=True)
dump_f = open(args.dump_text_log, "a", encoding="utf-8")
try:
for e in range(args.episodes):
obs, _ = env.reset()
tracker.reset()
done = False
score = 0.0
prev_ball: Optional[Tuple[Optional[float], Optional[float]]] = None
last_action = 0
step_i = 0
while not done:
rgb = env.render()
if rgb is None:
rgb = obs if isinstance(obs, np.ndarray) and obs.ndim == 3 else None
if rgb is not None and viewer is not None:
viewer.show_frame(rgb)
tracker.propose(rgb if rgb is not None else np.zeros((210, 160, 3), dtype=np.uint8))
t = tracker.telemetry
h0 = int(t.get("h") or 210)
w0 = int(t.get("w") or 160)
feat, prev_ball = _norm_xy(t, float(h0), float(w0), prev_ball)
state = np.reshape(feat, (1, state_dim))
if dump_f is not None:
rec = {
"episode": e,
"step": step_i,
"text": _telemetry_to_text(t, last_action, step_i),
"state": feat.tolist(),
}
dump_f.write(json.dumps(rec, ensure_ascii=False) + "\n")
action = agent.choose_action(state)
last_action = action
obs, reward, term, trunc, _ = env.step(action)
done = term or trunc
score += float(reward)
shaped = 0.1 if not done else -1.0
rgb2 = env.render()
if rgb2 is None:
rgb2 = obs if isinstance(obs, np.ndarray) and obs.ndim == 3 else None
tracker.propose(rgb2 if rgb2 is not None else np.zeros((210, 160, 3), dtype=np.uint8))
t2 = tracker.telemetry
h2 = int(t2.get("h") or 210)
w2 = int(t2.get("w") or 160)
next_feat, prev_ball = _norm_xy(t2, float(h2), float(w2), prev_ball)
next_state = np.reshape(next_feat, (1, state_dim))
agent.remember(state, action, shaped, next_state, done)
if len(agent.memory) >= agent.train_start and len(agent.memory) >= agent.batch_size:
agent.train_model()
step_i += 1
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(args.weights) or ".", exist_ok=True)
agent.model.save_weights(args.weights)
print(
f"ep {e:4d} | score {score:6.1f} | avg {score_avg:6.2f} | max_avg {max_score:6.2f} "
f"| mem {len(agent.memory):5d} | eps {agent.epsilon:.4f}",
flush=True,
)
scores.append(score_avg)
episodes_axis.append(e)
graph_ax.clear()
graph_ax.plot(episodes_axis, scores, "b")
graph_ax.set_xlabel("episode")
graph_ax.set_ylabel("avg score (ema)")
graph_fig.tight_layout()
graph_fig.savefig(args.graph)
finally:
if dump_f is not None:
dump_f.close()
env.close()
plt.close(graph_fig)
if viewer is not None:
viewer.close()
if app is not None:
app.quit()
if __name__ == "__main__":
main()