143 lines
5.1 KiB
Python
143 lines
5.1 KiB
Python
import gymnasium as gym
|
|
import torch
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt # 新增：用于画图
|
|
from agent import TRPOAgent
|
|
from utils import RolloutBuffer
|
|
|
|
# 新增：用于计算滑动平均，让曲线更平滑
|
|
def moving_average(data, window_size=5):
|
|
"""计算简单滑动平均"""
|
|
if len(data) < window_size:
|
|
return data
|
|
return np.convolve(data, np.ones(window_size)/window_size, mode='valid')
|
|
|
|
def plot_learning_curve(rewards, smoothed_rewards):
|
|
"""绘制学习曲线"""
|
|
plt.figure(figsize=(10, 6))
|
|
|
|
# 绘制原始奖励（浅色）
|
|
plt.plot(rewards, color='blue', alpha=0.3, label='Raw Average Reward')
|
|
|
|
# 绘制平滑后的奖励（深色，粗线）
|
|
# 注意：滑动平均后数据点会变少，需要调整 X 轴起始位置
|
|
if len(smoothed_rewards) > 0:
|
|
x_ticks = np.arange(len(rewards) - len(smoothed_rewards), len(rewards))
|
|
plt.plot(x_ticks, smoothed_rewards, color='red', linewidth=2, label='Smoothed Reward (MA-5)')
|
|
|
|
plt.title('TRPO Training Performance on Pendulum-v1')
|
|
plt.xlabel('Iteration')
|
|
plt.ylabel('Average Reward')
|
|
plt.grid(True, linestyle='--', alpha=0.5)
|
|
plt.legend()
|
|
|
|
# 保存图片
|
|
plt.savefig('trpo_training_curve.png')
|
|
print("\n训练曲线图已保存为 'trpo_training_curve.png'")
|
|
|
|
# 如果有 GUI 界面则显示
|
|
try:
|
|
plt.show()
|
|
except Exception:
|
|
print("无法显示图形界面（可能是无头服务器），已跳过 plt.show()。")
|
|
|
|
def main():
|
|
# 1. 初始化环境
|
|
#env_name = 'Pendulum-v1'
|
|
# env = gym.make('Pendulum-v1')
|
|
# 兼容 Gymnasium
|
|
try:
|
|
env = gym.make('Pendulum-v1', render_mode=None)
|
|
except Exception:
|
|
env = gym.make('Pendulum-v1')
|
|
|
|
state_dim = env.observation_space.shape[0] # type: ignore
|
|
action_dim = env.action_space.shape[0] # type: ignore
|
|
max_action = float(env.action_space.high[0]) # type: ignore
|
|
min_action = float(env.action_space.low[0]) # type: ignore
|
|
|
|
print(f"环境加载成功! 状态维度: {state_dim}, 动作维度: {action_dim}")
|
|
print(f"动作范围: [{min_action}, {max_action}]")
|
|
|
|
# 2. 初始化智能体和经验池
|
|
agent = TRPOAgent(state_dim, action_dim, max_kl=0.01)
|
|
buffer = RolloutBuffer()
|
|
|
|
# 3. 设置训练超参数
|
|
max_iterations = 300 # 训练迭代总轮数
|
|
batch_size = 2000 # 每次更新收集的步数
|
|
|
|
# 新增：用于记录绘图数据
|
|
history_rewards = []
|
|
|
|
# 4. 主训练循环
|
|
for iteration in range(max_iterations):
|
|
state = env.reset()
|
|
if isinstance(state, tuple): state = state[0]
|
|
|
|
episode_rewards = []
|
|
ep_reward = 0
|
|
steps = 0
|
|
done = False
|
|
|
|
while steps < batch_size:
|
|
state_tensor = torch.FloatTensor(state).unsqueeze(0)
|
|
|
|
with torch.no_grad():
|
|
action, log_prob = agent.actor.get_action(state_tensor)
|
|
value = agent.critic(state_tensor)
|
|
|
|
action_np = action.squeeze(0).numpy()
|
|
clipped_action = np.clip(action_np, min_action, max_action)
|
|
|
|
# env.step 兼容性
|
|
step_result = env.step(clipped_action)
|
|
if len(step_result) == 5:
|
|
next_state, reward, terminated, truncated, _ = step_result
|
|
done = terminated or truncated
|
|
else:
|
|
next_state, reward, done, _ = step_result
|
|
|
|
ep_reward += reward
|
|
buffer.add(state=state, action=action_np, reward=reward,
|
|
next_state=next_state, done=done,
|
|
log_prob=log_prob.item(), value=value.item())
|
|
|
|
state = next_state
|
|
steps += 1
|
|
|
|
if done:
|
|
episode_rewards.append(ep_reward)
|
|
state = env.reset()
|
|
if isinstance(state, tuple): state = state[0]
|
|
ep_reward = 0
|
|
|
|
print(f"正在更新参数 (Iteration {iteration + 1}/{max_iterations})...")
|
|
agent.update(buffer, state, done)
|
|
|
|
# 5. 记录和打印日志
|
|
if episode_rewards:
|
|
avg_reward = np.mean(episode_rewards)
|
|
history_rewards.append(avg_reward) # 新增：记录数据
|
|
print(f"Iteration: {iteration + 1} | Average Reward: {avg_reward:.2f} | Max Reward: {np.max(episode_rewards):.2f}")
|
|
print("-" * 50)
|
|
else:
|
|
# 如果 batch_size 刚好结束时没有完成任何 episode，
|
|
# 为了画图不中断，我们沿用上一次的奖励（或者简单处理）
|
|
if history_rewards:
|
|
history_rewards.append(history_rewards[-1])
|
|
else:
|
|
history_rewards.append(-2000) # 初始默认低分
|
|
|
|
print("训练结束！")
|
|
env.close()
|
|
|
|
# 6. 新增：绘图
|
|
print("正在生成训练曲线图...")
|
|
smoothed = moving_average(history_rewards, window_size=5)
|
|
plot_learning_curve(history_rewards, smoothed)
|
|
|
|
if __name__ == '__main__':
|
|
# 确保安装了 matplotlib: pip install matplotlib
|
|
main()
|