PyTorch实现REINFORCE时Softmax批推理输出求和不等于1问题求助

输出概率和不为1的问题完全来自Softmax层的维度参数配置错误，和批推理逻辑、模型初始化无关。

错配的具体逻辑 #

PyTorch批处理的张量默认维度规则是：模型输入的第0维是批大小（样本数量），后续维度是每个样本的特征维度。
对应你的模型计算流程：

• 输入批数据state_batch的形状为(3, 9)，即3个独立观测样本，每个样本9维观测值
• 经过最后一层线性层后，输出张量形状为(3, 2)，即每个样本对应2个动作的原始打分值
• 你当前设置的Softmax(dim=0)，是沿着**第0维（样本数量方向）**做归一化：也就是把所有样本的同一个动作的打分值放在一起算Softmax，完全没有在单个样本的维度上做概率归一化

两种推理场景的结果差异来源 #

• 单样本循环推理时，每次输入的单个样本形状是(9,)，经过最后一层线性层输出形状为(2,)的一维动作打分向量，这时候dim=0正好是沿着这个动作向量做归一化，自然单个样本的两个动作概率和为1，结果符合预期
• 整批推理时，输出是(3,2)的二维矩阵，dim=0的Softmax会跨3个样本做列方向的归一化，算出来的数值根本不是单个样本的动作概率分布，自然每行求和远不等于1

修正方案 #

把Softmax的维度参数改成dim=1，明确沿着每个样本的动作输出维度做归一化即可，修正后的模型定义如下：

model = torch.nn.Sequential(
        torch.nn.Linear(obs_size, HIDDEN_SIZE),
        torch.nn.ReLU(),
        torch.nn.Linear(HIDDEN_SIZE, n_actions),
        torch.nn.Softmax(dim=1)
    )

修改后整批推理的输出，在浮点精度范围内每个样本的概率和都会严格等于1，和单样本推理结果完全一致。

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