使用Stable Baselines3训练的A2C CartPole模型推理奖励偏低求助

我用Stable Baselines3在CartPole-v1环境实现A2C算法，训练后用evaluate_policy评估能拿到500满分，但自己写循环测试时奖励却很差，代码及运行情况如下：

初始化与训练代码 #

import os
import gymnasium as gym
import numpy as np

from stable_baselines3 import A2C
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
from stable_baselines3.common.env_util import make_vec_env

env_id = "CartPole-v1"
env = gym.make(env_id)

s_size = env.observation_space.shape
a_size = env.action_space

print("_____OBSERVATION SPACE_____ \n")
print("The State Space is: ", s_size)
print("Sample observation", env.observation_space.sample()) # Get a random observation

OBSERVATION SPACE
The State Space is: (4,) Sample observation [-2.3014314e+00
4.4097112e+37 -4.1089469e-01 2.7118910e+38]

envs = make_vec_env(env_id,seed=1, n_envs=4)
envs = VecNormalize(envs, norm_obs=True, norm_reward=True, clip_obs=10.)
model = A2C(policy = "MlpPolicy",env = envs, verbose=1)
model.learn(15_000)

# 保存模型与环境归一化参数
model.save("a2c-"+env_id)
envs.save("vec_normalize.pkl")

官方评估工具结果 #

# Load the saved statistics
eval_env = DummyVecEnv([lambda: gym.make(env_id)])
eval_env = VecNormalize.load("vec_normalize.pkl", eval_env)

# We need to override the render_mode
eval_env.render_mode = "rgb_array"

#  do not update them at test time
eval_env.training = False
# reward normalization is not needed at test time
eval_env.norm_reward = False

# Load the agent
model = A2C.load("a2c-"+env_id)
mean_reward, std_reward = evaluate_policy(model, eval_env)

print(f"Mean reward = {mean_reward:.2f} +/- {std_reward:.2f}")

Mean reward = 500.00 +/- 0.00

自行测试结果 #

for epi in range(10):
    score = 0
    state = env.reset()[0]
    done = False
    while not done:
        
        a , _ = model.predict(state)
        state_, r, done , _, _  = env.step(a)
        score += r
        state = np.copy(state_)
        env.render()
    print(f"Episode {epi} score {score}")

env.close()

Episode 0 score 71.0
Episode 1 score 70.0
Episode 2 score 83.0
Episode 3 score 62.0
Episode 4 score 63.0
Episode 5 score 59.0
Episode 6 score 52.0
Episode 7 score 54.0
Episode 8 score 60.0
Episode 9 score 69.0

问题原因 #

训练时模型是在经过VecNormalize归一化处理的环境中学习的，模型的输入期望是归一化后的状态数据。但自行测试时直接用原始环境的未归一化状态输入模型，导致模型无法做出正确预测，最终奖励暴跌。

修正后的测试代码 #

# 加载归一化环境
eval_env = DummyVecEnv([lambda: gym.make(env_id, render_mode="human")])
eval_env = VecNormalize.load("vec_normalize.pkl", eval_env)
eval_env.training = False
eval_env.norm_reward = False

for epi in range(10):
    score = 0
    # 从归一化环境获取初始状态
    state = eval_env.reset()
    done = False
    while not done:
        a, _ = model.predict(state)
        # 用归一化环境执行step，自动处理状态归一化
        state_, r, done, _ = eval_env.step(a)
        score += r
        state = state_
    print(f"Episode {epi} score {score[0]}")

eval_env.close()

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