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

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.

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.actor = PolicyNetwork(state_dim, action_dim, action_bound, hidden_dim).to(self.device)
-        self.critic = ValueNetwork(state_dim, hidden_dim)
+        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():
+        state_tensor = self._to_tensor(state)
-            return self.critic(torch.FloatTensor(state).unsqueeze(0)).squeeze().item()
+        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 = []
+        advantages = torch.zeros_like(rewards, device=self.device)
-        gae = 0
+        gae = torch.zeros(rewards.size(1), dtype=torch.float32, device=self.device)
-        for i in reversed(range(len(rewards))):
+
-            delta = rewards[i] + self.gamma * values[i + 1] * masks[i] - values[i]
+        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])
+            advantages[i] = gae
-        return returns
+
+        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]))
+        states = self._to_tensor(np.asarray(memory["states"]))
-        actions = torch.FloatTensor(np.array([m[1] for m in memory]))
+        actions = self._to_tensor(np.asarray(memory["actions"]))
-        rewards = [m[2] for m in memory]
+        rewards = self._to_tensor(np.asarray(memory["rewards"]))
-        next_states = torch.FloatTensor(np.array([m[3] for m in memory]))
+        next_states = self._to_tensor(np.asarray(memory["next_states"]))
-        masks = [m[4] for m in memory]
+        masks = self._to_tensor(np.asarray(memory["masks"]))
 
         # ---------- 2. GAE 优势估计 ----------
         with torch.no_grad():
-            values = self.critic(states).squeeze().numpy().tolist()
+            rollout_steps, num_envs = rewards.shape
-            next_value = self.critic(next_states[-1].unsqueeze(0)).squeeze().item()
+            flat_states = states.reshape(rollout_steps * num_envs, -1)
-            values.append(next_value)
+            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, advantages = self._compute_advantages(rewards, values, next_values, masks)
-        returns = torch.FloatTensor(returns)
-        values = torch.FloatTensor(values[:-1])
-        advantages = returns - values
         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

@@ -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.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.actor = PolicyNetwork(state_dim, action_dim, action_bound, hidden_dim).to(self.device)
-        self.critic = ValueNetwork(state_dim, hidden_dim)
+        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():
+        state_tensor = self._to_tensor(state)
-            return self.critic(torch.FloatTensor(state).unsqueeze(0)).squeeze().item()
+        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 = []
+        advantages = torch.zeros_like(rewards, device=self.device)
-        gae = 0
+        gae = torch.zeros(rewards.size(1), dtype=torch.float32, device=self.device)
-        for i in reversed(range(len(rewards))):
+
-            delta = rewards[i] + self.gamma * values[i + 1] * masks[i] - values[i]
+        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])
+            advantages[i] = gae
-        return returns
+
+        returns = advantages + values
+        return returns, advantages
 
     # --- 关键修复 1：将固定的 old_dist 作为参数传入 ---
     def _hessian_vector_product(self, states, old_dist, vector, damping=0.1):
@@ -88,30 +123,39 @@ class TRPOAgent:
         return x
 
     def update(self, memory):
-        states = torch.FloatTensor(np.array([m[0] for m in memory]))
+        if not memory:
-        actions = torch.FloatTensor(np.array([m[1] for m in memory]))
+            return {}
-        rewards = [m[2] for m in memory]
+
-        next_states = torch.FloatTensor(np.array([m[3] for m in memory]))
+        states = self._to_tensor(np.asarray(memory["states"]))
-        masks = [m[4] for m in memory]
+        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()
+            rollout_steps, num_envs = rewards.shape
-            next_value = self.critic(next_states[-1].unsqueeze(0)).squeeze().item()
+            flat_states = states.reshape(rollout_steps * num_envs, -1)
-            values.append(next_value)
+            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, advantages = self._compute_advantages(rewards, values, next_values, masks)
-        returns = torch.FloatTensor(returns)
-        values = torch.FloatTensor(values[:-1])
-        advantages = returns - values
         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)
@@ -145,6 +189,8 @@ class TRPOAgent:
         # 线性搜索
         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())
@@ -156,9 +202,20 @@ class TRPOAgent:
             # 【修复】：增加极小的浮点数宽容度，防止在极小提升时被误判失败而拒绝更新
             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

@@ -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

@@ -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):
+def parse_args():
-    """通用训练函数，适用于 PPO 和 TRPO"""
+    cpu_count = os.cpu_count() or 1
-    env = gym.make(env_name)
+    default_envs = max(1, min(4, cpu_count))
-    episode_rewards = []
+    default_threads = max(1, min(8, cpu_count))
 
-    state, _ = env.reset()
+    parser = argparse.ArgumentParser(description="使用 GPU / 多核并行训练 PPO 与 TRPO，并输出对比图。")
-    memory = []
+    parser.add_argument("--env-name", type=str, default="Pendulum-v1", help="Gymnasium 环境名称。")
-    current_ep_reward = 0
+    parser.add_argument("--num-episodes", type=int, default=500, help="每个算法训练的 episode 数。")
-    episodes_completed = 0
+    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:
+    progress = tqdm(total=num_episodes, desc=desc or agent.__class__.__name__, dynamic_ncols=True)
-        action = agent.get_action(state)
-        next_state, reward, terminated, truncated, _ = env.step(action)
-        done = terminated or truncated
 
-        mask = 0.0 if done else 1.0
+    try:
-        reward_store = reward
+        while len(episode_rewards) < num_episodes:
-        if truncated and not terminated:
+            actions = agent.get_action(states)
-            reward_store = reward + agent.gamma * agent.get_value(next_state)
+            next_states, rewards, terminated, truncated, _ = env.step(actions)
+            done = np.logical_or(terminated, truncated)
 
-        memory.append([state, action, reward_store, next_state, mask])
+            masks = (~done).astype(np.float32)
-        state = next_state
+            reward_store = rewards.astype(np.float32).copy()
-        current_ep_reward += reward
+            truncated_only = np.logical_and(truncated, ~terminated)
-        step_count += 1
+            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:
+            memory["states"].append(np.asarray(states, dtype=np.float32).copy())
-            episode_rewards.append(current_ep_reward)
+            memory["actions"].append(np.asarray(actions, dtype=np.float32).copy())
-            episodes_completed += 1
+            memory["rewards"].append(np.asarray(reward_store, dtype=np.float32).copy())
-            state, _ = env.reset()
+            memory["next_states"].append(np.asarray(next_states, dtype=np.float32).copy())
-            current_ep_reward = 0
+            memory["masks"].append(np.asarray(masks, dtype=np.float32).copy())
 
-            if episodes_completed % 10 == 0:
+            running_rewards += rewards
-                avg_reward = np.mean(episode_rewards[-10:])
+            step_count += num_envs
-                print(f"  Episode: {episodes_completed}, 平均奖励 (最近10轮): {avg_reward:.2f}")
+            total_env_steps += num_envs
 
-        if step_count >= batch_size:
+            if np.any(done):
-            agent.update(memory)
+                done_indices = np.flatnonzero(done)
-            memory.clear()
+                new_episode_rewards = running_rewards[done_indices].tolist()
-            step_count = 0
+                remaining = num_episodes - len(episode_rewards)
+                accepted = new_episode_rewards[:remaining]
+                if accepted:
+                    episode_rewards.extend(accepted)
+                    progress.update(len(accepted))
 
-    env.close()
+                running_rewards[done_indices] = 0.0
-    return episode_rewards
+                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]))
+        smoothed.append(float(np.mean(rewards[start:i + 1])))
-    return smoothed
+    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'
+    args = parse_args()
-    env = gym.make(env_name)
+    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 ---
+    ppo_agent = PPOAgent(
-    print("=" * 50)
+        state_dim=state_dim,
-    print("开始训练 PPO 智能体...")
+        action_dim=action_dim,
-    print("=" * 50)
+        action_bound=action_bound,
-    ppo_agent = PPOAgent(state_dim=state_dim, action_dim=action_dim, action_bound=action_bound)
+        hidden_dim=args.hidden_dim,
-    ppo_rewards = train_agent(ppo_agent, env_name, num_episodes)
+        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 ---
+    trpo_agent = TRPOAgent(
-    print("=" * 50)
+        state_dim=state_dim,
-    print("开始训练 TRPO 智能体...")
+        action_dim=action_dim,
-    print("=" * 50)
+        action_bound=action_bound,
-    trpo_agent = TRPOAgent(state_dim=state_dim, action_dim=action_dim, action_bound=action_bound)
+        hidden_dim=args.hidden_dim,
-    trpo_rewards = train_agent(trpo_agent, env_name, num_episodes)
+        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, figure_path, metrics_path = save_artifacts(ppo_rewards, trpo_rewards, output_dir, args.env_name)
-    fig, axes = plt.subplots(1, 2, figsize=(16, 6))
 
-    # 左图：原始奖励曲线
+    tqdm.write(f"对比图已保存至 {figure_path}")
-    axes[0].plot(ppo_rewards, alpha=0.3, color='blue', label='PPO (raw)')
+    tqdm.write(f"训练曲线原始数据已保存至 {metrics_path}")
-    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)
 
-    # 右图：滑动平均对比（更清晰）
+    if args.no_show:
-    axes[1].plot(smooth(ppo_rewards, 20), color='blue', linewidth=2, label='PPO')
+        plt.close(fig)
-    axes[1].plot(smooth(trpo_rewards, 20), color='red', linewidth=2, label='TRPO')
+    else:
-    axes[1].set_title('PPO vs TRPO (Smoothed, window=20)')
+        plt.show()
-    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 __name__ == '__main__':
+if __name__ == "__main__":
     main()

requirements.txt

@@ -2,3 +2,4 @@ numpy>=1.24
 matplotlib>=3.7
 torch>=2.1
 gymnasium[classic-control]>=0.29
+tqdm>=4.66