我处于强化学习环境中，我的环境动作空间取决于状态。因此，在对行为动作进行采样时，我会执行以下过程：
(1) 为所有可能的动作生成概率逻辑
(2) 在这些 logits 上计算 softmax
(3) 屏蔽在此状态下无效的动作（通过乘以一个由 0 和 1 组成的向量），从而将一些概率归零
(4) 重新归一化有效动作概率
(5) 来自该分布的样本
这对于生成动作非常有效。但是，当我需要计算策略梯度更新时遇到了问题。通常执行以下操作：
tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=A)
其中 logits 是概率 logits，A 是采样动作。但是，由于我在执行 softmax之后执行此掩码/重新归一化，因此在我的情况下，上面的代码片段不是正确的交叉熵。我想知道是否有一种优雅的方式来处理这种情况。我的理解是，应该始终使用 tensorflow 的交叉熵计算来实现数值稳定性，但是我不确定如何正确结合这种掩蔽/重整化。