Tony Wang's blogs

背景

在强化学习解决问题的场景中，动作是体现学习效果最直接的因素，直接影响了智能体下一步的走向和对环境状态的改变。在应用强化学习解决实际问题时，往往不同于gym库中倒立摆那样的情况，而是存在很多的约束。例如，在t时刻智能体可选的动作为1,2,3，但是在t+1时刻只能选1,2.3处于不可用的状态。在这种情况下，就需要借助掩码mask来对智能体的动作进行处理。

有人会疑问：就不能制定相应的奖励函数使得智能体学习到这种约束吗？这样做是可以的，但是付出的训练代价很大，并且极其容易导致模型发散。因此，在大多数RL落地的场景下，都会使用MASK掩码方法解决动作约束的问题。

MASK的方法

Mask的核心就是在输出的动作或者值函数的向量上戴个“面具”，点乘一个{0,1}或者{−∞,1}的行向量，以规范化输出。这样智能体选出的动作就可以进行简单的规范化。

MASK的两个关键点

由于强化学习，尤其是深度强化学习，学的最后还是分布，因此只是单单的不让智能体选择不符合规则的动作并不能加速模型的收敛。

因此，MASK一般加在选择动作前的值函数向量或者其他数据向量上，并且会将MASK后的值传入神经网络训练。
两个关键点分别是：

1-mask分布

2-回传训练

具体做法

以openai中MASK星际争霸智能体的动作为例：首先是环境部分self.env，使用的是为每个agent提供一个available的动作集合，可以随时调用这个方法以获取agent此时的可执行动作：

然后在agent的动作选择阶段，使用inf代替不符合要求的部分，使得softmax选择的动作合理：

最后在policy学习更新的部分，同样利用-9999999作为不合理动作的替换，使得反向传播的概率分布与采样一致：

在星际争霸游戏中，任何时刻，整个动作空间中只有一小部分子集的动作可以执行。为了防止 AI 在某些时刻选取当前时刻无法执行的动作，需要对动作空间进行 mask。具体操作时，如果选择了当前时刻不可用的动作，就会执行 no-op（no operation，即不操作）

实现

第一步，自定义环境：

class MyParamActionEnv(gym.Env):
    def __init__(self, max_avail_actions):
        self.action_space = Discrete(max_avail_actions)
        self.observation_space = Dict({
            "action_mask": Box(0, 1, shape=(max_avail_actions, )), # 添加action_mask 尺寸与action_space一致
            "avail_actions": Box(-1, 1, shape=(max_avail_actions, action_embedding_sz)),
            "real_obs": ...,
        })

第二步，自定义网络：

class ParametricActionsModel(TFModelV2):
    def __init__(self,
                 obs_space,
                 action_space,
                 num_outputs,
                 model_config,
                 name,
                 true_obs_shape=(4,),
                 action_embed_size=2):
        super(ParametricActionsModel, self).__init__(
            obs_space, action_space, num_outputs, model_config, name)
        self.action_embed_model = FullyConnectedNetwork(...)

    def forward(self, input_dict, state, seq_lens):
        # Extract the available actions tensor from the observation.
        avail_actions = input_dict["obs"]["avail_actions"]
        action_mask = input_dict["obs"]["action_mask"]

        # Compute the predicted action embedding
        action_embed, _ = self.action_embed_model({
            "obs": input_dict["obs"]["cart"]
        })

        # Expand the model output to [BATCH, 1, EMBED_SIZE]. Note that the
        # avail actions tensor is of shape [BATCH, MAX_ACTIONS, EMBED_SIZE].
        intent_vector = tf.expand_dims(action_embed, 1)

        # Batch dot product => shape of logits is [BATCH, MAX_ACTIONS].
        action_logits = tf.reduce_sum(avail_actions * intent_vector, axis=2)

        # Mask out invalid actions (use tf.float32.min for stability)
        inf_mask = tf.maximum(tf.log(action_mask), tf.float32.min)
        return action_logits + inf_mask, state

参考例子：

第一步：自定义环境：

class ActionMaskEnv(RandomEnv):
    """A randomly acting environment that publishes an action-mask each step."""

    def __init__(self, config):
        super().__init__(config)
        # Masking only works for Discrete actions.
        assert isinstance(self.action_space, Discrete)
        # Add action_mask to observations.
        self.observation_space = Dict(
            {
                "action_mask": Box(0.0, 1.0, shape=(self.action_space.n,)),
                "observations": self.observation_space,
            }
        )
        self.valid_actions = None

    def reset(self, *, seed=None, options=None):
        obs, info = super().reset()
        self._fix_action_mask(obs)
        return obs, info

    def step(self, action):
        # Check whether action is valid.
        if not self.valid_actions[action]:
            raise ValueError(
                f"Invalid action sent to env! " f"valid_actions={self.valid_actions}"
            )
        obs, rew, done, truncated, info = super().step(action)
        self._fix_action_mask(obs)
        return obs, rew, done, truncated, info

    def _fix_action_mask(self, obs):
        # Fix action-mask: Everything larger 0.5 is 1.0, everything else 0.0.
        self.valid_actions = np.round(obs["action_mask"])
        obs["action_mask"] = self.valid_actions

第二步：自定义网络：

class TorchActionMaskModel(TorchModelV2, nn.Module):
    """PyTorch version of above ActionMaskingModel."""

    def __init__(
        self,
        obs_space,
        action_space,
        num_outputs,
        model_config,
        name,
        **kwargs,
    ):
        orig_space = getattr(obs_space, "original_space", obs_space)
        assert (
            isinstance(orig_space, Dict)
            and "action_mask" in orig_space.spaces
            and "observations" in orig_space.spaces
        )

        TorchModelV2.__init__(
            self, obs_space, action_space, num_outputs, model_config, name, **kwargs
        )
        nn.Module.__init__(self)

        self.internal_model = TorchFC(
            orig_space["observations"],
            action_space,
            num_outputs,
            model_config,
            name + "_internal",
        )

        # disable action masking --> will likely lead to invalid actions
        self.no_masking = False
        if "no_masking" in model_config["custom_model_config"]:
            self.no_masking = model_config["custom_model_config"]["no_masking"]

    def forward(self, input_dict, state, seq_lens):
        # Extract the available actions tensor from the observation.
        action_mask = input_dict["obs"]["action_mask"]

        # Compute the unmasked logits.
        logits, _ = self.internal_model({"obs": input_dict["obs"]["observations"]})

        # If action masking is disabled, directly return unmasked logits
        if self.no_masking:
            return logits, state

        # Convert action_mask into a [0.0 || -inf]-type mask.
        inf_mask = torch.clamp(torch.log(action_mask), min=FLOAT_MIN)
        masked_logits = logits + inf_mask

        # Return masked logits.
        return masked_logits, state

    def value_function(self):
        return self.internal_model.value_function()

torch.clamp 将输入input张量每个元素的夹紧到区间 [min,max]

inf_mask趋近于负无穷，使用inf代替不符合要求的部分，使得softmax选择的动作合理

forward中包含一个batchsize内的所有数据的输入，Discrete(100)时，包含0-99的每个action取值的概率。