From 428e6f7f8104e190fbc27fc8d23d4a52e9b22f0a Mon Sep 17 00:00:00 2001 From: lhr Date: Thu, 2 Apr 2026 09:48:59 +0000 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E5=B9=B6=E8=A1=8C=E9=87=87?= =?UTF-8?q?=E6=A0=B7=E5=B9=B6=E5=AE=8C=E5=96=84=E8=AE=AD=E7=BB=83=E6=96=87?= =?UTF-8?q?=E6=A1=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 516 +++++++++++++++++++++++++++++++++++++++++++++++ agent/ppo.py | 99 ++++++--- agent/trpo.py | 127 ++++++++---- environment.yml | 1 + main.py | 340 +++++++++++++++++++++++-------- requirements.txt | 1 + 6 files changed, 943 insertions(+), 141 deletions(-) create mode 100644 README.md diff --git a/README.md b/README.md new file mode 100644 index 0000000..afd82c9 --- /dev/null +++ b/README.md @@ -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. \ No newline at end of file diff --git a/agent/ppo.py b/agent/ppo.py index a1ecd39..978b33c 100644 --- a/agent/ppo.py +++ b/agent/ppo.py @@ -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, + } diff --git a/agent/trpo.py b/agent/trpo.py index 1d52f21..db27449 100644 --- a/agent/trpo.py +++ b/agent/trpo.py @@ -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), + } diff --git a/environment.yml b/environment.yml index 316162f..2463a76 100644 --- a/environment.yml +++ b/environment.yml @@ -6,6 +6,7 @@ dependencies: - pip - numpy>=1.24 - matplotlib>=3.7 + - tqdm>=4.66 - pip: - torch>=2.1 - gymnasium[classic-control]>=0.29 diff --git a/main.py b/main.py index 9b9440c..c7bfb7b 100644 --- a/main.py +++ b/main.py @@ -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() diff --git a/requirements.txt b/requirements.txt index a3ed8f6..d07931f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,3 +2,4 @@ numpy>=1.24 matplotlib>=3.7 torch>=2.1 gymnasium[classic-control]>=0.29 +tqdm>=4.66