예제 소스
문랜딩 커스텀 환경: 2D 물리와 16가지 분사 행동 · 파일 6개
train.py
"""DQN 학습 스크립트.
사용법
------
python train.py # 기본 (1500 에피소드)
python train.py --episodes 2000
python train.py --render # 평가 시 GUI 렌더링
학습 결과
---------
checkpoints/best.pt # 평가 평균 reward 가 가장 높았던 모델
checkpoints/last.pt # 학습 종료 시점 모델
plots/training.png # 보상/성공률/ε/loss 곡선
logs/train.csv # 에피소드별 로그
"""
from __future__ import annotations
import argparse
import csv
import os
import time
from collections import deque
from pathlib import Path
import numpy as np
import torch
from dqn_agent import DQNAgent, DQNConfig
from lander_env import MoonLanderEnv
def make_dirs(root: Path) -> dict:
paths = {
"ckpt": root / "checkpoints",
"plots": root / "plots",
"logs": root / "logs",
}
for p in paths.values():
p.mkdir(parents=True, exist_ok=True)
return paths
# =====================================================================
# 평가
# =====================================================================
def evaluate(
agent: DQNAgent,
n_episodes: int = 10,
seed_base: int = 9000,
render: bool = False,
) -> dict:
"""ε=0 정책으로 n_episodes 만큼 실행.
return:
mean_reward, success_rate, mean_steps, mean_fuel_used
"""
renderer = None
env = MoonLanderEnv(seed=seed_base, randomize=True)
if render:
from render import LanderRenderer
renderer = LanderRenderer(env)
rewards = []
successes = 0
steps_list = []
fuel_used = []
for i in range(n_episodes):
obs = env.reset(seed=seed_base + i)
ep_r = 0.0
steps = 0
while True:
a = agent.act_greedy(obs)
obs, r, done, info = env.step(a)
ep_r += r
steps += 1
if renderer is not None:
if not renderer.pump_events():
render = False
renderer.close()
renderer = None
else:
renderer.draw({"reward": ep_r, "result": info.result or "..."})
renderer.tick(60)
if done:
if info.result == "landed":
successes += 1
break
rewards.append(ep_r)
steps_list.append(steps)
fuel_used.append(env.INIT_FUEL - env.state.fuel)
if renderer is not None:
renderer.close()
return {
"mean_reward": float(np.mean(rewards)),
"success_rate": successes / n_episodes,
"mean_steps": float(np.mean(steps_list)),
"mean_fuel": float(np.mean(fuel_used)),
}
# =====================================================================
# 학습
# =====================================================================
def train(args: argparse.Namespace) -> None:
root = Path(__file__).resolve().parent
paths = make_dirs(root)
torch.manual_seed(args.seed)
np.random.seed(args.seed)
env = MoonLanderEnv(seed=args.seed, randomize=True)
cfg = DQNConfig(
state_dim=env.STATE_DIM,
num_actions=env.NUM_ACTIONS,
hidden=args.hidden,
lr=args.lr,
gamma=args.gamma,
buffer_size=args.buffer_size,
batch_size=args.batch_size,
learn_start=args.learn_start,
target_tau=args.target_tau,
eps_start=args.eps_start,
eps_end=args.eps_end,
eps_decay_steps=args.eps_decay_steps,
)
agent = DQNAgent(cfg)
print(f"[device] {agent.device}")
print(f"[config] {cfg}")
# 로그 준비
csv_path = paths["logs"] / "train.csv"
csv_file = open(csv_path, "w", newline="")
writer = csv.writer(csv_file)
writer.writerow([
"episode", "env_steps", "reward", "length", "result",
"epsilon", "loss", "eval_reward", "eval_success",
])
history = {
"episode": [],
"reward": [],
"loss": [],
"epsilon": [],
"eval_reward": [],
"eval_success": [],
"eval_episode": [],
}
best_eval_reward = -1e9
recent_rewards: deque[float] = deque(maxlen=50)
recent_success: deque[int] = deque(maxlen=50)
last_eval_reward: float | None = None
last_eval_success: float | None = None
t_start = time.time()
for ep in range(1, args.episodes + 1):
obs = env.reset()
ep_reward = 0.0
ep_len = 0
last_loss = 0.0
result = None
while True:
action = agent.act(obs)
next_obs, reward, done, info = env.step(action)
agent.remember(obs, action, reward, next_obs, done)
loss = agent.learn()
if loss is not None:
last_loss = loss
obs = next_obs
ep_reward += reward
ep_len += 1
if done:
result = info.result
break
recent_rewards.append(ep_reward)
recent_success.append(1 if result == "landed" else 0)
history["episode"].append(ep)
history["reward"].append(ep_reward)
history["loss"].append(agent.avg_loss())
history["epsilon"].append(agent.epsilon())
# 주기적 평가
if ep % args.eval_every == 0:
eval_stats = evaluate(agent, n_episodes=args.eval_episodes,
seed_base=20000 + ep, render=False)
last_eval_reward = eval_stats["mean_reward"]
last_eval_success = eval_stats["success_rate"]
history["eval_reward"].append(last_eval_reward)
history["eval_success"].append(last_eval_success)
history["eval_episode"].append(ep)
elapsed = time.time() - t_start
print(
f"[ep {ep:5d}] "
f"R={np.mean(recent_rewards):7.1f} "
f"succ50={np.mean(recent_success):.2f} "
f"eps={agent.epsilon():.3f} "
f"loss={agent.avg_loss():.4f} "
f"| eval R={last_eval_reward:7.1f} "
f"succ={last_eval_success:.2f} "
f"| {elapsed:6.1f}s",
flush=True,
)
if last_eval_reward > best_eval_reward:
best_eval_reward = last_eval_reward
agent.save(str(paths["ckpt"] / "best.pt"))
writer.writerow([
ep, agent.total_env_steps, ep_reward, ep_len, result,
agent.epsilon(), agent.avg_loss(),
last_eval_reward, last_eval_success,
])
csv_file.flush()
# 마지막 모델 저장
agent.save(str(paths["ckpt"] / "last.pt"))
csv_file.close()
# 학습 곡선 저장
plot_history(history, str(paths["plots"] / "training.png"))
print(f"\n[done] best eval reward = {best_eval_reward:.2f}")
print(f"checkpoints: {paths['ckpt']}")
print(f"plots: {paths['plots']}")
print(f"logs: {paths['logs']}")
# 옵션: 최종 평가 GUI 렌더링
if args.render:
print("[render] best.pt 로 평가 5 에피소드 렌더링")
agent.load(str(paths["ckpt"] / "best.pt"))
evaluate(agent, n_episodes=5, seed_base=99000, render=True)
def plot_history(history: dict, save_path: str) -> None:
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
fig, axes = plt.subplots(2, 2, figsize=(12, 8))
ep = history["episode"]
# 1) reward
ax = axes[0, 0]
ax.plot(ep, history["reward"], color="#5b8def", alpha=0.4, label="episode")
if len(history["reward"]) >= 20:
# 이동 평균
w = 30
r = np.array(history["reward"])
ma = np.convolve(r, np.ones(w) / w, mode="valid")
ax.plot(ep[w - 1:], ma, color="#1f4dbf", label=f"MA{w}")
ax.set_title("Episode reward")
ax.set_xlabel("episode")
ax.legend()
ax.grid(alpha=0.3)
# 2) eval reward + success
ax = axes[0, 1]
if history["eval_episode"]:
ax.plot(history["eval_episode"], history["eval_reward"],
marker="o", color="#27a65a", label="eval reward")
ax.set_ylabel("eval reward", color="#27a65a")
ax2 = ax.twinx()
ax2.plot(history["eval_episode"], history["eval_success"],
marker="s", color="#e07b39", label="success rate")
ax2.set_ylabel("success rate", color="#e07b39")
ax2.set_ylim(-0.05, 1.05)
ax.set_title("Evaluation")
ax.set_xlabel("episode")
ax.grid(alpha=0.3)
# 3) epsilon
ax = axes[1, 0]
ax.plot(ep, history["epsilon"], color="#8e3fbf")
ax.set_title("Exploration ε")
ax.set_xlabel("episode")
ax.grid(alpha=0.3)
# 4) loss
ax = axes[1, 1]
ax.plot(ep, history["loss"], color="#bf3f5b")
ax.set_title("Avg loss (last 100)")
ax.set_xlabel("episode")
ax.set_yscale("log")
ax.grid(alpha=0.3)
fig.tight_layout()
fig.savefig(save_path, dpi=140)
plt.close(fig)
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser()
p.add_argument("--episodes", type=int, default=1500)
p.add_argument("--seed", type=int, default=42)
p.add_argument("--hidden", type=int, default=128)
p.add_argument("--lr", type=float, default=5e-4)
p.add_argument("--gamma", type=float, default=0.99)
p.add_argument("--buffer-size", type=int, default=100_000)
p.add_argument("--batch-size", type=int, default=64)
p.add_argument("--learn-start", type=int, default=2_000)
p.add_argument("--target-tau", type=float, default=0.005)
p.add_argument("--eps-start", type=float, default=1.0)
p.add_argument("--eps-end", type=float, default=0.05)
p.add_argument("--eps-decay-steps", type=int, default=80_000)
p.add_argument("--eval-every", type=int, default=50)
p.add_argument("--eval-episodes", type=int, default=10)
p.add_argument("--render", action="store_true",
help="학습 종료 후 best.pt 로 GUI 평가")
return p.parse_args()
if __name__ == "__main__":
train(parse_args())