예제 소스
문랜딩 프로그램으로 이해하는 강화학습 · 파일 6개
dqn_agent.py
"""DQN Agent (Dueling Double DQN, PyTorch).
핵심 구성요소
- QNetwork : Dueling 구조 (V(s) + A(s,a) - mean(A))
- ReplayBuffer : numpy 기반 원형 버퍼
- DQNAgent : ε-greedy 행동 선택, target network soft update,
학습 1 step 수행 (Double DQN 타깃 사용)
"""
from __future__ import annotations
import random
from collections import deque
from dataclasses import dataclass
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
# =====================================================================
# 모델
# =====================================================================
class QNetwork(nn.Module):
"""Dueling DQN 네트워크."""
def __init__(self, state_dim: int, num_actions: int, hidden: int = 128):
super().__init__()
self.feature = nn.Sequential(
nn.Linear(state_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
)
self.value_head = nn.Sequential(
nn.Linear(hidden, hidden // 2),
nn.ReLU(),
nn.Linear(hidden // 2, 1),
)
self.adv_head = nn.Sequential(
nn.Linear(hidden, hidden // 2),
nn.ReLU(),
nn.Linear(hidden // 2, num_actions),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
h = self.feature(x)
v = self.value_head(h)
a = self.adv_head(h)
# Dueling: Q = V + (A - mean(A))
return v + a - a.mean(dim=1, keepdim=True)
# =====================================================================
# Replay buffer
# =====================================================================
class ReplayBuffer:
def __init__(self, capacity: int, state_dim: int):
self.capacity = capacity
self.size = 0
self.idx = 0
self.states = np.zeros((capacity, state_dim), dtype=np.float32)
self.actions = np.zeros((capacity,), dtype=np.int64)
self.rewards = np.zeros((capacity,), dtype=np.float32)
self.next_states = np.zeros((capacity, state_dim), dtype=np.float32)
self.dones = np.zeros((capacity,), dtype=np.float32)
def push(self, s, a, r, ns, d):
i = self.idx
self.states[i] = s
self.actions[i] = a
self.rewards[i] = r
self.next_states[i] = ns
self.dones[i] = float(d)
self.idx = (self.idx + 1) % self.capacity
self.size = min(self.size + 1, self.capacity)
def sample(self, batch: int):
idx = np.random.randint(0, self.size, size=batch)
return (
self.states[idx],
self.actions[idx],
self.rewards[idx],
self.next_states[idx],
self.dones[idx],
)
def __len__(self) -> int:
return self.size
# =====================================================================
# Agent
# =====================================================================
@dataclass
class DQNConfig:
state_dim: int = 7
num_actions: int = 16
hidden: int = 128
lr: float = 5e-4
gamma: float = 0.99
buffer_size: int = 100_000
batch_size: int = 64
learn_start: int = 1_000 # 학습 시작 전 채울 transition 수
target_tau: float = 0.005 # soft update 비율
eps_start: float = 1.0
eps_end: float = 0.05
eps_decay_steps: int = 80_000 # 환경 step 기준
grad_clip: float = 5.0
class DQNAgent:
def __init__(self, cfg: DQNConfig | None = None, device: str | None = None):
self.cfg = cfg or DQNConfig()
self.device = torch.device(
device or ("cuda" if torch.cuda.is_available() else
"mps" if torch.backends.mps.is_available() else
"cpu")
)
self.q = QNetwork(self.cfg.state_dim, self.cfg.num_actions, self.cfg.hidden).to(self.device)
self.q_target = QNetwork(self.cfg.state_dim, self.cfg.num_actions, self.cfg.hidden).to(self.device)
self.q_target.load_state_dict(self.q.state_dict())
for p in self.q_target.parameters():
p.requires_grad_(False)
self.optimizer = torch.optim.Adam(self.q.parameters(), lr=self.cfg.lr)
self.buffer = ReplayBuffer(self.cfg.buffer_size, self.cfg.state_dim)
self.total_env_steps = 0
self.recent_losses: deque[float] = deque(maxlen=100)
# ---------------- ε ----------------
def epsilon(self) -> float:
c = self.cfg
frac = min(1.0, self.total_env_steps / c.eps_decay_steps)
return c.eps_start + frac * (c.eps_end - c.eps_start)
# ---------------- 행동 선택 ----------------
@torch.no_grad()
def act(self, state: np.ndarray, eps: float | None = None) -> int:
if eps is None:
eps = self.epsilon()
if random.random() < eps:
return random.randint(0, self.cfg.num_actions - 1)
s = torch.from_numpy(state).float().unsqueeze(0).to(self.device)
q = self.q(s)
return int(q.argmax(dim=1).item())
@torch.no_grad()
def act_greedy(self, state: np.ndarray) -> int:
return self.act(state, eps=0.0)
# ---------------- 저장 ----------------
def remember(self, s, a, r, ns, d) -> None:
self.buffer.push(s, a, r, ns, d)
self.total_env_steps += 1
# ---------------- 학습 ----------------
def learn(self) -> float | None:
c = self.cfg
if len(self.buffer) < max(c.learn_start, c.batch_size):
return None
s, a, r, ns, d = self.buffer.sample(c.batch_size)
s_t = torch.from_numpy(s).to(self.device)
a_t = torch.from_numpy(a).to(self.device)
r_t = torch.from_numpy(r).to(self.device)
ns_t = torch.from_numpy(ns).to(self.device)
d_t = torch.from_numpy(d).to(self.device)
# 현재 Q
q_pred = self.q(s_t).gather(1, a_t.unsqueeze(1)).squeeze(1)
# Double DQN 타깃
with torch.no_grad():
next_actions = self.q(ns_t).argmax(dim=1, keepdim=True)
next_q = self.q_target(ns_t).gather(1, next_actions).squeeze(1)
target = r_t + c.gamma * (1.0 - d_t) * next_q
loss = F.smooth_l1_loss(q_pred, target)
self.optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(self.q.parameters(), c.grad_clip)
self.optimizer.step()
# Target soft update
with torch.no_grad():
for p, p_t in zip(self.q.parameters(), self.q_target.parameters()):
p_t.data.mul_(1.0 - c.target_tau).add_(c.target_tau * p.data)
self.recent_losses.append(loss.item())
return loss.item()
def avg_loss(self) -> float:
if not self.recent_losses:
return 0.0
return float(np.mean(self.recent_losses))
# ---------------- 체크포인트 ----------------
def save(self, path: str) -> None:
torch.save({
"q": self.q.state_dict(),
"q_target": self.q_target.state_dict(),
"optimizer": self.optimizer.state_dict(),
"cfg": self.cfg.__dict__,
"total_env_steps": self.total_env_steps,
}, path)
def load(self, path: str, map_location: str | None = None) -> None:
ckpt = torch.load(path, map_location=map_location or self.device)
self.q.load_state_dict(ckpt["q"])
if "q_target" in ckpt:
self.q_target.load_state_dict(ckpt["q_target"])
if "optimizer" in ckpt:
try:
self.optimizer.load_state_dict(ckpt["optimizer"])
except Exception:
pass
self.total_env_steps = ckpt.get("total_env_steps", 0)