添加 PPO 算法实现及相关配置，更新训练入口以支持 SAC 和 PPO 模式
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user