299 lines
11 KiB
Python
299 lines
11 KiB
Python
import argparse
|
|
import os
|
|
import random
|
|
from pathlib import Path
|
|
|
|
os.environ.setdefault("MKL_THREADING_LAYER", "GNU")
|
|
|
|
import gymnasium as gym
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import torch
|
|
from tqdm.auto import tqdm
|
|
|
|
from agent.ppo import PPOAgent
|
|
from agent.trpo import TRPOAgent
|
|
|
|
|
|
def parse_args():
|
|
cpu_count = os.cpu_count() or 1
|
|
default_envs = max(1, min(4, cpu_count))
|
|
default_threads = max(1, min(8, cpu_count))
|
|
|
|
parser = argparse.ArgumentParser(description="使用 GPU / 多核并行训练 PPO 与 TRPO，并输出对比图。")
|
|
parser.add_argument("--env-name", type=str, default="Pendulum-v1", help="Gymnasium 环境名称。")
|
|
parser.add_argument("--num-episodes", type=int, default=500, help="每个算法训练的 episode 数。")
|
|
parser.add_argument("--batch-size", type=int, default=2000, help="每次策略更新前收集的环境步数。")
|
|
parser.add_argument("--num-envs", type=int, default=default_envs, help="并行采样环境数。大于 1 时可显著提高吞吐。")
|
|
parser.add_argument("--vector-mode", choices=("sync", "async"), default="async", help="向量环境模式。async 会启动多进程，更适合多核 CPU。")
|
|
parser.add_argument("--cpu-threads", type=int, default=default_threads, help="PyTorch CPU 线程数。GPU 训练时主要影响 CPU 侧数据准备。")
|
|
parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto", help="训练设备。auto 会优先选择 CUDA。")
|
|
parser.add_argument("--seed", type=int, default=42, help="全局随机种子。")
|
|
parser.add_argument("--hidden-dim", type=int, default=128, help="Actor/Critic 隐层宽度。")
|
|
parser.add_argument("--output-dir", type=str, default="outputs", help="图像与训练数据的输出目录。")
|
|
parser.add_argument("--no-show", action="store_true", help="仅保存图像，不弹出 matplotlib 窗口。")
|
|
return parser.parse_args()
|
|
|
|
|
|
def resolve_device(device_name):
|
|
if device_name == "auto":
|
|
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
if device_name == "cuda" and not torch.cuda.is_available():
|
|
raise RuntimeError("请求使用 CUDA，但当前环境不可用。请检查 GPU 驱动和 PyTorch CUDA 版本。")
|
|
|
|
return torch.device(device_name)
|
|
|
|
|
|
def configure_runtime(cpu_threads, seed):
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
if cpu_threads is not None:
|
|
torch.set_num_threads(max(1, cpu_threads))
|
|
if hasattr(torch, "set_num_interop_threads"):
|
|
torch.set_num_interop_threads(max(1, min(cpu_threads, 4)))
|
|
|
|
if hasattr(torch, "set_float32_matmul_precision"):
|
|
torch.set_float32_matmul_precision("high")
|
|
|
|
|
|
def make_env(env_name, seed, worker_idx):
|
|
def _thunk():
|
|
env = gym.make(env_name)
|
|
env.reset(seed=seed + worker_idx)
|
|
env.action_space.seed(seed + worker_idx)
|
|
return env
|
|
|
|
return _thunk
|
|
|
|
|
|
def create_vector_env(env_name, num_envs, seed, vector_mode):
|
|
env_fns = [make_env(env_name, seed, idx) for idx in range(num_envs)]
|
|
if num_envs == 1 or vector_mode == "sync":
|
|
return gym.vector.SyncVectorEnv(env_fns)
|
|
return gym.vector.AsyncVectorEnv(env_fns, context="spawn")
|
|
|
|
|
|
def format_postfix(episode_rewards, total_env_steps, updates, last_update_stats):
|
|
postfix = {
|
|
"avg20": f"{np.mean(episode_rewards[-20:]):.1f}" if episode_rewards else "n/a",
|
|
"steps": total_env_steps,
|
|
"updates": updates,
|
|
}
|
|
for key, value in last_update_stats.items():
|
|
if isinstance(value, (int, float, np.floating)):
|
|
postfix[key] = f"{float(value):.4f}"
|
|
return postfix
|
|
|
|
|
|
def train_agent(agent, env_name, num_episodes=500, batch_size=2000, num_envs=1, seed=42, vector_mode="async", desc=None):
|
|
"""通用训练函数，适用于 PPO 和 TRPO，支持 GPU 和向量化并行采样。"""
|
|
env = create_vector_env(env_name, num_envs, seed, vector_mode)
|
|
episode_rewards = []
|
|
memory = {
|
|
"states": [],
|
|
"actions": [],
|
|
"rewards": [],
|
|
"next_states": [],
|
|
"masks": [],
|
|
}
|
|
total_env_steps = 0
|
|
updates = 0
|
|
last_update_stats = {}
|
|
|
|
states, _ = env.reset(seed=seed)
|
|
running_rewards = np.zeros(num_envs, dtype=np.float64)
|
|
step_count = 0
|
|
|
|
progress = tqdm(total=num_episodes, desc=desc or agent.__class__.__name__, dynamic_ncols=True)
|
|
|
|
try:
|
|
while len(episode_rewards) < num_episodes:
|
|
actions = agent.get_action(states)
|
|
next_states, rewards, terminated, truncated, _ = env.step(actions)
|
|
done = np.logical_or(terminated, truncated)
|
|
|
|
masks = (~done).astype(np.float32)
|
|
reward_store = rewards.astype(np.float32).copy()
|
|
truncated_only = np.logical_and(truncated, ~terminated)
|
|
if np.any(truncated_only):
|
|
bootstrap_values = np.asarray(agent.get_value(next_states[truncated_only]), dtype=np.float32)
|
|
reward_store[truncated_only] += agent.gamma * bootstrap_values
|
|
|
|
memory["states"].append(np.asarray(states, dtype=np.float32).copy())
|
|
memory["actions"].append(np.asarray(actions, dtype=np.float32).copy())
|
|
memory["rewards"].append(np.asarray(reward_store, dtype=np.float32).copy())
|
|
memory["next_states"].append(np.asarray(next_states, dtype=np.float32).copy())
|
|
memory["masks"].append(np.asarray(masks, dtype=np.float32).copy())
|
|
|
|
running_rewards += rewards
|
|
step_count += num_envs
|
|
total_env_steps += num_envs
|
|
|
|
if np.any(done):
|
|
done_indices = np.flatnonzero(done)
|
|
new_episode_rewards = running_rewards[done_indices].tolist()
|
|
remaining = num_episodes - len(episode_rewards)
|
|
accepted = new_episode_rewards[:remaining]
|
|
if accepted:
|
|
episode_rewards.extend(accepted)
|
|
progress.update(len(accepted))
|
|
|
|
running_rewards[done_indices] = 0.0
|
|
progress.set_postfix(format_postfix(episode_rewards, total_env_steps, updates, last_update_stats))
|
|
|
|
if len(episode_rewards) >= num_episodes:
|
|
break
|
|
|
|
reset_mask = np.zeros(num_envs, dtype=bool)
|
|
reset_mask[done_indices] = True
|
|
reset_states, _ = env.reset(options={"reset_mask": reset_mask})
|
|
next_states[done_indices] = reset_states[done_indices]
|
|
|
|
states = next_states
|
|
|
|
if step_count >= batch_size:
|
|
last_update_stats = agent.update(memory)
|
|
for key in memory:
|
|
memory[key].clear()
|
|
step_count = 0
|
|
updates += 1
|
|
progress.set_postfix(format_postfix(episode_rewards, total_env_steps, updates, last_update_stats))
|
|
|
|
if memory["states"]:
|
|
last_update_stats = agent.update(memory)
|
|
updates += 1
|
|
progress.set_postfix(format_postfix(episode_rewards, total_env_steps, updates, last_update_stats))
|
|
finally:
|
|
progress.close()
|
|
env.close()
|
|
|
|
summary = {
|
|
"episodes": len(episode_rewards),
|
|
"total_env_steps": total_env_steps,
|
|
"updates": updates,
|
|
"final_avg20": float(np.mean(episode_rewards[-20:])) if episode_rewards else float("nan"),
|
|
}
|
|
summary.update(last_update_stats)
|
|
return episode_rewards, summary
|
|
|
|
|
|
def smooth(rewards, window=10):
|
|
"""滑动平均平滑曲线。"""
|
|
rewards = np.asarray(rewards, dtype=np.float32)
|
|
if rewards.size == 0:
|
|
return rewards
|
|
|
|
smoothed = []
|
|
for i in range(rewards.size):
|
|
start = max(0, i - window + 1)
|
|
smoothed.append(float(np.mean(rewards[start:i + 1])))
|
|
return np.asarray(smoothed, dtype=np.float32)
|
|
|
|
|
|
def save_artifacts(ppo_rewards, trpo_rewards, output_dir, env_name):
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
fig, axes = plt.subplots(1, 2, figsize=(16, 6))
|
|
|
|
axes[0].plot(ppo_rewards, alpha=0.3, color="blue", label="PPO (raw)")
|
|
axes[0].plot(trpo_rewards, alpha=0.3, color="red", label="TRPO (raw)")
|
|
axes[0].plot(smooth(ppo_rewards, 20), color="blue", linewidth=2, label="PPO (smooth)")
|
|
axes[0].plot(smooth(trpo_rewards, 20), color="red", linewidth=2, label="TRPO (smooth)")
|
|
axes[0].set_title(f"PPO vs TRPO on {env_name}")
|
|
axes[0].set_xlabel("Episode")
|
|
axes[0].set_ylabel("Total Reward")
|
|
axes[0].legend()
|
|
axes[0].grid(True)
|
|
|
|
axes[1].plot(smooth(ppo_rewards, 20), color="blue", linewidth=2, label="PPO")
|
|
axes[1].plot(smooth(trpo_rewards, 20), color="red", linewidth=2, label="TRPO")
|
|
axes[1].set_title("PPO vs TRPO (Smoothed, window=20)")
|
|
axes[1].set_xlabel("Episode")
|
|
axes[1].set_ylabel("Total Reward")
|
|
axes[1].legend()
|
|
axes[1].grid(True)
|
|
|
|
plt.tight_layout()
|
|
|
|
figure_path = output_dir / "ppo_vs_trpo_comparison.png"
|
|
metrics_path = output_dir / "training_metrics.npz"
|
|
fig.savefig(figure_path, dpi=150)
|
|
np.savez(metrics_path, ppo_rewards=np.asarray(ppo_rewards), trpo_rewards=np.asarray(trpo_rewards))
|
|
|
|
return fig, figure_path, metrics_path
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
configure_runtime(args.cpu_threads, args.seed)
|
|
device = resolve_device(args.device)
|
|
output_dir = Path(args.output_dir)
|
|
|
|
env = gym.make(args.env_name)
|
|
state_dim = env.observation_space.shape[0]
|
|
action_dim = env.action_space.shape[0]
|
|
action_bound = float(env.action_space.high[0])
|
|
env.close()
|
|
|
|
tqdm.write(
|
|
f"配置: env={args.env_name}, device={device}, num_envs={args.num_envs}, "
|
|
f"vector_mode={args.vector_mode}, cpu_threads={args.cpu_threads}, batch_size={args.batch_size}"
|
|
)
|
|
|
|
ppo_agent = PPOAgent(
|
|
state_dim=state_dim,
|
|
action_dim=action_dim,
|
|
action_bound=action_bound,
|
|
hidden_dim=args.hidden_dim,
|
|
device=device,
|
|
)
|
|
ppo_rewards, ppo_summary = train_agent(
|
|
ppo_agent,
|
|
args.env_name,
|
|
num_episodes=args.num_episodes,
|
|
batch_size=args.batch_size,
|
|
num_envs=args.num_envs,
|
|
seed=args.seed,
|
|
vector_mode=args.vector_mode,
|
|
desc="PPO",
|
|
)
|
|
tqdm.write(f"PPO 完成: {ppo_summary}")
|
|
|
|
trpo_agent = TRPOAgent(
|
|
state_dim=state_dim,
|
|
action_dim=action_dim,
|
|
action_bound=action_bound,
|
|
hidden_dim=args.hidden_dim,
|
|
device=device,
|
|
)
|
|
trpo_rewards, trpo_summary = train_agent(
|
|
trpo_agent,
|
|
args.env_name,
|
|
num_episodes=args.num_episodes,
|
|
batch_size=args.batch_size,
|
|
num_envs=args.num_envs,
|
|
seed=args.seed + 10_000,
|
|
vector_mode=args.vector_mode,
|
|
desc="TRPO",
|
|
)
|
|
tqdm.write(f"TRPO 完成: {trpo_summary}")
|
|
|
|
fig, figure_path, metrics_path = save_artifacts(ppo_rewards, trpo_rewards, output_dir, args.env_name)
|
|
|
|
tqdm.write(f"对比图已保存至 {figure_path}")
|
|
tqdm.write(f"训练曲线原始数据已保存至 {metrics_path}")
|
|
|
|
if args.no_show:
|
|
plt.close(fig)
|
|
else:
|
|
plt.show()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|