본문 바로가기

> Reinforcement Learning/> Basic

[강화학습 기초지식] Part 3: Intro to Policy Optimization (1)

본 포스트는 OpenAI에서 공개한 강화학습 교육자료인 Spinning Up의 도움을 받아, 강화학습(Reinforcement Learning, RL)에 대한 기초개념을 정리해보고자 제작한 시리즈의 일환입니다.

1)기초 지식 2)주요 논문, 3)최신 논문 의 순서로 시리즈를 정리하고자 합니다.

아래 링크에서 더 상세한 내용을 찾아볼 수 있습니다.

Part 3: Intro to Policy Optimization - Spinning Up documentation
In this section, we'll discuss the mathematical foundations of policy optimization algorithms, and connect the material to sample code. We will cover three key results in the theory of policy gradients: In the end, we'll tie those results together and describe the advantage-based expression for the policy gradient-the version we use in our Vanilla Policy Gradient implementation.
https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html

Deriving the Simplest Policy Gradient

Policy optimization의 근간이 될 policy gradient

이번 포스트에선, 지난번에 소개했던 policy optimization 알고리즘을 이해하기 위한 수학적 배경에 대해 다루고자 합니다. Sample code도 함께 리뷰해볼거구요.

먼저 policy의 조건을 정해보겠습니다. Stochastic, parameterized policy πθ\pi_\theta로 한정짓겠습니다. Expected return J(πθ)=Eτπθ[R(τ)]J(\pi_\theta)=\underset{\tau\sim\pi_\theta}{\mathbb{E}}[R(\tau)]를 최대화시키는 것이 우리의 목표이기에, finite-horizon undiscounted return R(τ)R(\tau)를 얻습니다. 만약 R(τ)R(\tau)를 infinite-horizon discounted return으로 설정한다해도 expected return의 정의는 동일합니다. 우리는 gradient ascent를 통해 policy를 최적화시킬 것이며, 이를 식으로 나타내면 아래와 같습니다.

θk+1=θk+αθJ(πθ)θk\theta_{k+1}=\theta_k+\alpha\nabla_\theta J(\pi_\theta)|_{\theta_k}

θJ(πθ)\nabla_\theta J(\pi_\theta)를 우리는 policy gradient라고 합니다. 그리고 이를 활용한 최적화방법을 policy gradient algorithm이라고 합니다. θJ(πθ)\nabla_\theta J(\pi_\theta)의 정확한 값을 계산하는 것이 물론 이상적입니다. 하지만 그러지 못하는 경우 agent와 environment 사이의 상호작용에서 획득한 많은 양의 data를 이용해 평균값을 추정할 수 있지요.

θJ(πθ)\nabla_\theta J(\pi_\theta)의 계산과정에 필요한 사전정보 5가지를 간단하게 요약해보겠습니다.

1. Probability of a Trajectory: Policy πθ\pi_\theta를 통해 얻어진 trajectory τ=(s0,a0,...,sT+1)\tau=(s_0, a_0,...,s_{T+1})가 벌어질 확률은 어떻게 될까요? 확률의 chain rule을 통해 아래와 같은 확률을 얻게 됩니다.

P(τθ)=ρ0(s0)t=0TP(st+1st,at)πθ(atst)P(\tau|\theta)=\rho_0(s_0)\prod_{t=0}^TP(s_{t+1}|s_t, a_t)\pi_\theta(a_t|s_t)

2. The Log-Derivative Trick: ddxlogx=1/x{d \over dx}\log x = 1/x 라는 기본 calculus 지식과 미분의 chain rule을 결합하면 아래와 같습니다.

θP(τθ)=P(τθ)θlogP(τθ)\nabla_\theta P(\tau|\theta)=P(\tau|\theta)\nabla_\theta \log{P(\tau|\theta)}

3. Log-Probability of a Trajectory: 1.에서 얻은 P(τθ)P(\tau|\theta)식의 양변에 log를 취하면 아래와 같습니다.

logP(τθ)=logρ0(s0)+t=0T(logP(st+1st,at)+logπθ(atst))\log P(\tau|\theta)=\log \rho_0(s_0)+\sum_{t=0}^T \Big( \log P(s_{t+1}|s_t, a_t)+\log\pi_\theta(a_t|s_t)\Big)

4. Gradients of Environment Functions: Policy라는 개념은 agent가 내리는 결정에 대한 기준이므로, environment는 policy 및 그의 θ\theta와 완전히 무관합니다. 그러므로 아래와 같은 결과를 낳습니다.

θρ0(s0)=θP(st+1st,at)=θR(τ)=0\nabla_\theta \rho_0(s_0)=\nabla_\theta P(s_{t+1}|s_t,a_t)=\nabla_\theta R(\tau)=0

5. Grad-Log-Prob of a Trajectory: 결국 θlogP(τθ)\nabla_\theta\log P(\tau|\theta)를 계산하면 아래와 같습니다.

θlogP(τθ)=θlogρ0(s0)+t=0T(θlogP(st+1st,at)+θlogπθ(atst))=t=0Tθlogπθ(atst)\begin{aligned} \nabla_\theta\log P(\tau|\theta) &= \cancel{\nabla_\theta\log \rho_0(s_0)} + \sum_{t=0}^T \Big( \cancel{\nabla_\theta\log P(s_{t+1}|s_t, a_t)}+\nabla_\theta\log\pi_\theta(a _t|s_t)\Big) \\ &= \sum_{t=0}^T \nabla_\theta\log\pi_\theta(a_t|s_t) \end{aligned}

위의 아이디어를 이용한 θJ(πθ)\nabla_\theta J(\pi_\theta)의 계산과정은 아래와 같습니다. 넷째 식으로 전개하는 과정에 log-derivative trick이 사용되었음을 알 수 있습니다. 마지막 식으로 전개하는 과정에 grad-log-prob of a trajectory가 사용되었습니다.

θJ(πθ)=θEτπθ[R(τ)]=θτP(τθ)R(τ)=τθP(τθ)R(τ)=τP(τθ)θlogP(τθ)R(τ)=Eτπθ[θlogP(τθ)R(τ)]θJ(πθ)=Eτπθ[t=0Tθlogπθ(atst)R(τ)]\begin{aligned} \nabla_\theta J(\pi_\theta) &= \nabla_\theta \underset{\tau\sim\pi_\theta}{\mathbb{E}}[R(\tau)] \\&= \nabla_\theta\int_\tau P(\tau|\theta)R(\tau) \\&= \int_\tau \nabla_\theta P(\tau|\theta)R(\tau) \\&= \int_\tau P(\tau|\theta)\nabla_\theta \log{P(\tau|\theta)} \cdot R(\tau) \\&= \underset{\tau\sim\pi_\theta}{\mathbb{E}}[\nabla_\theta \log{P(\tau|\theta)} \cdot R(\tau)] \end{aligned} \\ \therefore \nabla_\theta J(\pi_\theta)=\underset{\tau\sim\pi_\theta}{\mathbb{E}} \left[ \sum_{t=0}^T \nabla_\theta \log\pi_\theta(a_t|s_t) \cdot R(\tau) \right]

하지만 우리는 현실적으로 πθ\pi_\theta하에 생길 수 있는 '모든' 가짓수의 trajectory를 파악할 수 없습니다. 때문에 우리는 sample을 획득하여 평균의 계산에 활용할 것입니다. 이에 πθ\pi_\theta하에 생길 수 있는 N(=D)N(=\left| \mathcal{D} \right|)개의 sample trajectories D={τi}i=1,...,N\mathcal{D} = \{\tau_i\}_{i=1,...,N}를 획득합니다. 그 후, 각 τi\tau_i에 따라 결정되는 t=0Tθlogπθ(atst)R(τ)\sum_{t=0}^T \nabla_\theta\log\pi_\theta(a_t|s_t)R(\tau)값들의 산술평균을 구하여 간접적인 policy gradient를 얻습니다. 이는 아래의 식으로 나타낼 수 있습니다.

g^=1DτDt=0Tθlogπθ(atst)R(τ)\hat{g}= {1\over{\left| \mathcal{D} \right|}} \sum_{\tau\in\mathcal{D}} \sum_{t=0}^T \nabla_\theta\log\pi_\theta(a_t|s_t)R(\tau)

위의 방법으로 policy gradient를 계산하려면 θlogπθ(atst)\nabla_\theta\log\pi_\theta(a_t|s_t)을 계산할 수 있도록 πθ\pi_\theta의 정보를 알아야 하고, trajectory의 dataset을 온전히 계산에 활용할 수 있도록 해야합니다.

Implementing the Simplest Policy Gradient

PyTorch를 이용해 위에서 익힌 초간단 version의 policy gradient를 구현해보겠습니다. 이는 Spinning Up Github에서 상세히 볼 수 있습니다. 코드가 길어 toggle로 숨겨놓았으니 하단의 삼각형을 눌러주세요.

  • 1. Making the Policy Network:
    # 1_simple_pg.py, line 30:
    
    # make core of policy network
    logits_net = mlp(sizes=[obs_dim]+hidden_sizes+[n_acts])
    
    # make function to compute action distribution
    def get_policy(obs):
        logits = logits_net(obs)
        return Categorical(logits=logits)
    
    # make action selection function (outputs int actions, sampled from policy)
    def get_action(obs):
        return get_policy(obs).sample().item()

    이 부분에서 feedforward neural network, categorical policy를 이용한 logits_net module과 function들을 만들어줍니다. Categorical policy는 이전 포스트, Part 1에서 다룬 적이 있습니다. 함께 보시면 이해가 빠를 것입니다. mlp는 multi-layer perceptron을 의미합니다. logits_net module은 observation-to-action에 대한 확률분포와 log-prob을 구하기 위해 쓰입니다. get_action은 현재 logits_net라는 확률분포에 의해 얻어지는 action을 추출(sample)해냅니다.

    Line 36의 Categorical object는 PyTorch에서 제공하는 Distribution object의 일종으로, 미리 정의해놓은 확률분포를 line 40에서 action sampling에 사용하듯 손쉽게 쓸 수 있게 해줍니다.

  • 2. Making the Loss Function:
    # 1_simple_pg.py, line 42:
    
    # make loss function whose gradient, for the right data, is policy gradient
    def compute_loss(obs, act, weights):
        logp = get_policy(obs).log_prob(act)
        return -(logp * weights).mean()

    이 부분에서 우리는 policy gradient algorithm의 loss를 계산합니다. 여기서 구해진 loss에 gradient를 취해주면 위에서 공부한 θJ(πθ)=Eτπθ[θlogP(τθ)R(τ)]\nabla_\theta J(\pi_\theta) = \underset{\tau\sim\pi_\theta}{\mathbb{E}}[\nabla_\theta \log{P(\tau|\theta)} \cdot R(\tau)]를 구할 수 있게 되는 것이죠. 식과 함께 비교해보세요. 우리는 현재 policy에 의해 움직이며 획득된 (state, action, weight)의 tuple set을 data로 넣어줍니다. Weight은 처음 들어보신다구요? 현재 우리가 쓰고 있는 weight은 state-action pair에 의해 획득된 return을 의미하지만, 사용자에 따라 다른 값을 대입할 수도 있습니다.

    여기서 주의할 점은 통상 machine learning, 특히 supervised learning에서 이용되는 loss의 개념과 동일하게 생각하면 안된다는 점입니다. Supervised learning의 경우 ground truth(정해진 답)에 해당하는 data가 loss 계산에 쓰일 것이고, 이 data는 parameter의 분포에 무관하겠죠. 하지만 강화학습의 경우 training에 사용되는 data는 ground truth과 무관한, 일개 agent의 발자취에 불과합니다. 이 때문에 loss를 성능지표로 사용할 수 없습니다. 단지 '현재 policy'에 대해 fitting하는 것이 agent의 성능과 무슨 상관이 있을까요? 우리가 신경써야 할 것은 오직 average return뿐입니다.

  • 3. Running One Epoch of Training:
    # 1_simple_pg.py, line 42:
    
    # for training policy
    def train_one_epoch():
        # make some empty lists for logging.
        batch_obs = []          # for observations
        batch_acts = []         # for actions
        batch_weights = []      # for R(tau) weighting in policy gradient
        batch_rets = []         # for measuring episode returns
        batch_lens = []         # for measuring episode lengths
    
        # reset episode-specific variables
        obs = env.reset()       # first obs comes from starting distribution
        done = False            # signal from environment that episode is over
        ep_rews = []            # list for rewards accrued throughout ep
    
        # render first episode of each epoch
        finished_rendering_this_epoch = False
    
        # collect experience by acting in the environment with current policy
        while True:
    
            # rendering
            if (not finished_rendering_this_epoch) and render:
                env.render()
    
            # save obs
            batch_obs.append(obs.copy())
    
            # act in the environment
            act = get_action(torch.as_tensor(obs, dtype=torch.float32))
            obs, rew, done, _ = env.step(act)
    
            # save action, reward
            batch_acts.append(act)
            ep_rews.append(rew)
    
            if done:
                # if episode is over, record info about episode
                ep_ret, ep_len = sum(ep_rews), len(ep_rews)
                batch_rets.append(ep_ret)
                batch_lens.append(ep_len)
    
                # the weight for each logprob(a|s) is R(tau)
                batch_weights += [ep_ret] * ep_len
    
                # reset episode-specific variables
                obs, done, ep_rews = env.reset(), False, []
    
                # won't render again this epoch
                finished_rendering_this_epoch = True
    
                # end experience loop if we have enough of it
                if len(batch_obs) > batch_size:
                    break
    
        # take a single policy gradient update step
        optimizer.zero_grad()
        batch_loss = compute_loss(obs=torch.as_tensor(batch_obs, dtype=torch.float32),
                                  act=torch.as_tensor(batch_acts, dtype=torch.int32),
                                  weights=torch.as_tensor(batch_weights, dtype=torch.float32)
                                  )
        batch_loss.backward()
        optimizer.step()
        return batch_loss, batch_rets, batch_lens

    train_one_epoch()은 1 epoch의 policy gradient를 수행합니다. 이 1 epoch동안, agent는 '현재 policy'에 의거하여 몇 차례의 episode를 시행하며 batch data를 모읍니다(line 67~102). 그 후 policy gradient algorithm에 batch data를 적용시켜 policy를 update합니다(line 104~111).

남은 부분은 다음 포스트에서 다뤄보겠습니다! :)

Reference

[1] R. S. Sutton, A. G. Barto, et al., Introduction to reinforcement learning, vol. 135. MIT press Cambridge, 1998.

[2] OpenAI. Spinning Up, [Online]. Available: https://spinningup.openai.com/