[Bug]: evaluate_policy called multiple times vor vectorized environments

[LukasFehring]

🐛 Bug

When calling

from stable_baselines3.common.evaluation import evaluate_policy
def custom_callback(locals, globals):
    pass

evaluate_policy(callback=custom_callback)

with a vecenv, then the callback gets executed for each of the environments separately. However, the locals dict contains the aggregated results. Therefore you have to manually check for which environment the callback was called, or only execute it every n_envs time.

To Reproduce

import gym
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.evaluation import evaluate_policy

# Define a simple callback function
def callback(_locals, _globals):
    pass

# Function to create multiple environments
def make_env():
    return gym.make('CartPole-v1')

# Number of environments
num_envs = 4
envs = [make_env for _ in range(num_envs)]

# Create vectorized environment
vec_env = DummyVecEnv(envs)

# Create a model
model = PPO("MlpPolicy", vec_env, verbose=1)

# Train the model
model.learn(total_timesteps=5000)

# Evaluate the policy
mean_reward, std_reward = evaluate_policy(model, vec_env, n_eval_episodes=10, callback=callback)

print("Mean reward:", mean_reward, "STD reward:", std_reward)

Relevant log output / Error message

No response

System Info

• OS: Linux-3.10.0-1160.2.1.el7.x86_64-x86_64-with-glibc2.17 # 1 SMP Tue Oct 20 15:39:03 UTC 2020
• Python: 3.8.19
• Stable-Baselines3: 2.4.0a0
• PyTorch: 2.3.0+cu121
• GPU Enabled: False
• Numpy: 1.24.4
• Cloudpickle: 3.0.0
• Gymnasium: 0.29.1
• OpenAI Gym: 0.23.0

Checklist

• My issue does not relate to a custom gym environment. (Use the custom gym env template instead)
• I have checked that there is no similar issue in the repo
• I have read the documentation
• I have provided a minimal and working example to reproduce the bug
• I've used the markdown code blocks for both code and stack traces.

[araffin]

Hello,
what is your usecase/expected behavior?

the for loop also decompose the info per env:

stable-baselines3/stable_baselines3/common/evaluation.py

Lines 99 to 106 in 35eccaf

	# unpack values so that the callback can access the local variables
	reward = rewards[i]
	done = dones[i]
	info = infos[i]
	episode_starts[i] = done
	if callback is not None:
	callback(locals(), globals())

[LukasFehring]

How so? Both the globals and locals contain information on every environment in the vectorized environment. How am I supposed to determine for which env the callback is called?

[araffin]

there is the local variable "i"

[LukasFehring]

Ah ok sorry then.
A documentation of locals and globals would probably help to find that! :)

[araffin]

A documentation of locals and globals would probably help to find that! :)

feel free to open a PR that updates the doc ;)