Stable Baselines3中train/ep_rew_mean曲线异常，如何改为每步平均奖励？

我使用Stable Baselines3的SAC算法训练自定义强化学习环境，环境设置max_steps_per_episode=500，若智能体未在该步数内达成目标，回合会被截断并重置。

观察TensorBoard的train/ep_rew_mean图表时发现异常：曲线初始奖励值较高，随后下降并收敛至较低值，但正常情况下该曲线应先上升后收敛，以此体现智能体性能提升。

我分析问题根源在于：train/ep_rew_mean统计的是每回合累计奖励。随着训练推进，智能体学会了更高效的解决方案，能以更少步数达成目标（这一点可从train/ep_len_mean曲线的下降后收敛趋势验证：回合数增加时，每回合步数逐渐减少）。因此后期回合的累计奖励降低，但智能体的实际性能是在提升的。

图表趋势说明 #

• 左侧episode_length_mean图表：总时间步增加时，每回合步数逐渐减少并收敛，符合智能体学习高效解法的预期
• 右侧episode_reward_mean图表：呈现下降后收敛的趋势，与预期的上升平稳趋势相反，核心原因是该指标统计的是每回合累计奖励，而非每步平均奖励

环境step函数代码 #

def step(self, action):
        # Execute one time step within the environment
        terminated = False
        truncated = False    
        
        new_state, reward, self.done, info = self._take_action(action)
        
        if not self.best_episode_state:
            self.best_episode_state[reward] = new_state.get_points()
        elif reward > list(self.best_episode_state.keys())[0]:
            self.best_episode_state.clear()
            self.best_episode_state[reward] = new_state.get_points()
        
        if self.done == True:
            terminated = True
            print("**********(Terminated)**********")
           
        if self.current_step > 500 and not terminated:
            truncated = True
            print("***********(Truncated)***********")
        
        self.current_step += 1
        self.total_steps += 1
        self.state = new_state
        return np.array(new_state.get_points()), reward, terminated, truncated, info

需求 #

我希望图表能呈现上升趋势，体现智能体学习最大化奖励的过程，最终趋于平稳，准确反映智能体的性能提升，而非单纯展示每回合累计奖励。想了解是否可以通过修改图表配置或代码，让其显示每步平均奖励来实现更合理的可视化。

有三种可行方式实现每步平均奖励的可视化，按推荐优先级排序：

方法1：TensorBoard界面自定义计算图表（无需修改代码） #

TensorBoard支持在界面内基于现有指标创建自定义计算图表，快速验证你的分析：

1. 打开TensorBoard界面，点击顶部的「Add Chart」按钮
2. 选择「Scalar」类型，在表达式输入框中填写：train/ep_rew_mean / train/ep_len_mean
3. 将图表命名为avg_step_reward，即可直接查看每步平均奖励的趋势

这种方式零代码修改，适合快速验证你的结论。

方法2：自定义TensorBoard回调，记录每步平均奖励 #

通过Stable Baselines3的自定义回调函数，主动计算并记录每步平均奖励到TensorBoard，属于更规范的长期解决方案：

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

class StepRewardCallback(BaseCallback):
    def __init__(self, verbose=0):
        super().__init__(verbose)
        self.episode_rewards = []
        self.episode_step_count = 0

    def _on_step(self) -> bool:
        # 累计当前回合的奖励与步数
        self.episode_rewards.append(self.locals["rewards"][0])
        self.episode_step_count += 1
        
        # 当回合结束时（terminated或truncated）
        if self.locals["dones"][0]:
            # 计算每步平均奖励
            avg_step_reward = np.sum(self.episode_rewards) / self.episode_step_count
            # 写入TensorBoard，自定义指标名称
            self.logger.record("train/avg_step_rew_mean", avg_step_reward)
            # 重置回合缓存
            self.episode_rewards = []
            self.episode_step_count = 0
        return True

训练时添加该回调：

from stable_baselines3 import SAC

# 初始化你的自定义环境
env = YourCustomEnv()
model = SAC("MlpPolicy", env, tensorboard_log="./sac_tensorboard/")

# 加载自定义回调并开始训练
callback = StepRewardCallback()
model.learn(total_timesteps=100000, callback=callback)

完成后，TensorBoard中会新增train/avg_step_rew_mean指标，该曲线会呈现预期的上升后收敛趋势，准确反映智能体的性能提升。

方法3：调整奖励函数（需谨慎） #

如果你希望直接修改train/ep_rew_mean的统计逻辑，可以调整奖励的计算方式，将每步奖励设置为累计奖励的均分结果。但这种方式会改变智能体的学习信号，可能影响训练效果，需谨慎测试：

例如，在_take_action函数中修改奖励返回逻辑：

def _take_action(self, action):
    # 原有动作执行、状态更新逻辑
    # ...
    total_reward = 0  # 假设这是当前回合的累计奖励
    if self.done:
        # 回合完成时，将累计奖励按步数均分作为每步奖励
        reward = total_reward / self.current_step
    else:
        reward = 0  # 或保留原有每步奖励逻辑
    return new_state, reward, self.done, info

注意事项 #

• 优先使用方法1或方法2，避免修改奖励函数影响训练逻辑
• 自定义回调时，确保处理好多环境并行训练的情况（若使用make_vec_env，需调整回调中的数据获取逻辑）

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