96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
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')
|