为何我的REINFORCE算法在CartPole环境中无法正常学习？

我在CartPole-v0环境中训练REINFORCE算法，考虑到这个环境特性简单，原本预期它能快速完成学习，但实际训练效果并未达到预期。

算法核心代码 #

for i in range(episodes):
    print("i = ", i)
    state = env.reset()
    done = False
    transitions = []

    tot_rewards = 0
    while not done:

        act_proba = model(torch.from_numpy(state))
        action = np.random.choice(np.array([0,1]), p = act_proba.data.numpy())
        next_state, reward, done, info = env.step(action)
        tot_rewards += 1
        transitions.append((state, action, tot_rewards))
        state = next_state


    if i%50==0:
        print("i = ", i, ",reward = ", tot_rewards)
    score.append(tot_rewards)
    reward_batch = torch.Tensor([r for (s,a,r) in transitions])
    disc_rewards = discount_rewards(reward_batch)
    nrml_disc_rewards = normalize_rewards(disc_rewards)
    state_batch = torch.Tensor([s for (s,a,r) in transitions])
    action_batch = torch.Tensor([a for (s,a,r) in transitions])
    pred_batch = model(state_batch)
    prob_batch = pred_batch.gather(dim=1, index=action_batch.long().view(-1, 1)).squeeze()
    loss = -(torch.sum(torch.log(prob_batch)*nrml_disc_rewards))
    opt.zero_grad()
    loss.backward()
    opt.step()

完整算法代码 #

import numpy as np
import gym
import torch
from torch import nn

env = gym.make('CartPole-v0')
learning_rate = 0.0001
episodes = 10000

def discount_rewards(reward, gamma = 0.99):
    return torch.pow(gamma, torch.arange(len(reward)))*reward
def normalize_rewards(disc_reward):
    return disc_reward/(disc_reward.max())

class NeuralNetwork(nn.Module):
    def __init__(self, state_size, action_size):
        super(NeuralNetwork, self).__init__()
        self.state_size = state_size
        self.action_size = action_size
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(state_size, 300),
            nn.ReLU(),
            nn.Linear(300, 128),
            nn.ReLU(),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Linear(128, action_size),
            nn.Softmax()
        )

    def forward(self,x):
        x = self.linear_relu_stack(x)
        return x

model = NeuralNetwork(env.observation_space.shape[0], env.action_space.n)
opt = torch.optim.Adam(params = model.parameters(), lr = learning_rate)
score = []
for i in range(episodes):
    print("i = ", i)
    state = env.reset()
    done = False
    transitions = []

    tot_rewards = 0
    while not done:

        act_proba = model(torch.from_numpy(state))
        action = np.random.choice(np.array([0,1]), p = act_proba.data.numpy())
        next_state, reward, done, info = env.step(action)
        tot_rewards += 1
        transitions.append((state, action, tot_rewards))
        state = next_state


    if i%50==0:
        print("i = ", i, ",reward = ", tot_rewards)
    score.append(tot_rewards)
    reward_batch = torch.Tensor([r for (s,a,r) in transitions])
    disc_rewards = discount_rewards(reward_batch)
    nrml_disc_rewards = normalize_rewards(disc_rewards)
    state_batch = torch.Tensor([s for (s,a,r) in transitions])
    action_batch = torch.Tensor([a for (s,a,r) in transitions])
    pred_batch = model(state_batch)
    prob_batch = pred_batch.gather(dim=1, index=action_batch.long().view(-1, 1)).squeeze()
    loss = -(torch.sum(torch.log(prob_batch)*nrml_disc_rewards))
    opt.zero_grad()
    loss.backward()
    opt.step()

问题分析与修改建议 #

你的代码存在几个关键问题，导致REINFORCE算法在CartPole环境中学习效率低下：

1. 折扣奖励计算逻辑错误
   当前discount_rewards函数用gamma^t * r_t直接计算，这不符合折扣奖励的定义——正确的折扣奖励应该从后往前累加，每一步的奖励是当前值加上未来折扣奖励的总和。当前实现会让早期步骤的收益信号被过度衰减，无法正确传递长期奖励。
   修改后的函数：

   def discount_rewards(rewards, gamma=0.99):
       discounted = torch.zeros_like(rewards)
       running_sum = 0
       for t in reversed(range(len(rewards))):
           running_sum = running_sum * gamma + rewards[t]
           discounted[t] = running_sum
       return discounted
   

2. 奖励归一化方式不合理
   用disc_reward/(disc_reward.max())的归一化方式会受极值影响严重，且无法将奖励中心化到0附近，容易导致训练震荡。正确做法是做标准化处理：

   def normalize_rewards(disc_rewards):
       mean = disc_rewards.mean()
       std = disc_rewards.std()
       # 防止除以0，添加极小值epsilon
       return (disc_rewards - mean) / (std + 1e-9)
   

3. 网络结构过于复杂
   CartPole是低维度简单环境，3层隐藏层（300→128→128）会增加训练难度和收敛时间。建议简化网络：

   self.linear_relu_stack = nn.Sequential(
       nn.Linear(state_size, 64),
       nn.ReLU(),
       nn.Linear(64, action_size),
       nn.Softmax(dim=-1)
   )
   

   同时Softmax层需指定dim=-1，避免批量输入时维度错误。

4. 学习率设置过低
   0.0001的学习率对于简单环境来说太小，会大幅拖慢收敛速度。建议调整到0.001或0.0005，Adam优化器在这个区间能更快收敛。

5. 奖励存储逻辑错误
   当前把累计步数tot_rewards作为奖励存入transitions，但CartPole每步实际奖励是1，直到结束。错误的奖励信号会误导算法学习，应改为存储环境返回的reward：

   transitions.append((state, action, reward))  # 替换原有的tot_rewards
   

   tot_rewards计算保持tot_rewards += reward即可（结果和之前一致，但逻辑正确）。

6. 数值稳定性问题
   当prob_batch接近0时，torch.log(prob_batch)会趋近负无穷，引发数值不稳定。可添加小epsilon缓解：

   loss = -(torch.sum(torch.log(prob_batch + 1e-9) * nrml_disc_rewards))
   

按以上修改后，REINFORCE算法应该能在几百到一千个episode内收敛到CartPole的最大奖励（200步）。

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