Ray RLLib中带动作掩码的PPO策略梯度掩码化更新方案咨询

问题 #

在自定义环境中使用Ray RLLib训练PPO智能体，因动作空间存在大量非法动作，按照Ray文档实现了自定义ActionMaskModel来屏蔽非法动作，代码如下：

class ActionMaskModel(TFModelV2):
    """Model that handles simple discrete action masking.
    This assumes the outputs are logits for a single Categorical action dist.
    Getting this to work with a more complex output (e.g., if the action space
    is a tuple of several distributions) is also possible but left as an
    exercise to the reader.
    """

    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
        )

        super().__init__(obs_space, action_space, num_outputs, model_config, name)

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

        # disable action masking --> will likely lead to invalid actions
        self.no_masking = model_config["custom_model_config"].get("no_masking", False)

    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 = tf.maximum(tf.math.log(action_mask), tf.float32.min)
        masked_logits = logits + inf_mask

        # Return masked logits.
        return masked_logits, state

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

模型能成功屏蔽非法动作，但训练时出现KL散度非有限值的警告，尝试警告中的方法无效。查阅Huang等人2020年的论文后，确认问题根源是：动作采样使用掩码后的概率，但策略梯度更新却使用未掩码的概率，导致KL散度激增。需要解决如何确保策略梯度基于掩码后的概率更新。

解决方案 #

要让策略梯度完全基于掩码后的概率计算，需确保策略网络的前向传播全程使用掩码后的logits，同时PPO损失计算中，新旧策略的概率都基于掩码后的分布。具体修改如下：

1. 强化掩码逻辑的覆盖范围 #

当前代码仅在forward方法返回掩码后的logits，但需确保RLLib所有计算策略概率的环节（旧策略概率回溯、损失计算）都使用该结果。可以在forward中完成掩码后，不再保留未掩码logits的引用，强制后续操作依赖掩码后的输出。

2. 自定义支持掩码的动作分布（推荐） #

默认Categorical分布可能未正确处理掩码后的logits，自定义分布可确保所有概率计算都考虑掩码：

from ray.rllib.models.tf.tf_action_dist import Categorical
import tensorflow as tf

class MaskedCategorical(Categorical):
    def __init__(self, logits, action_mask=None, **kwargs):
        # 初始化时直接应用掩码到logits
        if action_mask is not None:
            inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min)
            logits = logits + inf_mask
        super().__init__(logits, **kwargs)

    def logp(self, x):
        # 确保对数概率基于掩码后的分布计算
        return super().logp(x)

    def kl(self, other):
        # 自定义KL散度计算，确保双方均基于掩码后的分布
        # 需保证other也是MaskedCategorical实例且使用相同action_mask
        return super().kl(other)

然后在ActionMaskModel的__init__方法中指定使用该分布：

self.action_dist_class = MaskedCategorical

3. 调整PPO配置确保掩码生效 #

在PPO训练配置中，明确指定自定义模型和分布，关闭可能绕过掩码的选项：

from ray.rllib.algorithms.ppo import PPOConfig
from ray.rllib.models import ModelCatalog

# 注册自定义模型和分布
ModelCatalog.register_custom_model("action_mask_model", ActionMaskModel)
ModelCatalog.register_custom_action_dist("masked_categorical", MaskedCategorical)

config = PPOConfig()
config.training(
    kl_coeff=0.2,
    kl_target=0.01,
)
config.model(
    custom_model="action_mask_model",
    custom_model_config={"no_masking": False},
    custom_action_dist="masked_categorical",
)

4. 验证掩码效果 #

训练时可添加临时日志，检查掩码前后的logits、概率值，确认KL散度计算基于掩码后的分布：

# 在ActionMaskModel的forward方法中添加
tf.print("Unmasked logits sample:", logits[0])
tf.print("Masked logits sample:", masked_logits[0])
tf.print("Action mask sample:", action_mask[0])

观察输出是否符合预期，确保掩码应用到所有概率计算环节。

内容的提问来源于stack exchange，提问作者Jakob Sejten