Advantage Actor Critic (A2C)

Actor-Critic

• 행동을 결정하는 Policy Net을 Actor라 칭하며
• 그 행위를 평가하는 Value Net을 Critic이라 칭함

두 신경망 모두 advantage \(A_t = V_{target} - V_{pred}\)를 가지고 업데이트를 하면서 최적의 행동을 고르는 actor와 정확하게 가치를 추정하는 critic을 학습하는 방법을 actor-critic이라고 한다. 여기서 \(V_{target}\)은 한 스텝만 가보고 얻은 target이고 \(V_{pred}\)는 critic의 예측값이다. 정확한 식은 아래에서 푼다.

 

근데 여기까지만 따지면 REINFORCE에 baseline으로 업데이트 하는 것도 actor-critic이라고 부를수 있지 않은가?

 

actor-critic이라고 부를려면 monte carlo처럼 에피소드를 끝까지 가보고 난 다음 return을 업데이트하면 인정 안해준다고 한다.(이유는 나도 모른다) 무튼 그래서 TD처럼 1 step만 가보고 return을 업데이트 한다면 그것은 actor-critic이라고 부를 수 있다.

Advantage 업데이트와 A2C

일단 TD를 다시 떠올려본다. TD에서 정확히 1 step에서 \(V(s_t)\)를 업데이트하는 식은 아래와 같았다.

\[ V(S_t) \leftarrow V(S_t) + \alpha \big[ R_{t+1} + \gamma V(S_{t+1}) - V(S_t) \big] \]

저기서 \(R_{t+1} + \gamma V(S_{t+1})\)은 return값(정확힌 임시 정답) 이였고 \(V(S_t)\)가 추정값이였다. 이제 이걸 advantage에 대입하는데 value function만 value network로 치환해서 대입하면 아래와 같다.

\[ A_t = \big( r + \gamma V_\phi(s') \big) - V_\phi(s_t) \]

이 \(r + \gamma V_\phi(s') - V_\phi(s_t)\)가 바로 TD error다. 실제로 받은 결과가 critic의 예측보다 얼마나 좋았는지를 나타낸다.

actor는 이 advantage 방향으로 policy를 업데이트한다.

\[ \nabla_\theta J(\theta) = \nabla_\theta \log \pi_\theta(a_t \mid s_t) \big( r + \gamma V_\phi(s') - V_\phi(s_t) \big) \]

\[ \theta \leftarrow \theta + \alpha \nabla_\theta J(\theta) \]

critic은 자기 예측 \(V_\phi(s_t)\)가 target \(r + \gamma V_\phi(s')\)에 가까워지도록 MSE로 업데이트한다.

\[ J(\phi) = \big( r + \gamma V_\phi(s') - V_\phi(s_t) \big)^2 \]

\[ \phi \leftarrow \phi - \alpha \nabla_\phi J(\phi) \]

actor는 ascent로 올리고 critic은 descent로 내린다. 매 스텝 이 둘을 동시에 한다.
이 TD기반의 actor-critic방식을 A2C라고 부른다.(Monte carlo기반은 REINFORCE)

Advantage A2C

A2C라는 이름에 advantage가 들어간 이유는 policy gradient를 advantage로 계산하기 때문이다. advantage는 원래 \(A(s,a) = Q(s,a) - V(s)\)로 정의된다. \(Q(s,a)\)는 state \(s\)에서 action \(a\)를 했을 때의 기대 return이고 \(V(s)\)는 그 state의 평균 기대 return이라 이 행동이 평균보다 얼마나 좋은지를 나타낸다.

여기서 \(Q(s,a)\)를 직접 구하긴 어려우니 \(Q(s,a) \approx r + \gamma V(s')\)로 근사하면

\[ A(s,a) = Q(s,a) - V(s) \approx r + \gamma V(s') - V(s) \]

위에서 본 식과 똑같다. 결국 A2C는 advantage로 "평균보다 좋은 행동은 강화하고 나쁜 행동은 억제"하는 actor-critic이다.

Actor-Critic 알고리즘

정리하면 전체 루프는 다음과 같다.

1. actor network 파라미터 \(\theta\)와 critic network 파라미터 \(\phi\)를 초기화한다.
2. 에피소드를 시작하고 매 스텝 \(t = 0, 1, \ldots, T-1\) 마다 3~6번을 반복한다.
3. 현재 policy로 action을 고른다 \(a_t \sim \pi_\theta(s_t)\). action을 실행해 reward \(r\)과 다음 state \(s'\)를 얻는다.
4. advantage를 계산한다 \(A = r + \gamma V_\phi(s') - V_\phi(s_t)\).
5. actor를 gradient ascent로 업데이트한다.
   \[ \theta \leftarrow \theta + \alpha \nabla_\theta \log \pi_\theta(a_t \mid s_t)\, A \]
6. critic을 gradient descent로 업데이트한다. loss는 \(A^2\) 즉 \(\big( r + \gamma V_\phi(s') - V_\phi(s_t) \big)^2\) 다.

매 스텝 바로바로 업데이트하니 에피소드가 끝나길 기다리는 Monte Carlo보다 빠르고 critic이 분산도 줄여준다.

A2C와 REINFORCE(baseline) code로 차이 분석

앞에서 말한 "네트워크는 같고 train logic만 다르다"를 코드로 직접 확인해보자. 먼저 두 방법이 똑같이 쓰는 Actor와 Critic이다.

class Actor(nn.Module):
    def __init__(self, state_dim, num_actions, hidden=128):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(state_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, num_actions),
        )

    def forward(self, state):
        return Categorical(logits=self.mlp(state))

class Critic(nn.Module):
    def __init__(self, state_dim, hidden=128):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(state_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, 1),
        )

    def forward(self, state):
        return self.mlp(state).squeeze(-1)

 

• Actor 즉 policy net은 state를 입력받아 action에 대한 distribution을 출력하는 network
• Critic 즉 value net은 state를 입력받아 그 state가 얼마나 좋은지를 점수 하나 즉 scalar로 출력하는 network

두 알고리즘 모두 위 Actor와 Critic을 그대로 쓴다. 차이는 오직 train loop다.

REINFORCE + baseline은 에피소드를 끝까지 굴려서 모은 다음 future_return을 뒤에서부터 누적해 만들고 루프 밖에서 한 번 업데이트한다.

def train_reinforce(env, actor, critic, actor_opt, critic_opt, gamma):
    state, _ = env.reset()
    log_probs, values, rewards = [], [], []
    done = False
    while not done:
        s = torch.as_tensor(state, dtype=torch.float32)
        dist = actor(s)
        action = dist.sample()
        next_state, reward, terminated, truncated, _ = env.step(action.item())
        log_probs.append(dist.log_prob(action))
        values.append(critic(s))
        rewards.append(reward)
        state = next_state
        done = terminated or truncated

    future_return = 0.0
    returns = []
    for reward in reversed(rewards):
        future_return = reward + gamma * future_return
        returns.insert(0, future_return)
    returns = torch.as_tensor(returns, dtype=torch.float32)

    values = torch.stack(values)
    log_probs = torch.stack(log_probs)
    advantage = returns - values.detach()

    actor_loss = -(log_probs * advantage).mean()
    critic_loss = (returns - values).pow(2).mean()

    actor_opt.zero_grad(); actor_loss.backward(); actor_opt.step()
    critic_opt.zero_grad(); critic_loss.backward(); critic_opt.step()

A2C는 루프 안에서 매 스텝 td_target을 critic으로 bootstrap해서 바로 업데이트한다.

def train_a2c(env, actor, critic, actor_opt, critic_opt, gamma):
    state, _ = env.reset()
    done = False
    while not done:
        s = torch.as_tensor(state, dtype=torch.float32)
        dist = actor(s)
        action = dist.sample()
        next_state, reward, terminated, truncated, _ = env.step(action.item())
        done = terminated or truncated
        s_next = torch.as_tensor(next_state, dtype=torch.float32)

        value = critic(s)
        td_target = reward + gamma * critic(s_next).detach() * (1 - done)
        advantage = td_target - value

        actor_loss = -dist.log_prob(action) * advantage.detach()
        critic_loss = advantage.pow(2)

        actor_opt.zero_grad(); actor_loss.backward(); actor_opt.step()
        critic_opt.zero_grad(); critic_loss.backward(); critic_opt.step()

        state = next_state

두 코드를 나란히 보면 차이가 딱 세 군데다.

• target 계산: REINFORCE는 에피소드가 끝난 뒤 future_return을 뒤에서부터 누적한 returns(Monte Carlo)를 쓰고, A2C는 매 스텝 td_target = reward + gamma * critic(s_next)로 critic을 bootstrap한다.
• 업데이트 위치: REINFORCE는 while 루프가 끝난 밖에서 한 번, A2C는 while 루프 안에서 매 스텝.
• 저장 범위: REINFORCE는 에피소드 전체의 log_prob와 reward를 모아둬야 하고, A2C는 지금 스텝의 transition 하나면 된다.

결국 Actor와 Critic 코드는 손도 안 대고 train loop에서 이 세 군데만 바뀌는 게 둘의 전부다.

참고로 A2C를 worker 여러 개로 병렬로 돌리며 각자 계산한 gradient를 global network에 비동기로 모으는 A3C라는 확장판도 있다.