策略梯度算法

Hands-on-RL - 这篇文章属于一个选集。

§ 6: 本文

§ 7: Actor-Critic 算法

§ 8: TRPO 算法

§ 9: PPO 算法

本系列是学习《动手学强化学习》过程中做的摘抄。

6.1 value-based 和 policy-based
#

基于价值的方法，如 Q-learning、DQN 及 DQN 改进算法，主要是学习值函数，然后根据值函数导出一个策略，学习过程中并不存在一个显式的策略。

基于策略的方法则是显式地学习一个目标策略，如 REINFORCE 算法。

6.2 策略梯度
#

基于策略的方法首先需要将策略参数化：假设目标策略 $\pi_{\theta}$ 是一个随机性策略，并且处处可微，其中 $\theta$ 是对应的参数。目标是寻找一个最优策略并最大化这个策略在环境中的期望回报（一个状态的期望回报被称为这个状态的价值）。可以将策略学习的目标函数定义为：

$$ J(\theta) = \text{E}_{s_{0}} \left[ V^{\pi_{\theta}}(s_{0}) \right] $$

其中，$s_0$ 表示初始状态。有了目标函数后就可以对策略 $\theta$ 求导，得到导数后就可以用梯度上升方法来最大化这个目标函数，从而得到最优策略。

$$ \nabla_{\theta} J(\theta) \propto \sum_{s\in \mathcal{S}}V^{\pi_{\theta}}(s) \sum_{a \in \mathcal{A}} Q^{\pi_{\theta}}(s,a) \nabla_{\theta} \pi_{\theta} (a|s) \\ = \text{E}_{\pi_{\theta}} \left[ Q^{\pi_{\theta}}(s,a) \nabla_{\theta} \log \pi_{\theta} (a|s) \right] $$

因为上式中期望 $\text{E}$ 的下标是 $\pi_{\theta}$，所以策略梯度算法为在线策略（on-policy）算法，即必须使用当前策略 $\pi_{\theta}$ 采样得到的数据来计算梯度。

直观地理解策略梯度：梯度的修改是让策略更多地采样到带来较高 $Q$ 值的动作，更少地采样到带来较低 $Q$ 值的动作。

6.3 REINFORCE
#

在计算策略梯度的公式中，需要用到 $Q^{\pi_{\theta}}(s,a)$，REINFORCE 算法采用蒙特卡洛方法来估计 $Q^{\pi_{\theta}}(s,a)$，对于一个有限步数的环境来说，REINFORCE 算法中的策略梯度为：

$$ \nabla_{\theta} J(\theta) = \text{E}_{\pi_{\theta}} \left[ \sum_{t=0}^{T} \left( \sum_{t'=t}^{T} \gamma^{t'-t} r_t \right) \nabla_{\theta} \log \pi_{\theta} (a_t|s_t) \right] $$

其中，$T$ 是和环境交互的最大步数。

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


class PolicyNet(nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim):
        super().__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, action_dim)

    def forward(self, x) -> torch.Tensor:
        # 输入是某个状态, 输出则是该状态下的动作概率分布
        x = F.relu(self.fc1(x))
        return F.softmax(self.fc2(x), dim=1)


class REINFORCE:
    def __init__(self, state_dim, hidden_dim, action_dim, lr: float, gamma: float, device):
        self.policy_net = PolicyNet(state_dim, hidden_dim, action_dim).to(device)
        self.optimizer = torch.optim.Adam(self.policy_net.parameters(), lr=lr)
        self.gamma = gamma  # 折扣因子
        self.device = device

    def tack_action(self, state):
        """根据动作概率分布随机采样"""
        state = torch.tensor([state], dtype=torch.float).to(self.device)
        probs = self.policy_net(state)
        action_dist = torch.distributions.Categorical(probs)
        action = action_dist.sample()
        return action.item()

    def update(self, transition_dict: dict):
        reward_list = transition_dict["rewards"]
        state_list = transition_dict["states"]
        action_list = transition_dict["actions"]

        G = 0
        self.optimizer.zero_grad()
        for i in reversed(range(len(reward_list))):
            reward = reward_list[i]
            state = torch.tensor([state_list[i]], dtype=torch.float).to(self.device)
            action = torch.tensor([action_list[i]]).view(-1, 1).to(self.device)
            log_prob = torch.log(self.policy_net(state).gather(1, action))
            G = self.gamma * G + reward
            loss = -log_prob * G
            loss.backward()
        self.optimizer.step()