Annonated Algorithm Visualization

View code on GitHub

Proximal Policy Optimization (PPO) 算法建模离散动作空间的入门示例（PyTorch版）
Related Link

PPO 是最经典常用的强化学习算法之一（策略梯度类）。它结合了经典的 Actor-Critic 范式和信赖域策略优化方法，并把相关工作整合为一个简洁而有效的算法。和之前传统的策略梯度类强化学习算法相比（例如 REINFORCE 和 A2C），PPO 可以更稳定高效地提升智能体策略，通过如下所示的截断式优化目标不断强化智能体：

J(\theta) = \min(\frac{\pi_{\theta}(a_{t}|s_{t})}{\pi_{\theta_k}(a_{t}|s_{t})}A^{\theta_k}(s_{t},a_{t}),\text{clip}(\frac{\pi_{\theta}(a_{t}|s_{t})}{\pi_{\theta_k}(a_{t}|s_{t})}, 1-\epsilon,1+\epsilon)A^{\theta_k}(s_{t},a_{t}))

这个截断式的优化目标是原始未截断版本的一个下界（即一种悲观的约束）。通过公式中的 min 操作，忽略掉一些对于策略提升较大的重要性采样系数（IS），但又在策略优化变得糟糕时保留足够的 IS，从而让整个优化过程更稳定。
详细的数学符号定义可以在符号表中找到 Related Link 。

离散动作空间是最常见的动作空间之一，例如像超级马里奥，雅达利（Atari）这样的视频游戏就通过这种动作空间来决策。具体来说，它包含一组可选的离散动作选项，每次决策需要从其中选择一个。离散动作空间常常被建模成广义伯努利分布来优化（类似一个分类问题）。

本文档将主要分为三个部分，读者可以从这些样例代码中一步一步进行学习，也可以将其中一些代码片段用到自己的程序中：
- 策略神经网络架构
- 动作采样函数
- 主函数（测试函数）
更多的可视化结果和实际应用样例，可以参考这个链接 Related Link

DiscretePolicyNetwork 定义概述
定义 PPO 中所使用的离散动作策略网络，其主要包含两部分：编码器（encoder）和决策输出头（head）


x
 
1
from typing import List
2
import torch
3
import torch.nn as nn
4
5
6
class DiscretePolicyNetwork(nn.Module):
7
8
    def __init__(self, obs_shape: int, action_shape: int) -> None:

继承 PyTorch 神经网络类所必需的操作，自定义的神经网络必须是 nn.Module 的子类


xxxxxxxxxx
9
 
9
        super(DiscretePolicyNetwork, self).__init__()

定义编码器模块，将原始的状态映射为特征向量。对于不同的环境，可能状态信息的模态不同，需要根据情况选择适当的编码器神经网络，例如对于图片状态信息就常常使用卷积神经网络
这里我们用一个简单的单层 MLP 作为例子，即:
$y = max(Wx+b, 0)$


xxxxxxxxxx
13
 
10
        self.encoder = nn.Sequential(
11
            nn.Linear(obs_shape, 32),
12
            nn.ReLU(),
13
        )

定义离散动作的决策输出头，一般仅仅一层全连接层即可，即:
$y=Wx+b$


xxxxxxxxxx
15
 
14
        self.head = nn.Linear(32, action_shape)
15

forward 函数功能概述
描述 PPO 中所使用的离散动作策略网络的前向计算图
x -> encoder -> head -> logit .


xxxxxxxxxx
16
 
16
    def forward(self, x: torch.Tensor) -> torch.Tensor:

将原始的状态信息转换为特征向量，维度变化为: $(B, *) -> (B, N)$


xxxxxxxxxx
17
 
17
        x = self.encoder(x)

为每个可能的离散动作选项，计算相应的 logit，维度变化为: $(B, N) -> (B, A)$


xxxxxxxxxx
21
 
18
        logit = self.head(x)
19
        return logit
20
21

MultiDiscretePolicyNetwork 定义概述
定义 PPO 中所使用的多维离散动作策略网络，其主要包含两部分：编码器（encoder）和多维决策输出头（head）


xxxxxxxxxx
24
 
22
class MultiDiscretePolicyNetwork(nn.Module):
23
24
    def __init__(self, obs_shape: int, action_shape: List[int]) -> None:

继承 PyTorch 神经网络类所必需的操作，自定义的神经网络必须是 nn.Module 的子类


xxxxxxxxxx
25
 
25
        super(MultiDiscretePolicyNetwork, self).__init__()

定义编码器模块，将原始的状态映射为特征向量。对于不同的环境，可能状态信息的模态不同，需要根据情况选择适当的编码器神经网络，例如对于图片状态信息就常常使用卷积神经网络
这里我们用一个简单的单层 MLP 作为例子，即:
$y = max(Wx+b, 0)$


xxxxxxxxxx
29
 
26
        self.encoder = nn.Sequential(
27
            nn.Linear(obs_shape, 32),
28
            nn.ReLU(),
29
        )

定义多维离散动作的决策输出头，即创建多个离散动作的预测器，整体用 nn.ModuleList 进行管理


xxxxxxxxxx
33
 
30
        self.head = nn.ModuleList()
31
        for size in action_shape:
32
            self.head.append(nn.Linear(32, size))
33

Overview
The computation graph of discrete action policy network used in PPO.
x -> encoder -> multiple head -> multiple logit .


xxxxxxxxxx
34
 
34
    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:

将原始的状态信息转换为特征向量，维度变化为: $(B, *) -> (B, N)$


xxxxxxxxxx
35
 
35
        x = self.encoder(x)

为每个可能的离散动作选项，计算相应的 logit，维度变化为: $(B, N) -> [(B, A_1), ..., (B, A_N)]$


xxxxxxxxxx
39
 
36
        logit = [h(x) for h in self.head]
37
        return logit
38
39

sample_action 函数功能概述
输入 logit 采样获得离散动作，输入维度为 (B, action_shape) 输出维度为 output shape = (B, )
在这个示例中，课程中提到的 distributions 工具库的三个维度分别为
batch_shape = (B, ), event_shape = (), sample_shape = ()


xxxxxxxxxx
40
 
40
def sample_action(logit: torch.Tensor) -> torch.Tensor:

将 logit 转化为概率（logit 一般指神经网络的原始输出，不经过激活函数，例如最后一层全连接层的输出）
$\text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}$


xxxxxxxxxx
41
 
41
    prob = torch.softmax(logit, dim=-1)

构建广义伯努利分布。它的概率质量函数为: $f(x=i|\boldsymbol{p})=p_i$
Related Link


xxxxxxxxxx
42
 
42
    dist = torch.distributions.Categorical(probs=prob)

为一个 batch 里的每个样本采样一个离散动作，并返回它


xxxxxxxxxx
45
 
43
    return dist.sample()
44
45

test_sample_discrete_action 函数功能概述
离散动作空间的主函数，构建一个简单的策略网络，执行前向计算过程，并采样得到一组离散动作


xxxxxxxxxx
46
 
46
def test_sample_discrete_action():

设置相关参数 batch_size = 4, obs_shape = 10, action_shape = 6.


xxxxxxxxxx
47
 
47
    B, obs_shape, action_shape = 4, 10, 6

从0-1 的均匀分布中生成状态数据


xxxxxxxxxx
48
 
48
    state = torch.rand(B, obs_shape)

定义策略网络


xxxxxxxxxx
49
 
49
    policy_network = DiscretePolicyNetwork(obs_shape, action_shape)

策略网络执行前向计算，即输入状态输出 logit
$logit = \pi(a|s)$


xxxxxxxxxx
51
 
50
    logit = policy_network(state)
51
    assert logit.shape == (B, action_shape)

根据 logit 采样得到最终的离散动作


xxxxxxxxxx
55
 
52
    action = sample_action(logit)
53
    assert action.shape == (B, )
54
55

test_sample_multi_discrete_action 函数功能概述
多维离散动作空间的主函数，构建一个简单的策略网络，执行前向计算过程，并采样得到一组多维离散动作


xxxxxxxxxx
56
 
56
def test_sample_multi_discrete_action():

Set batch_size = 4, obs_shape = 10, action_shape = [4, 5, 6].


xxxxxxxxxx
57
 
57
    B, obs_shape, action_shape = 4, 10, [4, 5, 6]

从0-1 的均匀分布中生成状态数据


xxxxxxxxxx
58
 
58
    state = torch.rand(B, obs_shape)

定义策略网络


xxxxxxxxxx
59
 
59
    policy_network = MultiDiscretePolicyNetwork(obs_shape, action_shape)

策略网络执行前向计算，即输入状态输出多个 logit
$logit = \pi(a|s)$


xxxxxxxxxx
62
 
60
    logit = policy_network(state)
61
    for i in range(len(logit)):
62
        assert logit[i].shape == (B, action_shape[i])

对于多个 logit，依次调用采样函数，根据相应的 logit 采样得到最终的离散动作


xxxxxxxxxx
67
 
63
    for i in range(len(logit)):
64
        action_i = sample_action(logit[i])
65
        assert action_i.shape == (B, )
66
67

如果读者关于本文档有任何问题和建议，可以在 GitHub 提 issue 或是直接发邮件给我们 (opendilab@pjlab.org.cn) 。