如何在TF-Agents中修改MountainCarContinuous-v0的奖励函数？

我用suite_gym()加载MountainCarContinuous-v0环境训练时陷入局部最优，这个环境默认奖励函数会惩罚最优解需要的大幅动作，所以想更换奖励函数优化训练效果。

在原生Gymnasium里，我可以用TransformReward包装器实现奖励函数修改：

import gymnasium as gym

env = gym.make("MountainCarContinuous-v0")
wrapped_env = gym.wrappers.TransformReward(env, lambda r: 0 if r <= 0 else 1)
state = wrapped_env.reset()
state, reward, done = wrapped_env.step([action])
# 现在奖励只会是0或1，取决于是否到达目标

但在TF-Agents里不知道怎么实现同样操作：

env = suite_gym.load("MountainCarContinuous-v0")
env = None  # 怎么修改底层环境？

我试过修改env._env，结果出现reset() got an unexpected keyword argument 'seed'错误；也试过下面的方式，但不知道怎么给需要初始化参数的Gym包装器传参：

test = suite_gym.load("MountainCarContinuous-v0", gym_env_wrappers=[gym.wrappers.TransformReward])

请问获取并修改TF-Agents底层Gym环境的正确方法是什么？

解决方案 #

方法1：自定义包装器类，通过gym_env_wrappers传递参数 #

suite_gym.load的gym_env_wrappers参数只接受包装器类，而TransformReward需要初始化时传入奖励转换函数，因此可以自定义一个继承自TransformReward的类，提前封装好转换逻辑：

import gymnasium as gym
from tf_agents.environments import suite_gym

# 自定义带预设奖励转换逻辑的包装器
class CustomMountainCarReward(gym.wrappers.TransformReward):
    def __init__(self, env):
        super().__init__(env, lambda r: 0 if r <= 0 else 1)

# 加载环境时传入自定义包装器
env = suite_gym.load(
    "MountainCarContinuous-v0",
    gym_env_wrappers=[CustomMountainCarReward]
)

方法2：手动获取底层环境并重新封装 #

直接修改env._env出错是因为TF-Agents的环境类对底层Gym环境的方法做了参数兼容封装，所以可以手动取出底层环境，包装后再重新封装成TF-Agents兼容的环境：

import gymnasium as gym
from tf_agents.environments import suite_gym, py_environment

# 加载原始TF-Agents环境
tf_env = suite_gym.load("MountainCarContinuous-v0")
# 获取底层Gym环境实例
gym_env = tf_env._env
# 用TransformReward包装底层环境
wrapped_gym_env = gym.wrappers.TransformReward(gym_env, lambda r: 0 if r <= 0 else 1)
# 重新封装为TF-Agents兼容的PyEnvironment
tf_wrapped_env = py_environment.PyEnvironment(wrapped_gym_env)

注意事项 #

• 不要直接修改env._env后直接使用，TF-Agents的环境类会处理seed等参数，直接修改会导致参数不兼容
• 如果需要更复杂的奖励逻辑（比如结合状态、动作计算奖励），可以在自定义包装器中扩展step方法实现

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