4. PPO的批量大小与学习率调优

任务目标

在Walker2d-v4环境中，以PPO算法为基线，分别调整batch size（512、2048、8192）与学习率（1e-4、3e-4、1e-3），通过固定其他超参数进行对照实验，绘制学习曲线并分析两个超参数对训练稳定性和最终回报的影响。

环境与依赖

  • Python 3.9+，PyTorch 2.0+，gymnasium 0.29+，mujoco 2.3.7+
  • 安装：pip install gymnasium[mujoco] torch tensorboard
  • Walker2d-v4环境自带，无需额外数据文件

核心实现

以下代码实现了一个最小可复现的PPO训练器，支持通过命令行参数切换batch size与学习率。策略网络与价值网络共享一个两层MLP骨干，使用标准 clipped surrogate 目标与 GAE 优势估计。每次实验固定运行步数，记录每个评估点的平均回报，输出为CSV格式供后续绘图。

import argparse, csv, os
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import gymnasium as gym
from torch.distributions import Normal

class PolicyNet(nn.Module):
    def __init__(self, obs_dim, act_dim, hidden=256):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(obs_dim, hidden), nn.Tanh(),
            nn.Linear(hidden, hidden), nn.Tanh())
        self.mean_head = nn.Linear(hidden, act_dim)
        self.log_std = nn.Parameter(torch.zeros(act_dim))

    def forward(self, obs):
        h = self.fc(obs)
        mean = self.mean_head(h)
        std = torch.exp(self.log_std.clamp(-20, 2))
        return mean, std

    def get_dist(self, obs):
        mean, std = self.forward(obs)
        return Normal(mean, std)

class ValueNet(nn.Module):
    def __init__(self, obs_dim, hidden=256):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(obs_dim, hidden), nn.Tanh(),
            nn.Linear(hidden, hidden), nn.Tanh(),
            nn.Linear(hidden, 1))

    def forward(self, obs):
        return self.fc(obs).squeeze(-1)

def compute_gae(rewards, dones, values, next_value, gamma=0.99, lam=0.95):
    advantages = []
    gae = 0.0
    values = values + [next_value]
    for t in reversed(range(len(rewards))):
        delta = rewards[t] + gamma * values[t+1] * (1 - dones[t]) - values[t]
        gae = delta + gamma * lam * (1 - dones[t]) * gae
        advantages.insert(0, gae)
    returns = [adv + val for adv, val in zip(advantages, values[:-1])]
    return advantages, returns

def train_ppo(env, policy, value, policy_opt, value_opt, batch_size, lr, total_steps=1_000_000, eval_interval=50_000):
    obs_dim = env.observation_space.shape[0]
    act_dim = env.action_space.shape[0]
    obs, _ = env.reset(seed=0)
    ep_rewards = []
    ep_len = 0
    ep_ret = 0.0
    rollout = []
    global_step = 0
    eval_results = []

    while global_step < total_steps:
        obs_t = torch.FloatTensor(obs)
        with torch.no_grad():
            dist = policy.get_dist(obs_t)
            value_t = value(obs_t)
            action = dist.sample()
            log_prob = dist.log_prob(action).sum(-1)
        action_np = action.numpy()
        next_obs, reward, terminated, truncated, _ = env.step(action_np)
        done = terminated or truncated
        rollout.append((obs, action_np, reward, done, log_prob.item(), value_t.item()))
        ep_ret += reward
        ep_len += 1
        obs = next_obs
        global_step += 1

        if done:
            ep_rewards.append(ep_ret)
            obs, _ = env.reset()
            ep_ret = 0.0
            ep_len = 0

        if len(rollout) == batch_size:
            obs_b, act_b, rew_b, done_b, lp_b, val_b = map(list, zip(*rollout))
            obs_b = torch.FloatTensor(np.array(obs_b))
            act_b = torch.FloatTensor(np.array(act_b))
            rew_b = torch.FloatTensor(rew_b)
            done_b = torch.FloatTensor(done_b)
            lp_b = torch.FloatTensor(lp_b)
            val_b = torch.FloatTensor(val_b)

            with torch.no_grad():
                next_val = value(torch.FloatTensor(obs)).item()
            adv_b, ret_b = compute_gae(rew_b.tolist(), done_b.tolist(), val_b.tolist(), next_val)
            adv_b = torch.FloatTensor(adv_b)
            ret_b = torch.FloatTensor(ret_b)
            adv_b = (adv_b - adv_b.mean()) / (adv_b.std() + 1e-8)

            for _ in range(10):
                dist = policy.get_dist(obs_b)
                new_lp = dist.log_prob(act_b).sum(-1)
                ratio = torch.exp(new_lp - lp_b)
                surr1 = ratio * adv_b
                surr2 = torch.clamp(ratio, 0.8, 1.2) * adv_b
                policy_loss = -torch.min(surr1, surr2).mean()
                value_loss = nn.MSELoss()(value(obs_b), ret_b)

                policy_opt.zero_grad()
                policy_loss.backward()
                nn.utils.clip_grad_norm_(policy.parameters(), 0.5)
                policy_opt.step()

                value_opt.zero_grad()
                value_loss.backward()
                nn.utils.clip_grad_norm_(value.parameters(), 0.5)
                value_opt.step()

            rollout = []

        if global_step % eval_interval == 0:
            eval_rewards = []
            for _ in range(5):
                e_obs, _ = env.reset()
                e_ret = 0.0
                for _ in range(1000):
                    with torch.no_grad():
                        e_dist = policy.get_dist(torch.FloatTensor(e_obs))
                        e_act = e_dist.mean
                    e_obs, e_rew, e_done, e_trunc, _ = env.step(e_act.numpy())
                    e_ret += e_rew
                    if e_done or e_trunc:
                        break
                eval_rewards.append(e_ret)
            eval_results.append((global_step, np.mean(eval_rewards)))
            print(f"Step {global_step}: avg eval reward = {np.mean(eval_rewards):.1f}")

    return eval_results

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch_size", type=int, default=2048)
    parser.add_argument("--lr", type=float, default=3e-4)
    parser.add_argument("--tag", type=str, default="default")
    args = parser.parse_args()

    env = gym.make("Walker2d-v4")
    obs_dim = env.observation_space.shape[0]
    act_dim = env.action_space.shape[0]

    policy = PolicyNet(obs_dim, act_dim)
    value = ValueNet(obs_dim)
    policy_opt = optim.Adam(policy.parameters(), lr=args.lr)
    value_opt = optim.Adam(value.parameters(), lr=args.lr)

    results = train_ppo(env, policy, value, policy_opt, value_opt,
                        batch_size=args.batch_size, lr=args.lr)

    os.makedirs("results", exist_ok=True)
    csv_path = f"results/ppo_walker2d_bs{args.batch_size}_lr{args.lr}.csv"
    with open(csv_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["step", "avg_eval_reward"])
        writer.writerows(results)
    print(f"Saved to {csv_path}")

运行与实验设计

执行以下命令进行6组对照实验（每组约需20～40分钟，视GPU/CPU而定）：

# 固定lr=3e-4，调整batch size
python ppo_walker2d.py --batch_size 512  --lr 3e-4 --tag bs512
python ppo_walker2d.py --batch_size 2048 --lr 3e-4 --tag bs2048
python ppo_walker2d.py --batch_size 8192 --lr 3e-4 --tag bs8192

# 固定batch_size=2048，调整学习率
python ppo_walker2d.py --batch_size 2048 --lr 1e-4 --tag lr1e4
python ppo_walker2d.py --batch_size 2048 --lr 3e-4 --tag lr3e4
python ppo_walker2d.py --batch_size 2048 --lr 1e-3 --tag lr1e3

结果可视化与敏感性分析

运行以下脚本绘制学习曲线并输出敏感性汇总：

import pandas as pd, matplotlib.pyplot as plt, glob

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# 左图：batch size对比
for bs in [512, 2048, 8192]:
    df = pd.read_csv(f"results/ppo_walker2d_bs{bs}_lr0.0003.csv")
    axes[0].plot(df["step"], df["avg_eval_reward"], label=f"batch={bs}")
axes[0].set_xlabel("Steps"); axes[0].set_ylabel("Avg Eval Reward")
axes[0].set_title("Batch Size Sensitivity (lr=3e-4)")
axes[0].legend(); axes[0].grid(alpha=0.3)

# 右图：学习率对比
for lr in [1e-4, 3e-4, 1e-3]:
    df = pd.read_csv(f"results/ppo_walker2d_bs2048_lr{lr}.csv")
    axes[1].plot(df["step"], df["avg_eval_reward"], label=f"lr={lr}")
axes[1].set_xlabel("Steps"); axes[1].set_ylabel("Avg Eval Reward")
axes[1].set_title("Learning Rate Sensitivity (batch=2048)")
axes[1].legend(); axes[1].grid(alpha=0.3)

plt.tight_layout()
plt.savefig("ppo_sensitivity_curves.png", dpi=150)
plt.show()

# 输出最终性能汇总
print("\n=== Final Performance (last eval point) ===")
for f in sorted(glob.glob("results/*.csv")):
    df = pd.read_csv(f)
    final = df["avg_eval_reward"].iloc[-1]
    print(f"{f.split('/')[-1]}: {final:.1f}")

验证方法与判定标准

  1. 运行验证：训练脚本正常输出每个评估点的平均回报，并在results/目录生成对应CSV文件，说明训练流程闭环成功。
  2. 曲线有效性：绘制出的学习曲线应呈现上升趋势（Walker2d任务最终回报通常可达2000～4000+），若曲线完全平坦或持续下降，说明超参数组合失效或实现有误。
  3. 敏感性判定：
    • batch size：较小的batch（512）通常波动更大但前期提升快；较大的batch（8192）更新更稳定但样本效率降低。对比最终回报与曲线波动幅度即可判断敏感性。
    • 学习率：lr=1e-3可能导致训练不稳定（曲线剧烈震荡或发散）；lr=1e-4收敛偏慢；lr=3e-4通常为稳定基线。若三组最终性能差异超过20%，则判定为敏感超参数。
  4. 可复现性：固定seed=0，同一超参数组合重复运行两次，最终回报波动应在±15%以内（PPO本身存在随机性，该容差用于区分超参数效应与随机噪声）。
PPO超参数敏感性实验设计 固定 lr = 3e-4 batch = 512 batch = 2048 batch = 8192 观察要点 • 收敛速度差异 • 最终回报高低 • 曲线波动幅度 固定 batch = 2048 lr = 1e-4 lr = 3e-4 lr = 1e-3 观察要点 • 是否发散/震荡 • 收敛速度差异 • 最终性能对比 对比 → 判定敏感性

常见问题

  • MuJoCo导入错误：需安装系统级依赖，Ubuntu执行apt install libosmesa6-dev libgl1-mesa-dev，macOS需brew install mujoco。
  • 训练不收敛：先确认回报计算正确（Walker2d的terminated为True时不应加额外惩罚），再检查GAE中dones是否在终止步正确置1。
  • 显存不足：batch=8192时若GPU显存不够，可将rollout存储在CPU上，仅在更新时转移至GPU。
  • CSV文件缺失：确认results目录存在且训练脚本完整执行完毕（未提前中断）。
← 上一章 📖 返回目录 下一章 →