Hongru
6b8156eceb
- 新增 RL_Algothrithms 模块，包含 A2C、QAC 智能体 - 添加 SAC 章节笔记和 C10 笔记 - 上传训练结果图片 - 完善 README 与 .gitignore
72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
import gymnasium as gym
|
|
import torch
|
|
from agents.qac import QACAgent
|
|
from agents.a2c import A2CAgent
|
|
from utils import plot_comparison
|
|
|
|
def train_agent(env_name, agent_class, device, num_episodes=500):
|
|
"""
|
|
通用的训练循环函数
|
|
"""
|
|
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
|
|
|
|
# QAC 属于 Sarsa 类，需要提前采样第一个动作
|
|
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 = train_agent(ENV_NAME, QACAgent, device, EPISODES)
|
|
|
|
print("\n--- 开始训练 A2C ---")
|
|
a2c_rewards = train_agent(ENV_NAME, A2CAgent, device, EPISODES)
|
|
|
|
# 收集结果并绘图对比
|
|
results = {
|
|
'QAC (High Variance)': qac_rewards,
|
|
'A2C (Low Variance)': a2c_rewards
|
|
}
|
|
|
|
plot_comparison(results, window=50, save_path='qac_vs_a2c_gpu.png')
|