63 lines
2.5 KiB
Python
63 lines
2.5 KiB
Python
# ==============================================================================
|
|
# Author: Hongru Liu
|
|
# Affiliation: School of Power and Energy, Northwestern Polytechnical University
|
|
# Version: 1.0
|
|
# Contact: hongruliu@mail.nwpu.edu.cn
|
|
# ==============================================================================
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import os
|
|
|
|
class Logger(object):
|
|
def __init__(self):
|
|
"""
|
|
初始化日志记录器，用于暂存训练过程中的各项指标
|
|
"""
|
|
self.episode_rewards = []
|
|
|
|
def record(self, reward):
|
|
"""
|
|
记录每个 Episode 的总奖励
|
|
"""
|
|
self.episode_rewards.append(reward)
|
|
|
|
def plot_learning_curve(self, save_dir="."):
|
|
"""
|
|
绘制并保存学习曲线 (Learning Curve)
|
|
"""
|
|
# 确保保存目录存在
|
|
os.makedirs(save_dir, exist_ok=True)
|
|
|
|
# 创建画布，严格设置白色背景
|
|
fig, ax = plt.subplots(figsize=(10, 6), facecolor='white')
|
|
ax.set_facecolor('white')
|
|
|
|
# 绘制奖励曲线，使用加粗线条以满足论文发表的视觉要求
|
|
ax.plot(self.episode_rewards, linewidth=2.5, color='#1f77b4', label='Episode Reward')
|
|
|
|
# 计算并绘制 10 个 Episode 的滑动平均线，让趋势更清晰
|
|
if len(self.episode_rewards) >= 10:
|
|
moving_avg = np.convolve(self.episode_rewards, np.ones(10)/10, mode='valid')
|
|
ax.plot(range(9, len(self.episode_rewards)), moving_avg,
|
|
linewidth=2.5, color='#ff7f0e', label='10-Episode Moving Average')
|
|
|
|
# 设置全英文的坐标轴标签和图例，调整字体大小
|
|
ax.set_xlabel('Episodes', fontsize=14, fontweight='bold')
|
|
ax.set_ylabel('Total Reward', fontsize=14, fontweight='bold')
|
|
ax.set_title('Training Learning Curve (Pendulum-v1)', fontsize=16, fontweight='bold')
|
|
|
|
# 设置刻度字体大小
|
|
ax.tick_params(axis='both', which='major', labelsize=12)
|
|
|
|
# 增加网格线并设置图例
|
|
ax.grid(True, linestyle='--', alpha=0.7)
|
|
ax.legend(fontsize=12, loc='lower right')
|
|
|
|
# 紧凑布局并保存
|
|
plt.tight_layout()
|
|
save_path = os.path.join(save_dir, "learning_curve.png")
|
|
plt.savefig(save_path, dpi=300, facecolor=fig.get_facecolor(), edgecolor='none')
|
|
plt.close()
|
|
print(f"[*] Learning curve saved to: {save_path}")
|