强化学习 4.策略梯度算法（实验部分）

发表于 2026-03-26 更新于 2026-04-10

前言

从零开始学习ai文章系列计划是个人在《动手学深度学习》和《磨菇书》两本书的学习中的个人笔记，文章也会以课本中的章节分开，即每个章节一片笔记。我会尽量的把主要内容以及遇到的难点进行记录与解决，如果哪里有错误的欢迎指正。或者不清晰的可以直接查看原文部分。

《蘑菇书》原文（课本）：https://datawhalechina.github.io/easy-rl/#/

（由于有时候公式太多，可能会直接贴图片）

蘑菇书的文章结构不会跟之前《动手学深度学习》按照原文章节进行，个人会适当调节。

1. REINFORCE：蒙特卡洛策略梯度

1.1 环境介绍

我们使用的环境是CartPole-v1。

状态空间：[ 小车位置，小车速度，杆子的角度，杆子的角速度 ]，4个连续值

动作空间： [左，右]，2个离散值

接下来我们创建环境并设置需要的超参数：

env = gym.make("CartPole-v1")
state_dim = env.observation_space.shape[0]  # 4
action_dim = env.action_space.n              # 2

policy = PolicyNet(state_dim, action_dim).to(device)

gamma = 0.99
num_episodes = 1000

接下来是我们的策略网络：

class PolicyNet(nn.Module):
    def __init__(self, state_dim, action_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, 128),
            nn.ReLU(),
            nn.Linear(128, action_dim)
        )
        #self.optimizer = optim.Adam(self.net.parameters(), lr=1e-3)
        self.optimizer = optim.Adam(policy.parameters(), lr=1e-3, weight_decay=1e-4) # 加上权重衰减后效果更稳定更好了

    def forward(self, x):
        logits = self.net(x)
        return torch.softmax(logits, dim=-1)

    def update(self,gamma,log_probs):
        def compute_returns(rewards, gamma):
            self.rewards = rewards
            G = 0
            returns = []
            for r in reversed(rewards):
                G = r + gamma * G
                returns.insert(0, G)

            self.returns = returns
            return torch.tensor(returns, dtype=torch.float32)

        def calculate_loss(log_probs, returns):
            loss = 0
            for log_prob, G in zip(log_probs, returns):
                loss += -log_prob * G
            return loss
        
        # 计算 G_t
        returns = compute_returns(self.rewards, gamma).to(device)
        # （可选但常用）回报归一化，降低方差
        # 高于平均则增加该动作概率，低于平均则减少该动作概率
        returns = (returns - returns.mean()) / (returns.std() + 1e-8)

        # REINFORCE loss
        loss = calculate_loss(log_probs,returns)

        torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

由于网络更新本身也是网络的功能之一，因此我写在了网络里面。

1.2 算法

流程代码：

def train():
    for episode in range(num_episodes):
        state, _ = env.reset()
        log_probs = []
        rewards = []

        # 1. 收集我们需要的log πθ(a|s)  和  rewards
        collect_data(policy,state, log_probs, rewards, env)

        # 2. 网络更新
        policy.update(gamma,log_probs,rewards)

        total_reward = sum(rewards)
        if episode % 50 == 0:
            print(f"Episode {episode}, Total Reward: {total_reward}")

    env.close()
    # 保存模型参数
    torch.save(policy.state_dict(), "policy_cartpole.pth")
    print("Policy saved to policy_cartpole.pth")

    return policy

简单来说，我们收集一条轨迹链的 log πθ(a|s) 和 rewards后，就利用这些数据进行一次网络的更新。

1.3 收集数据

我们传入2个需要的list：log_probs 和 rewards。对其直接的修改能影响到实际传入的list，所以不需要return。

def collect_data(policy,init_state,log_probs,rewards,env):
    state = init_state
    done = False
    while not done:
        state_tensor = torch.tensor(state, dtype=torch.float32).to(device)

        # πθ(a|s)  举例有3个动作，则probs=[0.2, 0.3, 0.5]
        probs = policy(state_tensor)
        # 按策略采样动作
        dist = torch.distributions.Categorical(probs)
        action = dist.sample()  # 举例，如果出来的是动作3，则 action = 2 (从0算起)
        # log πθ(a|s)   还是以动作3为例，传入action，会自动计算log(0.5)
        log_prob = dist.log_prob(action) 

        next_state, reward, terminated, truncated, _ = env.step(action.item())
        done = terminated or truncated

        log_probs.append(log_prob)
        rewards.append(reward)

        state = next_state

1.4 网络更新

其实也就是之前我们网络定义里面update的部分

def update(self,gamma,log_probs):
    def compute_returns(rewards, gamma):
        self.rewards = rewards
        G = 0
        returns = []
        for r in reversed(rewards):
            G = r + gamma * G
            returns.insert(0, G)

        self.returns = returns
        return torch.tensor(returns, dtype=torch.float32)

    def calculate_loss(log_probs, returns):
        loss = 0
        for log_prob, G in zip(log_probs, returns):
            loss += -log_prob * G
        return loss

    # 计算 G_t
    returns = compute_returns(self.rewards, gamma).to(device)
    # （可选但常用）回报归一化，降低方差
    # 高于平均则增加该动作概率，低于平均则减少该动作概率
    returns = (returns - returns.mean()) / (returns.std() + 1e-8)

    # REINFORCE loss
    loss = calculate_loss(log_probs,returns)


    self.optimizer.zero_grad()
    loss.backward()
            torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0) # 测试后发现裁剪没什么用，比原来效果还差了
    self.optimizer.step()

1.5 查看效果

main 函数和展示函数

def load_policy(model_path):
    state_dim = 4   # CartPole 固定
    action_dim = 2

    policy = PolicyNet(state_dim, action_dim)
    policy.load_state_dict(torch.load(model_path))
    policy.to(device)
    policy.eval()   # 重要：进入评估模式

    return policy

def render_policy(policy, episodes=10):
    env = gym.make("CartPole-v1", render_mode="human")

    for ep in range(episodes):
        state, _ = env.reset()
        done = False
        total_reward = 0

        while not done:
            state_tensor = torch.as_tensor(state, dtype=torch.float32, device=device)

            with torch.no_grad():
                probs = policy(state_tensor)
                action = torch.argmax(probs).item()

            state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            total_reward += reward

        print(f"Episode {ep}, Total Reward: {total_reward}")

    env.close()
    
if __name__ == "__main__":
    # ① 训练并保存
    policy = train()

    # ② 加载并展示
    # policy = load_policy("policy_cartpole.pth")
    # render_policy(policy)

对于权重衰减与梯度裁剪，可以参考我以前《动手学深度学习》的文章。

不启用权重衰减和梯度裁剪：

我们发现从650开始才完成目标

启用权重衰减：

我们发现从400轮开始就已经达到了目标

梯度裁剪

800轮才第一次成功，效果很难评价