深度强化学习中的动作屏蔽（Action Masking）

RLlib中的example有一个代码是action_masking，很感兴趣，所以学习了一下
主要功能是：

• “动作屏蔽”允许代理根据当前观察选择动作。这在许多实际场景中非常有用，在这些场景中，不同的时间步长可以执行不同的操作。
• 解释动作屏蔽的博客文章：https://boring-guy.sh/posts/masking-rl/ RLlib
• 支持动作屏蔽，即通过稍微调整环境和模型来禁止这些动作，如本示例所示。 在这里，ActionMaskEnv 包装了一个底层环境（这里是 RandomEnv），根据环境的观察仅将所有操作的子集定义为有效。如果选择了无效的操作，环境会引发错误 - 这绝不能发生！环境构造 Dict 观察，其中 obs[“observations”] 保存原始观察，obs[“action_mask”] 保存有效操作。
• 为了避免选择无效的操作，使用了ActionMaskModel。该模型采用原始观察结果，计算相应操作的逻辑，然后将所有无效操作的逻辑设置为零，从而禁用它们。
• 这仅适用于离散操作。

博客原文

简介

当我开始深度强化学习时，我工作的环境中的每个时间步都无法执行特定操作。

让我们具体说明不可能或不可用动作的概念：假设您想开发一个代理来玩马里奥赛车。接下来，假设代理有空库存（没有香蕉或任何东西）。代理无法执行“使用库存中的对象”操作。将代理限制为有意义的操作选择将使其能够以更智能的方式进行探索并输出更好的策略。

现在您了解了不可能或不可用操作的概念，自然的问题是：“我如何管理不可能的操作？” 我实施的第一个解决方案是，如果智能体采取不可能的操作，则分配负奖励。它的表现比不限制动作的选择，但我对这种方法不满意，因为它不能阻止代理选择不可能的动作。

然后我决定使用动作屏蔽（Action Masking）。这种方法实现起来简单且优雅，因为它限制代理只采取“有意义”的动作。

在我的深度强化学习实践中，我了解到有很多方法可以使用Masking。Masking可用于神经网络中的任何级别并用于不同的任务。不幸的是，除了 Costa Huang 的这篇精彩文章 [7] 之外，很少有强化学习的Masking实现可用。

这篇博文的范围是解释Masking的概念并通过图形和代码进行说明。事实上，这些Masking可以对我们在阅读这篇博文时看到的许多约束进行建模。请注意，整个过程是完全可微的。简而言之，Masking是为了简化您的生活。

要求

• 马尔可夫决策过程 (MDP) 的概念
• 策略梯度和 Q 学习算法的概念
• PyTorch 的一些知识或 numpy 的基础知识
• 自注意力的概念。

如果您想了解这个概念是什么，我邀请您阅读这篇解释 Transformer 的精彩文章 [6]

动作方面 Action level

概念：

深度强化学习中Masking的主要功能是过滤掉不可能或不可用的动作。例如，在《星际争霸 II》和《Dota 2》中，每个时间步的动作总数分别为$10^{26}$ 和$1,837,080$ 。然而，每个时间步的可能操作空间仅占可用操作空间的一小部分。因此，使用Masking有两个优点：

• 第一个是避免给环境带来无效的行为。
• 第二个是它是一种简单的方法，可以通过减少行动来管理广阔的空间。

图1说明了动作屏蔽的原理。其背后的想法很简单，它包括替换不可能的操作相关的 logits为$-\infty$。

那么，为什么应用这个掩码可以防止选择不可能的动作呢？
1. 基于价值的算法（Q-Learning）：
在基于价值的方法中，我们选择动作价值函数的最高估计值$Q(s,a)$：
$$a=\underset{a\in A}{\mathrm{argmax}}Q(s,.)$$
通过应用掩码，与不可能动作相关的 Q 值将等于$-\infty$ ，因此它们永远不会是最高值，因此永远不会被选择。

2. 基于策略的算法（策略梯度）：
在基于策略的方法中，我们根据模型输出的概率分布对动作进行采样：
$$a={\pi }_{\theta }(.|s)$$
因此，有必要将与不可能动作相关的概率设置为0。在我们使用Masking的时候不可能的动作是$-\infty$。我们使用 softmax 函数从 logits 转移到概率域：
$$\text{Softmax}(z{)}_i=\frac{e^{z_i}}{{\sum }_{j=1}^K{e}^{z_j}}\text{for}i=1,\dots ,K\text{and}z=(z_1,\dots ,z_K)\in {\mathbb{R}}^K.$$
考虑到我们已将与不可能动作相关的 logits 值设置为$-\infty$ ，对这些动作进行采样的概率等于0。

实现

现在让我们练习并实现离散动作空间和基于策略的算法的动作屏蔽。我使用 Costa Huang 的论文和动作屏蔽代码 [7] 作为起点。想法很简单，我们继承 PyTorch 的 Categorical 类并添加一个可选的 mask 参数。
当我们应用Masking时，我们会替换不可能动作的 logits。
然而，由于我们使用 float32，因此我们需要以 32 位表示的最小值。在 PyTorch 中，我们通过运行 torch.finfo(torch.float.dtype).min 来获取它，即 -3.40e+38。
最后，对于一些基于策略的方法，例如近端策略优化（PPO）[12]，有必要计算模型输出的概率分布熵。在我们的例子中，我们将仅计算可用操作的熵。

from typing import Optional

import torch
from torch.distributions.categorical import Categorical
from torch import einsum
from einops import  reduce


class CategoricalMasked(Categorical):
    def __init__(self, logits: torch.Tensor, mask: Optional[torch.Tensor] = None):
        self.mask = mask
        self.batch, self.nb_action = logits.size()
        if mask is None:
            super(CategoricalMasked, self).__init__(logits=logits)
        else:
            self.mask_value = torch.tensor(
                torch.finfo(logits.dtype).min, dtype=logits.dtype
            )
            logits = torch.where(self.mask, logits, self.mask_value)
            super(CategoricalMasked, self).__init__(logits=logits)

    def entropy(self):
        if self.mask is None:
            return super().entropy()
        # Elementwise multiplication
        p_log_p = einsum("ij,ij->ij", self.logits, self.probs)
        # Compute the entropy with possible action only
        p_log_p = torch.where(
            self.mask,
            p_log_p,
            torch.tensor(0, dtype=p_log_p.dtype, device=p_log_p.device),
        )
        return -reduce(p_log_p, "b a -> b", "sum", b=self.batch, a=self.nb_action)

以下代码块的目的是向您展示如何使用操作掩码。首先，我们创建虚拟逻辑和具有相同形状的虚拟蒙版。

logits_or_qvalues = torch.randn((2, 3), requires_grad=True) # batch size, nb action
print(logits_or_qvalues) 
# tensor([[-1.8222,  1.0769, -0.6567],
#         [-0.6729,  0.1665, -1.7856]])

mask = torch.zeros((2, 3), dtype=torch.bool) # batch size, nb action
mask[0][2] = True
mask[1][0] = True
mask[1][1] = True
print(mask) # False -> mask action 
# tensor([[False, False,  True],
#         [ True,  True, False]])

然后我们比较有和没有遮蔽的动作。

head = CategoricalMasked(logits=logits_or_qvalues)
print(head.probs) # Impossible action are not masked
# tensor([[0.0447, 0.8119, 0.1434], There remain 3 actions available
#         [0.2745, 0.6353, 0.0902]]) There remain 3 actions available

head_masked = CategoricalMasked(logits=logits_or_qvalues, mask=mask)
print(head_masked.probs) # Impossible action are  masked
# tensor([[0.0000, 0.0000, 1.0000], There remain 1 actions available
#         [0.3017, 0.6983, 0.0000]]) There remain 2 actions available

print(head.entropy())
# tensor([0.5867, 0.8601])

print(head_masked.entropy())
# tensor([-0.0000, 0.6123])

我们可以观察到，当我们应用掩码时，与不可能的动作相关的概率等于0 。因此，我们的智能体永远不会选择不可能的动作。
最后，当我们在熵计算中不包括不可能的动作时，我们就得到了一致的值。这种校正后的熵计算使代理能够仅在有效动作上最大化其探索。
这么酷的把戏！