72 lines
2.0 KiB
Python
72 lines
2.0 KiB
Python
import gymnasium as gym
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
from agent.ppo import PPOAgent
|
|
|
|
def main():
|
|
env = gym.make('Pendulum-v1')
|
|
state_dim = env.observation_space.shape[0]
|
|
action_dim = env.action_space.shape[0]
|
|
action_bound = float(env.action_space.high[0])
|
|
|
|
agent = PPOAgent(state_dim=state_dim, action_dim=action_dim, action_bound=action_bound)
|
|
|
|
num_episodes = 200
|
|
batch_size = 2000
|
|
|
|
episode_rewards = []
|
|
|
|
state, _ = env.reset()
|
|
memory = []
|
|
current_ep_reward = 0
|
|
episodes_completed = 0
|
|
|
|
print("开始训练 PPO 智能体...")
|
|
|
|
step_count = 0
|
|
while episodes_completed < num_episodes:
|
|
action = agent.get_action(state)
|
|
|
|
# 交互
|
|
next_state, reward, terminated, truncated, _ = env.step(action)
|
|
|
|
# 【极其关键的修复】：只有真正死亡 (terminated) 才清零未来价值
|
|
# 绝对不能把时间截断 (truncated) 算作 mask=0
|
|
mask = 0.0 if terminated else 1.0
|
|
|
|
done = terminated or truncated
|
|
memory.append([state, action, reward, next_state, mask])
|
|
|
|
state = next_state
|
|
current_ep_reward += reward
|
|
step_count += 1
|
|
|
|
if done:
|
|
episode_rewards.append(current_ep_reward)
|
|
episodes_completed += 1
|
|
state, _ = env.reset()
|
|
current_ep_reward = 0
|
|
|
|
if episodes_completed % 10 == 0:
|
|
avg_reward = np.mean(episode_rewards[-10:])
|
|
print(f"Episode: {episodes_completed}, 平均奖励 (最近10轮): {avg_reward:.2f}")
|
|
|
|
if step_count >= batch_size:
|
|
agent.update(memory)
|
|
memory.clear()
|
|
step_count = 0
|
|
|
|
env.close()
|
|
|
|
plt.figure(figsize=(10, 5))
|
|
plt.plot(episode_rewards)
|
|
plt.title('PPO Learning Curve on Pendulum-v1')
|
|
plt.xlabel('Episode')
|
|
plt.ylabel('Total Reward')
|
|
plt.grid(True)
|
|
plt.savefig('ppo_learning_curve_final.png')
|
|
plt.show()
|
|
|
|
if __name__ == '__main__':
|
|
main()
|