PyTorch实现REINFORCE算法时遇原地操作梯度计算错误求助

问题描述 #

在使用PyTorch实现REINFORCE强化学习算法时，运行以下代码出现错误：

import torch
import torch.optim as optim
from model_utils import Reinforce
from torch.distributions import Categorical


# train method
def reinforce(environment):
    learning_rate = 1e-3
    gamma = 0.99
    num_steps = 30
    max_episodes = 3000
    num_state_features = 21
    num_actions = 1000

    net = Reinforce(num_state_features, num_actions)
    optimizer = optim.Adam(net.parameters(), lr=learning_rate)
    net.train()

    for episode in range(max_episodes):
        state = torch.from_numpy(environment.reset())
        log_prob_seq = []
        reward_seq = []

        # generate trajectory
        for step in range(num_steps):
            policy_distro = net.forward(state)

            distro = Categorical(policy_distro)
            action = distro.sample()
            c = action.item()
            log_prob_seq.append(distro.log_prob(action))

            # compute reward, go to next state
            reward, new_state = environment.step(c)
            reward_seq.append(reward)
            new_state = torch.from_numpy(new_state)
            state = new_state

        # compute the return and loss
        loss = []
        returns = reward_seq.copy()
        for step in reversed(range(num_steps)):
            if step != num_steps - 1:
                returns[step] += gamma * returns[step + 1]
            loss.append(- (gamma ** step) * returns[step] * log_prob_seq[step])

        # update policy model parameters
        for step in range(num_steps):
            optimizer.zero_grad()
            loss[step].backward()
            optimizer.step()

错误提示：

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [128, 115]], which is output 0 of AsStridedBackward0, is at version 2; expected version 1 instead.

错误原因 #

核心问题是逐步更新模型参数的方式破坏了计算图的依赖关系：

• 每个loss[step]都依赖于当前episode开始时的模型参数状态
• 循环中每调用一次optimizer.step()会直接修改模型参数，导致后续loss[step]对应的计算图引用的参数版本不匹配，触发原地操作错误
• REINFORCE算法的标准逻辑是对整个轨迹的损失求和后，一次性完成反向传播与参数更新，而非逐步更新

解决方案 #

1. 合并损失，一次性更新参数 #

将所有时间步的损失累加为总损失，仅执行一次反向传播和优化器更新，保证计算图的连续性。修改后的核心代码如下：

# ... 其余代码保持不变 ...

# compute the return and loss
loss = []
returns = reward_seq.copy()
for step in reversed(range(num_steps)):
    if step != num_steps - 1:
        returns[step] += gamma * returns[step + 1]
    loss.append(- (gamma ** step) * returns[step] * log_prob_seq[step])

# 合并所有损失，一次性完成参数更新
optimizer.zero_grad()
total_loss = torch.stack(loss).sum()
total_loss.backward()
optimizer.step()

2. 排查模型中的原地操作 #

检查Reinforce模型的实现，避免使用带下划线的原地操作（如x.add_()、x.sigmoid_()、x += y），这类操作会直接修改张量、破坏计算图。替换为非原地操作：

• 把x += y改为x = x + y
• 把x.sigmoid_()改为x = x.sigmoid()

3. 统一设备（可选） #

如果模型运行在GPU上，确保状态张量也转移到对应设备，避免设备不匹配导致的隐性问题：

# 初始化模型后记录设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = Reinforce(num_state_features, num_actions).to(device)

# 加载状态时转移到对应设备
state = torch.from_numpy(environment.reset()).to(device)

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