Hongru
771eba8607
- PPO: 改为 Actor/Critic 联合小批量训练，新增梯度裁剪 (max_grad_norm)， 分离 actor_lr/critic_lr，添加 get_value()，GAE 部分补充论文公式注释 - TRPO: 添加 get_value()，调整 tau 从 0.97 到 0.95 - Networks: 移除 PolicyNet 输出层的 tanh，初始化 log_std=0 以增强探索 - Main: 抽取 train_agent() 通用训练函数，新增 TRPO 训练和 PPO vs TRPO 对比曲线图（原始曲线 + 滑动平均平滑曲线）
117 lines
3.7 KiB
Python
117 lines
3.7 KiB
Python
import gymnasium as gym
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
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 = []
|
|
|
|
state, _ = env.reset()
|
|
memory = []
|
|
current_ep_reward = 0
|
|
episodes_completed = 0
|
|
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
|
|
|
|
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)
|
|
|
|
memory.append([state, action, reward_store, next_state, mask])
|
|
state = next_state
|
|
current_ep_reward += reward
|
|
step_count += 1
|
|
|
|
if done:
|
|
episode_rewards.append(current_ep_reward)
|
|
episodes_completed += 1
|
|
state, _ = env.reset()
|
|
current_ep_reward = 0
|
|
|
|
if episodes_completed % 10 == 0:
|
|
avg_reward = np.mean(episode_rewards[-10:])
|
|
print(f" Episode: {episodes_completed}, 平均奖励 (最近10轮): {avg_reward:.2f}")
|
|
|
|
if step_count >= batch_size:
|
|
agent.update(memory)
|
|
memory.clear()
|
|
step_count = 0
|
|
|
|
env.close()
|
|
return episode_rewards
|
|
|
|
|
|
def smooth(rewards, window=10):
|
|
"""滑动平均平滑曲线"""
|
|
smoothed = []
|
|
for i in range(len(rewards)):
|
|
start = max(0, i - window + 1)
|
|
smoothed.append(np.mean(rewards[start:i + 1]))
|
|
return smoothed
|
|
|
|
|
|
def main():
|
|
env_name = 'Pendulum-v1'
|
|
env = gym.make(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
|
|
|
|
# --- 训练 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)
|
|
|
|
# --- 训练 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)
|
|
|
|
# --- 对比画图 ---
|
|
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('PPO vs TRPO on Pendulum-v1')
|
|
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()
|
|
plt.savefig('ppo_vs_trpo_comparison.png', dpi=150)
|
|
plt.show()
|
|
print("对比图已保存至 ppo_vs_trpo_comparison.png")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|