如何让Stable Baselines3的PPO回调函数在每一步执行任务？

Stable Baselines3（SB3）中PPO的BaseCallback._on_step默认是每完成一批次（对应n_steps步）的经验收集后才会被调用，而非每一次环境交互的单步。这是因为PPO是基于批次更新的算法，回调的触发逻辑和经验收集的批次周期绑定，和你设置的n_steps=1024直接相关。

下面提供两种可靠的实现方式，可根据需求选择：

方案一：自定义环境包装器（推荐） #

通过包装环境，在每一次环境的step方法执行后触发任务，完全不干扰PPO的训练逻辑，是最直接的实现方式。

from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from gymnasium import Wrapper

class StepWrapper(Wrapper):
    def step(self, action):
        # 执行原环境的step逻辑
        obs, reward, terminated, truncated, info = super().step(action)
        # 在这里写入每一步需要执行的任务
        print("执行单步任务")
        return obs, reward, terminated, truncated, info

# 创建并包装环境
env = make_vec_env("CartPole-v1", n_envs=1, wrapper_class=StepWrapper)
model = PPO("MlpPolicy", env, n_steps=1024)
model.learn(total_timesteps=25000)

如果使用多环境（n_envs>1），每个环境的每一步交互都会触发任务，符合并行训练的逻辑。

方案二：在回调中遍历批次内的所有步骤 #

如果必须通过回调实现，可以在_on_step中访问当前批次收集的所有经验数据，遍历每一步并执行任务。这种方法利用PPO每批次结束触发回调的特性，一次性处理批次内的所有单步。

from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.callbacks import BaseCallback

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

    def _on_step(self):
        # 获取当前批次的经验缓存
        rollout_buffer = self.model.rollout_buffer
        # 遍历批次内的每一步（buffer_size = n_steps * n_envs）
        for step_idx in range(rollout_buffer.buffer_size):
            # 可按需获取该步的观测、动作、奖励等数据
            current_obs = rollout_buffer.observations[step_idx]
            current_action = rollout_buffer.actions[step_idx]
            # 执行单步任务
            print(f"处理第{self.num_timesteps - rollout_buffer.buffer_size + step_idx + 1}步")
        return True

env = make_vec_env("CartPole-v1", n_envs=1)
model = PPO("MlpPolicy", env, n_steps=1024)
model.learn(
    total_timesteps=25000,
    callback=CustomCallback(),
)

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