RL_TRPO

这个项目实现了连续动作空间上的 PPO 和 TRPO,并针对训练工程侧做了三类增强:

  • 自动选择 GPU 进行前向与反向传播
  • 使用 Gymnasium 向量环境进行并行采样,支持多核 CPU
  • 使用 tqdm 替代周期性 print,实时展示 episode 进度、滚动奖励和优化指标

当前默认实验环境是 Pendulum-v1,用来对比 PPO 和 TRPO 在同一任务上的训练稳定性、优化代价和收敛表现。

1. 项目结构

RL_TRPO/
├── agent/
│   ├── ppo.py
│   └── trpo.py
├── environment.yml
├── main.py
├── networks.py
├── requirements.txt
└── README.md

核心职责如下:

  • main.py: 训练入口、设备选择、并行环境创建、tqdm 日志、结果绘图与指标导出
  • agent/ppo.py: PPO-Clip 算法实现
  • agent/trpo.py: TRPO 算法实现,包括共轭梯度和线搜索
  • networks.py: Actor/Critic 网络定义

2. 功能特性

2.1 GPU 训练

如果机器上有可用 CUDA,程序默认会自动选择 GPU:

python main.py --device auto

也可以显式指定:

python main.py --device cuda
python main.py --device cpu

注意:环境仿真本身仍然主要发生在 CPU 上,GPU 主要加速以下部分:

  • Actor 前向推理
  • Critic 前向推理
  • PPO / TRPO 更新中的反向传播
  • TRPO 中的 Hessian-vector product 等二阶近似计算

因此,单环境串行采样时 GPU 利用率不会特别高。真正想把吞吐拉起来,需要同时打开并行采样。

2.2 多核并行采样

训练入口支持 Gymnasium 的向量环境:

  • SyncVectorEnv: 单进程同步向量化
  • AsyncVectorEnv: 多进程异步采样,更适合多核 CPU

例如,开启 8 个并行环境:

python main.py --num-envs 8 --vector-mode async

这会把环境交互压力分散到多个 CPU 核心上,让策略更新前的数据收集更快。

这里有一个很关键的语义:

  • batch-size 表示总环境步数,不是单个子环境的步数
  • 例如 num-envs=8 且 batch-size=2000,意味着每次 update 前大约只会为每个子环境收集 250 步

如果你想保持和单环境训练接近的数据长度,可以按 num-envs 近似线性放大 batch-size。例如原来单环境用 2000,那么 8 环境时可以先尝试 16000。

2.3 tqdm 进度日志

训练时不再每 10 个 episode 打印一次,而是通过进度条持续显示:

  • 当前完成 episode 数
  • 最近 20 个 episode 的平均回报
  • 已采样环境步数
  • 已做的 update 次数
  • PPO 的策略损失、价值损失、熵
  • TRPO 的 surrogate loss、KL、critic loss、line search 是否成功

3. 安装

3.1 使用 Conda

conda env create -f environment.yml
conda activate rl_trpo

3.2 使用 pip

pip install -r requirements.txt

3.3 GPU 版本 PyTorch

requirements.txt 里只写了通用的 torch 依赖。若你想确保安装 CUDA 版本,请优先按照 PyTorch 官方安装命令安装。例如在 Linux + CUDA 12.x 场景下,通常应该使用官方提供的对应索引源或 Conda 渠道。

安装完成后可以验证:

python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

如果输出里的 torch.cuda.is_available() 为 True,就说明代码里的 GPU 路径会生效。

4. 快速开始

4.1 默认运行

python main.py

默认行为:

  • 环境: Pendulum-v1
  • episode 数: 500
  • batch size: 2000
  • 设备: auto
  • 并行环境数: min(4, CPU 核数)
  • 向量环境模式: async

4.2 典型高吞吐训练命令

python main.py \
  --device cuda \
  --num-envs 8 \
  --vector-mode async \
  --cpu-threads 8 \
  --num-episodes 500 \
  --batch-size 16000 \
  --output-dir outputs_gpu

建议理解:

  • num-envs 决定并行采样宽度
  • vector-mode async 会启动多进程,通常更能吃满多核 CPU
  • cpu-threads 控制训练进程里 PyTorch 的 CPU 线程数
  • batch-size 是总环境步数。并行环境变多后,如果不相应增大 batch-size,每个子环境能提供的时间长度会变短,优势估计会更依赖 bootstrap
  • batch-size 越大,单次更新更稳定,但每次更新等待的采样时间也更长

4.3 纯 CPU 基线

python main.py --device cpu --num-envs 1 --vector-mode sync

4.4 不弹图,只保存结果

python main.py --no-show --output-dir outputs

输出文件包括:

  • outputs/ppo_vs_trpo_comparison.png
  • outputs/training_metrics.npz

5. 命令行参数

参数 说明 默认值
--env-name Gymnasium 环境名称 Pendulum-v1
--num-episodes 每个算法训练的 episode 数 500
--batch-size 每次更新前采样的总环境步数 2000
--num-envs 并行环境数量 min(4, CPU 核数)
--vector-mode sync 或 async async
--cpu-threads PyTorch CPU 线程数 min(8, CPU 核数)
--device auto / cpu / cuda auto
--seed 随机种子 42
--hidden-dim 隐层维度 128
--output-dir 输出目录 outputs
--no-show 只保存图,不弹窗 False

6. 算法背景

6.1 强化学习目标

我们考虑折扣回报最大化问题:


J(\pi_\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\left[\sum_{t=0}^{T} \gamma^t r_t\right]

其中:

  • \pi_\theta(a \mid s) 是参数化策略
  • \tau = (s_0, a_0, s_1, a_1, \dots) 是轨迹
  • \gamma \in (0, 1) 是折扣因子

策略梯度定理给出:


\nabla_\theta J(\pi_\theta) = \mathbb{E}_{s_t, a_t \sim \pi_\theta}\left[\nabla_\theta \log \pi_\theta(a_t \mid s_t) \cdot A^{\pi}(s_t, a_t)\right]

这里的关键就是如何稳定地估计优势函数 $A^{\pi}(s_t, a_t)$,以及如何限制策略更新幅度,避免一步走太猛。

6.2 GAE: 广义优势估计

项目中的 PPO 和 TRPO 都使用 GAE


\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)

\hat{A}_t = \delta_t + (\gamma \lambda)\delta_{t+1} + (\gamma \lambda)^2\delta_{t+2} + \cdots

也可以写成:


\hat{A}_t = \sum_{l=0}^{\infty}(\gamma \lambda)^l \delta_{t+l}

对应回报目标为:


\hat{R}_t = \hat{A}_t + V(s_t)

GAE 在偏差和方差之间做折中:

  • \lambda 越大,估计越接近 Monte Carlo,方差更大
  • \lambda 越小,估计更依赖 bootstrap,偏差更大

7. PPO 原理

7.1 PPO 的核心思想

PPO 希望在不显式求解复杂约束优化问题的前提下,限制新旧策略差异。它定义重要性采样比率:


r_t(\theta) = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{\theta_{\text{old}}}(a_t \mid s_t)}

如果直接最大化 $r_t(\theta)\hat{A}_t$,更新可能过大。PPO-Clip 使用裁剪目标:


L^{\text{CLIP}}(\theta) = \mathbb{E}_t\left[
\min\left(
r_t(\theta)\hat{A}_t,
\operatorname{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon)\hat{A}_t
\right)
\right]

它的含义是:

  • 当更新幅度还在可信范围内时,正常优化
  • 当比率偏离过大时,直接截断收益,防止目标函数鼓励过激更新

7.2 PPO 的工程特点

  • 一阶方法,易实现
  • 可配合 Adam 稳定训练
  • 可对同一批数据做多轮 epoch 更新
  • 常常是工业与研究中最常用的 on-policy baseline 之一

本项目中的 PPO 实现还包括:

  • mini-batch 更新
  • 梯度裁剪
  • advantage 标准化

8. TRPO 原理

8.1 TRPO 的核心思想

TRPO 直接把策略更新写成一个带 KL 约束的优化问题:


\max_\theta \; \mathbb{E}_t\left[
\frac{\pi_\theta(a_t \mid s_t)}{\pi_{\theta_{\text{old}}}(a_t \mid s_t)} \hat{A}_t
\right]

满足:


\mathbb{E}_t\left[D_{\mathrm{KL}}\left(\pi_{\theta_{\text{old}}}(\cdot \mid s_t) \;\|\; \pi_\theta(\cdot \mid s_t)\right)\right] \leq \delta

其中 \delta 是 trust region 半径。

8.2 二阶近似与自然梯度

TRPO 会对目标函数做一阶近似,对 KL 约束做二阶近似。最终可得到类似自然梯度方向的更新:


\theta_{k+1} = \theta_k + \alpha x

其中方向 x 通过求解线性系统获得:


Hx = g

这里:

  • g 是 surrogate objective 对参数的梯度
  • H 是 KL 关于参数的 Hessian,实践中通常用 Fisher 信息矩阵近似

由于 H 很大,TRPO 不直接显式构造 Hessian,而是通过 Hessian-vector product 配合共轭梯度法近似求解。

8.3 线搜索

得到方向后,TRPO 还会做回溯线搜索,保证:

  • surrogate objective 没有退化
  • KL 没有超过阈值

这也是 TRPO 稳定但实现复杂、单次更新较重的关键原因。

9. PPO 与 TRPO 的数学比较

9.1 相同点

  • 都属于 on-policy policy gradient 方法
  • 都使用优势函数估计来降低方差
  • 都试图限制策略更新幅度,避免性能崩掉

9.2 不同点概览

维度 PPO TRPO
更新思想 用 clip 近似 trust region 显式 KL 约束 trust region
优化阶数 一阶优化 二阶近似优化
目标函数 裁剪 surrogate objective 约束优化 surrogate objective
约束实现 软限制,靠 clip 硬一些,靠 KL 上界
求解代价
实现复杂度
常见实践 更普遍 更偏研究或教学

9.3 数学层面对比

PPO: 近似约束

PPO 并没有真正解下面这个约束问题:


\max_\theta \; L(\theta) \quad \text{s.t.} \quad D_{\mathrm{KL}}(\pi_{\theta_{old}}, \pi_\theta) \le \delta

它是通过剪裁比率:


r_t(\theta) \in [1-\epsilon, 1+\epsilon]

间接限制策略变化。这个做法更简单,但不是严格的 trust region。

TRPO: 显式约束

TRPO 则直接围绕 KL 散度建立局部二次约束:


\max_x \; g^T x \quad \text{s.t.} \quad \frac{1}{2} x^T H x \le \delta

这个问题有解析形式的最优方向:


x^* \propto H^{-1} g

也就是自然梯度方向。

因此,从数学角度讲:

  • PPO 更像是对 trust region 的廉价近似
  • TRPO 更接近真正的受约束自然梯度更新

9.4 实践层面对比

PPO 往往更常用,因为:

  • 写起来简单
  • 调参更直接
  • 可以自然接入 mini-batch 和多 epoch
  • 和现代深度学习工具链更兼容

TRPO 的价值主要体现在:

  • 作为 trust region policy optimization 的经典代表
  • 数学解释更强
  • 在一些场景下更新更保守、更稳定

但代价是:

  • 代码更复杂
  • 单次更新更慢
  • 在 GPU 上也并不像 PPO 那样容易把吞吐做大

10. 当前实现中的工程细节

10.1 设备路径

Actor 和 Critic 都会被显式移动到指定 device:

  • PPOAgent(..., device=device)
  • TRPOAgent(..., device=device)

rollout 收集到的数据在 update 阶段会统一搬到对应设备上,避免反复构造 CPU tensor 再隐式迁移。

10.2 向量环境与部分 reset

并行环境中,不同 worker 的 episode 结束时间不一致。实现里会:

  • 用 reset_mask 只重置已完成的子环境
  • 为每个子环境分别维护 running episode reward
  • 在一个 batch 中把多个环境的 transition 扁平化后统一交给 PPO/TRPO 更新

10.3 截断 bootstrap

对于 truncated 但未真正 terminated 的样本,代码会执行:


r_t^{\text{store}} = r_t + \gamma V(s_{t+1})

这样可以在时间截断时保留 bootstrap 信息,减小价值目标偏差。

11. 结果解读建议

建议重点观察:

  • 原始 reward 曲线波动大小
  • 平滑曲线提升速度
  • 相同 episode 数下谁更稳定
  • tqdm 中 KL / policy_loss / critic_loss 的变化趋势

一个常见现象是:

  • PPO 更新轻、吞吐高,通常更快得到可用结果
  • TRPO 单次更新更贵,但策略变化通常更保守

12. 常见问题

12.1 GPU 可用但速度不快

这是正常的,因为强化学习瓶颈常常在环境采样而不是网络计算。可以尝试:

  • 提高 num-envs
  • 使用 async 向量环境
  • 增大 batch-size
  • 在更复杂网络或更大 batch 下训练

12.2 AsyncVectorEnv 报多进程相关错误

可以先退回同步模式:

python main.py --vector-mode sync --num-envs 4

12.3 Matplotlib 保存图片时报 MKL / OpenMP 冲突

如果你看到类似下面的错误:

libmkl_intel_thread.so.2: undefined symbol: __kmpc_global_thread_num

或者:

libmkl_intel_thread.so.2: undefined symbol: omp_get_num_procs

通常是 MKL 和 OpenMP 运行库的兼容性问题。本项目已经在 main.py 顶部默认设置了:

MKL_THREADING_LAYER=GNU

如果你是自己写外部脚本调用这些模块,也建议在导入 matplotlib / numpy 之前设置同样的环境变量。

12.4 想让实验更稳定

可以尝试:

  • 固定 seed
  • 增大 batch-size
  • 调整 hidden-dim
  • 增减 PPO 的 clip 系数或 TRPO 的 KL 阈值

13. 后续可扩展方向

  • 增加 TensorBoard 或 Weights & Biases 记录
  • 保存 checkpoint 与最佳模型
  • 扩展到 MountainCarContinuous-v0、BipedalWalker-v3 等连续控制任务
  • 为 PPO / TRPO 增加更完整的超参数命令行入口
  • 统计 wall-clock time、samples/sec、updates/sec 等性能指标

14. 参考文献

  1. Schulman, J., Levine, S., Abbeel, P., Jordan, M., Moritz, P. Trust Region Policy Optimization. ICML 2015.
  2. Schulman, J., Wolski, F., Dhariwal, P., Radford, A., Klimov, O. Proximal Policy Optimization Algorithms. arXiv 2017.
  3. Schulman, J., Moritz, P., Levine, S., Jordan, M., Abbeel, P. High-Dimensional Continuous Control Using Generalized Advantage Estimation. arXiv 2015.
S
Description
TRPO与PPO一类Policy Gradient算法
Readme
56 KiB
Languages
Python 100%