强化学习——PyTorch 实现 Advantage Actor-Critic (A2C)

前言

  本博客的理论知识来自王树森老师《深度强化学习》，这本书写得简直太好了，强烈推荐，只是现在还在校对没出版，可能有些小瑕疵，但并不影响阅读和学习。

Advantage Actor-Critic (A2C)

  本次的 A2C 的原理我们从带基线的策略梯度开始，在对带基线的策略梯度做蒙特卡洛近似，得到策略梯度的一个无偏估计：
$$\mathbf{g}(s,a,;\mathbf{\theta })=\left[Q_{\pi }(s,a)-V_{\pi }(s)\cdot \nabla \ln \pi (a|s;\theta )\right]\text{\hspace{1em}(1)}$$
公式中的 $Q_{\pi }-V_{\pi }$ 被称作优势函数 (Advantage Function)。因此，基于上面公式得到的 Actor-Critic 方法被称为 Advantage Actor-Critic，缩写 A2C。A2C 属于 Actor-Critic 方法。有一个策略网络 $\pi (a|s;\theta )$，相当于演员，用于控制智能体运动。还有一个价值网络 $v(s;w)$，相当于评委，他的评分可以帮助策略网络（演员）改进技术。

算法推导

  训练价值网络： 训练价值网络 $v(s;w)$ 的算法是从贝尔曼公式来的：
$$V_{\pi }(s_t)={\mathbb{E}}_{A_t\sim \pi (\cdot |s_t;\theta )}[{\mathbb{E}}_{S_{t+1}\sim p(\cdot |s_t,A_t)}[R_t+\gamma \cdot V_{\pi }(S_{t+1})]]$$
对贝尔曼方程左右两边做近似：

• 方程左边的 $V_{\pi }(s_t)$ 可以近似成 $v(s_t;w)$。$v(s_t;w)$ 是价值网络在 $t$ 时刻对 $V_{\pi }(s_t)$ 做出的估计；
• 方程右边的期望是关于当前时刻动作 $A_t$ 与下一时刻状态 $S_{t+1}$ 求的。给定当前状态 $s_t$，智能体执行动作 $a_t$，环境会给出奖励 $r_t$ 和新的状态 $s_{t+1}$。用观测到的 $r_t$、$s_{t+1}$ 对期望做蒙特卡洛近似，得到：
  $$r_t+\gamma \cdot V_{\pi }(s_{t+1})\text{\hspace{1em}(2)}$$
• 公式 (1) 中的 $V_{\pi }(s_{t+1})$ 近似成 $v(s_{t+1};\mathbf{w})$，得到：
  $$\widehat{y_t}:=r_t+\gamma \cdot v(s_{t+1};\mathbf{w})$$
  把它称作 TD 目标。它是价值网络在 $t+1$ 时刻对 $V_{\pi }(s_t)$ 做出的估计。

$v(s_t;w)$ 和 ${\hat{y}}_t$ 都是对动作价值 $V{\pi }_{(}{s}_t)$ 的估计。由于 ${\hat{y}}_t$ 部分基于真实观测到的奖励 $r_t$，我们认为 ${\hat{y}}_t$ 比 $v(s_t;\mathbf{w})$ 更可靠。所以把 ${\hat{y}}_t$ 固定住，更新 $\mathbf{w}$，使得 $v(s_t;\mathbf{w})$ 更接近 ${\hat{y}}_t$。具体一点，定义损失函数：
$$L(\mathbf{w}):=\frac{1}{2}[v(s_t;\mathbf{w})-{\hat{y}}_t{]}^2$$
设 ${\hat{v}}_t:=v(s_t;\mathbf{w})$，损失函数的梯度为：
$${\nabla }_{\mathbf{w}}L(\mathbf{w})=\underset{TD误差{\delta }_t}{({\hat{v}}_t-{\hat{y}}_t)}\cdot {\nabla }_{\mathbf{w}}v(s_t;\mathbf{w})$$
定义 TD 误差为 $\delta :={\hat{v}}_t-{\hat{y}}_t$。做一轮梯度下降更新 $\mathbf{w}$：
$$\mathbf{w}\leftarrow \mathbf{w}-\alpha \cdot {\delta }_t\cdot {\nabla }_{\mathbf{w}}v(s_t;\mathbf{w})$$
这样可以让价值网络的预测 $v(s_t;\mathbf{w})$ 更接近 ${\hat{y}}_t$。

  训练策略网络：A2C 从公式 (1) 出发，对 $\mathbf{g}(s,a,;\mathbf{\theta })$ 做近似，记作 $\tilde{\mathbf{g}}$，然后用 $\tilde{\mathbf{g}}$ 更新策略网络参数 $\mathbf{\theta }$。下面我们做数学推导。回忆一下贝尔曼公式：
$$Q_{\pi }(s_t,a_t)={\mathbb{E}}_{S_{t+1}\sim p(\cdot |s_t,A_t)}[R_t+\gamma \cdot V_{\pi }(S_{t+1})]$$
把近似策略梯度 $\mathbf{g}(s,a,;\mathbf{\theta })$ 中的 $Q_{\pi }(s_t,a_t)$ 替换成上面的期望，得到：
$$\begin{array}{rl}\mathbf{g}(s,a,;\mathbf{\theta })& =[Q_{\pi }(s_t,a_t)-V_{\pi }(s_t)]\cdot {\nabla }_{\mathbf{\theta }}\ln \pi (a_t|s_t;\mathbf{\theta })\\ & =\left[{\mathbb{E}}_{S_{t+1}}[R_t+\gamma \cdot V_{\pi }(S_{t+1})]-V_{\pi }(S_t)\right]\cdot {\nabla }_{\mathbf{\theta }}\ln \pi (a_t|s_t;\mathbf{\theta })\end{array}$$
当智能体执行动作 $a_t$ 之后，环境给出新的状态 $s_{t+1}$ 和奖励 $r_t$；利用 $s_{t+1}$ 和 $r_t$ 对上面的期望做蒙特卡洛近似，得到：
$$\mathbf{g}(s,a,;\mathbf{\theta })\approx [r_t+\gamma \cdot V_{\pi }(s_{t+1})-V_{\pi }(s_t)]\cdot \nabla \ln \pi (a|s;\theta )$$
进一步把状态价值函数 $V_{\pi }(s)$ 替换成价值网络 $v(s;\mathbf{w})$，得到：
$$\tilde{\mathbf{g}}(s,a,;\mathbf{\theta }):=[\underset{TD目标{\hat{y}}_t}{r_t+\gamma \cdot v(s;\mathbf{w})}-v(s;\mathbf{w})\cdot \nabla \ln \pi (a|s;\theta )$$
前面定义了 TD 目标和 TD 误差，$\widehat{y_t}:=r_t+\gamma \cdot v(s_{t+1};\mathbf{w})$ 和 ${\delta }_t:=v(s_t;\mathbf{w})-{\hat{y}}_t$。因此，可以将 $$\mathbf{g}(s,a,;\mathbf{\theta }):=-{\delta }_t\cdot \nabla \ln \pi (a|s;\theta )$$
$\tilde{\mathbf{g}}$ 是 $\mathbf{g}$ 的近似，所以也是策略梯度 ${\nabla }_{\mathbf{\theta }}J(\mathbf{\theta })$ 的近似。用 $\tilde{\mathbf{g}}$更新策略网络参数 $\mathbf{\theta }$：
$$\mathbf{\theta }\leftarrow \mathbf{\theta }+\beta \cdot \mathbf{g}(s,a,;\mathbf{\theta })$$
这样可以让目标函数 $J(\mathbf{\theta })$ 变大。

训练流程

  下面概括 A2C 训练流程。设当前策略网络参数是 ${\mathbf{\theta }}_{\mathbf{now}}$，价值网络参数是 ${\mathbf{w}}_{\mathbf{now}}$。执行下面的步骤，将参数更新成 ${\mathbf{\theta }}_{\mathbf{new}}$ 和 ${\mathbf{w}}_{\mathbf{new}}$：

1. 观测到当前状态 $s_t$，根据策略网络做决策：$a_t\sim \pi (\cdot |s_t;{\mathbf{\theta }}_{\mathbf{now}})$，并让智能体执行动作 $a_t$；
2. 从环境中观测到奖励 $r_t$ 和新的状态 $s_{t+1}$；
3. 让价值网络打分：${\hat{v}}_t=v(s_t;{\mathbf{w}}_{\mathbf{now}})$ 和 ${\hat{v}}_{t+1}=v(s_{s+1};{\mathbf{w}}_{\mathbf{now}})$；
4. 计算 TD 目标和 TD 误差：$\widehat{y_t}=r_t+\gamma \cdot {\hat{v}}_{t+1}$ 和 ${\delta }_t={\hat{v}}_t-{\hat{y}}_t$；
5. 更新价值网络：${\mathbf{w}}_{\mathbf{new}}\leftarrow {\mathbf{w}}_{\mathbf{now}}-\alpha \cdot {\delta }_t\cdot {\nabla }_{\mathbf{w}}v(s_t;{\mathbf{w}}_{\mathbf{now}})$；
6. 更新策略网络：${\mathbf{\theta }}_{\mathbf{new}}\leftarrow {\mathbf{\theta }}_{\mathbf{now}}-\beta \cdot {\delta }_t\cdot {\nabla }_{\mathbf{\theta }}\ln \pi (a_t|s_t;{\mathbf{\theta }}_{\mathbf{now}})$.

用目标网络改进训练

  上述训练价值网络的算法存在自举——即用价值网络自己的估值 ${\hat{v}}_{t+1}$ 去更新价值网络自己。为了缓解自举造成的偏差，可以使用目标网络 (Target Network) 计算 TD 目标。把目标网络记作 $v(s;{\mathbf{w}}^{-})$，它的结构与价值网络的结构相同，但是参数不同。使用目标网络计算 TD 目标，那么 A2C 的训练就变成了：

1. 观测到当前状态 $s_t$，根据策略网络做决策：$a_t\sim \pi (\cdot |s_t;{\mathbf{\theta }}_{\mathbf{now}})$，并让智能体执行动作 $a_t$；
2. 从环境中观测到奖励 $r_t$ 和新的状态 $s_{t+1}$
3. 让价值网络给 $s_t$ 打分：${\hat{v}}_t=v(s_t;{\mathbf{w}}_{\mathbf{now}})$；
4. 让目标网络给 $s_{t+1}$ 打分：${\hat{v}}_{t+1}=v(s_{s+1};{\mathbf{w}}_{\mathbf{now}}^{-})$；
5. 计算 TD 目标和 TD 误差：$\widehat{y_t^{-}}=r_t+\gamma \cdot {\hat{v}}_{t+1}^{-}$ 和 ${\delta }_t={\hat{v}}_t-{\hat{y}}_t^{-}$；
6. 更新价值网络：${\mathbf{w}}_{\mathbf{new}}\leftarrow {\mathbf{w}}_{\mathbf{now}}-\alpha \cdot {\delta }_t\cdot {\nabla }_{\mathbf{w}}v(s_t;{\mathbf{w}}_{\mathbf{now}})$；
7. 更新策略网络：${\mathbf{\theta }}_{\mathbf{new}}\leftarrow {\mathbf{\theta }}_{\mathbf{now}}-\beta \cdot {\delta }_t\cdot {\nabla }_{\mathbf{\theta }}\ln \pi (a_t|s_t;{\mathbf{\theta }}_{\mathbf{now}})$.
8. 设 $\tau \in (0,1)$ 是需要手动调的参数，做加权平均更新目标网络的参数：${\mathbf{w}}_{\mathbf{new}}^{-}\leftarrow \tau \cdot {\mathbf{w}}_{\mathbf{new}}+(1-\tau )\cdot {\mathbf{w}}_{\mathbf{now}}^{-}$.

PyTorch 实现

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import gym
import numpy as np

LR_ACTOR = 0.01     # 策略网络的学习率
LR_CRITIC = 0.001   # 价值网络的学习率
GAMMA = 0.9         # 奖励的折扣因子
EPSILON = 0.9       # ϵ-greedy 策略的概率
TARGET_REPLACE_ITER = 100                 # 目标网络更新的频率
env = gym.make('CartPole-v0')             # 加载游戏环境
N_ACTIONS = env.action_space.n            # 动作数
N_SPACES = env.observation_space.shape[0] # 状态数量
env = env.unwrapped

# 网络参数初始化，采用均值为 0，方差为 0.1 的高斯分布
def init_weights(m) :
    if isinstance(m, nn.Linear) :
        nn.init.normal_(m.weight, mean = 0, std = 0.1)

# 策略网络
class Actor(nn.Module) :
    def __init__(self):
        super(Actor, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(N_SPACES, 50),
            nn.ReLU(),
            nn.Linear(50, N_ACTIONS) # 输出为各个动作的概率，维度为 3
        )

    def forward(self, s):
        output = self.net(s)
        output = F.softmax(output, dim = -1) # 概率归一化
        return output

# 价值网络
class Critic(nn.Module) :
    def __init__(self):
        super(Critic, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(N_SPACES, 20),
            nn.ReLU(),
            nn.Linear(20, 1) # 输出值是对当前状态的打分，维度为 1
        )

    def forward(self, s):
        output = self.net(s)
        return output

# A2C 的主体函数
class A2C :
    def __init__(self):
        # 初始化策略网络，价值网络和目标网络。价值网络和目标网络使用同一个网络
        self.actor_net, self.critic_net, self.target_net = Actor().apply(init_weights), Critic().apply(init_weights), Critic().apply(init_weights)
        self.learn_step_counter = 0 # 学习步数
        self.optimizer_actor = optim.Adam(self.actor_net.parameters(), lr = LR_ACTOR)    # 策略网络优化器
        self.optimizer_critic = optim.Adam(self.critic_net.parameters(), lr = LR_CRITIC) # 价值网络优化器
        self.criterion_critic = nn.MSELoss() # 价值网络损失函数

    def choose_action(self, s):
        s = torch.unsqueeze(torch.FloatTensor(s), dim = 0) # 增加维度
        if np.random.uniform() < EPSILON :                 # ϵ-greedy 策略对动作进行采取
            action_value = self.actor_net(s)
            action = torch.max(action_value, dim = 1)[1].item()
        else :
            action = np.random.randint(0, N_ACTIONS)

        return action

    def learn(self, s, a, r, s_):
        if self.learn_step_counter % TARGET_REPLACE_ITER == 0 :          # 更新目标网络
            self.target_net.load_state_dict(self.critic_net.state_dict())

        self.learn_step_counter += 1

        s = torch.FloatTensor(s)
        s_ = torch.FloatTensor(s_)

        q_actor = self.actor_net(s)               # 策略网络
        q_critic = self.critic_net(s)             # 价值对当前状态进行打分
        q_next = self.target_net(s_).detach()     # 目标网络对下一个状态进行打分
        q_target = r + GAMMA * q_next             # 更新 TD 目标
        td_error = (q_critic - q_target).detach() # TD 误差

        # 更新价值网络
        loss_critic = self.criterion_critic(q_critic, q_target)
        self.optimizer_critic.zero_grad()
        loss_critic.backward()
        self.optimizer_critic.step()

        # 更新策略网络
        log_q_actor = torch.log(q_actor)
        actor_loss = log_q_actor[a] * td_error
        self.optimizer_actor.zero_grad()
        actor_loss.backward()
        self.optimizer_actor.step()

a2c = A2C()

for epoch in range(10000) :
    s = env.reset()
    ep_r = 0
    while True :
        env.render()
        a = a2c.choose_action(s)       # 选择动作

        s_, r, done, info = env.step(a)# 执行动作

        x, x_dot, theta, theta_dot = s_
        # 修改奖励，为了更快收敛
        r1 = (env.x_threshold - abs(x)) / env.x_threshold - 0.8
        r2 = (env.theta_threshold_radians - abs(theta)) / env.theta_threshold_radians - 0.5
        r = r1 + r2

        ep_r += r

        # 学习
        a2c.learn(s, a, r, s_)

        if done :
            break

        s = s_
    print(f'Ep: {epoch} | Ep_r: {round(ep_r, 2)}')