添加 DDPG、DPAC、Off-PAC 算法实现及连续动作空间训练代码

2026-03-19 10:29:22 +00:00
parent 6b8156eceb
commit 0cff2adcf3
8 changed files with 593 additions and 0 deletions
@@ -0,0 +1,95 @@
+import gymnasium as gym
+import torch
+import os
+import pickle
+from agents.qac import QACAgent
+from agents.a2c import A2CAgent
+from agents.off_pac import OffPACAgent
+from utils import plot_comparison
+
+RESULTS_DIR = 'training_results'
+
+def get_result_path(agent_class, env_name, num_episodes):
+    os.makedirs(RESULTS_DIR, exist_ok=True)
+    return os.path.join(RESULTS_DIR, f'{agent_class.__name__}_{env_name}_{num_episodes}.pkl')
+
+def load_or_train(env_name, agent_class, device, num_episodes):
+    """检查是否有保存的结果，有则加载，没有则训练并保存。"""
+    result_path = get_result_path(agent_class, env_name, num_episodes)
+    if os.path.exists(result_path):
+        with open(result_path, 'rb') as f:
+            rewards = pickle.load(f)
+        print(f"[{agent_class.__name__}] 已找到训练结果文件，直接加载: {result_path}")
+        return rewards
+
+    # 没有结果文件，开始训练
+    print(f"[{agent_class.__name__}] 未找到结果文件，开始训练...")
+    rewards = _train_agent(env_name, agent_class, device, num_episodes)
+
+    with open(result_path, 'wb') as f:
+        pickle.dump(rewards, f)
+    print(f"[{agent_class.__name__}] 训练完成，结果已保存至: {result_path}")
+    return rewards
+
+def _train_agent(env_name, agent_class, device, num_episodes):
+    """实际的训练循环逻辑"""
+    env = gym.make(env_name)
+    state_dim = env.observation_space.shape[0] # type: ignore
+    action_dim = int(env.action_space.n)       # type: ignore
+
+    agent = agent_class(state_dim, action_dim, device)
+    rewards_history = []
+
+    for episode in range(num_episodes):
+        state, _ = env.reset()
+        episode_reward = 0
+
+        action = agent.select_action(state)
+
+        while True:
+            next_state, reward, terminated, truncated, _ = env.step(action)
+            done = terminated or truncated
+
+            next_action = agent.select_action(next_state)
+            agent.update(state, action, reward, next_state, next_action, done)
+
+            state = next_state
+            action = next_action
+            episode_reward += reward # type: ignore
+
+            if done:
+                break
+
+        rewards_history.append(episode_reward)
+        if (episode + 1) % 100 == 0:
+            print(f"[{agent_class.__name__}] 回合 {episode+1}/{num_episodes}, 近100回合均分: {sum(rewards_history[-100:])/100:.2f}")
+
+    env.close()
+    return rewards_history
+
+if __name__ == "__main__":
+    # 检测 GPU
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"正在使用的计算设备: {device}")
+    
+    ENV_NAME = 'CartPole-v1'
+    EPISODES = 600
+    
+    print("\n--- 开始训练 QAC ---")
+    qac_rewards = load_or_train(ENV_NAME, QACAgent, device, EPISODES)
+
+    print("\n--- 开始训练 A2C ---")
+    a2c_rewards = load_or_train(ENV_NAME, A2CAgent, device, EPISODES)
+
+    print("\n--- 开始训练 Off-Policy A2C ---")
+    off_pac_rewards = load_or_train(ENV_NAME, OffPACAgent, device, EPISODES)
+    
+    # 将其加入字典一起画图对比
+    results = {
+        'QAC': qac_rewards,
+        'A2C': a2c_rewards,
+        'Off-Policy A2C': off_pac_rewards
+    }
+
+    
+    plot_comparison(results, window=50, save_path='qac_vs_a2c_gpu.png')