如何优雅地向TensorBoard添加PettingZoo环境自定义指标？

问题背景 #

• 现有训练流程使用TensorBoard监控PettingZoo多智能体环境的强化学习训练过程，目前已可正常观测训练过程中奖励值随迭代逐步上升的趋势
• 需求为新增更多环境专属自定义指标，在TensorBoard中生成对应指标随训练迭代变化的趋势图表
• 已验证修改SB3（Stable Baselines3）源码的方案可生效：在OnPolicyAlgorithm类的learn方法中插入指标统计代码后，可成功在TensorBoard面板生成目标自定义指标图表
• 该方案存在明显工程缺陷：直接修改第三方依赖源码不符合开发规范，自定义逻辑需要全部收敛在自有业务代码范围内

无侵入实现方案 #

不需要修改任何第三方依赖源码，以下三种方案均可实现需求，可根据业务场景选择：

方案1：基于SB3原生回调机制实现（最推荐） #

SB3原生提供了完整的训练生命周期回调钩子，是官方预留的自定义扩展入口，完全无侵入：

• 实现逻辑：
  1. 自定义继承BaseCallback的指标统计类，在对应生命周期方法中拉取环境侧的自定义指标
  2. 调用回调自带的日志记录方法写入指标，SB3会自动同步写入TensorBoard日志
  3. 训练时将自定义回调实例传入model.learn()的callback参数即可生效
• 最小实现代码参考：

from stable_baselines3.common.callbacks import BaseCallback
import numpy as np

class EnvCustomMetricCallback(BaseCallback):
    def __init__(self, verbose=0):
        super().__init__(verbose)

    def _on_step(self) -> bool:
        # 单步执行逻辑：可在此处做逐帧指标累计，返回True继续训练
        return True

    def _on_rollout_end(self) -> None:
        # 每个采样轮次结束时统一上报指标，避免逐step写入的性能损耗
        # 单环境场景直接读取unwrapped环境的自定义属性
        if self.training_env.num_envs == 1:
            coop_action_rate = self.training_env.envs[0].unwrapped.coop_action_count / max(self.training_env.envs[0].unwrapped.step_count, 1)
            invalid_action_rate = self.training_env.envs[0].unwrapped.invalid_action_count / max(self.training_env.envs[0].unwrapped.step_count, 1)
        # 多并行环境场景拉取所有环境实例的指标做聚合
        else:
            coop_counts = self.training_env.get_attr("coop_action_count")
            invalid_counts = self.training_env.get_attr("invalid_action_count")
            step_counts = self.training_env.get_attr("step_count")
            coop_action_rate = np.mean([c/max(s,1) for c,s in zip(coop_counts, step_counts)])
            invalid_action_rate = np.mean([i/max(s,1) for i,s in zip(invalid_counts, step_counts)])
        
        # 指标写入TensorBoard，会自动归到对应分组下
        self.logger.record("env_metrics/coop_action_rate", coop_action_rate)
        self.logger.record("env_metrics/invalid_action_rate", invalid_action_rate)

# 训练时传入回调即可
# model.learn(total_timesteps=2_000_000, callback=EnvCustomMetricCallback())

• 适配提示：PettingZoo环境的自定义指标只需要在环境类的step方法中做好计数，存为实例属性，就能通过上述回调方法读取到，不需要对环境做额外改造。

方案2：环境Wrapper封装指标逻辑（低耦合场景适用） #

如果需要让指标统计逻辑和训练框架完全解耦，方便后续切换不同RL训练库复用：

• 自定义继承gym.Wrapper的PettingZoo环境包装类，在包装类的step、reset方法中完成所有自定义指标的累计、重置逻辑，对外暴露指标读取接口
• 依然使用上述的自定义回调，从wrapper实例中读取聚合后的指标值上报即可
• 该方案下环境侧的指标统计逻辑完全独立，不管用SB3、RLlib还是其他训练框架，都可以直接复用包装后的环境，不需要重写指标统计代码。

方案3：独立日志写入（多来源指标场景适用） #

如果自定义指标不止来自环境交互，还包含离线评估、数据预处理、模型权重监控等多来源数据：

• 在自有业务代码中单独初始化TensorBoard的SummaryWriter实例，日志存储路径和SB3的日志路径保持一致
• 在任何需要记录指标的业务节点，直接调用writer.add_scalar()方法写入指标值和对应的全局训练步数即可
• 启动TensorBoard时指向统一的日志根目录，会自动加载SB3原生指标和自定义写入的指标，统一在面板展示
• 注意点：需要手动对齐SB3训练的全局步数，避免横轴错位导致趋势图异常。

内容的提问来源于stack exchange，提问作者Ram Rachum