516 lines
13 KiB
Markdown
516 lines
13 KiB
Markdown
# RL_TRPO
|
||
|
||
这个项目实现了连续动作空间上的 PPO 和 TRPO，并针对训练工程侧做了三类增强：
|
||
|
||
- 自动选择 GPU 进行前向与反向传播
|
||
- 使用 Gymnasium 向量环境进行并行采样，支持多核 CPU
|
||
- 使用 tqdm 替代周期性 print，实时展示 episode 进度、滚动奖励和优化指标
|
||
|
||
当前默认实验环境是 Pendulum-v1，用来对比 PPO 和 TRPO 在同一任务上的训练稳定性、优化代价和收敛表现。
|
||
|
||
## 1. 项目结构
|
||
|
||
```text
|
||
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：
|
||
|
||
```bash
|
||
python main.py --device auto
|
||
```
|
||
|
||
也可以显式指定：
|
||
|
||
```bash
|
||
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 个并行环境：
|
||
|
||
```bash
|
||
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
|
||
|
||
```bash
|
||
conda env create -f environment.yml
|
||
conda activate rl_trpo
|
||
```
|
||
|
||
### 3.2 使用 pip
|
||
|
||
```bash
|
||
pip install -r requirements.txt
|
||
```
|
||
|
||
### 3.3 GPU 版本 PyTorch
|
||
|
||
requirements.txt 里只写了通用的 torch 依赖。若你想确保安装 CUDA 版本，请优先按照 PyTorch 官方安装命令安装。例如在 Linux + CUDA 12.x 场景下，通常应该使用官方提供的对应索引源或 Conda 渠道。
|
||
|
||
安装完成后可以验证：
|
||
|
||
```bash
|
||
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"
|
||
```
|
||
|
||
如果输出里的 torch.cuda.is_available() 为 True，就说明代码里的 GPU 路径会生效。
|
||
|
||
## 4. 快速开始
|
||
|
||
### 4.1 默认运行
|
||
|
||
```bash
|
||
python main.py
|
||
```
|
||
|
||
默认行为：
|
||
|
||
- 环境: Pendulum-v1
|
||
- episode 数: 500
|
||
- batch size: 2000
|
||
- 设备: auto
|
||
- 并行环境数: min(4, CPU 核数)
|
||
- 向量环境模式: async
|
||
|
||
### 4.2 典型高吞吐训练命令
|
||
|
||
```bash
|
||
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 基线
|
||
|
||
```bash
|
||
python main.py --device cpu --num-envs 1 --vector-mode sync
|
||
```
|
||
|
||
### 4.4 不弹图，只保存结果
|
||
|
||
```bash
|
||
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 报多进程相关错误
|
||
|
||
可以先退回同步模式：
|
||
|
||
```bash
|
||
python main.py --vector-mode sync --num-envs 4
|
||
```
|
||
|
||
### 12.3 Matplotlib 保存图片时报 MKL / OpenMP 冲突
|
||
|
||
如果你看到类似下面的错误：
|
||
|
||
```text
|
||
libmkl_intel_thread.so.2: undefined symbol: __kmpc_global_thread_num
|
||
```
|
||
|
||
或者：
|
||
|
||
```text
|
||
libmkl_intel_thread.so.2: undefined symbol: omp_get_num_procs
|
||
```
|
||
|
||
通常是 MKL 和 OpenMP 运行库的兼容性问题。本项目已经在 main.py 顶部默认设置了：
|
||
|
||
```bash
|
||
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. |