REINFORCE实现中Policy Network批量与单状态输入输出不一致问题

你的问题核心是Policy Network中Softmax层的dim参数配置错误。

具体原因解释 #

• 逐次输入单个状态时：单状态输入张量形状为(4,)（CartPole状态维度为4），经过输出层Linear后得到形状为(2,)的张量，此时Softmax(dim=0)会对这2个值做归一化，刚好对应两个动作的概率，结果符合预期。
• 批量输入状态时：批量张量形状为(N, 4)（N为单episode的步数），经过输出层Linear后得到形状为(N, 2)的张量，此时Softmax(dim=0)会对**批次维度（即所有样本的第一个动作值、所有样本的第二个动作值分别）**做归一化，而非对每个样本的两个动作值做归一化，因此结果和逐次输入差异极大。

只需要把Policy Network中Softmax层的dim参数改为-1（表示对张量的最后一个维度做归一化）即可：

self.model = nn.Sequential(
    nn.Linear(state_dim, 64),
    nn.ReLU(),
    nn.Linear(64, n_actions),
    nn.Softmax(dim=-1) # 把原dim=0修改为dim=-1
).float()

修改后两种输入方式得到的动作概率会完全一致。

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