新增 TRPO 算法实现，包括核心数学引擎、智能体、网络结构及训练入口，完善环境交互与数据处理功能

2026-03-25 15:38:08 +08:00
parent e53486fece
commit 7f9d7b2ee6
5 changed files with 660 additions and 11 deletions
@@ -0,0 +1,142 @@
+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()