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. 参考文献
- Schulman, J., Levine, S., Abbeel, P., Jordan, M., Moritz, P. Trust Region Policy Optimization. ICML 2015.
- Schulman, J., Wolski, F., Dhariwal, P., Radford, A., Klimov, O. Proximal Policy Optimization Algorithms. arXiv 2017.
- Schulman, J., Moritz, P., Levine, S., Jordan, M., Abbeel, P. High-Dimensional Continuous Control Using Generalized Advantage Estimation. arXiv 2015.