如何在Stable Baselines3自定义环境中添加奖励到TensorBoard日志

基础内置奖励记录（无需额外定制代码） #

你已经在PPO初始化时正确配置了tensorboard_log参数，仅需要在调用learn()方法时指定tb_log_name参数，Stable Baselines3就会自动记录每个episode的平均奖励、episode长度等核心指标到TensorBoard。
示例代码：

# 训练时添加tb_log_name参数指定实验标识
model.learn(
    total_timesteps=100000, # 按你的训练需求调整总步数
    tb_log_name="ppo_custom_env_exp1" # TensorBoard中显示的实验名称
)

启动TensorBoard查看日志的命令：

tensorboard --logdir ./tensorboard/

运行命令后访问终端输出的本地地址（默认是http://localhost:6006），即可在scalars面板查看内置的奖励曲线，对应指标名为rollout/ep_rew_mean。 #

自定义奖励分量记录 #

如果你的总奖励由多个子项组成，需要单独记录各个子项的变化趋势，需要通过自定义回调函数实现：

步骤1：自定义环境返回奖励分量 #

在你自定义环境的step()方法中，将需要单独记录的奖励分量存入返回的info字典中，示例：

def step(self, action):
    # 你的原有环境逻辑代码
    survival_reward = 0.1
    target_reward = 10 if reach_target else 0
    action_penalty = 0.05 * abs(action)
    total_reward = survival_reward + target_reward - action_penalty

    # 将自定义奖励分量存入info
    info = {
        "survival_reward": survival_reward,
        "target_reward": target_reward,
        "action_penalty": action_penalty
    }
    return obs, total_reward, terminated, truncated, info

步骤2：实现自定义回调类 #

继承SB3的BaseCallback类，在每步训练时提取info中的奖励分量写入日志：

from stable_baselines3.common.callbacks import BaseCallback

class CustomRewardLogCallback(BaseCallback):
    def __init__(self, verbose=0):
        super().__init__(verbose)
    
    def _on_step(self) -> bool:
        # 兼容单环境/向量化多环境场景
        for info in self.locals["infos"]:
            # 仅在episode结束时记录聚合指标，不需要聚合可去掉该判断
            if "episode" in info:
                # 记录episode总奖励（和内置指标重复，仅作示例）
                self.logger.record("custom/ep_total_reward", info["episode"]["r"])
                # 记录自定义奖励分量
                if "survival_reward" in info:
                    self.logger.record("custom/survival_reward", info["survival_reward"])
                if "target_reward" in info:
                    self.logger.record("custom/target_reward", info["target_reward"])
                if "action_penalty" in info:
                    self.logger.record("custom/action_penalty", info["action_penalty"])
        return True

步骤3：训练时传入回调 #

在调用learn()方法时传入实例化的回调对象即可：

reward_callback = CustomRewardLogCallback()
model.learn(
    total_timesteps=100000,
    tb_log_name="ppo_custom_env_exp1",
    callback=reward_callback
)

自定义的奖励指标会出现在TensorBoard的custom分类下，可和内置指标同时查看对比。

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