Hongru/RL_TRPO
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

修复并行采样并完善训练文档

Browse Source
This commit is contained in:
Hongru
2026-04-02 09:48:59 +00:00
parent 771eba8607
commit 428e6f7f81
6 changed files with 943 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+516
View File
@@ -0,0 +1,516 @@
# 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.
agent/ppo.py
+72 -27
View File
@@ -2,7 +2,7 @@ import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal
from networks import PolicyNetwork, ValueNetwork
@@ -32,6 +32,7 @@ class PPOAgent:
k_epochs=10,
minibatch_size=64,
max_grad_norm=0.5,
device="cpu",
):
self.gamma = gamma
self.tau = tau
@@ -39,37 +40,58 @@ class PPOAgent:
self.k_epochs = k_epochs # 每次更新对同一批数据的训练轮数
self.minibatch_size = minibatch_size
self.max_grad_norm = max_grad_norm # 梯度裁剪阈值
self.device = torch.device(device)
self.actor = PolicyNetwork(state_dim, action_dim, action_bound, hidden_dim)
self.critic = ValueNetwork(state_dim, hidden_dim)
self.actor = PolicyNetwork(state_dim, action_dim, action_bound, hidden_dim).to(self.device)
self.critic = ValueNetwork(state_dim, hidden_dim).to(self.device)
self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)
self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=critic_lr)
def _to_tensor(self, array_like):
return torch.as_tensor(array_like, dtype=torch.float32, device=self.device)
def get_action(self, state):
state_tensor = torch.FloatTensor(state).unsqueeze(0)
state_tensor = self._to_tensor(state)
single_state = state_tensor.ndim == 1
if single_state:
state_tensor = state_tensor.unsqueeze(0)
with torch.no_grad():
dist = self.actor.evaluate(state_tensor)
action = dist.sample()
return action.squeeze(0).numpy()
action_np = action.detach().cpu().numpy()
return action_np[0] if single_state else action_np
def get_value(self, state):
with torch.no_grad():
return self.critic(torch.FloatTensor(state).unsqueeze(0)).squeeze().item()
state_tensor = self._to_tensor(state)
single_state = state_tensor.ndim == 1
if single_state:
state_tensor = state_tensor.unsqueeze(0)
def _compute_advantages(self, rewards, values, masks):
with torch.no_grad():
values = self.critic(state_tensor).squeeze(-1)
values_np = values.detach().cpu().numpy()
return float(values_np[0]) if single_state else values_np
def _compute_advantages(self, rewards, values, next_values, masks):
"""
GAE (Generalized Advantage Estimation)
Eq.(11): Â_t = δ_t + (γλ)δ_{t+1} + ... + (γλ)^{T-t+1} δ_{T-1}
Eq.(12): δ_t = r_t + γV(s_{t+1}) - V(s_t)
"""
returns = []
gae = 0
for i in reversed(range(len(rewards))):
delta = rewards[i] + self.gamma * values[i + 1] * masks[i] - values[i]
advantages = torch.zeros_like(rewards, device=self.device)
gae = torch.zeros(rewards.size(1), dtype=torch.float32, device=self.device)
for i in reversed(range(rewards.size(0))):
delta = rewards[i] + self.gamma * next_values[i] * masks[i] - values[i]
gae = delta + self.gamma * self.tau * masks[i] * gae
returns.insert(0, gae + values[i])
return returns
advantages[i] = gae
returns = advantages + values
return returns, advantages
def update(self, memory):
"""
@@ -80,25 +102,32 @@ class PPOAgent:
L = L^CLIP - c1 * L^VF (Eq.9，c2=0 不加熵)
2. 梯度裁剪 (max_grad_norm)
"""
if not memory:
return {}
# ---------- 1. 准备数据 ----------
states = torch.FloatTensor(np.array([m[0] for m in memory]))
actions = torch.FloatTensor(np.array([m[1] for m in memory]))
rewards = [m[2] for m in memory]
next_states = torch.FloatTensor(np.array([m[3] for m in memory]))
masks = [m[4] for m in memory]
states = self._to_tensor(np.asarray(memory["states"]))
actions = self._to_tensor(np.asarray(memory["actions"]))
rewards = self._to_tensor(np.asarray(memory["rewards"]))
next_states = self._to_tensor(np.asarray(memory["next_states"]))
masks = self._to_tensor(np.asarray(memory["masks"]))
# ---------- 2. GAE 优势估计 ----------
with torch.no_grad():
values = self.critic(states).squeeze().numpy().tolist()
next_value = self.critic(next_states[-1].unsqueeze(0)).squeeze().item()
values.append(next_value)
rollout_steps, num_envs = rewards.shape
flat_states = states.reshape(rollout_steps * num_envs, -1)
flat_next_states = next_states.reshape(rollout_steps * num_envs, -1)
values = self.critic(flat_states).squeeze(-1).reshape(rollout_steps, num_envs)
next_values = self.critic(flat_next_states).squeeze(-1).reshape(rollout_steps, num_envs)
returns = self._compute_advantages(rewards, values, masks)
returns = torch.FloatTensor(returns)
values = torch.FloatTensor(values[:-1])
advantages = returns - values
returns, advantages = self._compute_advantages(rewards, values, next_values, masks)
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
states = flat_states
actions = actions.reshape(rollout_steps * num_envs, -1)
returns = returns.reshape(rollout_steps * num_envs)
advantages = advantages.reshape(rollout_steps * num_envs)
# ---------- 3. 记录旧策略 π_θold ----------
with torch.no_grad():
old_dist = self.actor.evaluate(states)
@@ -107,13 +136,17 @@ class PPOAgent:
# ---------- 4. K epochs × minibatch 联合训练 ----------
dataset_size = states.size(0)
indices = np.arange(dataset_size)
policy_loss_value = 0.0
critic_loss_value = 0.0
entropy_value = 0.0
minibatch_updates = 0
for _ in range(self.k_epochs):
np.random.shuffle(indices)
for start in range(0, dataset_size, self.minibatch_size):
end = start + self.minibatch_size
mb_idx = indices[start:end]
mb_idx = torch.as_tensor(indices[start:end], device=self.device, dtype=torch.long)
mb_states = states[mb_idx]
mb_actions = actions[mb_idx]
@@ -144,3 +177,15 @@ class PPOAgent:
self.critic_optimizer.zero_grad()
critic_loss.backward()
self.critic_optimizer.step()
policy_loss_value += float(policy_loss.detach().item())
critic_loss_value += float(critic_loss.detach().item())
entropy_value += float(dist.entropy().sum(dim=1).mean().detach().item())
minibatch_updates += 1
divisor = max(1, minibatch_updates)
return {
"policy_loss": policy_loss_value / divisor,
"critic_loss": critic_loss_value / divisor,
"entropy": entropy_value / divisor,
}
agent/trpo.py
+92 -35
View File
@@ -3,44 +3,79 @@ import torch
import torch.nn.functional as F
from torch.distributions import Normal
from torch.nn.utils import parameters_to_vector, vector_to_parameters
from networks import PolicyNetwork, ValueNetwork
class TRPOAgent:
def __init__(self, state_dim, action_dim, action_bound, hidden_dim=128, kl_margin=0.01, gamma=0.99, tau=0.95, cg_iters=10):
def __init__(
self,
state_dim,
action_dim,
action_bound,
hidden_dim=128,
kl_margin=0.01,
gamma=0.99,
tau=0.95,
cg_iters=10,
critic_epochs=40,
device="cpu",
):
self.gamma = gamma
self.tau = tau
self.kl_margin = kl_margin
self.kl_margin = kl_margin
self.cg_iters = cg_iters
self.critic_epochs = critic_epochs
self.device = torch.device(device)
self.actor = PolicyNetwork(state_dim, action_dim, action_bound, hidden_dim)
self.critic = ValueNetwork(state_dim, hidden_dim)
self.actor = PolicyNetwork(state_dim, action_dim, action_bound, hidden_dim).to(self.device)
self.critic = ValueNetwork(state_dim, hidden_dim).to(self.device)
self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=1e-3)
def _to_tensor(self, array_like):
return torch.as_tensor(array_like, dtype=torch.float32, device=self.device)
def get_action(self, state):
"""
根据当前状态采样动作
"""
state_tensor = torch.FloatTensor(state).unsqueeze(0)
state_tensor = self._to_tensor(state)
single_state = state_tensor.ndim == 1
if single_state:
state_tensor = state_tensor.unsqueeze(0)
with torch.no_grad():
dist = self.actor.evaluate(state_tensor)
action = dist.sample()
return action.squeeze(0).numpy()
action_np = action.detach().cpu().numpy()
return action_np[0] if single_state else action_np
def get_value(self, state):
with torch.no_grad():
return self.critic(torch.FloatTensor(state).unsqueeze(0)).squeeze().item()
state_tensor = self._to_tensor(state)
single_state = state_tensor.ndim == 1
if single_state:
state_tensor = state_tensor.unsqueeze(0)
def _compute_advantages(self, rewards, values, masks):
with torch.no_grad():
values = self.critic(state_tensor).squeeze(-1)
values_np = values.detach().cpu().numpy()
return float(values_np[0]) if single_state else values_np
def _compute_advantages(self, rewards, values, next_values, masks):
"""
使用广义优势估计 (GAE) 计算优势函数
"""
returns = []
gae = 0
for i in reversed(range(len(rewards))):
delta = rewards[i] + self.gamma * values[i + 1] * masks[i] - values[i]
advantages = torch.zeros_like(rewards, device=self.device)
gae = torch.zeros(rewards.size(1), dtype=torch.float32, device=self.device)
for i in reversed(range(rewards.size(0))):
delta = rewards[i] + self.gamma * next_values[i] * masks[i] - values[i]
gae = delta + self.gamma * self.tau * masks[i] * gae
returns.insert(0, gae + values[i])
return returns
advantages[i] = gae
returns = advantages + values
return returns, advantages
# --- 关键修复 1：将固定的 old_dist 作为参数传入 ---
def _hessian_vector_product(self, states, old_dist, vector, damping=0.1):
@@ -58,7 +93,7 @@ class TRPOAgent:
# 与给定向量点乘
kl_v = (flat_grad_kl * vector).sum()
# 二阶导数
grads = torch.autograd.grad(kl_v, self.actor.parameters())
flat_grad_grad_kl = torch.cat([grad.contiguous().view(-1) for grad in grads])
@@ -88,30 +123,39 @@ class TRPOAgent:
return x
def update(self, memory):
states = torch.FloatTensor(np.array([m[0] for m in memory]))
actions = torch.FloatTensor(np.array([m[1] for m in memory]))
rewards = [m[2] for m in memory]
next_states = torch.FloatTensor(np.array([m[3] for m in memory]))
masks = [m[4] for m in memory]
if not memory:
return {}
states = self._to_tensor(np.asarray(memory["states"]))
actions = self._to_tensor(np.asarray(memory["actions"]))
rewards = self._to_tensor(np.asarray(memory["rewards"]))
next_states = self._to_tensor(np.asarray(memory["next_states"]))
masks = self._to_tensor(np.asarray(memory["masks"]))
with torch.no_grad():
values = self.critic(states).squeeze().numpy().tolist()
next_value = self.critic(next_states[-1].unsqueeze(0)).squeeze().item()
values.append(next_value)
returns = self._compute_advantages(rewards, values, masks)
returns = torch.FloatTensor(returns)
values = torch.FloatTensor(values[:-1])
advantages = returns - values
rollout_steps, num_envs = rewards.shape
flat_states = states.reshape(rollout_steps * num_envs, -1)
flat_next_states = next_states.reshape(rollout_steps * num_envs, -1)
values = self.critic(flat_states).squeeze(-1).reshape(rollout_steps, num_envs)
next_values = self.critic(flat_next_states).squeeze(-1).reshape(rollout_steps, num_envs)
returns, advantages = self._compute_advantages(rewards, values, next_values, masks)
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
states = flat_states
actions = actions.reshape(rollout_steps * num_envs, -1)
returns = returns.reshape(rollout_steps * num_envs)
advantages = advantages.reshape(rollout_steps * num_envs)
# 【修复】：加大 Critic 的训练力度，从 10 提升到 40 Epochs
# 确保裁判的眼光足够准确
for _ in range(40):
critic_loss_value = 0.0
for _ in range(self.critic_epochs):
critic_loss = F.mse_loss(self.critic(states).squeeze(), returns)
self.critic_optimizer.zero_grad()
critic_loss.backward()
self.critic_optimizer.step()
critic_loss_value += float(critic_loss.detach().item())
with torch.no_grad():
old_mean, old_std = self.actor(states)
@@ -139,26 +183,39 @@ class TRPOAgent:
lm = torch.sqrt(shs / self.kl_margin)
fullstep = step_dir / lm
old_params = parameters_to_vector(self.actor.parameters())
# 线性搜索
success = False
step_size = 1.0
kl_value = 0.0
surrogate_value = float(surrogate_loss.detach().item())
for _ in range(10):
new_params = old_params + step_size * fullstep
vector_to_parameters(new_params, self.actor.parameters())
with torch.no_grad():
new_surrogate_loss, new_dist = compute_surrogate_loss()
kl = torch.distributions.kl_divergence(old_dist, new_dist).mean()
# 【修复】：增加极小的浮点数宽容度，防止在极小提升时被误判失败而拒绝更新
if new_surrogate_loss >= surrogate_loss - 1e-8 and kl <= self.kl_margin * 1.5:
success = True
surrogate_value = float(new_surrogate_loss.item())
kl_value = float(kl.item())
break
step_size *= 0.5
kl_value = float(kl.item())
if not success:
vector_to_parameters(old_params, self.actor.parameters())
return {
"critic_loss": critic_loss_value / max(1, self.critic_epochs),
"surrogate_loss": surrogate_value,
"kl": kl_value,
"line_search_success": float(success),
"step_scale": float(step_size if success else 0.0),
}
environment.yml
+1
View File
@@ -6,6 +6,7 @@ dependencies:
- pip
- numpy>=1.24
- matplotlib>=3.7
- tqdm>=4.66
- pip:
- torch>=2.1
- gymnasium[classic-control]>=0.29
main.py
+261 -79
View File
@@ -1,116 +1,298 @@
import argparse
import os
import random
from pathlib import Path
os.environ.setdefault("MKL_THREADING_LAYER", "GNU")
import gymnasium as gym
import matplotlib.pyplot as plt
import numpy as np
import torch
from tqdm.auto import tqdm
from agent.ppo import PPOAgent
from agent.trpo import TRPOAgent
def train_agent(agent, env_name, num_episodes=500, batch_size=2000):
"""通用训练函数，适用于 PPO 和 TRPO"""
env = gym.make(env_name)
episode_rewards = []
def parse_args():
cpu_count = os.cpu_count() or 1
default_envs = max(1, min(4, cpu_count))
default_threads = max(1, min(8, cpu_count))
state, _ = env.reset()
memory = []
current_ep_reward = 0
episodes_completed = 0
parser = argparse.ArgumentParser(description="使用 GPU / 多核并行训练 PPO 与 TRPO，并输出对比图。")
parser.add_argument("--env-name", type=str, default="Pendulum-v1", help="Gymnasium 环境名称。")
parser.add_argument("--num-episodes", type=int, default=500, help="每个算法训练的 episode 数。")
parser.add_argument("--batch-size", type=int, default=2000, help="每次策略更新前收集的环境步数。")
parser.add_argument("--num-envs", type=int, default=default_envs, help="并行采样环境数。大于 1 时可显著提高吞吐。")
parser.add_argument("--vector-mode", choices=("sync", "async"), default="async", help="向量环境模式。async 会启动多进程，更适合多核 CPU。")
parser.add_argument("--cpu-threads", type=int, default=default_threads, help="PyTorch CPU 线程数。GPU 训练时主要影响 CPU 侧数据准备。")
parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto", help="训练设备。auto 会优先选择 CUDA。")
parser.add_argument("--seed", type=int, default=42, help="全局随机种子。")
parser.add_argument("--hidden-dim", type=int, default=128, help="Actor/Critic 隐层宽度。")
parser.add_argument("--output-dir", type=str, default="outputs", help="图像与训练数据的输出目录。")
parser.add_argument("--no-show", action="store_true", help="仅保存图像，不弹出 matplotlib 窗口。")
return parser.parse_args()
def resolve_device(device_name):
if device_name == "auto":
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
if device_name == "cuda" and not torch.cuda.is_available():
raise RuntimeError("请求使用 CUDA，但当前环境不可用。请检查 GPU 驱动和 PyTorch CUDA 版本。")
return torch.device(device_name)
def configure_runtime(cpu_threads, seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
if cpu_threads is not None:
torch.set_num_threads(max(1, cpu_threads))
if hasattr(torch, "set_num_interop_threads"):
torch.set_num_interop_threads(max(1, min(cpu_threads, 4)))
if hasattr(torch, "set_float32_matmul_precision"):
torch.set_float32_matmul_precision("high")
def make_env(env_name, seed, worker_idx):
def _thunk():
env = gym.make(env_name)
env.reset(seed=seed + worker_idx)
env.action_space.seed(seed + worker_idx)
return env
return _thunk
def create_vector_env(env_name, num_envs, seed, vector_mode):
env_fns = [make_env(env_name, seed, idx) for idx in range(num_envs)]
if num_envs == 1 or vector_mode == "sync":
return gym.vector.SyncVectorEnv(env_fns)
return gym.vector.AsyncVectorEnv(env_fns, context="spawn")
def format_postfix(episode_rewards, total_env_steps, updates, last_update_stats):
postfix = {
"avg20": f"{np.mean(episode_rewards[-20:]):.1f}" if episode_rewards else "n/a",
"steps": total_env_steps,
"updates": updates,
}
for key, value in last_update_stats.items():
if isinstance(value, (int, float, np.floating)):
postfix[key] = f"{float(value):.4f}"
return postfix
def train_agent(agent, env_name, num_episodes=500, batch_size=2000, num_envs=1, seed=42, vector_mode="async", desc=None):
"""通用训练函数，适用于 PPO 和 TRPO，支持 GPU 和向量化并行采样。"""
env = create_vector_env(env_name, num_envs, seed, vector_mode)
episode_rewards = []
memory = {
"states": [],
"actions": [],
"rewards": [],
"next_states": [],
"masks": [],
}
total_env_steps = 0
updates = 0
last_update_stats = {}
states, _ = env.reset(seed=seed)
running_rewards = np.zeros(num_envs, dtype=np.float64)
step_count = 0
while episodes_completed < num_episodes:
action = agent.get_action(state)
next_state, reward, terminated, truncated, _ = env.step(action)
done = terminated or truncated
progress = tqdm(total=num_episodes, desc=desc or agent.__class__.__name__, dynamic_ncols=True)
mask = 0.0 if done else 1.0
reward_store = reward
if truncated and not terminated:
reward_store = reward + agent.gamma * agent.get_value(next_state)
try:
while len(episode_rewards) < num_episodes:
actions = agent.get_action(states)
next_states, rewards, terminated, truncated, _ = env.step(actions)
done = np.logical_or(terminated, truncated)
memory.append([state, action, reward_store, next_state, mask])
state = next_state
current_ep_reward += reward
step_count += 1
masks = (~done).astype(np.float32)
reward_store = rewards.astype(np.float32).copy()
truncated_only = np.logical_and(truncated, ~terminated)
if np.any(truncated_only):
bootstrap_values = np.asarray(agent.get_value(next_states[truncated_only]), dtype=np.float32)
reward_store[truncated_only] += agent.gamma * bootstrap_values
if done:
episode_rewards.append(current_ep_reward)
episodes_completed += 1
state, _ = env.reset()
current_ep_reward = 0
memory["states"].append(np.asarray(states, dtype=np.float32).copy())
memory["actions"].append(np.asarray(actions, dtype=np.float32).copy())
memory["rewards"].append(np.asarray(reward_store, dtype=np.float32).copy())
memory["next_states"].append(np.asarray(next_states, dtype=np.float32).copy())
memory["masks"].append(np.asarray(masks, dtype=np.float32).copy())
if episodes_completed % 10 == 0:
avg_reward = np.mean(episode_rewards[-10:])
print(f" Episode: {episodes_completed}, 平均奖励 (最近10轮): {avg_reward:.2f}")
running_rewards += rewards
step_count += num_envs
total_env_steps += num_envs
if step_count >= batch_size:
agent.update(memory)
memory.clear()
step_count = 0
if np.any(done):
done_indices = np.flatnonzero(done)
new_episode_rewards = running_rewards[done_indices].tolist()
remaining = num_episodes - len(episode_rewards)
accepted = new_episode_rewards[:remaining]
if accepted:
episode_rewards.extend(accepted)
progress.update(len(accepted))
env.close()
return episode_rewards
running_rewards[done_indices] = 0.0
progress.set_postfix(format_postfix(episode_rewards, total_env_steps, updates, last_update_stats))
if len(episode_rewards) >= num_episodes:
break
reset_mask = np.zeros(num_envs, dtype=bool)
reset_mask[done_indices] = True
reset_states, _ = env.reset(options={"reset_mask": reset_mask})
next_states[done_indices] = reset_states[done_indices]
states = next_states
if step_count >= batch_size:
last_update_stats = agent.update(memory)
for key in memory:
memory[key].clear()
step_count = 0
updates += 1
progress.set_postfix(format_postfix(episode_rewards, total_env_steps, updates, last_update_stats))
if memory["states"]:
last_update_stats = agent.update(memory)
updates += 1
progress.set_postfix(format_postfix(episode_rewards, total_env_steps, updates, last_update_stats))
finally:
progress.close()
env.close()
summary = {
"episodes": len(episode_rewards),
"total_env_steps": total_env_steps,
"updates": updates,
"final_avg20": float(np.mean(episode_rewards[-20:])) if episode_rewards else float("nan"),
}
summary.update(last_update_stats)
return episode_rewards, summary
def smooth(rewards, window=10):
"""滑动平均平滑曲线"""
"""滑动平均平滑曲线。"""
rewards = np.asarray(rewards, dtype=np.float32)
if rewards.size == 0:
return rewards
smoothed = []
for i in range(len(rewards)):
for i in range(rewards.size):
start = max(0, i - window + 1)
smoothed.append(np.mean(rewards[start:i + 1]))
return smoothed
smoothed.append(float(np.mean(rewards[start:i + 1])))
return np.asarray(smoothed, dtype=np.float32)
def save_artifacts(ppo_rewards, trpo_rewards, output_dir, env_name):
output_dir.mkdir(parents=True, exist_ok=True)
fig, axes = plt.subplots(1, 2, figsize=(16, 6))
axes[0].plot(ppo_rewards, alpha=0.3, color="blue", label="PPO (raw)")
axes[0].plot(trpo_rewards, alpha=0.3, color="red", label="TRPO (raw)")
axes[0].plot(smooth(ppo_rewards, 20), color="blue", linewidth=2, label="PPO (smooth)")
axes[0].plot(smooth(trpo_rewards, 20), color="red", linewidth=2, label="TRPO (smooth)")
axes[0].set_title(f"PPO vs TRPO on {env_name}")
axes[0].set_xlabel("Episode")
axes[0].set_ylabel("Total Reward")
axes[0].legend()
axes[0].grid(True)
axes[1].plot(smooth(ppo_rewards, 20), color="blue", linewidth=2, label="PPO")
axes[1].plot(smooth(trpo_rewards, 20), color="red", linewidth=2, label="TRPO")
axes[1].set_title("PPO vs TRPO (Smoothed, window=20)")
axes[1].set_xlabel("Episode")
axes[1].set_ylabel("Total Reward")
axes[1].legend()
axes[1].grid(True)
plt.tight_layout()
figure_path = output_dir / "ppo_vs_trpo_comparison.png"
metrics_path = output_dir / "training_metrics.npz"
fig.savefig(figure_path, dpi=150)
np.savez(metrics_path, ppo_rewards=np.asarray(ppo_rewards), trpo_rewards=np.asarray(trpo_rewards))
return fig, figure_path, metrics_path
def main():
env_name = 'Pendulum-v1'
env = gym.make(env_name)
args = parse_args()
configure_runtime(args.cpu_threads, args.seed)
device = resolve_device(args.device)
output_dir = Path(args.output_dir)
env = gym.make(args.env_name)
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.shape[0]
action_bound = float(env.action_space.high[0])
env.close()
num_episodes = 500
tqdm.write(
f"配置: env={args.env_name}, device={device}, num_envs={args.num_envs}, "
f"vector_mode={args.vector_mode}, cpu_threads={args.cpu_threads}, batch_size={args.batch_size}"
)
# --- 训练 PPO ---
print("=" * 50)
print("开始训练 PPO 智能体...")
print("=" * 50)
ppo_agent = PPOAgent(state_dim=state_dim, action_dim=action_dim, action_bound=action_bound)
ppo_rewards = train_agent(ppo_agent, env_name, num_episodes)
ppo_agent = PPOAgent(
state_dim=state_dim,
action_dim=action_dim,
action_bound=action_bound,
hidden_dim=args.hidden_dim,
device=device,
)
ppo_rewards, ppo_summary = train_agent(
ppo_agent,
args.env_name,
num_episodes=args.num_episodes,
batch_size=args.batch_size,
num_envs=args.num_envs,
seed=args.seed,
vector_mode=args.vector_mode,
desc="PPO",
)
tqdm.write(f"PPO 完成: {ppo_summary}")
# --- 训练 TRPO ---
print("=" * 50)
print("开始训练 TRPO 智能体...")
print("=" * 50)
trpo_agent = TRPOAgent(state_dim=state_dim, action_dim=action_dim, action_bound=action_bound)
trpo_rewards = train_agent(trpo_agent, env_name, num_episodes)
trpo_agent = TRPOAgent(
state_dim=state_dim,
action_dim=action_dim,
action_bound=action_bound,
hidden_dim=args.hidden_dim,
device=device,
)
trpo_rewards, trpo_summary = train_agent(
trpo_agent,
args.env_name,
num_episodes=args.num_episodes,
batch_size=args.batch_size,
num_envs=args.num_envs,
seed=args.seed + 10_000,
vector_mode=args.vector_mode,
desc="TRPO",
)
tqdm.write(f"TRPO 完成: {trpo_summary}")
# --- 对比画图 ---
fig, axes = plt.subplots(1, 2, figsize=(16, 6))
fig, figure_path, metrics_path = save_artifacts(ppo_rewards, trpo_rewards, output_dir, args.env_name)
# 左图：原始奖励曲线
axes[0].plot(ppo_rewards, alpha=0.3, color='blue', label='PPO (raw)')
axes[0].plot(trpo_rewards, alpha=0.3, color='red', label='TRPO (raw)')
axes[0].plot(smooth(ppo_rewards, 20), color='blue', linewidth=2, label='PPO (smooth)')
axes[0].plot(smooth(trpo_rewards, 20), color='red', linewidth=2, label='TRPO (smooth)')
axes[0].set_title('PPO vs TRPO on Pendulum-v1')
axes[0].set_xlabel('Episode')
axes[0].set_ylabel('Total Reward')
axes[0].legend()
axes[0].grid(True)
tqdm.write(f"对比图已保存至 {figure_path}")
tqdm.write(f"训练曲线原始数据已保存至 {metrics_path}")
# 右图：滑动平均对比（更清晰）
axes[1].plot(smooth(ppo_rewards, 20), color='blue', linewidth=2, label='PPO')
axes[1].plot(smooth(trpo_rewards, 20), color='red', linewidth=2, label='TRPO')
axes[1].set_title('PPO vs TRPO (Smoothed, window=20)')
axes[1].set_xlabel('Episode')
axes[1].set_ylabel('Total Reward')
axes[1].legend()
axes[1].grid(True)
plt.tight_layout()
plt.savefig('ppo_vs_trpo_comparison.png', dpi=150)
plt.show()
print("对比图已保存至 ppo_vs_trpo_comparison.png")
if args.no_show:
plt.close(fig)
else:
plt.show()
if __name__ == '__main__':
if __name__ == "__main__":
main()
requirements.txt
+1
View File
@@ -2,3 +2,4 @@ numpy>=1.24
matplotlib>=3.7
torch>=2.1
gymnasium[classic-control]>=0.29
tqdm>=4.66