예제 소스
테트리스를 배우는 강화학습 에이전트 만들기 · 파일 31개
train_qt_live.py
"""train_qt_live.py
학습(LocalTetrisEnv 헤드리스) + Qt 시각화를 동시에 실행.
백그라운드 스레드: LocalTetrisEnv로 DQN 학습 → model/live_latest.zip 주기 저장
메인 스레드 : Qt 앱에 TCP 연결, 현재 policy로 액션 전송, 파일 변경 시 자동 리로드
human_play_mode: 정책으로 플레이하지 않고, Qt 키 입력 전이만 리플레이 버퍼에 넣어 DQN.train으로 학습.
"""
import argparse
import json
import os
import re
import shutil
import socket
import threading
import time
from datetime import datetime
from pathlib import Path
import numpy as np
# ── 공통 유틸 ─────────────────────────────────────────────────────────────────
def _mtime(path: Path) -> float:
"""파일 수정 시각 반환 (없으면 0)."""
for p in (path, Path(str(path) + ".zip")):
try:
return os.path.getmtime(p)
except OSError:
pass
return 0.0
def _open_model(path: Path):
"""path 또는 path.zip 을 열어 파일 객체 반환."""
for p in (path, Path(str(path) + ".zip")):
if p.exists():
return open(p, "rb")
raise FileNotFoundError(f"모델 파일 없음: {path}")
def _make_vec_env_low_level(args):
"""휴먼 데모·저수준 DQN용: 항상 LocalTetrisEnv(Discrete(7))."""
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import DummyVecEnv
from local_tetris_env import LocalTetrisConfig, LocalTetrisEnv
from llm_reward_wrapper import LLMRewardConfig, LLMRewardWrapper
def make_env():
env = LocalTetrisEnv(
lib_path=args.env_id,
cfg=LocalTetrisConfig(
stack_height_penalty_scale=float(getattr(args, "stack_height_penalty_scale", 0.0)),
bumpiness_penalty_scale=float(getattr(args, "bumpiness_penalty_scale", 0.0)),
max_height_penalty_scale=float(getattr(args, "max_height_penalty_scale", 0.0)),
),
)
llm_cfg = getattr(args, "llm_reward_cfg", None)
if isinstance(llm_cfg, dict) and bool(llm_cfg.get("enabled", False)):
env = LLMRewardWrapper(env, LLMRewardConfig(**llm_cfg))
env.reset(seed=args.seed)
return Monitor(env)
return DummyVecEnv([make_env])
def _init_dqn(args, *, force_low_level: bool):
"""
DQN 생성 또는 로드.
force_low_level True: 휴먼 데모용 Discrete(7)만 (매크로 zip이면 새로 생성).
"""
from stable_baselines3 import DQN
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.preprocessing import is_image_space
from stable_baselines3.common.vec_env import DummyVecEnv
from local_tetris_env import LocalTetrisConfig, LocalTetrisEnv
from macro_planning_env import MacroPlanningTetrisEnv
from llm_reward_wrapper import LLMRewardConfig, LLMRewardWrapper
if force_low_level:
vec_env = _make_vec_env_low_level(args)
if args.load_model:
load_path = Path(args.load_model)
with _open_model(load_path) as f:
tmp = DQN.load(f)
try:
n = int(getattr(tmp.action_space, "n", 0))
except Exception:
n = 0
if n == 7:
print(f"[train] 휴먼용 모델 로드: {load_path}")
with _open_model(load_path) as f:
model = DQN.load(f, env=vec_env)
model.learning_rate = args.learning_rate
return model
print(
f"[train] 휴먼 모드는 Discrete(7)만 사용합니다 (로드된 액션 수: {n}). 새 DQN을 만듭니다."
)
return DQN(
policy="CnnPolicy" if is_image_space(vec_env.observation_space) else "MlpPolicy",
env=vec_env,
learning_rate=args.learning_rate,
buffer_size=args.buffer_size,
batch_size=args.batch_size,
train_freq=args.train_freq,
target_update_interval=args.target_update_interval,
exploration_fraction=args.exploration_fraction,
exploration_final_eps=args.exploration_final_eps,
verbose=1,
seed=args.seed,
)
def make_env():
cfg = LocalTetrisConfig(
stack_height_penalty_scale=float(getattr(args, "stack_height_penalty_scale", 0.0)),
bumpiness_penalty_scale=float(getattr(args, "bumpiness_penalty_scale", 0.0)),
max_height_penalty_scale=float(getattr(args, "max_height_penalty_scale", 0.0)),
)
if getattr(args, "macro", False):
env = MacroPlanningTetrisEnv(lib_path=args.env_id, cfg=cfg)
else:
env = LocalTetrisEnv(lib_path=args.env_id, cfg=cfg)
llm_cfg = getattr(args, "llm_reward_cfg", None)
if isinstance(llm_cfg, dict) and bool(llm_cfg.get("enabled", False)):
env = LLMRewardWrapper(env, LLMRewardConfig(**llm_cfg))
env.reset(seed=args.seed)
return Monitor(env)
if args.load_model:
load_path = Path(args.load_model)
print(f"[train] 기존 모델 로드: {load_path}")
with _open_model(load_path) as f:
tmp = DQN.load(f)
try:
n = int(getattr(tmp.action_space, "n", 0))
except Exception:
n = 0
if n == 40:
args.macro = True
elif n == 7:
args.macro = False
else:
pass
vec_env = DummyVecEnv([make_env])
with _open_model(load_path) as f:
model = DQN.load(f, env=vec_env)
model.learning_rate = args.learning_rate
return model
vec_env = DummyVecEnv([make_env])
return DQN(
policy="CnnPolicy" if is_image_space(vec_env.observation_space) else "MlpPolicy",
env=vec_env,
learning_rate=args.learning_rate,
buffer_size=args.buffer_size,
batch_size=args.batch_size,
train_freq=args.train_freq,
target_update_interval=args.target_update_interval,
exploration_fraction=args.exploration_fraction,
exploration_final_eps=args.exploration_final_eps,
verbose=1,
seed=args.seed,
)
# ── 학습 스레드 ───────────────────────────────────────────────────────────────
def _train_thread(args, live_path: Path, ready_event: threading.Event,
stop_event: threading.Event):
if getattr(args, "human_play_mode", False):
model = _init_dqn(args, force_low_level=True)
model.save(str(live_path))
ready_event.set()
print(
"[train] human_play_mode=True: 헤드리스 env.step 학습은 하지 않습니다. "
"메인 스레드가 Qt 휴먼 전이로 리플레이 버퍼·DQN을 갱신합니다."
)
while not stop_event.is_set():
time.sleep(0.5)
return
from stable_baselines3.common.callbacks import BaseCallback
model = _init_dqn(args, force_low_level=False)
class LiveCallback(BaseCallback):
def __init__(self, live_path, freq, ready_event, stop_event):
super().__init__()
self.live_path = live_path
self.freq = freq
self.ready_event = ready_event
self.stop_event = stop_event
def _on_step(self):
if self.stop_event.is_set():
return False
if self.n_calls % self.freq == 0:
self.model.save(str(self.live_path))
print(f"[train] 체크포인트 저장 ({self.n_calls:,} steps)")
if not self.ready_event.is_set():
self.ready_event.set()
return True
# 초기 체크포인트: 학습 전에 현재 모델 저장해서 viz가 바로 시작할 수 있게
model.save(str(live_path))
ready_event.set()
print(f"[train] 초기 모델 저장: {live_path}.zip")
round_num = 1
while not stop_event.is_set():
print(f"\n[train] ── 라운드 {round_num} 시작 ({args.timesteps:,} steps) ──")
callback = LiveCallback(live_path, args.checkpoint_freq, ready_event, stop_event)
model.learn(total_timesteps=args.timesteps, callback=callback,
progress_bar=False, reset_num_timesteps=False)
if stop_event.is_set():
break
# 라운드 완료마다 live_latest 갱신 (이전 파일은 .bak 1개만 유지)
live_zip = live_path.with_suffix(".zip")
if live_zip.exists():
backup_path = live_zip.with_name(f"{live_zip.name}.bak")
if backup_path.exists():
backup_path.unlink()
live_zip.replace(backup_path)
print(f"[train] 이전 체크포인트 백업: {backup_path}")
model.save(str(live_path))
print(f"[train] 라운드 {round_num} 저장: {live_zip}")
round_num += 1
# ── Qt 시각화 (메인 스레드) ───────────────────────────────────────────────────
def _ensure_sb3_logger(model, tb_log_name: str = "human_demo") -> None:
"""
SB3는 `_logger`를 저장하지 않으며 `learn()` → `_setup_learn`에서만 설정한다.
`DQN.load()` 직후 `train()`만 호출하면 AttributeError가 나므로 사전에 구성한다.
"""
if getattr(model, "_logger", None) is not None:
return
from stable_baselines3.common import utils
model._custom_logger = False
model._logger = utils.configure_logger(
model.verbose,
model.tensorboard_log,
tb_log_name,
reset_num_timesteps=False,
)
def _dqn_offline_step(model) -> None:
"""DQN.collect_rollouts 밖에서 target Q망을 주기적으로 맞춤."""
from stable_baselines3.common.utils import polyak_update
model._n_calls += 1
if model._n_calls % max(model.target_update_interval // model.n_envs, 1) == 0:
polyak_update(model.q_net.parameters(), model.q_net_target.parameters(), model.tau)
if len(model.batch_norm_stats) > 0:
polyak_update(model.batch_norm_stats, model.batch_norm_stats_target, 1.0)
def _run_qt_human(args, live_path: Path, stop_event: threading.Event) -> None:
"""human_play_mode: 정책 추론 없이 Qt 휴먼 전이만 리플레이 버퍼에 넣고 DQN.train 호출."""
from stable_baselines3 import DQN
def load_model():
ve = _make_vec_env_low_level(args)
with _open_model(live_path) as f:
return DQN.load(f, env=ve)
print("[viz] 휴먼 모드 — DQN 로드 중…")
model = load_model()
_ensure_sb3_logger(model)
teacher = None
llm_cfg = getattr(args, "llm_reward_cfg", None)
if isinstance(llm_cfg, dict) and bool(llm_cfg.get("enabled", False)):
try:
from llm_reward_wrapper import LLMRewardConfig
from llm_teacher import LocalGGUFTeacher, LLMTeacherConfig
from vlm_teacher import LocalQwenVLViaTransformersTeacher, VLMTeacherConfig
c = LLMRewardConfig(**llm_cfg)
if c.backend == "qwen2_5_vl":
teacher = LocalQwenVLViaTransformersTeacher(
VLMTeacherConfig(
enabled=True,
model_dir=c.vlm_model_dir,
call_every=c.call_every,
cell_px=c.vlm_cell_px,
temperature=c.temperature,
top_p=c.top_p,
max_new_tokens=c.max_tokens,
)
)
else:
teacher = LocalGGUFTeacher(
LLMTeacherConfig(
enabled=True,
model_path=c.model_path,
n_ctx=c.n_ctx,
n_threads=c.n_threads,
n_gpu_layers=c.n_gpu_layers,
temperature=c.temperature,
top_p=c.top_p,
max_tokens=c.max_tokens,
call_every=c.call_every,
reward_alpha=c.alpha,
)
)
except Exception as e:
print(f"[viz] teacher 초기화 실패(무시): {e}")
tf = args.train_freq
grad_steps = tf[0] if isinstance(tf, tuple) else int(tf)
grad_steps = max(1, min(grad_steps, 128))
sock = None
print(f"[viz] human 모드 — Qt 서버({args.host}:{args.port}) 연결 시도 중...")
deadline = time.time() + args.connect_timeout
while time.time() < deadline and not stop_event.is_set():
try:
sock = socket.create_connection((args.host, args.port), timeout=2.0)
break
except OSError:
time.sleep(0.5)
if sock is None:
print(f"[viz] Qt 서버 연결 실패 — tetris_rl --rl-port {args.port} 실행 여부 확인")
stop_event.set()
return
f_in = sock.makefile("r", encoding="utf-8", newline="\n")
f_out = sock.makefile("w", encoding="utf-8", newline="\n")
print("[viz] Qt 서버 연결됨 (human 명령 전송)")
# create_connection(timeout=...) sets socket timeout for subsequent reads too.
# We want a blocking stream for readline() based protocol.
sock.settimeout(None)
f_out.write("human\n")
f_out.flush()
last_seen_seq: int | None = None
last_obs: np.ndarray | None = None
n_added = 0
last_print = 0
last_loss: float | None = None
try:
while not stop_event.is_set():
try:
line = f_in.readline()
except TimeoutError:
continue
if not line:
break
try:
state = json.loads(line)
except json.JSONDecodeError:
continue
seq_c = int(state.get("step_seq", 0))
if last_seen_seq is not None and seq_c == last_seen_seq:
continue
la_c = int(state.get("last_action", -1))
w = int(state["w"])
vh = int(state["visible_h"])
obs_c = np.asarray(state["grid"], dtype=np.uint8).reshape(vh, w, 1)
go = bool(state.get("game_over", False))
r = float(state["reward"])
if last_obs is not None and la_c >= 0:
obs_b = last_obs[np.newaxis, ...]
next_b = obs_c[np.newaxis, ...]
act_b = np.array([[la_c]], dtype=np.float32)
rew_b = np.array([r], dtype=np.float32)
done_b = np.array([float(go)], dtype=np.float32)
model.replay_buffer.add(obs_b, next_b, act_b, rew_b, done_b, [{}])
model.num_timesteps += 1
n_added += 1
_dqn_offline_step(model)
min_buf = max(model.learning_starts, model.batch_size)
if model.replay_buffer.size() >= min_buf:
try:
model.train(gradient_steps=grad_steps, batch_size=args.batch_size)
# SB3 logger에서 가장 최근 loss를 꺼내서 찍기
try:
last_loss = float(model.logger.name_to_value.get("train/loss")) # type: ignore[attr-defined]
except Exception:
pass
except Exception as e:
print(f"[viz] train() 실패: {e}")
if n_added > 0 and n_added % args.checkpoint_freq == 0:
model.save(str(live_path))
print(f"[viz] human 체크포인트 저장 (전이 {n_added:,}개)")
# 진행 상황 출력(눈으로 학습 확인)
if n_added - last_print >= 100:
last_print = n_added
buf = model.replay_buffer.size()
if last_loss is None:
print(f"[viz] 전이 {n_added:,}개, replay {buf:,}/{model.replay_buffer.buffer_size:,} (학습 대기/진행 중)")
else:
print(
f"[viz] 전이 {n_added:,}개, replay {buf:,}/{model.replay_buffer.buffer_size:,}, "
f"loss {last_loss:.5f}"
)
# SB3 stdout logger dump(있으면 표 형태로도 출력)
try:
model.logger.dump(step=model.num_timesteps)
except Exception:
pass
last_seen_seq = seq_c
if go:
last_obs = None
f_out.write("restart\n")
f_out.flush()
else:
last_obs = obs_c.copy()
# HUD에 teacher 평가 표시 (옵션)
if teacher is not None and la_c >= 0:
try:
na, sc = teacher.judge_and_suggest(obs_c, la_c, state)
f_out.write(f"hud llm_score={sc:+.3f} llm_next={na}\n")
f_out.flush()
except Exception:
pass
except (KeyboardInterrupt, BrokenPipeError):
pass
finally:
try:
model.save(str(live_path))
print(f"[viz] 종료 시 저장 완료 (누적 전이 {n_added:,}개)")
except Exception as e:
print(f"[viz] 종료 저장 실패: {e}")
if sock is not None:
sock.close()
def _run_qt(args, live_path: Path, ready_event: threading.Event,
stop_event: threading.Event):
if getattr(args, "human_play_mode", False):
print("[viz] human_play_mode: 휴먼 플레이 전이로만 학습 (정책 추론 없음)")
ready_event.wait()
_run_qt_human(args, live_path, stop_event)
return
from stable_baselines3 import DQN
from macro_planning_env import next_micro_for_target
# 첫 체크포인트가 저장될 때까지 대기
print("[viz] 첫 체크포인트 대기 중...")
ready_event.wait()
def load():
with _open_model(live_path) as f:
return DQN.load(f)
model = load()
last_mtime = _mtime(live_path)
print(f"[viz] 모델 로드 완료")
teacher = None
llm_cfg = getattr(args, "llm_reward_cfg", None)
if isinstance(llm_cfg, dict) and bool(llm_cfg.get("enabled", False)):
try:
from llm_reward_wrapper import LLMRewardConfig
from llm_teacher import LocalGGUFTeacher, LLMTeacherConfig
from vlm_teacher import LocalQwenVLViaTransformersTeacher, VLMTeacherConfig
c = LLMRewardConfig(**llm_cfg)
if c.backend == "qwen2_5_vl":
teacher = LocalQwenVLViaTransformersTeacher(
VLMTeacherConfig(
enabled=True,
model_dir=c.vlm_model_dir,
call_every=c.call_every,
cell_px=c.vlm_cell_px,
temperature=c.temperature,
top_p=c.top_p,
max_new_tokens=c.max_tokens,
)
)
else:
teacher = LocalGGUFTeacher(
LLMTeacherConfig(
enabled=True,
model_path=c.model_path,
n_ctx=c.n_ctx,
n_threads=c.n_threads,
n_gpu_layers=c.n_gpu_layers,
temperature=c.temperature,
top_p=c.top_p,
max_tokens=c.max_tokens,
call_every=c.call_every,
reward_alpha=c.alpha,
)
)
except Exception as e:
print(f"[viz] teacher 초기화 실패(무시): {e}")
# Qt 서버 연결
sock = None
print(f"[viz] Qt 서버({args.host}:{args.port}) 연결 시도 중...")
deadline = time.time() + args.connect_timeout
while time.time() < deadline and not stop_event.is_set():
try:
sock = socket.create_connection((args.host, args.port), timeout=2.0)
break
except OSError:
time.sleep(0.5)
if sock is None:
print(f"[viz] Qt 서버 연결 실패 — tetris_rl --rl-port {args.port} 실행 여부 확인")
stop_event.set()
return
f_in = sock.makefile("r", encoding="utf-8", newline="\n")
f_out = sock.makefile("w", encoding="utf-8", newline="\n")
print(f"[viz] Qt 서버 연결됨")
sock.settimeout(None)
step = 0
state = None
try:
while not stop_event.is_set():
if state is None:
try:
line = f_in.readline()
except TimeoutError:
continue
if not line:
break
state = json.loads(line)
status = state.get("status", "running")
if status == "game_over":
f_out.write("restart\n")
f_out.flush()
state = None
cur = _mtime(live_path)
if cur > last_mtime:
try:
model = load()
last_mtime = cur
print(f"[viz] 모델 리로드")
except Exception as e:
print(f"[viz] 리로드 실패: {e}")
continue
if status == "paused":
f_out.write("resume\n")
f_out.flush()
try:
line = f_in.readline()
except TimeoutError:
continue
state = json.loads(line) if line else None
continue
w = int(state["w"])
vh = int(state["visible_h"])
grid = np.asarray(state["grid"], dtype=np.uint8).reshape(vh, w, 1)
if getattr(args, "macro", False):
for _k in ("piece_rot", "piece_x", "piece_type"):
if _k not in state:
raise RuntimeError(
"매크로 모드는 Qt `tetris_rl`이 JSON에 piece_type, piece_x, piece_y, piece_rot를 "
"보내야 합니다. 소스가 오래된 빌드입니다.\n"
" cd tetris && cmake --build build -j\n"
" ./build/tetris_rl --rl-port <포트>\n"
"또는 --macro 없이 저수준(7액션) 학습만 사용하세요."
)
macro_a, _ = model.predict(grid, deterministic=args.deterministic)
target_rot = int(macro_a) % 4
target_col = int(macro_a) // 4
start_pi = int(state["piece"])
while (
int(state["piece"]) == start_pi
and not state.get("game_over", False)
and not stop_event.is_set()
):
micro = next_micro_for_target(
int(state.get("piece_type", 0)),
int(state.get("piece_rot", 0)),
int(state.get("piece_x", 0)),
target_rot,
target_col,
)
f_out.write(f"{micro}\n")
f_out.flush()
line = f_in.readline()
if not line:
stop_event.set()
break
state = json.loads(line)
step += 1
else:
action, _ = model.predict(grid, deterministic=args.deterministic)
f_out.write(f"{int(action)}\n")
f_out.flush()
try:
line = f_in.readline()
except TimeoutError:
continue
state = json.loads(line) if line else None
step += 1
if teacher is not None and state is not None:
try:
na, sc = teacher.judge_and_suggest(grid, int(action), state)
f_out.write(f"hud llm_score={sc:+.3f} llm_next={na}\n")
f_out.flush()
except Exception:
pass
if args.reload_every > 0 and step % args.reload_every == 0:
cur = _mtime(live_path)
if cur > last_mtime:
try:
model = load()
last_mtime = cur
print(f"[viz] 모델 리로드 ({step} steps)")
except Exception as e:
print(f"[viz] 리로드 실패: {e}")
except (KeyboardInterrupt, BrokenPipeError, TimeoutError):
pass
finally:
sock.close()
# ── 메인 ──────────────────────────────────────────────────────────────────────
def main():
p = argparse.ArgumentParser(description="학습 + Qt 시각화 동시 실행 (config 기반)")
p.add_argument(
"--config",
default="config/rl_config.py",
help="설정 파일 경로 (기본: config/rl_config.py)",
)
args = p.parse_args()
# config 로드
import importlib.util
import sys
cfg_path = Path(args.config).resolve()
spec = importlib.util.spec_from_file_location("rl_config", str(cfg_path))
if spec is None or spec.loader is None:
raise RuntimeError(f"config 로드 실패: {cfg_path}")
mod = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = mod
spec.loader.exec_module(mod) # type: ignore[union-attr]
cfg = getattr(mod, "CONFIG")
# 기존 코드가 args.*를 참조하므로 간단히 namespace처럼 맞춰준다
class _Args:
pass
run = _Args()
run.env_id = cfg.env.env_id
run.seed = cfg.env.seed
run.macro = bool(cfg.env.macro)
run.human_play_mode = bool(cfg.env.human_play_mode)
run.llm_reward_cfg = cfg.llm_reward.__dict__ if hasattr(cfg, "llm_reward") else None
run.stack_height_penalty_scale = float(getattr(cfg.env, "stack_height_penalty_scale", 0.0))
run.bumpiness_penalty_scale = float(getattr(cfg.env, "bumpiness_penalty_scale", 0.0))
run.max_height_penalty_scale = float(getattr(cfg.env, "max_height_penalty_scale", 0.0))
run.timesteps = cfg.dqn.timesteps
run.learning_rate = cfg.dqn.learning_rate
run.buffer_size = cfg.dqn.buffer_size
run.batch_size = cfg.dqn.batch_size
run.train_freq = cfg.dqn.train_freq
run.target_update_interval = cfg.dqn.target_update_interval
run.exploration_fraction = cfg.dqn.exploration_fraction
run.exploration_final_eps = cfg.dqn.exploration_final_eps
run.model_dir = cfg.io.model_dir
run.load_model = cfg.io.load_model
run.host = cfg.qt.host
run.port = cfg.qt.port
run.connect_timeout = cfg.qt.connect_timeout
run.deterministic = bool(cfg.qt.deterministic)
run.checkpoint_freq = cfg.live.checkpoint_freq
run.reload_every = cfg.live.reload_every
Path(run.model_dir).mkdir(parents=True, exist_ok=True)
live_path = Path(run.model_dir) / "live_latest"
ready_event = threading.Event()
stop_event = threading.Event()
train_t = threading.Thread(
target=_train_thread,
args=(run, live_path, ready_event, stop_event),
daemon=True, name="train",
)
train_t.start()
try:
_run_qt(run, live_path, ready_event, stop_event)
except KeyboardInterrupt:
print("\n[main] 중단됨")
finally:
stop_event.set()
train_t.join(timeout=5)
print("[main] 종료")
if __name__ == "__main__":
main()