添加 PPO 算法实现及相关配置,更新训练入口以支持 SAC 和 PPO 模式
This commit is contained in:
+1
-1
@@ -18,7 +18,7 @@ env/
|
|||||||
*.pth
|
*.pth
|
||||||
|
|
||||||
# Images
|
# Images
|
||||||
step_response.png
|
*.png
|
||||||
|
|
||||||
# Jupyter
|
# Jupyter
|
||||||
.ipynb_checkpoints/
|
.ipynb_checkpoints/
|
||||||
|
|||||||
@@ -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),
|
||||||
|
}
|
||||||
@@ -94,7 +94,6 @@ class SAC(object):
|
|||||||
min_q_pi = torch.min(q1_pi, q2_pi)
|
min_q_pi = torch.min(q1_pi, q2_pi)
|
||||||
|
|
||||||
# 计算 Actor 的损失:最小化 (alpha * log_prob - min_Q)
|
# 计算 Actor 的损失:最小化 (alpha * log_prob - min_Q)
|
||||||
# 等价于最大化 (min_Q - alpha * log_prob) -> 既要 Q 值大,又要熵大(分布广)
|
|
||||||
actor_loss = (self.alpha * log_prob - min_q_pi).mean()
|
actor_loss = (self.alpha * log_prob - min_q_pi).mean()
|
||||||
|
|
||||||
# 优化 Actor 网络
|
# 优化 Actor 网络
|
||||||
|
|||||||
@@ -1,22 +1,19 @@
|
|||||||
# ==========================================
|
# ==========================================
|
||||||
# Soft Actor-Critic (SAC) - Pendulum-v1 配置
|
# Soft Actor-Critic (SAC)
|
||||||
|
# Pendulum-v1 配置
|
||||||
# ==========================================
|
# ==========================================
|
||||||
|
|
||||||
# --- 算法核心参数 ---
|
# --- 算法核心参数 ---
|
||||||
gamma: 0.99 # 折扣因子 (越接近1越看重长期收益)
|
gamma: 0.99 # 折扣因子
|
||||||
tau: 0.005 # 目标网络软更新系数 (EMA平滑系数,越小越稳定)
|
tau: 0.005 # 目标网络软更新系数 (EMA)
|
||||||
alpha: 0.2 # 熵温度系数 (控制探索力度,Pendulum中0.2比较合适,如果是复杂环境可能需要调整或自动学习)
|
alpha: 0.2 # 熵温度系数
|
||||||
lr: 0.0003 # 学习率 (Actor 和 Critic 保持一致,3e-4 是 Adam 优化器的万金油)
|
lr: 0.0003 # 学习率 (Adam, 3e-4)
|
||||||
|
|
||||||
# --- 经验回放池参数 ---
|
# --- 经验回放池参数 ---
|
||||||
buffer_size: 1000000 # 回放池最大容量 (100万条)
|
buffer_size: 1000000 # 回放池最大容量 (100万条)
|
||||||
batch_size: 256 # 每次梯度更新抽样的 batch 大小
|
batch_size: 256 # 每次梯度更新抽样的 batch 大小
|
||||||
|
|
||||||
# --- 训练循环控制 ---
|
# --- 训练循环控制 ---
|
||||||
max_episodes: 200 # 总共训练多少个回合 (Episode)
|
max_episodes: 200 # 总共训练多少个 episode
|
||||||
max_steps: 200 # 每个回合最多走多少步 (Gym Pendulum 默认 200 步截断)
|
max_steps: 200 # 每 episode 最多步数 (Pendulum-v1 默认 200 步截断)
|
||||||
start_steps: 10000 # 纯随机动作探索的步数 (用来快速填补经验池的高多样性数据)
|
start_steps: 10000 # 纯随机动作探索的步数
|
||||||
|
|
||||||
# --- 扩展参数 (为了后续画图和保存模型备用) ---
|
|
||||||
eval_freq: 10 # 每隔多少个回合评估一次策略
|
|
||||||
save_model: true # 是否保存最终模型权重
|
|
||||||
|
|||||||
@@ -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 步截断)
|
||||||
@@ -23,6 +23,11 @@ def evaluate_and_plot(model_path):
|
|||||||
config = yaml.safe_load(f)
|
config = yaml.safe_load(f)
|
||||||
|
|
||||||
env = gym.make('Pendulum-v1')
|
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]
|
state_dim = env.observation_space.shape[0]
|
||||||
action_dim = env.action_space.shape[0]
|
action_dim = env.action_space.shape[0]
|
||||||
max_action = float(env.action_space.high[0])
|
max_action = float(env.action_space.high[0])
|
||||||
@@ -47,9 +52,9 @@ def evaluate_and_plot(model_path):
|
|||||||
actions = []
|
actions = []
|
||||||
time_steps = []
|
time_steps = []
|
||||||
|
|
||||||
episode_reward = 0
|
episode_reward: float = 0.0
|
||||||
for step in range(config['max_steps']):
|
for step in range(config['max_steps']):
|
||||||
# 【重点】:设置 evaluate=True,让网络输出确定的均值动作,关闭随机探索
|
# 设置 evaluate=True,让网络输出确定的均值动作,关闭随机探索
|
||||||
action = agent.select_action(state, 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)
|
next_state, reward, terminated, truncated, _ = env.step(action)
|
||||||
state = next_state
|
state = next_state
|
||||||
episode_reward += reward
|
episode_reward += float(reward)
|
||||||
|
|
||||||
if terminated or truncated:
|
if terminated or truncated:
|
||||||
break
|
break
|
||||||
|
|||||||
@@ -96,3 +96,79 @@ class Actor(nn.Module):
|
|||||||
mean = torch.tanh(mean) * self.max_action
|
mean = torch.tanh(mean) * self.max_action
|
||||||
|
|
||||||
return action, log_prob, mean
|
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)
|
||||||
|
|||||||
@@ -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 numpy as np
|
||||||
import torch
|
import torch
|
||||||
|
import gymnasium as gym
|
||||||
|
from typing import cast
|
||||||
|
from gymnasium.spaces import Box
|
||||||
|
|
||||||
from algorithms.sac import SAC
|
from algorithms.sac import SAC
|
||||||
|
from algorithms.ppo import PPO
|
||||||
from utils.replay_buffer import ReplayBuffer
|
from utils.replay_buffer import ReplayBuffer
|
||||||
import yaml
|
from utils.logger import Logger
|
||||||
import os
|
|
||||||
|
|
||||||
# 读取 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}")
|
|
||||||
|
|
||||||
# ---------------------------------------------------------
|
def load_config(path: str) -> dict:
|
||||||
# 3. 实例化 SAC 代理与经验回放池
|
with open(path, "r", encoding="utf-8") as f:
|
||||||
# ---------------------------------------------------------
|
return yaml.safe_load(f)
|
||||||
agent = SAC(state_dim, action_dim, max_action, config)
|
|
||||||
|
|
||||||
|
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'])
|
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']):
|
for episode in range(config['max_episodes']):
|
||||||
# 重置环境,获取初始状态 (Gymnasium 返回 state 和 info)
|
|
||||||
state, _ = env.reset()
|
state, _ = env.reset()
|
||||||
episode_reward = 0
|
episode_reward: float = 0.0
|
||||||
|
|
||||||
for step in range(config['max_steps']):
|
for _ in range(config['max_steps']):
|
||||||
# --- a. 动作选择策略 ---
|
# 动作选择:预热期使用纯随机,之后用策略
|
||||||
# 强化学习工程技巧:在训练初期使用纯随机动作,收集高多样性的初始数据
|
|
||||||
if total_steps < config['start_steps']:
|
if total_steps < config['start_steps']:
|
||||||
action = env.action_space.sample()
|
action = env.action_space.sample()
|
||||||
else:
|
else:
|
||||||
# 预热结束后,交由 SAC 的策略网络进行带噪声的采样
|
|
||||||
action = agent.select_action(state, evaluate=False)
|
action = agent.select_action(state, evaluate=False)
|
||||||
|
|
||||||
# --- b. 与环境交互 ---
|
|
||||||
next_state, reward, terminated, truncated, _ = env.step(action)
|
next_state, reward, terminated, truncated, _ = env.step(action)
|
||||||
|
|
||||||
# 判断回合是否真正结束 (超时截断 truncated 不算做环境动力学意义上的 done)
|
|
||||||
done = float(terminated)
|
done = float(terminated)
|
||||||
|
|
||||||
# --- c. 存入经验池 ---
|
|
||||||
replay_buffer.add(state, action, reward, next_state, done)
|
replay_buffer.add(state, action, reward, next_state, done)
|
||||||
|
|
||||||
state = next_state
|
state = next_state
|
||||||
episode_reward += reward
|
episode_reward += float(reward)
|
||||||
total_steps += 1
|
total_steps += 1
|
||||||
|
|
||||||
# --- d. 核心学习逻辑 ---
|
|
||||||
# 只有当经验池里的数据量足够凑齐一个 Batch 时,才开始更新网络
|
|
||||||
if replay_buffer.size > config['batch_size']:
|
if replay_buffer.size > config['batch_size']:
|
||||||
agent.update(replay_buffer, config['batch_size'])
|
agent.update(replay_buffer, config['batch_size'])
|
||||||
|
|
||||||
# 如果提前倒地或撞毁,结束当前回合
|
|
||||||
if terminated or truncated:
|
if terminated or truncated:
|
||||||
break
|
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:
|
if (episode + 1) % 50 == 0:
|
||||||
torch.save(agent.actor.state_dict(), f"sac_actor_pendulum_ep{episode+1}.pth")
|
path = os.path.join(ROOT, f"sac_actor_pendulum_ep{episode+1}.pth")
|
||||||
print(f"[*] 已保存第 {episode+1} 回合的模型权重。")
|
torch.save(agent.actor.state_dict(), path)
|
||||||
|
print(f" [*] 模型已保存: {path}")
|
||||||
|
|
||||||
env.close()
|
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__":
|
if __name__ == "__main__":
|
||||||
main()
|
main()
|
||||||
|
|||||||
@@ -60,3 +60,78 @@ class Logger(object):
|
|||||||
plt.savefig(save_path, dpi=300, facecolor=fig.get_facecolor(), edgecolor='none')
|
plt.savefig(save_path, dpi=300, facecolor=fig.get_facecolor(), edgecolor='none')
|
||||||
plt.close()
|
plt.close()
|
||||||
print(f"[*] Learning curve saved to: {save_path}")
|
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}")
|
||||||
|
|||||||
@@ -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
|
||||||
Reference in New Issue
Block a user