REINFORCE算法实现求助：CartPole任务无性能提升

我正尝试按照Sutton与Barto所著《强化学习导论》中的内容实现REINFORCE算法，经过数小时调试仍无法定位问题，自我感觉实现逻辑正确。但在CartPole-v0环境中，智能体的每回合累计奖励未随训练回合数提升，恳请技术帮助。

#Import libraries
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import gym
import matplotlib.pyplot as plt
import numpy as np
from torch.distributions import Categorical

class Actor(nn.Module):
    
    """
    Use a ANN to construct a parameterized policy

    Input layer: observation space length
    First hidden layer: 128
    ReLU activation
    second hidden layer: 256
    ReLU activation
    Output layer: action space length
    Softmax activation
    """
    
    def __init__(self, state_size, action_size):
        super(Actor, self).__init__()
        self.state_size = state_size
        self.action_size = action_size
        self.linear1 = nn.Linear(self.state_size,64)
        self.linear2 = nn.Linear(64, self.action_size)
        
        
    def forward(self,state):
        x = F.relu(self.linear1(state))
        x = self.linear2(x)
        x = F.softmax(x, dim=-1)
        return x
    
    def select_action(self,state):
        
        #convert state to float tensor, add 1 dimension, allocate tensor on device
        state = torch.from_numpy(state).float().unsqueeze(0).to(device)
        action_probs = self.forward(state)
        #sample an action using the probability distribution
        dist = Categorical(action_probs)
        action = dist.sample()
        
        return action.item(), dist.log_prob(action)
    
    

device = "cuda" if torch.cuda.is_available() else "cpu" #Set device
episodes = 1000
step = 0
lr = 0.005
gamma = 0.9999
env = gym.make('CartPole-v0') # import cartpole environment
env.seed(0)
"""
The cartpole environment has an observation space of length 4
1) Cart Position, 2) Cart Velocity, 3) Pole Angle, 4) Pole Angular Velocity

The action space is of length two
1) Push cart to the left, Push cart to the right
"""
state_size = env.observation_space.shape[0]
action_size = env.action_space.n
actor = Actor(state_size, action_size).to(device)
optimizer_actor = optim.Adam(actor.parameters(), lr = lr) # import optimizer

for i in range(episodes):
    #Loop forever

    state = env.reset() # Initialize S
    cumulative_reward = 0
    done = False
    actions,states, rewards = [], [], []
    #print(1)
    while not done:
        # Generate episode S0, A0, R1,.. ,S_{T-1}, A_{T-1},RT, following pi(·|·, theta)
        
        #get action and log probability
        action, logp = actor.select_action(state)
        
        state, reward, done, _ = env.step(action) #take action in environment
        rewards.append(reward)
        actions.append(action)
        states.append(state)
        cumulative_reward +=reward
        #Store episode trajectory
        
    discounted_returns = []
    for t in range(len(rewards)):
        G = 0.0
        for k, r in enumerate(rewards[t:]):
            G += (gamma**k)*r
        discounted_returns.append(G)
    
    for state, action, G in zip(states, actions, discounted_returns):
        
        
        state = torch.from_numpy(state).float().unsqueeze(0).to(device)
        probs = actor(state)
        dist = Categorical(probs=probs)    
        action = torch.tensor(action, dtype=torch.int).to(device)
        logp = dist.log_prob(action)
        loss = - logp*G
        
        
        optimizer_actor.zero_grad()
        loss.backward()
        optimizer_actor.step()
        
    
    print('Episode {} Score: {:.2f}'.format(i, cumulative_reward))

你的代码存在几个关键问题，导致奖励无法提升：

1. 状态存储错误 #

回合循环中，你执行env.step(action)后才把状态存入states列表，但REINFORCE需要的是执行动作前的状态，而非动作后的状态。这会导致策略梯度计算时使用错误状态，完全偏离正确更新方向。

修复：在执行动作前保存当前状态

while not done:
    states.append(state)  # 先存动作前的状态
    action, logp = actor.select_action(state)
    state, reward, done, _ = env.step(action)
    rewards.append(reward)
    actions.append(action)
    cumulative_reward += reward

2. 梯度更新时机错误 #

当前对每个状态-动作对单独执行optimizer.step()，相当于单样本更新，训练稳定性极差。正确做法是累加整个回合的损失，再统一执行一次反向传播和优化。

修复后的更新逻辑：

total_loss = 0.0
for state, action, G in zip(states, actions, discounted_returns):
    state = torch.from_numpy(state).float().unsqueeze(0).to(device)
    probs = actor(state)
    dist = Categorical(probs=probs)    
    action = torch.tensor(action, dtype=torch.int64).to(device)  # 用int64避免类型错误
    logp = dist.log_prob(action)
    total_loss += -logp * G

optimizer_actor.zero_grad()
total_loss.backward()
optimizer_actor.step()

3. 超参数不合理 #

• lr=0.005过大，容易导致梯度震荡，建议调低至0.001或0.0005
• gamma=0.9999接近1，会让折扣回报的方差极大，建议改为0.99或0.95平衡偏差与方差

4. 折扣回报计算效率优化（非核心但建议） #

原嵌套循环计算折扣回报的时间复杂度为O(T²)，可改为反向遍历的O(T)方式：

discounted_returns = []
G = 0.0
for r in reversed(rewards):
    G = r + gamma * G
    discounted_returns.insert(0, G)

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import gym
import numpy as np
from torch.distributions import Categorical

class Actor(nn.Module):
    def __init__(self, state_size, action_size):
        super(Actor, self).__init__()
        self.state_size = state_size
        self.action_size = action_size
        self.linear1 = nn.Linear(self.state_size, 64)
        self.linear2 = nn.Linear(64, self.action_size)
        
    def forward(self, state):
        x = F.relu(self.linear1(state))
        x = self.linear2(x)
        x = F.softmax(x, dim=-1)
        return x
    
    def select_action(self, state):
        state = torch.from_numpy(state).float().unsqueeze(0).to(device)
        action_probs = self.forward(state)
        dist = Categorical(action_probs)
        action = dist.sample()
        return action.item(), dist.log_prob(action)

device = "cuda" if torch.cuda.is_available() else "cpu"
episodes = 1000
lr = 0.001  # 调低学习率
gamma = 0.99  # 调整gamma
env = gym.make('CartPole-v0')
env.seed(0)

state_size = env.observation_space.shape[0]
action_size = env.action_space.n
actor = Actor(state_size, action_size).to(device)
optimizer_actor = optim.Adam(actor.parameters(), lr=lr)

for i in range(episodes):
    state = env.reset()
    cumulative_reward = 0
    done = False
    actions, states, rewards = [], [], []
    
    while not done:
        states.append(state)  # 保存动作前的状态
        action, logp = actor.select_action(state)
        state, reward, done, _ = env.step(action)
        rewards.append(reward)
        actions.append(action)
        cumulative_reward += reward
    
    # 高效计算折扣回报
    discounted_returns = []
    G = 0.0
    for r in reversed(rewards):
        G = r + gamma * G
        discounted_returns.insert(0, G)
    
    # 累加整个回合的损失
    total_loss = 0.0
    for state, action, G in zip(states, actions, discounted_returns):
        state = torch.from_numpy(state).float().unsqueeze(0).to(device)
        probs = actor(state)
        dist = Categorical(probs=probs)
        action = torch.tensor(action, dtype=torch.int64).to(device)
        logp = dist.log_prob(action)
        total_loss += -logp * G
    
    # 统一更新梯度
    optimizer_actor.zero_grad()
    total_loss.backward()
    optimizer_actor.step()
    
    print('Episode {} Score: {:.2f}'.format(i, cumulative_reward))

修改后，智能体的回合奖励应随训练逐步提升，最终稳定达到CartPole-v0的最大奖励（200）。

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