修复并行采样并完善训练文档
This commit is contained in:
@@ -1,116 +1,298 @@
|
||||
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 train_agent(agent, env_name, num_episodes=500, batch_size=2000):
|
||||
"""通用训练函数，适用于 PPO 和 TRPO"""
|
||||
env = gym.make(env_name)
|
||||
episode_rewards = []
|
||||
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))
|
||||
|
||||
state, _ = env.reset()
|
||||
memory = []
|
||||
current_ep_reward = 0
|
||||
episodes_completed = 0
|
||||
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
|
||||
|
||||
while episodes_completed < num_episodes:
|
||||
action = agent.get_action(state)
|
||||
next_state, reward, terminated, truncated, _ = env.step(action)
|
||||
done = terminated or truncated
|
||||
progress = tqdm(total=num_episodes, desc=desc or agent.__class__.__name__, dynamic_ncols=True)
|
||||
|
||||
mask = 0.0 if done else 1.0
|
||||
reward_store = reward
|
||||
if truncated and not terminated:
|
||||
reward_store = reward + agent.gamma * agent.get_value(next_state)
|
||||
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)
|
||||
|
||||
memory.append([state, action, reward_store, next_state, mask])
|
||||
state = next_state
|
||||
current_ep_reward += reward
|
||||
step_count += 1
|
||||
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
|
||||
|
||||
if done:
|
||||
episode_rewards.append(current_ep_reward)
|
||||
episodes_completed += 1
|
||||
state, _ = env.reset()
|
||||
current_ep_reward = 0
|
||||
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())
|
||||
|
||||
if episodes_completed % 10 == 0:
|
||||
avg_reward = np.mean(episode_rewards[-10:])
|
||||
print(f" Episode: {episodes_completed}, 平均奖励 (最近10轮): {avg_reward:.2f}")
|
||||
running_rewards += rewards
|
||||
step_count += num_envs
|
||||
total_env_steps += num_envs
|
||||
|
||||
if step_count >= batch_size:
|
||||
agent.update(memory)
|
||||
memory.clear()
|
||||
step_count = 0
|
||||
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))
|
||||
|
||||
env.close()
|
||||
return episode_rewards
|
||||
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(len(rewards)):
|
||||
for i in range(rewards.size):
|
||||
start = max(0, i - window + 1)
|
||||
smoothed.append(np.mean(rewards[start:i + 1]))
|
||||
return smoothed
|
||||
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():
|
||||
env_name = 'Pendulum-v1'
|
||||
env = gym.make(env_name)
|
||||
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()
|
||||
|
||||
num_episodes = 500
|
||||
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 ---
|
||||
print("=" * 50)
|
||||
print("开始训练 PPO 智能体...")
|
||||
print("=" * 50)
|
||||
ppo_agent = PPOAgent(state_dim=state_dim, action_dim=action_dim, action_bound=action_bound)
|
||||
ppo_rewards = train_agent(ppo_agent, env_name, num_episodes)
|
||||
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 ---
|
||||
print("=" * 50)
|
||||
print("开始训练 TRPO 智能体...")
|
||||
print("=" * 50)
|
||||
trpo_agent = TRPOAgent(state_dim=state_dim, action_dim=action_dim, action_bound=action_bound)
|
||||
trpo_rewards = train_agent(trpo_agent, env_name, num_episodes)
|
||||
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, axes = plt.subplots(1, 2, figsize=(16, 6))
|
||||
fig, figure_path, metrics_path = save_artifacts(ppo_rewards, trpo_rewards, output_dir, args.env_name)
|
||||
|
||||
# 左图：原始奖励曲线
|
||||
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('PPO vs TRPO on Pendulum-v1')
|
||||
axes[0].set_xlabel('Episode')
|
||||
axes[0].set_ylabel('Total Reward')
|
||||
axes[0].legend()
|
||||
axes[0].grid(True)
|
||||
tqdm.write(f"对比图已保存至 {figure_path}")
|
||||
tqdm.write(f"训练曲线原始数据已保存至 {metrics_path}")
|
||||
|
||||
# 右图：滑动平均对比（更清晰）
|
||||
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()
|
||||
plt.savefig('ppo_vs_trpo_comparison.png', dpi=150)
|
||||
plt.show()
|
||||
print("对比图已保存至 ppo_vs_trpo_comparison.png")
|
||||
if args.no_show:
|
||||
plt.close(fig)
|
||||
else:
|
||||
plt.show()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user