예제 소스
문랜딩 커스텀 환경: 2D 물리와 16가지 분사 행동 · 파일 6개
play.py
"""학습된 DQN 모델을 실행 / 또는 무한 연속 학습.
3가지 모드:
1) 추론 (기본) : 학습된 모델로 ε=0 greedy 실행
2) 연속 학습 : --learn 플래그. 시각화하면서 무한 학습 + 자동 저장
3) 사람 조작 : --human 플래그. 키보드로 직접 조종 (환경 점검용)
사용 예시
---------
# 1) 학습된 모델 데모
python play.py --episodes 5
# 2) 연속 학습 (무한 실행, 창 닫으면 종료)
python play.py --learn
# 3) 기존 best.pt 에서 이어서 학습
python play.py --learn --init-from checkpoints/best.pt
# 4) 빠르게 학습하고 가끔만 화면 갱신 (10배속)
python play.py --learn --speed 10
# 5) 처음부터 새로 시작 (기존 live.pt 무시)
python play.py --learn --fresh
# 6) 사람 키보드 조작
python play.py --human
키 조작
-------
공통:
[Esc] 저장 후 종료
연속 학습 모드(--learn):
[+] / [=] 화면 1 프레임당 환경 step 수 증가 (학습 속도 ↑)
[-] 감소
[Space] 일시정지 / 재개
[S] 즉시 체크포인트 저장
사람 조작(--human):
[Space] 메인 엔진 [↑/W] 상부 RCS
[←/A] 좌측 RCS [→/D] 우측 RCS
[R] 에피소드 리셋
"""
from __future__ import annotations
import argparse
import time
from collections import deque
from pathlib import Path
import numpy as np
import pygame
from dqn_agent import DQNAgent, DQNConfig
from lander_env import MoonLanderEnv
# =====================================================================
# 모드 1: 추론 (greedy)
# =====================================================================
def run_dqn(args: argparse.Namespace) -> None:
env = MoonLanderEnv(seed=args.seed, randomize=not args.fixed_start)
agent = DQNAgent(DQNConfig(
state_dim=env.STATE_DIM,
num_actions=env.NUM_ACTIONS,
))
ckpt = Path(args.ckpt)
if not ckpt.exists():
raise FileNotFoundError(f"체크포인트 없음: {ckpt}")
agent.load(str(ckpt))
print(f"[load] {ckpt}")
renderer = None
if not args.no_render:
from render import LanderRenderer, show_terminal_banner
renderer = LanderRenderer(env, caption="Moon Lander DQN — play")
successes = 0
rewards = []
for ep in range(1, args.episodes + 1):
obs = env.reset(seed=args.seed + ep)
ep_r = 0.0
result = None
while True:
a = agent.act_greedy(obs)
obs, r, done, info = env.step(a)
ep_r += r
if renderer is not None:
if not renderer.pump_events():
renderer.close()
print("[user-quit]")
_print_summary(rewards, successes, ep - 1)
return
renderer.draw({
"episode": ep,
"reward": ep_r,
"result": info.result or "...",
})
renderer.tick(args.fps)
if done:
result = info.result
break
rewards.append(ep_r)
if result == "landed":
successes += 1
print(f" episode {ep:3d}: reward={ep_r:7.2f} result={result} steps={env.steps}")
if renderer is not None:
color = (120, 240, 140) if result == "landed" else (240, 120, 120)
show_terminal_banner(renderer.screen, renderer.font_big,
f"{result.upper()}", color)
pygame.display.flip()
pygame.time.delay(800)
_print_summary(rewards, successes, args.episodes)
if renderer is not None:
renderer.close()
def _print_summary(rewards, successes, total) -> None:
if total == 0:
return
arr = np.array(rewards)
print("\n" + "=" * 50)
print(f" episodes : {total}")
print(f" success rate : {successes}/{total} = {successes/total:.1%}")
print(f" mean reward : {arr.mean():.2f}")
print(f" std reward : {arr.std():.2f}")
print(f" min / max : {arr.min():.2f} / {arr.max():.2f}")
print("=" * 50)
# =====================================================================
# 모드 2: 연속 학습 (무한 실행, 시각화 + 학습)
# =====================================================================
def run_live_train(args: argparse.Namespace) -> None:
"""창을 닫을 때까지 무한히 환경을 돌리며 매 step 마다 학습한다."""
from render import LanderRenderer, show_terminal_banner
env = MoonLanderEnv(seed=args.seed, randomize=not args.fixed_start)
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,
eps_decay_steps=args.eps_decay_steps,
eps_end=args.eps_end,
)
agent = DQNAgent(cfg)
ckpt_path = Path(args.ckpt)
ckpt_path.parent.mkdir(parents=True, exist_ok=True)
best_path = ckpt_path.with_name(ckpt_path.stem + "_best.pt")
# 가중치 로드 우선순위: --fresh 면 새로 시작, 아니면
# 1) 지정한 ckpt 가 존재하면 로드 (이어 학습)
# 2) --init-from 으로 다른 ckpt 에서 가중치만 가져오기
if args.fresh:
print("[init] starting from scratch")
elif ckpt_path.exists():
agent.load(str(ckpt_path))
print(f"[load] resumed from {ckpt_path}, env_steps={agent.total_env_steps}")
elif args.init_from and Path(args.init_from).exists():
agent.load(args.init_from)
# 가중치만 가져왔으니 step 카운터는 0 부터 시작해서
# 약간의 추가 탐험 여지를 남긴다 (eps 가 처음부터 eps_end 가 되지 않도록).
agent.total_env_steps = max(0, int(args.eps_decay_steps * 0.7))
print(f"[load] initialized weights from {args.init_from}")
else:
print(f"[init] no existing ckpt at {ckpt_path}, starting fresh")
print(f"[device] {agent.device} [ckpt] {ckpt_path} [best] {best_path}")
print("[hint] [+]/[-] 속도 조절, [Space] 일시정지, [S] 즉시 저장, [Esc] 종료")
renderer = LanderRenderer(env, caption="Moon Lander — live training")
rewards_50: deque[float] = deque(maxlen=50)
success_50: deque[int] = deque(maxlen=50)
best_succ = 0.0
speed = max(1, int(args.speed))
paused = False
episode = 0
total_lifetime_steps = 0
t0 = time.time()
last_print = t0
def hud(ep_r: float, eps_now: float, status: str) -> dict:
succ = float(np.mean(success_50)) if success_50 else 0.0
return {
"ep": episode,
"reward": ep_r,
"eps": eps_now,
"succ50": succ,
"best_succ": best_succ,
"loss": agent.avg_loss(),
"env_steps": agent.total_env_steps,
"speed": f"x{speed}{' PAUSED' if paused else ''}",
"status": status,
}
def handle_events() -> bool:
"""반환 False 면 종료."""
nonlocal speed, paused
for event in pygame.event.get():
if event.type == pygame.QUIT:
return False
if event.type == pygame.KEYDOWN:
if event.key == pygame.K_ESCAPE:
return False
elif event.key in (pygame.K_PLUS, pygame.K_EQUALS, pygame.K_KP_PLUS):
speed = min(speed * 2, 256)
print(f"[speed] x{speed}")
elif event.key in (pygame.K_MINUS, pygame.K_KP_MINUS):
speed = max(speed // 2, 1)
print(f"[speed] x{speed}")
elif event.key == pygame.K_SPACE:
paused = not paused
print(f"[pause] {paused}")
elif event.key == pygame.K_s:
agent.save(str(ckpt_path))
print(f"[save] {ckpt_path} (manual)")
return True
try:
running = True
while running:
episode += 1
obs = env.reset()
ep_r = 0.0
steps_since_render = 0
result = None
while True:
# 일시정지 처리: 화면만 갱신하고 환경/학습은 멈춤
if paused:
if not handle_events():
running = False
break
renderer.draw(hud(ep_r, max(agent.epsilon(), args.eps_min), "paused"))
renderer.tick(30)
continue
# 행동 선택 (ε floor 적용)
eps_now = max(agent.epsilon(), args.eps_min)
a = agent.act(obs, eps=eps_now)
ns, r, done, info = env.step(a)
agent.remember(obs, a, r, ns, done)
agent.learn()
obs = ns
ep_r += r
total_lifetime_steps += 1
steps_since_render += 1
if done or steps_since_render >= speed:
steps_since_render = 0
if not handle_events():
running = False
break
renderer.draw(hud(ep_r, eps_now, info.result or "..."))
if not done:
# speed 가 클수록 fps lock 을 약하게 걸어 빠르게 진행
renderer.tick(args.fps if speed <= 4 else 240)
if done:
break
if not running:
break
result = info.result
rewards_50.append(ep_r)
success_50.append(1 if result == "landed" else 0)
# 베스트 success rate 갱신 시 별도 저장
if len(success_50) == success_50.maxlen:
cur_succ = float(np.mean(success_50))
if cur_succ > best_succ + 1e-6:
best_succ = cur_succ
agent.save(str(best_path))
# 콘솔 로그
now = time.time()
if now - last_print >= 5.0 or episode % args.print_every == 0:
last_print = now
elapsed = now - t0
sps = total_lifetime_steps / max(elapsed, 1e-6)
print(
f"[ep {episode:5d}] "
f"R={float(np.mean(rewards_50)):7.1f} "
f"succ50={float(np.mean(success_50)):.2f} "
f"best={best_succ:.2f} "
f"eps={eps_now:.3f} "
f"loss={agent.avg_loss():.4f} "
f"steps={agent.total_env_steps} "
f"({sps:.0f}/s)",
flush=True,
)
# 주기적 저장
if episode % args.save_every == 0:
agent.save(str(ckpt_path))
# 결과 배너 (일시 표시). speed 가 크면 짧게.
color = (120, 240, 140) if result == "landed" else (240, 120, 120)
show_terminal_banner(renderer.screen, renderer.font_big,
f"{(result or 'END').upper()}", color)
pygame.display.flip()
banner_ms = max(50, args.banner_ms // max(1, speed // 2))
pygame.time.delay(banner_ms)
except KeyboardInterrupt:
print("\n[interrupt] saving and exiting...")
finally:
try:
agent.save(str(ckpt_path))
print(f"[save] {ckpt_path} (final, episode={episode})")
print(f"[best] {best_path} succ50_best={best_succ:.2f}")
except Exception as e:
print(f"[warn] save failed: {e}")
renderer.close()
# =====================================================================
# 모드 3: 사람 조작 (환경 점검)
# =====================================================================
def run_human(args: argparse.Namespace) -> None:
from render import LanderRenderer, show_terminal_banner
env = MoonLanderEnv(seed=args.seed, randomize=not args.fixed_start)
renderer = LanderRenderer(env, caption="Moon Lander — human")
obs = env.reset()
ep_r = 0.0
running = True
while running:
for event in pygame.event.get():
if event.type == pygame.QUIT:
running = False
elif event.type == pygame.KEYDOWN:
if event.key == pygame.K_ESCAPE:
running = False
elif event.key == pygame.K_r:
obs = env.reset()
ep_r = 0.0
keys = pygame.key.get_pressed()
bits = 0
if keys[pygame.K_SPACE]:
bits |= 1 << 0
if keys[pygame.K_UP] or keys[pygame.K_w]:
bits |= 1 << 1
if keys[pygame.K_LEFT] or keys[pygame.K_a]:
bits |= 1 << 2
if keys[pygame.K_RIGHT] or keys[pygame.K_d]:
bits |= 1 << 3
obs, r, done, info = env.step(bits)
ep_r += r
renderer.draw({"reward": ep_r, "result": info.result or "..."})
renderer.tick(args.fps)
if done:
color = (120, 240, 140) if info.result == "landed" else (240, 120, 120)
show_terminal_banner(renderer.screen, renderer.font_big,
f"{info.result.upper()} R={ep_r:.1f}", color)
pygame.display.flip()
pygame.time.delay(1200)
obs = env.reset()
ep_r = 0.0
renderer.close()
# =====================================================================
# CLI
# =====================================================================
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser()
# 공통
p.add_argument("--ckpt", type=str, default=None,
help="체크포인트 경로. 추론 모드 기본값 checkpoints/best.pt, "
"--learn 모드 기본값 checkpoints/live.pt")
p.add_argument("--episodes", type=int, default=5)
p.add_argument("--seed", type=int, default=99000)
p.add_argument("--fps", type=int, default=60)
p.add_argument("--no-render", action="store_true",
help="GUI 없이 헤드리스 평가 (--learn 에는 적용 안 됨)")
p.add_argument("--fixed-start", action="store_true",
help="초기 위치/속도 고정")
# 모드 선택
p.add_argument("--learn", action="store_true",
help="연속 학습 모드 (무한 실행, 매 step 학습 + 자동 저장)")
p.add_argument("--human", action="store_true",
help="DQN 대신 키보드 조작")
# --learn 전용
p.add_argument("--init-from", type=str, default=None,
help="--learn 시 가중치 초기화에 사용할 다른 ckpt 경로")
p.add_argument("--fresh", action="store_true",
help="--learn 시 기존 ckpt 무시하고 새로 시작")
p.add_argument("--speed", type=int, default=1,
help="--learn 화면 한 프레임당 실행할 env step 수 (>=1)")
p.add_argument("--save-every", type=int, default=20,
help="--learn 체크포인트 저장 주기 (에피소드 단위)")
p.add_argument("--print-every", type=int, default=10,
help="--learn 콘솔 로그 주기 (에피소드 단위)")
p.add_argument("--banner-ms", type=int, default=400,
help="--learn 에피소드 결과 배너 표시 시간(ms)")
p.add_argument("--eps-min", type=float, default=0.05,
help="--learn 행동 ε 의 하한 (탐험 유지)")
# 학습 하이퍼파라미터 (대부분 train.py 기본값과 동일)
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("--eps-decay-steps", type=int, default=80_000)
p.add_argument("--eps-end", type=float, default=0.05)
args = p.parse_args()
# ckpt 기본값 모드별로 설정
if args.ckpt is None:
args.ckpt = "checkpoints/live.pt" if args.learn else "checkpoints/best.pt"
return args
if __name__ == "__main__":
args = parse_args()
if args.learn:
run_live_train(args)
elif args.human:
run_human(args)
else:
run_dqn(args)