From a2ce5073c5a9c50e14ed45b1e02202a4602e7995 Mon Sep 17 00:00:00 2001 From: Hongru Date: Sat, 4 Apr 2026 17:50:02 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=20PPO=20=E7=AE=97=E6=B3=95?= =?UTF-8?q?=E5=AE=9E=E7=8E=B0=E5=8F=8A=E7=9B=B8=E5=85=B3=E9=85=8D=E7=BD=AE?= =?UTF-8?q?=EF=BC=8C=E6=9B=B4=E6=96=B0=E8=AE=AD=E7=BB=83=E5=85=A5=E5=8F=A3?= =?UTF-8?q?=E4=BB=A5=E6=94=AF=E6=8C=81=20SAC=20=E5=92=8C=20PPO=20=E6=A8=A1?= =?UTF-8?q?=E5=BC=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 2 +- algorithms/ppo.py | 139 +++++++++++++++ algorithms/sac.py | 1 - configs/pendulum_config.yaml | 25 ++- configs/ppo_pendulum_config.yaml | 25 +++ eval.py | 11 +- models/networks.py | 76 ++++++++ train.py | 290 +++++++++++++++++++++++++------ utils/logger.py | 75 ++++++++ utils/rollout_buffer.py | 104 +++++++++++ 10 files changed, 675 insertions(+), 73 deletions(-) create mode 100644 algorithms/ppo.py create mode 100644 configs/ppo_pendulum_config.yaml create mode 100644 utils/rollout_buffer.py diff --git a/.gitignore b/.gitignore index 7d1c5e5..83d69b4 100644 --- a/.gitignore +++ b/.gitignore @@ -18,7 +18,7 @@ env/ *.pth # Images -step_response.png +*.png # Jupyter .ipynb_checkpoints/ diff --git a/algorithms/ppo.py b/algorithms/ppo.py new file mode 100644 index 0000000..0c372c2 --- /dev/null +++ b/algorithms/ppo.py @@ -0,0 +1,139 @@ +import torch +import torch.nn.functional as F +import torch.optim as optim +import numpy as np +from models.networks import PPOActor, ValueNet +from utils.rollout_buffer import RolloutBuffer + + +class PPO: + """ + Proximal Policy Optimization (PPO-Clip) + 参考论文：Schulman et al., 2017 (arXiv:1707.06347) + + 核心目标函数（论文公式9）： + L^{CLIP+VF+S} = E[ L^CLIP - c1 * L^VF + c2 * S[π](s) ] + + 使用独立的 Actor 和 Value 网络，GAE 优势估计（公式11/12）。 + """ + + def __init__(self, state_dim, action_dim, max_action, config): + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # 超参数 + self.gamma = config.get('gamma', 0.99) + self.gae_lambda = config.get('gae_lambda', 0.95) + self.clip_epsilon = config.get('clip_epsilon', 0.2) + self.n_epochs = config.get('n_epochs', 10) + self.batch_size = config.get('batch_size', 64) + self.vf_coef = config.get('vf_coef', 0.5) + self.entropy_coef = config.get('entropy_coef', 0.01) + self.max_grad_norm = config.get('max_grad_norm', 0.5) + steps_per_update = config.get('steps_per_update', 2048) + + # 策略网络 (Actor) + self.actor = PPOActor(state_dim, action_dim, max_action).to(self.device) + self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=config.get('lr', 3e-4)) + + # 价值网络 (V-net) + self.value_net = ValueNet(state_dim).to(self.device) + self.value_optimizer = optim.Adam(self.value_net.parameters(), lr=config.get('lr', 3e-4)) + + # On-policy 滚动缓冲区 + self.rollout = RolloutBuffer( + state_dim, action_dim, steps_per_update, + self.gamma, self.gae_lambda, self.device + ) + + # ------------------------------------------------------------------ + # 与环境交互 + # ------------------------------------------------------------------ + + @torch.no_grad() + def select_action(self, state): + """ + 采样动作并返回： + action (np.ndarray): clamp 后的实际动作 + action_unbounded (np.ndarray): 未裁剪的高斯采样值（存入 RolloutBuffer 用于 evaluate） + log_prob (float) + value (float): V(s) + """ + state_t = torch.FloatTensor(state).unsqueeze(0).to(self.device) + + # Actor：采样 → clamp + action, log_prob, action_unbounded = self.actor.sample(state_t) + + # Critic：价值估计 + value = self.value_net(state_t) + + return ( + action.cpu().numpy().flatten(), + action_unbounded.cpu().numpy().flatten(), + log_prob.cpu().item(), + value.cpu().item(), + ) + + @torch.no_grad() + def get_value(self, state): + """获取当前状态的 V(s)，用于 GAE 计算的 last_value。""" + state_t = torch.FloatTensor(state).unsqueeze(0).to(self.device) + return self.value_net(state_t).cpu().item() + + # ------------------------------------------------------------------ + # 核心更新 + # ------------------------------------------------------------------ + + def update(self): + """ + 使用 RolloutBuffer 中收集的数据，进行 K 轮 mini-batch 更新。 + 对应论文 Algorithm 1。 + 注意：调用前需已执行 rollout.compute_returns_and_advantages() + """ + actor_losses, value_losses, entropy_bonuses = [], [], [] + + for _ in range(self.n_epochs): + for states, actions_unbounded, old_log_probs, returns, advantages in \ + self.rollout.get_batches(self.batch_size): + + # ---- 重新计算当前策略的 log_prob 和熵 ---- + new_log_probs, entropy = self.actor.evaluate(states, actions_unbounded) + + # 概率比率 r_t(θ) = π_θ(a|s) / π_θ_old(a|s) + ratio = torch.exp(new_log_probs - old_log_probs) + + # ---- L^CLIP 目标（论文公式7）---- + surrogate1 = ratio * advantages + surrogate2 = torch.clamp(ratio, 1 - self.clip_epsilon, 1 + self.clip_epsilon) * advantages + actor_loss = -torch.min(surrogate1, surrogate2).mean() + + # ---- L^VF 价值损失 ---- + current_values = self.value_net(states) + value_loss = F.mse_loss(current_values, returns) + + # ---- 熵奖励 ---- + entropy_bonus = entropy.mean() + + # ---- 总损失（公式9）---- + loss = actor_loss + self.vf_coef * value_loss - self.entropy_coef * entropy_bonus + + # ---- 梯度更新 ---- + self.actor_optimizer.zero_grad() + self.value_optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(self.actor.parameters(), self.max_grad_norm) + torch.nn.utils.clip_grad_norm_(self.value_net.parameters(), self.max_grad_norm) + self.actor_optimizer.step() + self.value_optimizer.step() + + actor_losses.append(actor_loss.item()) + value_losses.append(value_loss.item()) + entropy_bonuses.append(entropy_bonus.item()) + + # 清空缓冲区，准备下一轮收集 + self.rollout.clear() + + return { + 'actor_loss': np.mean(actor_losses), + 'value_loss': np.mean(value_losses), + 'entropy': np.mean(entropy_bonuses), + } diff --git a/algorithms/sac.py b/algorithms/sac.py index 1b349b5..cf033af 100644 --- a/algorithms/sac.py +++ b/algorithms/sac.py @@ -94,7 +94,6 @@ class SAC(object): min_q_pi = torch.min(q1_pi, q2_pi) # 计算 Actor 的损失：最小化 (alpha * log_prob - min_Q) - # 等价于最大化 (min_Q - alpha * log_prob) -> 既要 Q 值大，又要熵大（分布广） actor_loss = (self.alpha * log_prob - min_q_pi).mean() # 优化 Actor 网络 diff --git a/configs/pendulum_config.yaml b/configs/pendulum_config.yaml index c627dce..e982b9c 100644 --- a/configs/pendulum_config.yaml +++ b/configs/pendulum_config.yaml @@ -1,22 +1,19 @@ # ========================================== -# Soft Actor-Critic (SAC) - Pendulum-v1 配置 +# Soft Actor-Critic (SAC) +# Pendulum-v1 配置 # ========================================== # --- 算法核心参数 --- -gamma: 0.99 # 折扣因子 (越接近1越看重长期收益) -tau: 0.005 # 目标网络软更新系数 (EMA平滑系数，越小越稳定) -alpha: 0.2 # 熵温度系数 (控制探索力度，Pendulum中0.2比较合适，如果是复杂环境可能需要调整或自动学习) -lr: 0.0003 # 学习率 (Actor 和 Critic 保持一致，3e-4 是 Adam 优化器的万金油) +gamma: 0.99 # 折扣因子 +tau: 0.005 # 目标网络软更新系数 (EMA) +alpha: 0.2 # 熵温度系数 +lr: 0.0003 # 学习率 (Adam, 3e-4) # --- 经验回放池参数 --- -buffer_size: 1000000 # 回放池最大容量 (100万条) -batch_size: 256 # 每次梯度更新抽样的 batch 大小 +buffer_size: 1000000 # 回放池最大容量 (100万条) +batch_size: 256 # 每次梯度更新抽样的 batch 大小 # --- 训练循环控制 --- -max_episodes: 200 # 总共训练多少个回合 (Episode) -max_steps: 200 # 每个回合最多走多少步 (Gym Pendulum 默认 200 步截断) -start_steps: 10000 # 纯随机动作探索的步数 (用来快速填补经验池的高多样性数据) - -# --- 扩展参数 (为了后续画图和保存模型备用) --- -eval_freq: 10 # 每隔多少个回合评估一次策略 -save_model: true # 是否保存最终模型权重 +max_episodes: 200 # 总共训练多少个 episode +max_steps: 200 # 每 episode 最多步数 (Pendulum-v1 默认 200 步截断) +start_steps: 10000 # 纯随机动作探索的步数 diff --git a/configs/ppo_pendulum_config.yaml b/configs/ppo_pendulum_config.yaml new file mode 100644 index 0000000..7109d1b --- /dev/null +++ b/configs/ppo_pendulum_config.yaml @@ -0,0 +1,25 @@ +# ========================================== +# Proximal Policy Optimization (PPO-Clip) +# Pendulum-v1 配置 +# 参考论文: Schulman et al., 2017 (arXiv:1707.06347) +# ========================================== + +# --- 算法核心参数 --- +gamma: 0.9 # 折扣因子 (Pendulum 短周期任务用 0.9 比 0.99 更易收敛) +gae_lambda: 0.95 # GAE λ (论文 Table 3) +clip_epsilon: 0.2 # 概率比率裁剪范围 [1-ε, 1+ε] +lr: 0.001 # 学习率 (PPO on-policy 更新少，适当提高 lr) + +# --- 网络更新参数 --- +n_epochs: 10 # 每轮收集后用同一批数据重复优化的 epoch 数 +batch_size: 64 # mini-batch 大小 +vf_coef: 0.5 # 价值损失系数 c1 (公式9) +entropy_coef: 0.0 # 熵奖励系数 c2 (Pendulum 简单任务，不需要额外探索奖励) +max_grad_norm: 0.5 # 梯度裁剪上限 + +# --- 数据收集参数 --- +steps_per_update: 1024 # 每次更新前收集的步数 (缩短到 5 个 episode 更新一次，加速学习) + +# --- 训练循环控制 --- +max_episodes: 500 # PPO 是 on-policy，需要更多 episode 才能收敛 +max_steps: 200 # 每 episode 最多步数 (Pendulum-v1 默认 200 步截断) diff --git a/eval.py b/eval.py index 69e4387..a474077 100644 --- a/eval.py +++ b/eval.py @@ -23,6 +23,11 @@ def evaluate_and_plot(model_path): config = yaml.safe_load(f) env = gym.make('Pendulum-v1') + if not isinstance(env.observation_space, gym.spaces.Box) or env.observation_space.shape is None: + raise TypeError("Pendulum-v1 observation_space must be a Box with a valid shape.") + if not isinstance(env.action_space, gym.spaces.Box) or env.action_space.shape is None: + raise TypeError("Pendulum-v1 action_space must be a Box with a valid shape.") + state_dim = env.observation_space.shape[0] action_dim = env.action_space.shape[0] max_action = float(env.action_space.high[0]) @@ -47,9 +52,9 @@ def evaluate_and_plot(model_path): actions = [] time_steps = [] - episode_reward = 0 + episode_reward: float = 0.0 for step in range(config['max_steps']): - # 【重点】：设置 evaluate=True，让网络输出确定的均值动作，关闭随机探索 + # 设置 evaluate=True，让网络输出确定的均值动作，关闭随机探索 action = agent.select_action(state, evaluate=True) # 记录当前步的数据 @@ -66,7 +71,7 @@ def evaluate_and_plot(model_path): # 与环境交互 next_state, reward, terminated, truncated, _ = env.step(action) state = next_state - episode_reward += reward + episode_reward += float(reward) if terminated or truncated: break diff --git a/models/networks.py b/models/networks.py index eedfc20..1d9334b 100644 --- a/models/networks.py +++ b/models/networks.py @@ -96,3 +96,79 @@ class Actor(nn.Module): mean = torch.tanh(mean) * self.max_action return action, log_prob, mean + + +# =========================================================================== +# PPO 专用网络 +# =========================================================================== + +class PPOActor(nn.Module): + """ + PPO 策略网络 — 标准高斯策略（不使用 tanh 压缩） + + 与 SAC Actor 的核心区别： + - SAC 需要 tanh squashing + log_prob Jacobian 修正来精确计算熵 + - PPO 直接使用高斯分布的 log_prob / entropy，再 clamp 到合法范围 + - log_std 是全局可学习参数（不依赖状态），更稳定 + """ + def __init__(self, state_dim, action_dim, max_action): + super(PPOActor, self).__init__() + self.max_action = max_action + + self.net = nn.Sequential( + nn.Linear(state_dim, 64), + nn.Tanh(), + nn.Linear(64, 64), + nn.Tanh(), + nn.Linear(64, action_dim), + ) + # 全局可学习的对数标准差，初始化为 0 → std=1.0（足够的初始探索） + self.log_std = nn.Parameter(torch.zeros(action_dim)) + + def forward(self, state): + mean = self.net(state) + std = self.log_std.exp().expand_as(mean) + return mean, std + + def get_dist(self, state): + mean, std = self.forward(state) + return Normal(mean, std) + + def sample(self, state): + """ + 采样动作，直接 clamp 到 [-max_action, max_action] + 返回: (action, log_prob) + """ + dist = self.get_dist(state) + action_unbounded = dist.rsample() + # 直接 clamp（不做 tanh，避免 log_prob 被 Jacobian 修正污染） + action = torch.clamp(action_unbounded, -self.max_action, self.max_action) + log_prob = dist.log_prob(action_unbounded).sum(1, keepdim=True) + return action, log_prob, action_unbounded + + def evaluate(self, state, action_unbounded): + """ + 给定之前保存的未裁剪动作，重新计算 log_prob 和熵（用于 PPO K 轮更新） + """ + dist = self.get_dist(state) + log_prob = dist.log_prob(action_unbounded).sum(1, keepdim=True) + entropy = dist.entropy().sum(1, keepdim=True) + return log_prob, entropy + + +class ValueNet(nn.Module): + """ + PPO 价值网络：估计状态价值函数 V(s) + """ + def __init__(self, state_dim): + super(ValueNet, self).__init__() + self.net = nn.Sequential( + nn.Linear(state_dim, 256), + nn.Tanh(), + nn.Linear(256, 256), + nn.Tanh(), + nn.Linear(256, 1), + ) + + def forward(self, state): + return self.net(state) diff --git a/train.py b/train.py index e0af75c..0719e32 100644 --- a/train.py +++ b/train.py @@ -1,86 +1,268 @@ -import gymnasium as gym +""" +训练入口 — 支持 SAC / PPO / 对比(compare) 三种模式 + +用法: + python train.py --algo sac # 仅训练 SAC（行为与原始版本一致） + python train.py --algo ppo # 仅训练 PPO + python train.py --algo compare # 依次训练 SAC 和 PPO，结束后输出对比曲线 + +可选参数: + --env Gymnasium 环境 ID（默认 Pendulum-v1） + --seed 随机种子（默认 0） +""" +import argparse +import os +import yaml import numpy as np import torch +import gymnasium as gym +from typing import cast +from gymnasium.spaces import Box + from algorithms.sac import SAC +from algorithms.ppo import PPO from utils.replay_buffer import ReplayBuffer -import yaml -import os +from utils.logger import Logger -# 读取 YAML 配置文件 -config_path = os.path.join(os.path.dirname(__file__), 'configs', 'pendulum_config.yaml') -with open(config_path, "r", encoding="utf-8") as f: - config = yaml.safe_load(f) -def main(): - # --------------------------------------------------------- - # 2. 实例化环境与获取维度信息 - # --------------------------------------------------------- - env = gym.make('Pendulum-v1') - - # 动态获取环境的维度信息，确保算法通用性 - state_dim = env.observation_space.shape[0] - action_dim = env.action_space.shape[0] - max_action = float(env.action_space.high[0]) - - print(f"环境已加载: 状态维度 {state_dim}, 动作维度 {action_dim}, 最大动作界限 {max_action}") +# =========================================================================== +# 配置加载 +# =========================================================================== - # --------------------------------------------------------- - # 3. 实例化 SAC 代理与经验回放池 - # --------------------------------------------------------- - agent = SAC(state_dim, action_dim, max_action, config) +def load_config(path: str) -> dict: + with open(path, "r", encoding="utf-8") as f: + return yaml.safe_load(f) + + +ROOT = os.path.dirname(os.path.abspath(__file__)) +SAC_CFG_PATH = os.path.join(ROOT, "configs", "pendulum_config.yaml") +PPO_CFG_PATH = os.path.join(ROOT, "configs", "ppo_pendulum_config.yaml") + + +# =========================================================================== +# 工具函数 +# =========================================================================== + +def set_seed(seed: int, env: gym.Env): + np.random.seed(seed) + torch.manual_seed(seed) + env.reset(seed=seed) + + +def make_env(env_id: str): + env = gym.make(env_id) + obs_space = cast(Box, env.observation_space) + act_space = cast(Box, env.action_space) + if obs_space.shape is None or act_space.shape is None: + raise ValueError("环境的 observation/action space 不支持 shape 维度读取。") + state_dim = obs_space.shape[0] + action_dim = act_space.shape[0] + max_action = float(act_space.high[0]) + return env, state_dim, action_dim, max_action + + +# =========================================================================== +# SAC 训练循环 (off-policy) +# =========================================================================== + +def train_sac(env_id: str, seed: int) -> list: + """ + 训练 SAC 并返回每个 episode 的总奖励列表。 + """ + config = load_config(SAC_CFG_PATH) + env, state_dim, action_dim, max_action = make_env(env_id) + set_seed(seed, env) + + print(f"\n{'='*60}") + print(f" Training SAC on {env_id}") + print(f" state_dim={state_dim}, action_dim={action_dim}, max_action={max_action}") + print(f"{'='*60}\n") + + agent = SAC(state_dim, action_dim, max_action, config) replay_buffer = ReplayBuffer(state_dim, action_dim, max_size=config['buffer_size']) + logger = Logger() - total_steps = 0 # 记录与环境交互的总步数 + total_steps = 0 - # --------------------------------------------------------- - # 4. 训练大循环 - # --------------------------------------------------------- for episode in range(config['max_episodes']): - # 重置环境，获取初始状态 (Gymnasium 返回 state 和 info) state, _ = env.reset() - episode_reward = 0 - - for step in range(config['max_steps']): - # --- a. 动作选择策略 --- - # 强化学习工程技巧：在训练初期使用纯随机动作，收集高多样性的初始数据 + episode_reward: float = 0.0 + + for _ in range(config['max_steps']): + # 动作选择：预热期使用纯随机，之后用策略 if total_steps < config['start_steps']: action = env.action_space.sample() else: - # 预热结束后，交由 SAC 的策略网络进行带噪声的采样 action = agent.select_action(state, evaluate=False) - - # --- b. 与环境交互 --- + next_state, reward, terminated, truncated, _ = env.step(action) - - # 判断回合是否真正结束 (超时截断 truncated 不算做环境动力学意义上的 done) done = float(terminated) - - # --- c. 存入经验池 --- + replay_buffer.add(state, action, reward, next_state, done) - state = next_state - episode_reward += reward + episode_reward += float(reward) total_steps += 1 - - # --- d. 核心学习逻辑 --- - # 只有当经验池里的数据量足够凑齐一个 Batch 时，才开始更新网络 + if replay_buffer.size > config['batch_size']: agent.update(replay_buffer, config['batch_size']) - - # 如果提前倒地或撞毁，结束当前回合 + if terminated or truncated: break - - # 打印当前回合的训练结果 - print(f"Episode: {episode+1:03d} | Total Steps: {total_steps:06d} | Reward: {episode_reward:.2f}") - # --- 阶段性保存模型 (可选) --- + logger.record(episode_reward) + print(f"[SAC] Episode: {episode+1:03d} | Steps: {total_steps:06d} | Reward: {episode_reward:.2f}") + if (episode + 1) % 50 == 0: - torch.save(agent.actor.state_dict(), f"sac_actor_pendulum_ep{episode+1}.pth") - print(f"[*] 已保存第 {episode+1} 回合的模型权重。") + path = os.path.join(ROOT, f"sac_actor_pendulum_ep{episode+1}.pth") + torch.save(agent.actor.state_dict(), path) + print(f" [*] 模型已保存: {path}") env.close() - print("训练结束！") + logger.plot_learning_curve(save_dir=ROOT) + print("\n[SAC] 训练完成！\n") + return logger.episode_rewards + + +# =========================================================================== +# PPO 训练循环 (on-policy) +# =========================================================================== + +def train_ppo(env_id: str, seed: int) -> list: + """ + 训练 PPO 并返回每个 episode 的总奖励列表。 + + PPO 是 on-policy 的：先收集固定 T 步数据（steps_per_update）， + 然后用这批数据做 K 轮 epoch 更新，再继续收集。 + """ + config = load_config(PPO_CFG_PATH) + env, state_dim, action_dim, max_action = make_env(env_id) + set_seed(seed, env) + + print(f"\n{'='*60}") + print(f" Training PPO on {env_id}") + print(f" state_dim={state_dim}, action_dim={action_dim}, max_action={max_action}") + print(f"{'='*60}\n") + + agent = PPO(state_dim, action_dim, max_action, config) + logger = Logger() + + steps_per_update = config.get('steps_per_update', 2048) + max_episodes = config['max_episodes'] + max_steps = config['max_steps'] + + total_steps = 0 + buffer_steps = 0 + episode = 0 + episode_reward: float = 0.0 + state, _ = env.reset() + + while episode < max_episodes: + action, action_raw, log_prob, value = agent.select_action(state) + next_state, reward, terminated, truncated, _ = env.step(action) + done = float(terminated) + + agent.rollout.add(state, action_raw, log_prob, reward, done, value) + state = next_state + episode_reward += float(reward) + total_steps += 1 + buffer_steps += 1 + + # ---- episode 结束 ---- + if terminated or truncated: + logger.record(episode_reward) + print(f"[PPO] Episode: {episode+1:03d} | Steps: {total_steps:06d} | Reward: {episode_reward:.2f}") + + # 阶段性保存（避免重复保存） + if (episode + 1) % 50 == 0: + path = os.path.join(ROOT, f"ppo_actor_pendulum_ep{episode+1}.pth") + torch.save(agent.actor.state_dict(), path) + print(f" [*] 模型已保存: {path}") + + episode += 1 + episode_reward = 0.0 + state, _ = env.reset() + + if episode >= max_episodes: + break + + # ---- 收集够 T 步 → 触发 PPO 更新 ---- + if buffer_steps >= steps_per_update: + last_value = 0.0 if done else agent.get_value(state) + agent.rollout.compute_returns_and_advantages(last_value) + info = agent.update() + buffer_steps = 0 + print(f" [PPO update] actor_loss={info['actor_loss']:.4f} | " + f"value_loss={info['value_loss']:.4f} | entropy={info['entropy']:.4f}") + + env.close() + + # 若缓冲区中还有剩余数据，做最后一次更新 + if buffer_steps > 0: + agent.rollout.compute_returns_and_advantages(0.0) + agent.update() + + logger.plot_learning_curve(save_dir=ROOT) + print("\n[PPO] 训练完成！\n") + return logger.episode_rewards + + +# =========================================================================== +# 对比模式：依次训练两者，输出对比曲线 +# =========================================================================== + +def train_compare(env_id: str, seed: int): + print("\n" + "="*60) + print(" Compare Mode: SAC vs PPO") + print("="*60) + + rewards_sac = train_sac(env_id, seed) + rewards_ppo = train_ppo(env_id, seed) + + Logger.plot_comparison( + rewards_dict={'SAC': rewards_sac, 'PPO': rewards_ppo}, + window=10, + save_dir=ROOT, + ) + print("\n[Compare] 对比训练完成！已生成 3 张图：") + print(" - comparison_curve.png (同 episode 范围对比)") + print(" - sac_learning_curve.png (SAC 独立完整曲线)") + print(" - ppo_learning_curve.png (PPO 独立完整曲线)") + + +# =========================================================================== +# 命令行入口 +# =========================================================================== + +def parse_args(): + parser = argparse.ArgumentParser( + description="SAC / PPO 强化学习训练脚本 (Pendulum-v1)" + ) + parser.add_argument( + "--algo", type=str, default="sac", + choices=["sac", "ppo", "compare"], + help="选择训练的算法: sac | ppo | compare (默认: sac)" + ) + parser.add_argument( + "--env", type=str, default="Pendulum-v1", + help="Gymnasium 环境 ID (默认: Pendulum-v1)" + ) + parser.add_argument( + "--seed", type=int, default=0, + help="随机种子 (默认: 0)" + ) + return parser.parse_args() + + +def main(): + args = parse_args() + + if args.algo == "sac": + train_sac(args.env, args.seed) + elif args.algo == "ppo": + train_ppo(args.env, args.seed) + elif args.algo == "compare": + train_compare(args.env, args.seed) + if __name__ == "__main__": main() diff --git a/utils/logger.py b/utils/logger.py index d5f7b29..0e24e43 100644 --- a/utils/logger.py +++ b/utils/logger.py @@ -60,3 +60,78 @@ class Logger(object): plt.savefig(save_path, dpi=300, facecolor=fig.get_facecolor(), edgecolor='none') plt.close() print(f"[*] Learning curve saved to: {save_path}") + + @staticmethod + def plot_comparison(rewards_dict, window=10, save_dir="."): + """ + 生成三张图： + 1. comparison_curve.png — 在同 episode 范围内对比（截取到最短算法的长度） + 2. sac_learning_curve.png — SAC 独立完整曲线 + 3. ppo_learning_curve.png — PPO 独立完整曲线 + + Args: + rewards_dict (dict): { 'SAC': [r1, r2, ...], 'PPO': [r1, r2, ...] } + window (int): 滑动平均窗口大小 + save_dir (str): 图片保存目录 + """ + os.makedirs(save_dir, exist_ok=True) + + palette = { + 'SAC': '#1f77b4', # 蓝色 + 'PPO': '#d62728', # 红色 + } + fallback = ['#2ca02c', '#9467bd', '#8c564b'] + + # ========== 图 1: 同 episode 对比（截取到最短长度） ========== + min_len = min(len(r) for r in rewards_dict.values()) + fig, ax = plt.subplots(figsize=(12, 7), facecolor='white') + ax.set_facecolor('white') + color_idx = 0 + for algo_name, rewards in rewards_dict.items(): + color = palette.get(algo_name, fallback[color_idx % len(fallback)]) + color_idx += 1 + r = rewards[:min_len] + eps = list(range(1, min_len + 1)) + ax.plot(eps, r, linewidth=1.0, color=color, alpha=0.3) + if min_len >= window: + avg = np.convolve(r, np.ones(window) / window, mode='valid') + ax.plot(range(window, min_len + 1), avg, + linewidth=2.5, color=color, + label=f'{algo_name} ({window}-ep avg)') + ax.set_xlabel('Episodes', fontsize=14, fontweight='bold') + ax.set_ylabel('Total Reward', fontsize=14, fontweight='bold') + ax.set_title(f'SAC vs PPO — Same Episode Range (1-{min_len})', fontsize=16, fontweight='bold') + ax.tick_params(axis='both', which='major', labelsize=12) + ax.grid(True, linestyle='--', alpha=0.7) + ax.legend(fontsize=13, loc='lower right') + plt.tight_layout() + p = os.path.join(save_dir, "comparison_curve.png") + plt.savefig(p, dpi=300, facecolor=fig.get_facecolor(), edgecolor='none') + plt.close() + print(f"[*] Comparison curve saved to: {p}") + + # ========== 图 2 & 3: 各算法独立完整曲线 ========== + for algo_name, rewards in rewards_dict.items(): + color = palette.get(algo_name, '#333333') + fig, ax = plt.subplots(figsize=(10, 6), facecolor='white') + ax.set_facecolor('white') + eps = list(range(1, len(rewards) + 1)) + ax.plot(eps, rewards, linewidth=1.0, color=color, alpha=0.3, label='Episode Reward') + if len(rewards) >= window: + avg = np.convolve(rewards, np.ones(window) / window, mode='valid') + ax.plot(range(window, len(rewards) + 1), avg, + linewidth=2.5, color=color, + label=f'{window}-Episode Moving Average') + ax.set_xlabel('Episodes', fontsize=14, fontweight='bold') + ax.set_ylabel('Total Reward', fontsize=14, fontweight='bold') + ax.set_title(f'{algo_name} Learning Curve — Pendulum-v1 ({len(rewards)} episodes)', + fontsize=16, fontweight='bold') + ax.tick_params(axis='both', which='major', labelsize=12) + ax.grid(True, linestyle='--', alpha=0.7) + ax.legend(fontsize=12, loc='lower right') + plt.tight_layout() + fname = f"{algo_name.lower()}_learning_curve.png" + p = os.path.join(save_dir, fname) + plt.savefig(p, dpi=300, facecolor=fig.get_facecolor(), edgecolor='none') + plt.close() + print(f"[*] {algo_name} learning curve saved to: {p}") diff --git a/utils/rollout_buffer.py b/utils/rollout_buffer.py new file mode 100644 index 0000000..930abb2 --- /dev/null +++ b/utils/rollout_buffer.py @@ -0,0 +1,104 @@ +import numpy as np +import torch + + +class RolloutBuffer: + """ + On-policy 滚动缓冲区，用于 PPO 算法。 + + 每次 collect() 收集 T 步数据后，调用 compute_returns_and_advantages() + 计算 GAE 优势估计，然后通过 get_batches() 将数据切分为 mini-batch + 供 K 轮 epoch 使用，最后 clear() 清空等待下一轮收集。 + """ + + def __init__(self, state_dim, action_dim, steps_per_update, gamma, gae_lambda, device): + self.steps = steps_per_update + self.gamma = gamma + self.gae_lambda = gae_lambda + self.device = device + + # 预分配存储空间 + self.states = np.zeros((steps_per_update, state_dim), dtype=np.float32) + self.actions_unbounded = np.zeros((steps_per_update, action_dim), dtype=np.float32) # clamp 之前的高斯采样值 + self.log_probs = np.zeros((steps_per_update, 1), dtype=np.float32) + self.rewards = np.zeros((steps_per_update, 1), dtype=np.float32) + self.dones = np.zeros((steps_per_update, 1), dtype=np.float32) + self.values = np.zeros((steps_per_update, 1), dtype=np.float32) + + # 计算后填充 + self.returns = np.zeros((steps_per_update, 1), dtype=np.float32) + self.advantages = np.zeros((steps_per_update, 1), dtype=np.float32) + + self.ptr = 0 + self.full = False + + def add(self, state, action_unbounded, log_prob, reward, done, value): + """ + 向缓冲区写入一步数据。 + action_unbounded: 未裁剪的高斯采样值（形状 [action_dim]） + log_prob: 该步的 log π(a|s)（标量） + value: V(s) 的估计值（标量） + """ + idx = self.ptr + self.states[idx] = state + self.actions_unbounded[idx] = action_unbounded + self.log_probs[idx] = log_prob + self.rewards[idx] = reward + self.dones[idx] = done + self.values[idx] = value + self.ptr += 1 + if self.ptr >= self.steps: + self.full = True + + def compute_returns_and_advantages(self, last_value): + """ + 反向遍历轨迹，计算 GAE 优势估计（论文公式11/12）。 + 只对实际填充的 self.ptr 步数据计算，避免无效数据参与。 + + Args: + last_value: V(s_{T+1})，下一个状态的价值估计（若 episode 结束则为 0） + """ + n = self.ptr # 实际有效数据条数 + last_gae = 0.0 + for t in reversed(range(n)): + if t == n - 1: + next_non_terminal = 1.0 - self.dones[t] + next_value = last_value + else: + next_non_terminal = 1.0 - self.dones[t] + next_value = self.values[t + 1] + + delta = self.rewards[t] + self.gamma * next_value * next_non_terminal - self.values[t] + last_gae = delta + self.gamma * self.gae_lambda * next_non_terminal * last_gae + self.advantages[t] = last_gae + + # 回报 G_t = Â_t + V(s_t) + self.returns[:n] = self.advantages[:n] + self.values[:n] + + # 优势归一化 + adv = self.advantages[:n] + self.advantages[:n] = (adv - adv.mean()) / (adv.std() + 1e-8) + + def get_batches(self, batch_size): + """ + 将实际填充的数据随机打乱后按 batch_size 切片，生成 mini-batch。 + + Yields: + (states, actions_unbounded, log_probs, returns, advantages) — 均为 Tensor + """ + n = self.ptr + indices = np.random.permutation(n) + for start in range(0, n, batch_size): + idx = indices[start: start + batch_size] + yield ( + torch.FloatTensor(self.states[idx]).to(self.device), + torch.FloatTensor(self.actions_unbounded[idx]).to(self.device), + torch.FloatTensor(self.log_probs[idx]).to(self.device), + torch.FloatTensor(self.returns[idx]).to(self.device), + torch.FloatTensor(self.advantages[idx]).to(self.device), + ) + + def clear(self): + """清空缓冲区，为下一轮收集做准备。""" + self.ptr = 0 + self.full = False