Advantage Actor-Critic (A2C)

[Actor-Critic 핵심 개념]

정책 기반과 가치 기반 방법을 결합
- Actor (정책 결정): 정책 $\pi_\theta(a|s)$ 를 파라미터화, 상태에서 행동 확률 분포 출력
- Critic (가치 평가): Value Function $V_\psi(s)$ 를 파라미터화, 현재 상태의 가치(기대 반환) 추정
장점
- 실시간 반응: 에피소드가 끝나지 않아도 업데이트 가능
  - 에피소드 단위로 Return 계산하지 않음
  - 환경 변화에 즉각적 대응 가능 > REINFORCE의 고분산 문제 완화

[Advantage Actor-Critic (A2C) 알고리즘]

Advantage 함수 사용
- $A(s, a) = Q(s, a) - V(s)$ : 특정 상태에서 행동을 선택했을 때 얻는 추가적인 이득
- Critic은 $V(s)$ 를 추정, Actor는 $A(s, a)$ 를 최대화하는 방향으로 학습
업데이트 방식
- Actor: $\theta \leftarrow \theta + \alpha \nabla_\theta \log \pi_\theta(a|s) A(s, a)$
  - Policy Gradient를 사용하여 정책 업데이트
  - $\alpha$ : 학습률
- Critic: $\psi \leftarrow \psi + \beta (r + \gamma V(s') - V(s)) \nabla_\psi V(s)$
  - TD 오차(Temporal-Difference error)를 최소화하는 방향으로 Value Function 업데이트
  - $\beta$ : 학습률
  - $r$ : 보상
  - $\gamma$ : 할인율
  - $s'$ : 다음 상태
A2C의 효능
- 안정적인 학습: Critic이 가치를 평가, Actor는 Critic의 평가를 기반으로 정책을 업데이트하므로, REINFORCE에 비해 안정적인 학습 가능
- 효율적인 탐색: Advantage 함수를 사용하여, 어떤 행동이 평균보다 더 나은 결과를 가져올 수 있는지 파악, 효율적인 탐색 가능

[A2C 구현]

Actor 네트워크와 Critic 네트워크 설계
- Actor 네트워크: $s \rightarrow \boxed{\text{FC} + \text{ReLU}} \rightarrow \boxed{\text{FC} + \text{Softmax}} \rightarrow \pi(a|s)$
  - 상태를 입력받아 행동에 대한 확률 분포 출력
  - REINFORCE의 정책 네트워크 구조와 동일
- Critic 네크워크: $s \rightarrow \boxed{\text{FC} + \text{ReLU}} \rightarrow \boxed{\text{FC}} \rightarrow V(s)$
  - 상태를 입력받아 해당 상태의 가치(Value)를 출력
  - DQN과 같은 MLP를 사용할 수 있지만, Critic의 출력은 상태별 Scalar $V(s)$ 이며 DQN의 행동별 Q-value 벡터와 다름
환경과의 상호작용을 통해 데이터 수집
수집된 데이터를 사용하여 Actor와 Critic 네트워크 업데이트

[A2C 한계 및 개선 방향]

On-policy 학습: 현재 정책에 따라서만 학습 데이터를 수집하므로, 샘플 효율성이 낮음
개선 방향
- Proximal Policy Optimization (PPO) (Week 9)
  - 정책 업데이트 시, 새로운 정책과 기존 정책의 차이를 제한하여 안정적인 학습 가능
- Soft Actor-Critic SAC (Week 10)
  - 엔트로피 개념을 도입하여 탐험을 장려하고, 연속적인 행동 공간에서도 효과적인 학습 가능

[실습: Advantage Actor-Critic(A2C)을 이용한 CartPole 실습]

1. 라이브러리 가져오기 및 하이퍼 파라미터 설정

필요한 라이브러리 (PyTorch, Gym, NumPy 등)를 가져옵니다.
A2C 알고리즘에 사용될 하이퍼파라미터들을 설정합니다.
- gamma: 할인율 (미래 보상의 중요도)
- learning_rate: 학습률
- hidden_size: 은닉층의 크기 (2층/3층 Neural Network 선택 가능)
- num_steps: 각 롤아웃(rollout)에서 수행할 단계 수
  - 롤아웃: Advantage 계산 단위 구간
- num_episodes: 총 학습 에피소드 수
- check_interval: 학습 진행 상황 확인 주기 (에피소드 단위)
- early_stop_threshold: 조기 종료 조건 (평균 보상 기준)
- early_stop_patience : 조기 종료 조건 유지 횟수
- render_gif_path: 렌더링 결과를 저장할 GIF 파일 경로

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical
import gymnasium as gym
import numpy as np
import imageio

# 하이퍼파라미터
gamma = 0.99
learning_rate = 0.003
hidden_size = 128  # 기본값 (2층 신경망)
num_steps = 10
num_episodes = 2000
progress_interval = 50
early_stop_threshold = 400 # CartPole-v1 대체로 200 이상이면 성공으로 분류
early_stop_patience = 10
render_gif_path = "cartpole_simulation_a2c.gif"

2. A2C Neural Network 모델 (Actor-Critic) 정의

__init__ 메서드에서 네트워크 구조를 정의
use_three_layers 인자를 통해 2층 또는 3층 Neural Network을 선택할 수 있음
fc1, fc2 (선택적 fc_h), fc_pi, fc_v 레이어를 정의
fc_pi는 정책(policy)을 나타내는 확률 분포를 출력
fc_v는 상태 가치(value)를 출력
forward 메서드에서 입력을 받아 정책과 상태 가치를 반환

class ActorCritic(nn.Module):
    def __init__(self, input_size, num_actions, hidden_size=128, use_three_layers=False):
        super(ActorCritic, self).__init__()
        self.use_three_layers = use_three_layers
        self.fc1 = nn.Linear(input_size, hidden_size)
        if self.use_three_layers:
            self.fc_h = nn.Linear(hidden_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc_pi = nn.Linear(hidden_size, num_actions)
        self.fc_v = nn.Linear(hidden_size, 1)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        if self.use_three_layers:
            x = F.relu(self.fc_h(x))
        x = F.relu(self.fc2(x))
        policy = F.softmax(self.fc_pi(x), dim=0)
        value = self.fc_v(x)
        return policy, value

3.학습 루프 함수 정의

정책/가치 예측 및 행동 선택

state_tensor = torch.from_numpy(state).float()
- 현재 상태(NumPy 배열)를 PyTorch 텐서로 변환
policy, value = model(state_tensor)
- 모델에 현재 상태를 입력해 정책과 상태 가치 얻기
action_dist = Categorical(policy)
- 정책(확률 분포)을 이용하여 행동을 샘플링하기 위한 Categorical 객체를 생성
action = action_dist.sample()
- Categorical 객체로부터 확률적으로 행동을 선택

선택한 행동을 적용해 환경과 상호작용

next_state, reward, done, _ = env.step(action.item())
- 선택한 행동을 환경에 적용하여 다음 상태, 보상, 종료 여부, 추가 정보 얻기
- action.item()은 텐서 형태의 행동을 정수 값으로 변환
total_reward += reward: 에피소드 동안의 누적 보상을 계산

경험 저장

memory.append((state_tensor, action, reward, value, done))
- 현재 상태, 선택한 행동, 받은 보상, 상태 가치, 종료 여부를 memory 리스트에 튜플 형태로 저장
- 이 데이터는 나중에 Advantage와 손실을 계산하는 데 사용

롤아웃 데이터로부터 Advantage 및 손실 계산

loss = 0: 손실을 초기화
_, last_value = model(torch.from_numpy(next_state).float())
- 롤아웃의 마지막 상태(다음 상태)에 대한 Value Function 값을 계산
- 이 값은 Returns를 계산하는 데 사용
returns = last_value.detach()
- 마지막 상태의 Value Function 값을 Returns의 초기값으로 설정
- .detach()는 이 값이 기울기 계산에 포함되지 않도록 함
for state_tensor, action, reward, value, done in reversed(memory)
- memory 리스트에 저장된 경험들을 역순으로 순회합니다
  - 롤아웃의 끝에서부터 시작하여 처음으로 거슬러 올라감
- returns = reward + gamma * returns * (1 - done)
  - Returns는 현재 보상과 할인된 미래 보상의 합
  - 에피소드가 끝났으면(done=1), 현재 보상만 Returns가 됨
- advantage = returns - value
  - Advantage는 Returns와 현재 상태의 Value Function 값의 차이
  - Advantage는 특정 행동이 평균적인 행동보다 얼마나 더 좋거나 나쁜지를 나타냄
- loss += ~: Loss Function을 계산
  - -Categorical(model(state_tensor)[0]).log_prob(action) * advantage.detach()
    - Actor (정책)의 손실
    - 정책 그레이디언트(Policy Gradient) 방법을 사용
    - Advantage가 높을수록 해당 행동을 선택할 확률을 높이고, Advantage가 낮을수록 확률을 낮춤
    - .detach()는 Advantage가 기울기 계산에 포함되지 않도록 함
  - F.smooth_l1_loss(value, returns.detach())
    - Critic (Value Function)의 손실
    - Huber Loss를 사용하여 Value Function 예측값(value)과 실제 Returns 간의 차이를 줄임
    - .detach()는 Returns가 기울기 계산에 포함되지 않도록 함

최적화:

optimizer.zero_grad()
- 옵티마이저에 저장된 이전 기울기 값을 0으로 초기화
loss.backward()
- Loss Function에 대한 Backpropagation를 수행하여 모델의 각 파라미터에 대한 기울기를 계산
optimizer.step()
- 계산된 기울기를 사용하여 모델의 파라미터를 업데이트 (경사 하강법 적용)
memory = []
- 다음 롤아웃을 위해 메모리 비우기

def train(model, optimizer, env, num_episodes, num_steps, gamma, check_interval, early_stop_threshold, early_stop_patience):
    scores = []
    for episode in range(num_episodes):
        state, _ = env.reset()
        done = False
        total_reward = 0
        memory = []

        while not done:
            for t in range(num_steps):
                state_tensor = torch.from_numpy(state).float()
                policy, value = model(state_tensor)
                action_dist = Categorical(policy)
                action = action_dist.sample()

                next_state, reward, terminated, truncated, _ = env.step(action.item())
                total_reward += reward
                done = terminated or truncated

                memory.append((state_tensor, action, reward, value, done))
                state = next_state

                if done:
                    break

            # 롤아웃 데이터로부터 Advantage 계산 및 손실 함수 계산
            loss = 0
            _, last_value = model(torch.from_numpy(next_state).float())
            returns = last_value.detach()

            for state_tensor, action, reward, value, done in reversed(memory):
                returns = reward + gamma * returns * (1 - done)
                advantage = returns - value
                loss += -Categorical(model(state_tensor)[0]).log_prob(action) * advantage.detach() + \
                        F.smooth_l1_loss(value, returns.detach()) # Huber loss

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            memory = [] # 메모리 초기화

        scores.append(total_reward)

        # 진행 상황 확인 및 조기 종료 검사
        if (episode+1) % check_interval == 0:
            avg_score = np.mean(scores[-check_interval:])
            print(f"Episode: {episode+1}, Average Reward: {avg_score:.2f}")

        if np.min(scores[-early_stop_patience:]) >= early_stop_threshold:
            print(f"Early stopping triggered! at Episode: {episode+1}")
            break

4. 학습 실행 및 결과 확인

Gym 환경을 생성합니다. (CartPole-v1)
ActorCritic 모델 인스턴스를 생성 (2층 또는 3층 선택)
optim.Adam 옵티마이저를 생성
train 함수를 호출하여 학습을 시작

# 환경 생성
env = gym.make('CartPole-v1', render_mode = 'rgb_array')
input_size = env.observation_space.shape[0]
num_actions = env.action_space.n

# 모델 생성 (3층 신경망 사용 예시)
model = ActorCritic(input_size, num_actions, hidden_size, use_three_layers=True)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 학습 실행
train(model, optimizer, env, num_episodes, num_steps, gamma, progress_interval, early_stop_threshold, early_stop_patience)

env.close()

Episode: 50, Average Reward: 16.78
Episode: 100, Average Reward: 24.02
Episode: 150, Average Reward: 18.66
Episode: 200, Average Reward: 37.02
Episode: 250, Average Reward: 39.84
Episode: 300, Average Reward: 31.26
Episode: 350, Average Reward: 55.00
Episode: 400, Average Reward: 104.56
Early stopping triggered! at Episode: 422

5. 학습된 모델 테스트 및 GIF 렌더링

학습이 완료된 모델을 사용하여 한 번의 에피소드를 실행
각 스텝의 화면을 수집
imageio 라이브러리를 사용하여 GIF를 생성
에피소드 동안 얻은 총 보상을 출력

state, _ = env.reset()
done = False
total_reward = 0
frames = []

with torch.no_grad():  # 기울기 계산 비활성화 (추론 모드)
    while not done:
        state_tensor = torch.from_numpy(state).float()
        policy, _ = model(state_tensor)
        action_dist = Categorical(policy)
        action = action_dist.sample()

        next_state, reward, terminated, truncated, _ = env.step(action.item())
        total_reward += reward
        state = next_state
        done = terminated or truncated

        # 렌더링 이미지를 프레임 리스트에 추가
        frames.append(env.render())

# GIF 파일로 저장
imageio.mimsave(render_gif_path, frames, duration=33)  # fps 조절 가능 duration = 1000/fps
print(f"Test Episode Reward: {total_reward}")
env.close()

Test Episode Reward: 500.0