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

添加 PPO 算法实现及相关配置，更新训练入口以支持 SAC 和 PPO 模式

Browse Source
This commit is contained in:
Hongru
2026-04-04 17:50:02 +08:00
parent f3d2d1a85f
commit a2ce5073c5
10 changed files with 675 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1 -1
View File
@@ -18,7 +18,7 @@ env/
*.pth
# Images
step_response.png
*.png
# Jupyter
.ipynb_checkpoints/
algorithms/ppo.py
+139
View File
@@ -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),
}
algorithms/sac.py
-1
View File
@@ -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 网络
configs/pendulum_config.yaml
+11 -14
View File
@@ -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 # 纯随机动作探索的步数
configs/ppo_pendulum_config.yaml
+25
View File
@@ -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 步截断)
eval.py
+8 -3
View File
@@ -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
models/networks.py
+76
View File
@@ -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)
train.py
+236 -54
View File
@@ -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()
utils/logger.py
+75
View File
@@ -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}")
utils/rollout_buffer.py
+104
View File
@@ -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