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 。

混合动作空间常常用于一些复杂的实践决策应用，例如《星际争霸2》和《王者荣耀》。它包含一系列可控制的决策变量，形式上可以用一棵树来表示，树的中间节点表示多级的离散选择，而叶节点则可以是任意的离散或连续动作空间。由于这种复杂的结构，混合动作空间需要更加复杂的算法设计和代码实现。

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

P.S, 如果需要安装 treetesor，你可以使用这样的命令
pip install DI-treetensor .

HybridPolicyNetwork 概述
定义 PPO 中所使用的混合动作空间（这里特指参数化动作空间）策略网络，其主要包含三部分：编码器（encoder），动作类型预测器（离散）和动作参数预测器（continuous）

from typing import Dict
import torch
import torch.nn as nn
import treetensor.torch as ttorch
from torch.distributions import Normal, Independent


class HybridPolicyNetwork(nn.Module):
    def __init__(self, obs_shape: int, action_shape: Dict[str, int]) -> None:

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

        super(HybridPolicyNetwork, self).__init__()

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

        self.encoder = nn.Sequential(
            nn.Linear(obs_shape, 16),
            nn.ReLU(),
            nn.Linear(16, 32),
            nn.ReLU(),
        )

定义离散动作类型预测器，输出离散动作的 logit
$$ y = Wx + b $$

        self.action_type_shape = action_shape['action_type_shape']
        self.action_type_head = nn.Linear(32, self.action_type_shape)

定义连续动作参数预测器，类似 PPO 在连续动作上的设计，输出相应的 mu 和 sigma
$$ \mu = Wx + b $$
$$\sigma = e^w$$

        self.action_args_shape = action_shape['action_args_shape']
        self.action_args_mu = nn.Linear(32, self.action_args_shape)
        self.action_args_log_sigma = nn.Parameter(torch.zeros(1, self.action_args_shape))

forward 函数功能概述
描述 PPO 中所使用的混合动作策略网络的前向计算图
x -> encoder -> action_type_head -> action_type_logit .
x -> encoder -> action_args_mu -> \mu .
action_args_log_sigma -> exp -> sigma .

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:

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

        x = self.encoder(x)

输出离散动作类型的 logit

        logit = self.action_type_head(x)

根据特征向量输出动作均值 mu

        mu = self.action_args_mu(x)

借助”广播“机制让对数空间标准差的维度和均值一致（在 batch 维度上复制
zeros_like 操作并不会传递梯度
Related Link

        log_sigma = self.action_args_log_sigma + torch.zeros_like(mu)

通过取指数操作得到最终的标准差 sigma
$$\sigma = e^w$$

        sigma = torch.exp(log_sigma)

返回 treetensor 类型的输出

        return ttorch.as_tensor({
            'action_type': logit,
            'action_args': {
                'mu': mu,
                'sigma': sigma
            }
        })

sample_hybrid_action 函数功能概述
输入 logit 采样得到混合动作，输入是一个 treetensor 包含 action_type 和 action_args 两部分

def sample_hybrid_action(logit: ttorch.Tensor) -> torch.Tensor:

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

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

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

    discrete_dist = torch.distributions.Categorical(probs=prob)

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

    action_type = discrete_dist.sample()

根据 mu 和 sigma 构建高斯分布
$$X \sim \mathcal{N}(\mu,\,\sigma^{2})$$
它的概率密度函数为: $$f(x) = \frac{1}{\sigma\sqrt{2\pi}} \exp\left( -\frac{1}{2}\left(\frac{x-\mu}{\sigma}\right)^{\!2}\,\right)$$
Related Link

    continuous_dist = Normal(logit.action_args.mu, logit.action_args.sigma)

将 action_shape 个高斯分布转义为一个有着对角协方差矩阵的多维高斯分布。
并保证高斯分布中，每一维之间都是互相独立的（因为协方差矩阵是对角矩阵）
Related Link

    continuous_dist = Independent(continuous_dist, 1)

为一个 batch 里的每个样本采样一个维度为 action_shape 的连续参数

    action_args = continuous_dist.sample()

返回最终的混合动作（参数化动作）

    return ttorch.as_tensor({
        'action_type': action_type,
        'action_args': action_args,
    })

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

def test_sample_hybrid_action():

设置相关参数 batch_size = 4, obs_shape = 10, action_shape 是一个字典，包含 3 种可选择的离散动作类型选项和 3 种对应的连续参数。动作类型和参数之间的关系将用下方的 mask 变量来表示

    B, obs_shape, action_shape = 4, 10, {'action_type_shape': 3, 'action_args_shape': 3}
    mask = [[0, 1, 0], [1, 0, 0], [0, 0, 1]]

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

    state = torch.rand(B, obs_shape)

定义策略网络，包含编码器，离散动作类型预测器和连续动作参数预测器

    policy_network = HybridPolicyNetwork(obs_shape, action_shape)

策略网络执行前向计算，即输入状态输出 treetensor 类型的 logit

    logit = policy_network(state)
    assert isinstance(logit, ttorch.Tensor)
    assert logit.action_type.shape == (B, action_shape['action_type_shape'])
    assert logit.action_args.mu.shape == (B, action_shape['action_args_shape'])
    assert logit.action_args.sigma.shape == (B, action_shape['action_args_shape'])

根据 logit 中的相关部分采样相应的动作部分

    action = sample_hybrid_action(logit)
    assert action.action_type.shape == (B, )
    assert action.action_args.shape == (B, action_shape['action_args_shape'])

通过动作类型查找获得每个数据样本具体所对应的 mask

    data_mask = torch.as_tensor([mask[i] for i in action.action_type]).bool()

用 mask 选择（过滤）出动作类型相对应的动作参数，并重新赋值

    filtered_action_args = ttorch.masked_select(action.action_args, data_mask)
    action.action_args = filtered_action_args
    assert action.action_args.shape == (B, )