PyTorch REINFORCE实现中autograd如何跨非张量反向传播？

核心疑问 #

研究PyTorch官方REINFORCE实现时，对policy_loss.backward()的反向传播存在疑惑：从状态输入到policy_loss的过程中，存在policy.saved_log_probs列表、policy.rewards列表及计算得到的returns，这些看似非张量的变量为何没有阻断autograd的反向传播，反而能让梯度成功传递到Policy类的affine1线性层权重？

关键代码片段 #

第75行核心反向传播代码

policy_loss.backward()

policy_loss计算逻辑

for log_prob, R in zip(policy.saved_log_probs, returns):
    policy_loss.append(-log_prob * R)

Policy类相关属性定义

self.saved_log_probs = []
self.rewards = []

returns计算逻辑

for r in policy.rewards[::-1]:
    R = r + args.gamma * R
    returns.appendleft(R)

Policy类初始化

class Policy(nn.Module):
    def __init__(self):
        super(Policy, self).__init__()
        self.affine1 = nn.Linear(4, 128)

原理解析 #

• saved_log_probs存储的是带梯度的张量
  saved_log_probs列表里保存的并非普通数值，而是Policy模型前向传播时生成的log_prob张量（Tensor）。这些张量保留了完整的计算图信息——它们是从模型的线性层（affine1）、激活层，到最终输出动作概率的对数这一完整计算流程的产物，因此和模型的可训练参数（如affine1的权重、偏置）存在直接的梯度连接。

• returns和rewards是常数系数，不阻断梯度链
  rewards是环境返回的标量奖励值，计算returns时会基于折扣因子累加，最终得到的returns元素无论是普通数值还是被转为张量，在和log_prob相乘时都属于常数项（对于当前反向传播步骤来说，returns是基于历史交互数据计算的固定值）。PyTorch中，张量与常数相乘只会缩放张量的梯度，不会破坏原有的计算图连接，因此-log_prob * R的结果依然保留着从log_prob到模型参数的梯度路径。

• policy_loss是带梯度张量的聚合
  代码中policy_loss列表的元素都是带梯度的张量，后续会通过sum操作聚合为一个标量张量。调用backward()时，autograd会沿着每个log_prob的计算图反向追溯，将梯度传递到模型的所有可训练参数，自然包括affine1线性层的权重。

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