张量流中交叉熵计算之前的掩蔽概率

Question

我处于强化学习环境中，我的环境动作空间取决于状态。因此，我在对行为动作进行采样时会经历以下过程：

(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 的交叉熵计算，但是我不确定如何正确地合并这种掩蔽/重整化。