dqn_agent.py

import random
from collections import deque
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

class DuelingDQN(nn.Module):
    """
    Dueling DQN 구조:
    가치(Value) 함수와 이점(Advantage) 함수를 분리하여 더 빠르고 안정적인 Q-학습 제공
    Q(s, a) = V(s) + (A(s, a) - mean(A(s)))
    """
    def __init__(self, state_dim=8, action_dim=4):
        super(DuelingDQN, self).__init__()
        
        # 공통 특징 추출층
        self.feature_network = nn.Sequential(
            nn.Linear(state_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 128),
            nn.ReLU()
        )
        
        # 상태 가치(State Value) 스트림 V(s)
        self.value_stream = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 1)
        )
        
        # 행동 이점(Action Advantage) 스트림 A(s, a)
        self.advantage_stream = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, action_dim)
        )

    def forward(self, state):
        features = self.feature_network(state)
        values = self.value_stream(features)
        advantages = self.advantage_stream(features)
        
        # Dueling 공식: Q = V + (A - mean(A))
        q_values = values + (advantages - advantages.mean(dim=-1, keepdim=True))
        return q_values


class ReplayBuffer:
    """경험 리플레이 버퍼"""
    def __init__(self, capacity=100000):
        self.buffer = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)
        return (
            np.array(states, dtype=np.float32),
            np.array(actions, dtype=np.int64),
            np.array(rewards, dtype=np.float32),
            np.array(next_states, dtype=np.float32),
            np.array(dones, dtype=np.float32),
        )

    def __len__(self):
        return len(self.buffer)


class DQNAgent:
    """
    LunarLander 착륙용 Dueling Double-DQN Agent
    - Double DQN: 오버에스티메이션 방지
    - Epsilon Decay: 100% -> 5% 점진적 감쇄
    - Soft Target Update (Polyak averaging)
    """
    def __init__(
        self,
        state_dim=8,
        action_dim=4,
        lr=5e-4,
        gamma=0.99,
        tau=0.005,
        buffer_size=100000,
        batch_size=64,
        epsilon_start=1.0,
        epsilon_end=0.05,
        total_episodes=1000
    ):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.gamma = gamma
        self.tau = tau
        self.batch_size = batch_size
        
        # 탐색률(Epsilon) 파라미터 (1.0 -> 0.05)
        self.epsilon_start = epsilon_start
        self.epsilon_end = epsilon_end
        self.total_episodes = total_episodes
        self.epsilon = epsilon_start
        
        # 1000 에피소드 중 약 75% 지점(750 ep)에서 epsilon_end(0.05)에 도달하도록 감쇄율 설정
        self.epsilon_decay = (epsilon_end / epsilon_start) ** (1.0 / (total_episodes * 0.75))

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # 정책망 & 타깃망 생성
        self.policy_net = DuelingDQN(state_dim, action_dim).to(self.device)
        self.target_net = DuelingDQN(state_dim, action_dim).to(self.device)
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()

        self.optimizer = optim.AdamW(self.policy_net.parameters(), lr=lr, weight_decay=1e-4)
        self.criterion = nn.SmoothL1Loss()  # Huber Loss (노이즈에 강건)
        self.memory = ReplayBuffer(buffer_size)

    def select_action(self, state, evaluate=False):
        """
        행동 선택 및 각 행동별 Q-Value 반환 (웹 대시보드 시각화용)
        evaluate=True 일 경우 순수 Greedy 행동 (간지 착륙 시연)
        """
        state_t = torch.FloatTensor(state).unsqueeze(0).to(self.device)
        
        with torch.no_grad():
            q_values = self.policy_net(state_t).cpu().numpy()[0]

        if not evaluate and random.random() < self.epsilon:
            action = random.randrange(self.action_dim)
        else:
            action = int(np.argmax(q_values))

        return action, q_values.tolist()

    def update(self):
        """Double DQN 기반 신경망 가중치 1스텝 업데이트"""
        if len(self.memory) < self.batch_size:
            return None

        states, actions, rewards, next_states, dones = self.memory.sample(self.batch_size)

        states_t = torch.FloatTensor(states).to(self.device)
        actions_t = torch.LongTensor(actions).unsqueeze(1).to(self.device)
        rewards_t = torch.FloatTensor(rewards).unsqueeze(1).to(self.device)
        next_states_t = torch.FloatTensor(next_states).to(self.device)
        dones_t = torch.FloatTensor(dones).unsqueeze(1).to(self.device)

        # 현재 상태의 Q-값 계산: Q(s, a)
        curr_q = self.policy_net(states_t).gather(1, actions_t)

        # Double DQN: Policy Net으로 최적 행동 선택 -> Target Net으로 해당 행동의 Q-값 평가
        with torch.no_grad():
            best_actions = self.policy_net(next_states_t).argmax(1, keepdim=True)
            next_q = self.target_net(next_states_t).gather(1, best_actions)
            target_q = rewards_t + (1.0 - dones_t) * self.gamma * next_q

        # Loss 계산 및 역전파
        loss = self.criterion(curr_q, target_q)

        self.optimizer.zero_grad()
        loss.backward()
        # 안정적인 학습을 위한 Gradient Clipping
        nn.utils.clip_grad_norm_(self.policy_net.parameters(), max_norm=10.0)
        self.optimizer.step()

        # 타깃망 Soft Update
        self.soft_update()

        return loss.item()

    def soft_update(self):
        """Polyak Averaging 타깃망 소프트 업데이트: θ_target = τ*θ_local + (1 - τ)*θ_target"""
        for target_param, policy_param in zip(self.target_net.parameters(), self.policy_net.parameters()):
            target_param.data.copy_(self.tau * policy_param.data + (1.0 - self.tau) * target_param.data)

    def decay_epsilon(self):
        """에피소드 종료 시 Epsilon 감쇄"""
        self.epsilon = max(self.epsilon_end, self.epsilon * self.epsilon_decay)

    def save(self, filepath="best_lunar_lander_dqn.pth"):
        torch.save({
            'policy_net': self.policy_net.state_dict(),
            'target_net': self.target_net.state_dict(),
            'optimizer': self.optimizer.state_dict(),
            'epsilon': self.epsilon
        }, filepath)

    def load(self, filepath="best_lunar_lander_dqn.pth"):
        checkpoint = torch.load(filepath, map_location=self.device)
        self.policy_net.load_state_dict(checkpoint['policy_net'])
        self.target_net.load_state_dict(checkpoint['target_net'])
        if 'optimizer' in checkpoint:
            self.optimizer.load_state_dict(checkpoint['optimizer'])
        if 'epsilon' in checkpoint:
            self.epsilon = checkpoint['epsilon']