1. PPO算法核心组件实现

PPO算法核心组件实现 学习卡片

任务目标

实现PPO(Proximal Policy Optimization)算法的三个核心组件:策略网络(Actor)、价值网络(Critic)和经验缓冲区(Replay Buffer),并通过一个简单的随机环境验证组件的前向传播、经验存储与采样的基本功能是否正常。

环境依赖

  • Python 3.8+
  • PyTorch 1.12+
  • NumPy

核心实现

import torch
import torch.nn as nn
import numpy as np
from collections import deque
import random

# ---------- 策略网络(Actor) ----------
class PolicyNetwork(nn.Module):
    """输出动作概率分布参数(连续动作:均值+标准差)"""
    def __init__(self, state_dim, action_dim, hidden_dim=64):
        super().__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.mean_head = nn.Linear(hidden_dim, action_dim)
        self.log_std = nn.Parameter(torch.zeros(action_dim))

    def forward(self, state):
        x = torch.relu(self.fc1(state))
        x = torch.relu(self.fc2(x))
        mean = self.mean_head(x)
        std = torch.exp(self.log_std)
        return mean, std

    def sample_action(self, state):
        """从高斯分布采样动作,返回动作和对数概率"""
        mean, std = self.forward(state)
        dist = torch.distributions.Normal(mean, std)
        action = dist.sample()
        log_prob = dist.log_prob(action).sum(dim=-1)
        return action, log_prob

    def evaluate(self, state, action):
        """计算给定动作的对数概率和熵(用于PPO更新)"""
        mean, std = self.forward(state)
        dist = torch.distributions.Normal(mean, std)
        log_prob = dist.log_prob(action).sum(dim=-1)
        entropy = dist.entropy().sum(dim=-1)
        return log_prob, entropy

# ---------- 价值网络(Critic) ----------
class ValueNetwork(nn.Module):
    """输出状态价值V(s)"""
    def __init__(self, state_dim, hidden_dim=64):
        super().__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.value_head = nn.Linear(hidden_dim, 1)

    def forward(self, state):
        x = torch.relu(self.fc1(state))
        x = torch.relu(self.fc2(x))
        return self.value_head(x)

# ---------- 经验缓冲区 ----------
class RolloutBuffer:
    """存储一条轨迹的经验,用于PPO按epoch更新"""
    def __init__(self, capacity=2048):
        self.capacity = capacity
        self.clear()

    def clear(self):
        self.states = []
        self.actions = []
        self.log_probs = []
        self.rewards = []
        self.dones = []
        self.values = []

    def store(self, state, action, log_prob, reward, done, value):
        """存储单步经验"""
        self.states.append(state)
        self.actions.append(action)
        self.log_probs.append(log_prob)
        self.rewards.append(reward)
        self.dones.append(done)
        self.values.append(value)

    def is_full(self):
        return len(self.states) >= self.capacity

    def get_batch(self):
        """返回整批经验并清空缓冲区"""
        batch = {
            'states': torch.stack(self.states),
            'actions': torch.stack(self.actions),
            'log_probs': torch.stack(self.log_probs),
            'rewards': torch.tensor(self.rewards, dtype=torch.float32),
            'dones': torch.tensor(self.dones, dtype=torch.float32),
            'values': torch.stack(self.values).squeeze(-1)
        }
        self.clear()
        return batch

# ---------- 功能验证 ----------
def verify_components():
    state_dim, action_dim = 4, 2
    actor = PolicyNetwork(state_dim, action_dim)
    critic = ValueNetwork(state_dim)
    buffer = RolloutBuffer(capacity=8)

    # 1. 验证前向传播
    test_state = torch.randn(3, state_dim)
    mean, std = actor(test_state)
    value = critic(test_state)
    assert mean.shape == (3, action_dim), f"均值形状错误: {mean.shape}"
    assert std.shape == (action_dim,), f"标准差形状错误: {std.shape}"
    assert value.shape == (3, 1), f"价值形状错误: {value.shape}"
    print("[PASS] 策略网络和价值网络前向传播正常")

    # 2. 验证采样与评估
    action, log_prob = actor.sample_action(test_state)
    assert action.shape == (3, action_dim)
    assert log_prob.shape == (3,)
    eval_log_prob, entropy = actor.evaluate(test_state, action)
    assert torch.allclose(log_prob, eval_log_prob, atol=1e-5), "采样与评估对数概率不一致"
    print("[PASS] 动作采样与对数概率计算一致")

    # 3. 验证经验缓冲区
    for _ in range(8):
        s = torch.randn(state_dim)
        a, lp = actor.sample_action(s.unsqueeze(0))
        v = critic(s.unsqueeze(0))
        buffer.store(s, a.squeeze(0), lp.squeeze(0), 1.0, False, v.squeeze(0))

    assert buffer.is_full(), "缓冲区应已满"
    batch = buffer.get_batch()
    assert batch['states'].shape == (8, state_dim)
    assert batch['actions'].shape == (8, action_dim)
    assert batch['rewards'].shape == (8,)
    assert len(buffer.states) == 0, "取批后缓冲区应清空"
    print("[PASS] 经验缓冲区存储与取批正常")

    # 4. 验证PPO更新所需数据完整性
    assert 'log_probs' in batch and 'values' in batch
    print("[PASS] 所有PPO更新所需数据字段完整")
    print("\n全部验证通过:PPO核心组件功能正常")

if __name__ == "__main__":
    verify_components()

运行方式

python ppo_components.py

验证方法与判定标准

  • 前向传播验证:策略网络输出动作均值形状为 (batch_size, action_dim),标准差形状为 (action_dim,);价值网络输出形状为 (batch_size, 1)。若形状不匹配,程序会抛出断言错误。
  • 采样一致性验证:通过 sample_action 采样得到的动作,再用 evaluate 计算对数概率,应与采样时的对数概率一致(误差 < 1e-5)。这确保PPO更新时能正确计算新旧策略的比率。
  • 缓冲区功能验证:存入8条经验后缓冲区应报告已满;调用 get_batch() 后返回包含 states、actions、log_probs、rewards、dones、values 六个字段的字典,且缓冲区被清空。
  • 通过条件:程序运行后输出 [PASS] 标记的四行信息,最后一行显示 全部验证通过:PPO核心组件功能正常,即表示所有组件实现正确。

关键说明

  • 本实现针对连续动作空间设计,策略网络输出高斯分布的均值和标准差;若处理离散动作,需将输出层改为 logits 并使用分类分布。
  • 经验缓冲区采用整批更新方式(而非随机小批量采样),这是PPO的标准做法——收集一条完整轨迹后,用同一批数据做多轮epoch更新。
  • 价值网络输出 V(s),用于计算优势函数(GAE),本案例未包含优势计算和策略更新循环,这些属于PPO训练主循环的内容,将在后续案例中实现。