[PyTorch로 시작하는 강화학습 입문] 7편: 정책기반 접근 살펴보기 – REINFORCE 알고리즘 구현

 

지금까지 다룬 DQN 계열 알고리즘은 Q값(Q(s,a))을 근사하고, 이를 바탕으로 최적 행동을 선택하는 가치기반(Value-based) 방식이었습니다. 반면, 정책기반(Policy-based) 방법은 Q함수를 명시적으로 다루지 않고, 정책(π(a|s))을 직접 파라미터화(파라미터 θ)하고 이를 최적화하는 접근을 사용합니다.

정책기반 접근의 장점:

  • 연속적이고 큰 행동 공간 처리 용이: Q테이블이나 Q함수를 모든 행동에 대해 근사하는 것이 어려운 상황에서 정책을 직접 근사하면 편리합니다.
  • 확률적 정책: 정책이 확률적으로 행동을 샘플링하기 때문에 탐색을 내장하고 있습니다.
  • 정책 개선의 직관성: 목표는 "정책의 기대 return을 최대화"하는 것이며, 이를 직접 최적화 가능합니다.

이번 글에서는 가장 기초적인 정책기반 알고리즘인 REINFORCE를 예제로 구현해 보겠습니다.

REINFORCE 알고리즘 개요:

  • 에피소드 단위로 환경을 플레이하고, 한 에피소드의 모든 (s,a,r) 데이터를 수집
  • 에피소드 종료 후, 각 행동에 대해 "Return(G)"을 구하고, 이 Return을 가중치로 정책 확률을 높이거나 낮추는 방향으로 파라미터 업데이트
  • 수식적으로:
    Update θ ← θ + α * Σ_t (∇_θ log π_θ(a_t|s_t) * G_t)
    여기서 G_t는 시점 t 이후의 누적보상

REINFORCE는 단순하나, 매 에피소드가 끝나야 업데이트할 수 있고, 고분산 업데이트 문제 등 단점도 있지만, 정책기반 접근의 기본 철학을 배우는 데 훌륭한 출발점입니다.

참고자료:

  • Sutton & Barto, "Reinforcement Learning: An Introduction" (REINFORCE 소개)

REINFORCE 알고리즘 구현 아이디어

  1. 에이전트는 파라미터 θ를 가진 정책 신경망 π_θ(a|s)를 갖습니다. 여기서 정책은 주로 소프트맥스 출력으로 행동의 확률분포를 나타냅니다.
  2. 한 에피소드를 실행하며 (s_t, a_t, r_t)를 모두 저장합니다.
  3. 에피소드가 끝나면, 각 타임스텝 t에 대해 G_t(그 시점 이후의 누적보상) 계산
  4. 그 후 ∇_θ log π_θ(a_t|s_t)*G_t를 모든 t에 대해 합산한 값을 θ에 적용해 파라미터 업데이트
  5. 반복하며 정책이 좋은 행동의 확률을 점차 높여나갑니다.

코드 예제 (REINFORCE 구현)

아래 코드는 CartPole 환경에서 REINFORCE를 구현한 예제입니다. 에피소드 단위로 데이터 수집 후 업데이트하는 구조를 보여줍니다.

import gym
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np

###################################
# 정책 신경망 정의
###################################
class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_size=64):
        super(PolicyNetwork, self).__init__()
        self.fc1 = nn.Linear(state_dim, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, action_dim)
    
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        logits = self.fc3(x)  # action_dim 출력
        # 소프트맥스로 행동 확률 계산
        return F.softmax(logits, dim=-1)

###################################
# REINFORCE 에이전트
###################################
class REINFORCEAgent:
    def __init__(self, state_dim, action_dim, gamma=0.99, lr=1e-3):
        self.gamma = gamma
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        self.policy = PolicyNetwork(state_dim, action_dim).to(self.device)
        self.optimizer = optim.Adam(self.policy.parameters(), lr=lr)
        
        # 한 에피소드 동안 저장할 (log_prob, reward)
        self.log_probs = []
        self.rewards = []

    def select_action(self, state):
        # 상태 넣어 행동 확률 얻기
        state_t = torch.FloatTensor(state).unsqueeze(0).to(self.device)
        action_probs = self.policy(state_t)
        # 확률로 행동 샘플
        action_dist = torch.distributions.Categorical(action_probs)
        action = action_dist.sample()
        self.log_probs.append(action_dist.log_prob(action))
        return action.item()

    def store_reward(self, reward):
        # 스텝마다 받은 보상 저장
        self.rewards.append(reward)

    def update(self):
        # 에피소드 끝난 뒤 호출
        # G_t 계산
        G = 0
        returns = []
        for r in reversed(self.rewards):
            G = r + self.gamma * G
            returns.insert(0, G)
        returns = torch.FloatTensor(returns).to(self.device)
        # 정규화(Optional): 학습 안정화를 위해 리턴을 정규화 가능
        returns = (returns - returns.mean()) / (returns.std() + 1e-9)
        
        # 정책 기울기 업데이트
        loss = 0
        for log_p, Gt in zip(self.log_probs, returns):
            loss += -log_p * Gt
        
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        
        # 다음 에피소드 위해 메모리 초기화
        self.log_probs = []
        self.rewards = []

def train_reinforce(env_name="CartPole-v1", max_episodes=300):
    env = gym.make(env_name)
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n
    
    agent = REINFORCEAgent(state_dim, action_dim)
    
    reward_history = []
    for ep in range(max_episodes):
        state = env.reset()
        total_reward = 0
        done = False
        while not done:
            action = agent.select_action(state)
            next_state, reward, done, info = env.step(action)
            agent.store_reward(reward)
            state = next_state
            total_reward += reward
        
        agent.update()
        reward_history.append(total_reward)
        
        if (ep+1) % 20 == 0:
            avg_reward = np.mean(reward_history[-20:])
            print(f"Episode {ep+1}, Avg Reward(last 20): {avg_reward:.2f}")

    env.close()

if __name__ == "__main__":
    train_reinforce()

코드 해설

  • PolicyNetwork: 상태를 입력받아 행동 확률분포를 출력하는 신경망. 소프트맥스 통해 확률 얻음.
  • REINFORCEAgent:
    • select_action: 정책에 따라 행동 샘플, log_prob 저장 (나중에 정책 그래디언트 계산)
    • store_reward: 스텝마다 받은 보상을 저장
    • update: 에피소드 끝난 뒤 누적보상(G)을 계산하고, ∇_θ log π(a|s)*G 합산해 정책 파라미터 업데이트
  • REINFORCE는 에피소드 단위로 업데이트하므로, 초기 학습이 느리고 변동성 클 수 있음. 하지만 정책기반 접근의 기본 원리를 익히는 데 유용.

마무리

이번 글에서는 DQN 계열과 달리 정책을 직접 파라미터화하고 업데이트하는 정책기반 접근, 그리고 그 중 가장 기본적인 REINFORCE 알고리즘을 구현해보았습니다. 정책기반 방법은 연속형 행동공간, 고차원 문제 등에서 유용하며, 앞으로 배우게 될 Actor-Critic 방법론(PPO, A2C, A3C 등)의 기초가 됩니다.

다음 글에서는 Actor-Critic 알고리즘을 도입해 REINFORCE의 고분산 문제를 완화하고, 더 효율적으로 정책을 업데이트하는 방법을 배워볼 수 있습니다.

반응형