01搭建强化学习绘制工具

Vita2026/1/2大约 6 分钟

搭建强化学习科研绘图工具

一、为什么要搭建这个工具

TensorBoard 是给“程序员”看的监控器，而论文图是给“审稿人”看的成果展示图。

在强化学习训练过程中，为了解决 TensorBoard 图片无法直接高质量下载、难以满足论文排版规范 的问题，我们搭建了一个基于 Python 的强化学习曲线绘制工具。

二、解决思路

第一步：提取数据

从 TensorBoard 日志中提取训练曲线数据。
将原始二进制数据转换为通用的 CSV 表格。

第二步：绘制图片

读取 CSV 数据，进行平滑处理（Exponential Moving Average）。
双层渲染：同时绘制原始波动曲线（Raw）和平滑后的趋势曲线（Smoothed）。
学术定制：自定义字体（衬线体）、颜色、高 DPI 以及图例。
自动化管理：按实验名称自动创建文件夹，输出高清 PNG 图片，满足论文、汇报和答辩需求。

三、搭建工具结构

创建一个 drawing_tool 文件夹，其目录结构如下：

drawing_tool/
├── image_env/              # 虚拟环境（存放安装的库）
├── scripts/                # 存放 Python 源代码
│   ├── extract_data.py     # 脚本1：数据提取
│   └── plot_results.py     # 脚本2：精准绘图
├── data/                   # 存放导出的 CSV 原始数据
└── figures/                # 存放生成的 PNG 图片
    └── [实验日期_名称]/      # 自动按实验生成的子文件夹

四、环境配置

在 drawing_tool 目录下执行以下命令，建立独立的绘图环境：

1. 创建并激活虚拟环境

python3 -m venv image_env

source image_env/bin/activate

2. 安装核心依赖

pip install tbparse pandas matplotlib seaborn

📚 核心库功能简介

1. tbparse —— 数据“翻译官”

作用：专门用于解析 TensorBoard 的二进制事件文件（.tfevents）。
为什么用它：TensorBoard 的原始数据是协议缓冲区（Protobuf）格式，人类无法直观读取。tbparse 能把这些复杂的二进制流一键转换成我们熟悉的 Pandas 表格（DataFrame）。

2. pandas —— 数据“管家”

作用：Python 中最强大的数据处理工具，类似于“代码版的 Excel”。
为什么用它：它负责存储从日志里提取出来的成千上万个数据点。我们可以利用它轻松完成数据清洗、按步数排序、指标筛选以及计算平滑平均值。

3. matplotlib —— 绘图“画板”

作用：Python 绘图界的“鼻祖”和基石。
为什么用它：它提供了极高的自由度。无论是调整坐标轴的字体（如 Times New Roman）、设置图片的分辨率（DPI），还是精确控制线条的粗细，都离不开它。它是我们生成 .png 或 .pdf 文件的核心引擎。

4. seaborn —— 视觉“美妆师”

作用：基于 matplotlib 开发的高级绘图库。
为什么用它：它自带了一套非常优雅的学术配色方案和绘图风格（如 ticks 风格）。它能让我们只用几行代码，就画出带有阴影区间或美观图例的专业曲线，让图片瞬间告别“廉价感”。

五、具体功能实现

1. 创建数据提取脚本 (`scripts/extract_data.py`)

该脚本负责将 TensorBoard 的二进制日志转换成 CSV。它会自动识别路径末尾的文件夹名作为文件名。

import os
import pandas as pd
from tbparse import SummaryReader

# --- 配置区 ---
# 修改为你需要提取的 TensorBoard 日志路径
LOG_DIR = "/home/isst/IsaacLab/logs/skrl/reach_ur10/2026-04-02_20-12-40_ppo_torch"

# 自动化命名：提取路径最后一部分作为文件名
exp_name = os.path.basename(LOG_DIR.rstrip("/"))
OUTPUT_CSV = os.path.join("../data", f"{exp_name}.csv")
# --------------

def main():
    if not os.path.exists(LOG_DIR):
        print(f"❌ 路径不存在：{LOG_DIR}")
        return

    print(f"🚀 正在提取实验数据: {exp_name}")
    
    # pivot=True 会自动将指标（Reward, Loss等）转换为列
    reader = SummaryReader(LOG_DIR, pivot=True)
    df = reader.scalars

    if df.empty:
        print("❓ 警告：该路径下未发现有效的 TensorBoard 数据。")
    else:
        os.makedirs("../data", exist_ok=True)
        df.to_csv(OUTPUT_CSV, index=False)
        print(f"✅ 提取成功！CSV 已存至: {OUTPUT_CSV}")

if __name__ == "__main__":
    main()

2. 创建精准绘图脚本 (`scripts/plot_results.py`)

该脚本读取 CSV，并为每个指标生成“背景噪声+主趋势线”的高清图，且自动分类存放。

import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import seaborn as sns
import os

# --- 核心配置区 ---
# 填入你需要画图的实验名称（需与 data 文件夹下的 CSV 文件名一致）
EXP_NAME = "2026-04-02_20-12-40_ppo_torch"

DATA_PATH = f"../data/{EXP_NAME}.csv"
OUTPUT_DIR = f"../figures/{EXP_NAME}"
SMOOTH_SPAN = 5  # 平滑系数，越小越还原波动细节
# -----------------

def plot_all_metrics():
    if not os.path.exists(DATA_PATH):
        print(f"❌ 找不到数据文件: {DATA_PATH}")
        return

    # 加载并排序数据
    df = pd.read_csv(DATA_PATH).sort_values('step')

    # 设置学术绘图风格
    sns.set_theme(style="ticks") 
    plt.rcParams.update({
        "font.family": "serif",      # 使用衬线字体（模拟 Times New Roman）
        "font.size": 11,
        "axes.grid": True,           # 开启网格
        "grid.linestyle": "--",
        "grid.alpha": 0.6
    })

    # 自动创建实验对应的子文件夹
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # 自动获取所有指标列
    metrics = [c for c in df.columns if c not in ['step', 'wall_time', 'dir_name']]
    
    print(f"🚀 正在为实验 [{EXP_NAME}] 生成 {len(metrics)} 张高清图片...")

    for m in metrics:
        fig = plt.figure(figsize=(8, 5), dpi=300)
        ax = fig.add_subplot(111)

        # 1. 绘制原始数据 (Raw) - 灰色透明背景，不省略任何数据细节
        ax.plot(df['step'], df[m], color='gray', alpha=0.3, label='Raw', linewidth=0.8)
        
        # 2. 绘制平滑数据 (Smoothed) - 蓝色主趋势线
        smoothed = df[m].ewm(span=SMOOTH_SPAN).mean()
        ax.plot(df['step'], smoothed, color='#1f77b4', linewidth=2, label='Smoothed')

        # 3. 细节美化与坐标轴处理
        ax.set_xlabel("Training Steps", fontweight='bold')
        ax.set_ylabel(m.split('/')[-1].strip(), fontweight='bold')
        ax.set_title(f"{EXP_NAME}\n{m.strip()}", fontsize=10, pad=10)
        
        # 坐标轴单位转换（如 10000 变为 10k）
        ax.xaxis.set_major_formatter(ticker.EngFormatter())
        ax.legend(loc='lower right', frameon=True)

        # 4. 清理文件名并保存
        safe_name = m.strip().replace('/', '_').replace(' ', '_').replace('(', '').replace(')', '').replace('__', '_')
        plt.tight_layout()
        plt.savefig(os.path.join(OUTPUT_DIR, f"{safe_name}.png"), bbox_inches='tight')
        plt.close(fig) 
        print(f" ✅ 已生成: {safe_name}.png")

    print(f"\n🎉 任务完成！所有图片已放入: {OUTPUT_DIR}")

if __name__ == "__main__":
    plot_all_metrics()