CS 285 Online Course at UC Berkeley의 Assingment3에서 다룬 SAC(Soft Actor Critic) 알고리즘에 대한 내용입니다.

해결하고자 하는 문제

문제

SAC 알고리즘은 기존의 model-free deep RL 알고리즘이 가지고 있는 문제를 해결하기 위해 설계되었습니다. 특히 다음과 같은 두 가지 주요 문제입니다.

• Sample complexity 샘플 복잡성
• Hyperparameter sensitivity 하이퍼파라미터 민감도

Sample complexity

model-free deep RL, 특히 on-policy learning 알고리즘은 매 gradient update마다 새로운 sample을 수집해야 하므로, task가 복잡할수록 많은 데이터가 필요합니다. 이는 매우 비효율적이며, 실제 환경에서는 이러한 샘플을 얻는 것이 더욱 어렵습니다.

Hyperparameter sensitivity

off-policy 알고리즘은 replay-buffer를 활용해 과거의 경험을 재사용하는 것을 목표로 합니다. Q-learning 기반 방법에서 이를 간단하게 구현할 수 있지만, off-policy learning과 Neural Network와 같은 high-dimension nonlinear function approximation이 결합하게 되면 알고리즘의 안정성과 수렴에 큰 어려움이 생깁니다. 이는 문제 설정에 맞게 하이퍼파라미터를 신중하게 조정해야 하며, 많은 시도와 오류를 요구하게 됩니다.

해결방법

그렇다면 SAC 알고리즘은 이 문제들을 어떻게 해결하고 continuous state and action space에서 효율적이고 안정적인 model-free deep RL 알고리즘으로 설계할 수 있었을까요? 그 답은 논문의 제목에 있습니다. SAC 알고리즘을 소개하고 있는 논문의 제목은

Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor

입니다. 몇 가지 키워드가 눈에 들어옵니다. Soft, Off-policy, Entropy, Stochastic Actor. 논문의 Related Work에 다음과 같은 문장이 있습니다.

Our method instead combines off-policy actor-critic training with a stochastic actor, and further aims to maximize the entropy of this actor with an entropy maximization objective.

이 문장에는 크게 두 가지 해결방법이 적혀 있습니다.

• combine off-policy actor-critic training with a stochastic actor
• maximize the entropy of this actor with an entropy maximization objective

Related work

Preliminaries

논문의 이해를 높이기 위해 알아두어야 할 개념이 있다. 바로 Maximum Entropy Reinforcement Learning. entropy라 하면 기계공학을 전공한 나에겐 익숙한 단어이다. 열역학에서도 등장하는 개념으로 한 단어로 표현하면 불확실성이다. 이는 당연하게 강화학습에서도 같은 개념으로 사용된다.

Entropy

엔트로피(Entropy)는 확률 분포의 불확실성을 측정하는 척도입니다. 강화 학습에서는 정책의 무작위성을 높이기 위해 엔트로피 보너스를 추가하여 더 많은 탐색을 유도할 수 있습니다. 즉, 엔트로피가 높은 정책은 가능한 행동들을 더 균등하게 선택하게 되어 탐색 성향이 강해집니다.

엔트로피(Entropy)는 확률 분포의 불확실성을 측정하는 척도입니다. 강화 학습에서는 정책의 무작위성을 높이기 위해 엔트로피 보너스를 추가하여 더 많은 탐색을 유도할 수 있습니다. 엔트로피가 높은 정책은 가능한 행동들을 더 균등하게 선택하게 되어 탐색 성향이 강해집니다.

확률 분포 $P$의 엔트로피 $H(P)$는 다음과 같이 정의됩니다:

여기서 $X$는 가능한 결과들의 집합입니다.

예를 들어, 두 가지 행동 $a_1$과 $a_2$가 있는 경우, 정책 $\pi$가 두 행동을 각각 0.5의 확률로 선택한다고 가정해보겠습니다. 이 정책의 엔트로피는 다음과 같습니다.

Entropy bonous

강화 학습에서 엔트로피 보너스를 추가하는 이유는 정책이 가능한 행동들을 더 균등하게 선택하도록 유도하여 더 많은 탐색을 촉진하는 것입니다. 엔트로피 보너스는 다음과 같이 목표 함수에 추가됩니다.

여기서,

• $Q(s, \mu_{\theta}(s) + \sigma_{\theta}(s)\epsilon)$는 상태 $s$와 행동 $a$에 대한 Q-값입니다. 이때 $\mu_{\theta}(s) + \sigma_{\theta}(s)\epsilon$은 행동 a로, distribution에서 re-parameterize한 action입니다.
• $\beta$는 엔트로피 보너스의 가중치입니다.
• $\mathcal{H}(\pi(a|s))$는 정책 $\pi$의 엔트로피로, $\mathcal{H}(\pi(a|s)) = \mathbb{E}_{a \sim \pi}[-\log \pi(a|s)]$

policy entropy를 구현한 함수는 다음과 같다.

def entropy(self, action_distribution: torch.distributions.Distribution):
    # reparameterization sampling. 
    action_sample = action_distribution.rsample()

    # Compute log probabilities of the sampled actions
    log_probs = action_distribution.log_prob(action_sample)

    # Calculate the negative log probabilities
    negative_log_probs = -log_probs

    # Estimate the entropy as the mean of the negative log probabilities
    approximated_entropy = negative_log_probs.mean()

    return approximated_entropy

엔트로피 보너스를 사용한 학습 목표

엔트로피 보너스를 사용한 target value는 다음과 같이 수정됩니다.

예를 들어, $\beta = 0.1$이고, 정책 $\pi$가 두 행동 $a_1$과 $a_2$를 각각 0.8과 0.2의 확률로 선택한다고 가정해봅시다. 이 정책의 엔트로피는 다음과 같습니다.

이 엔트로피 보너스를 학습 목표에 추가하면, 정책이 더 다양한 행동을 선택하도록 유도하여 탐색 성향을 높이고, 최적의 정책을 더 빨리 찾을 수 있습니다. 이와 같이 엔트로피는 강화 학습에서 정책의 무작위성을 조절하고 탐색을 촉진하는 중요한 역할을 합니다.

From Soft Policy Iteration to Soft Actor-Critic

앞서 "off-policy learning과 high-dimension nonlinear function approximation이 결합되면 알고리즘의 안정성과 수렴에 큰 어려움이 생깁니다."라는 말을 했습니다. 이는 non-tabular case인 경우 optimal에 수렴하지 않는다는 의미입니다. 반면, tabular case인 Soft Policy Iteration은 optimal policy로 수렴한다는 것이 보장됩니다. off-policy SAC 알고리즘은 maximum entropy variant of the policy iteration 방법으로부터 도출할 수 있습니다.

Soft Policy Iteration 프레임워크는 Soft Policy Evaluation와 Soft Policy Improvement의 두 가지 기본적인 단계를 반복적으로 수행합니다.

Soft Policy Evaluation

소프트 정책 평가 단계에서는 주어진 정책 $\pi$에 대해 soft Q-함수 $Q^\pi(s_t, a_t)$를 계산합니다. 이를 위해 Bellman backup operator $\mathcal{T}^\pi$를 사용하여 soft Q-함수를 반복적으로 계산하면, 이 값이 특정 정책 $\pi$에 대해 수렴함을 보장합니다.

여기서 $V(s_t)$는 soft state function으로 다음과 같이 정의됩니다.

이 과정은 다음과 같은 순서로 이루어집니다.

1. 초기화: 임의의 $Q^0 : S \times A \to \mathbb{R}$ 함수로 시작합니다.
2. 반복적 적용: $\mathcal{T}^\pi$ 연산자를 반복적으로 적용하여 $Q^{k+1} = \mathcal{T}^\pi Q^k$를 계산합니다.
3. 수렴: $k \to \infty$일 때, $Q^k$는 정책 $\pi$에 대한 소프트 Q-함수에 수렴합니다.

이 과정에서 $V(s_t)$는 상태 $s_t$에서 가능한 모든 행동 $a_t$에 대한 Q-값과 해당 행동의 로그 확률을 결합한 값의 기대값입니다. 이를 통해 정책의 탐색 성향을 유지하면서도 최적의 Q-값을 학습할 수 있습니다.

Soft Policy Improvement

소프트 정책 개선 단계에서는 현재의 소프트 Q-함수를 기반으로 정책을 개선합니다. 새로운 정책 $\pi_{\text{new}}$는 다음과 같이 정의됩니다.

여기서 $D_{KL}$은 쿨백-라이블러 발산(Kullback-Leibler divergence)이고, $Z^{\pi_{\text{old}}}(s_t)$는 정규화 함수입니다.

Soft policy improvement 과정은 다음의 순서로 이루어진다.

1. 소프트 Q-함수 계산: 정책 $\pi_{\text{old}}$에 대한 소프트 Q-함수 $Q^{\pi_{\text{old}}}(s_t, a_t)$를 계산합니다.
2. 정보 투영: 새로운 정책 $\pi_{\text{new}}$는 $\pi_{\text{old}}$에서 계산된 소프트 Q-함수를 사용하여 KL-Divergence을 최소화하는 방식으로 개선됩니다.
3. 최적화: 개선된 정책 $\pi_{\text{new}}$는 다음과 같이 KL-Divergence을 최소화하여 정의됩니다.
   $\pi_{\text{new}} = \arg \min_{\pi' \in \Pi} D_{KL} \left( \pi'(\cdot|s_t) \bigg\| \frac{\exp \left( \frac{1}{\alpha} Q^{\pi_{\text{old}}}(s_t, \cdot) \right)}{Z^{\pi_{\text{old}}}(s_t)} \right)$

KL-Divergence을 최소화하는 정보 투영은 새로운 정책이 이전 정책보다 더 높은 소프트 Q-값을 가지도록 보장합니다. 이 과정은 소프트 Q-함수와 엔트로피 보너스를 모두 고려하여 정책을 개선합니다. 새로운 정책 $\pi_{\text{new}}$는 $\pi_{\text{old}}$보다 더 높은 값을 가지며, 이는 최적의 최대 엔트로피 정책에 수렴하는 과정에서 중요한 역할을 합니다.

Soft Actor-Critic

Continuous domain에서 soft policy iteration의 근사를 도출하기 위해서 Q-함수와 policy에 대해 function approximation을 사용하고, evaluation과 improvement를 수렴할 때까지 반복하는 대신 두 네트워크를 최적화하는 과정을 SGD 방법으로 번갈아가며 실행합니다.

매개변수화된 state function $V_{\psi}(s_t)$, soft Q-function $Q_{\theta}(s_t, a_t)$, 그리고 tractable policy $\pi_{\phi}(a_t|s_t)$을 구합니다.

State value & Soft Q-function (critic) update

SAC는 두 개의 소프트 Q-함수를 학습합니다. 두 Q-함수는 동일한 타겟을 사용하여 업데이트되며, 이중 Q-러닝(Double Q-Learning)과 유사하게 동작하여 Q-함수의 과대평가 문제를 완화합니다. 두 Q-함수의 업데이트는 다음과 같이 이루어집니다:

여기서 타겟 값 $y_t$는 다음과 같이 정의됩니다:

이때 $\bar{\theta}_j$는 Polyak averaging (exponential moving average) 매개변수를 나타냅니다. 또한, $min Q$는 Clipped double-Q를 나타낸다.

def update_critic(
    self,
    obs: torch.Tensor,
    action: torch.Tensor,
    reward: torch.Tensor,
    next_obs: torch.Tensor,
    done: torch.Tensor,
):
    # Compute target values
    with torch.no_grad():
        next_action_distribution: torch.distributions.Distribution = self.actor(next_obs)
        next_action = next_action_distribution.sample()
        next_qs = self.target_critic(next_obs, next_action)
        next_qs = self.q_backup_strategy(next_qs)
        if self.use_entropy_bonus and self.backup_entropy:
            next_action_entropy = self.entropy(next_action_distribution)
            next_qs += self.temperature * next_action_entropy
        target_values: torch.Tensor = reward + self.discount * next_qs * (1 - done.float())
    q_values = self.critic(obs, action)
    loss: torch.Tensor = self.critic_loss(q_values, target_values)
    self.critic_optimizer.zero_grad()
    loss.backward()
    self.critic_optimizer.step()

이때 다음 action을 새롭게 sampling 하는 것을 알 수 있다. 이는 target value가 represent the value of the latest policy하기 위함이다. 즉, 현재 policy에서 sampling함으로써 다음 state가 현재 정책에 의해 결정될 다음 action의 가치를 반영할 수 있는 것이다.

그리고 다음 코드는 Polyak averaging (exponential moving average)을 구현한 코드이다.

def update_target_critic(self):
    self.soft_update_target_critic(1.0)

def soft_update_target_critic(self, tau):
    for target_critic, critic in zip(self.target_critics, self.critics):
        for target_param, param in zip(
            target_critic.parameters(), critic.parameters()
        ):
            target_param.data.copy_(
                target_param.data * (1.0 - tau) + param.data * tau
            )

Policy (actor) update

정책 $\pi_{\phi}$는 현재 Q-함수에 따라 업데이트됩니다. 정책 업데이트는 주어진 상태에서의 행동을 샘플링하고, 엔트로피 보너스를 포함한 목적 함수를 최대화하는 방식으로 이루어집니다.

여기서 $\beta$는 엔트로피 보너스의 가중치를 조절하는 온도 매개변수입니다. 정책은 주어진 상태에서 높은 Q-값을 갖는 행동을 선택하는 동시에 엔트로피를 최대화하도록 학습됩니다.

이때 target value가 신경망으로 표현되는 Q 함수이며, 그래디언트를 역전파할 수 있으므로 재매개변수화 트릭을 적용하는 것이 편리하며, 이는 더 낮은 분산 추정기로 이어집니다. 이를 위해 정책을 신경망 변환을 사용하여 재매개변수화합니다.

def actor_loss_reparametrize(self, obs: torch.Tensor):
    action_distribution: torch.distributions.Distribution = self.actor(obs)
    action = action_distribution.rsample()
    q_values = self.critic(obs, action)
    loss = -q_values.mean()
    return loss, torch.mean(self.entropy(action_distribution))