如何在Ray Tune中为自定义Gym环境使用OpenAI Gym Wrappers

我基于Ray 2.1.0（Python 3.9.15）用Ray Tune的tune.run()结合PPO训练自定义强化学习环境CustomEnv（类似OpenAI Gym的CartPoleEnv），训练流程正常但学习速度慢。我已经实现了ObservationWrapper、RewardWrapper、ActionWrapper来缩放观测、奖励和动作以加速学习，这些Wrapper在本地传统训练循环里能正常工作，但把包装后的环境传给tune.run()的env参数时会报错——这个参数只接受环境类或已注册的Gym环境ID。我想找无需注册自定义环境的解决方案，其他可行方案也接受。

相关代码示例 #

1. 初始训练代码 #

import ray
from ray import tune
tune.run(
        "PPO",                         # 使用PPO算法
        config={"env": CustomEnv,      # 传入自定义环境类
            "framework": "tf2",
            "evaluation_interval": 100, 
            "evaluation_duration": 100,
            },
        checkpoint_freq = 100,             # 每100轮评估保存一次检查点
        local_dir=checkpoint_dir,          # 结果保存目录
        stop={"episode_reward_mean": 250},  # 停止训练的条件
        )

2. Wrapper实现代码 #

import gym

class ObservationWrapper(gym.ObservationWrapper):
    def __init__(self, env):
        super().__init__(env)
        self.o_min = 0.
        self.o_max = 5000.

    def observation(self, ob):
        # 归一化观测值
        ob = (ob - self.o_min)/(self.o_max - self.o_min)
        return ob

class RewardWrapper(gym.RewardWrapper):
    def __init__(self, env):
        super().__init__(env)
        self.r_min = -500
        self.r_max = 100

    def reward(self, reward):
        # 缩放奖励值
        reward = reward/(self.r_max - self.r_min)
        return reward

class ActionWrapper(gym.ActionWrapper):
    def __init__(self, env):
        super().__init__(env)

    def action(self, action):
        # 缩放动作值
        action = action/10
        return action

3. 本地训练循环代码 #

from my_file import CustomEnv

env = CustomEnv()
wrapped_env = ObservationWrapper(RewardWrapper(ActionWrapper(env)))
episodes = 10

for episode in range(1,episodes+1):
    obs = wrapped_env.reset()
    done = False
    score = 0
    
    while not done:
        action = wrapped_env.action_space.sample()
        obs, reward, done, info = wrapped_env.step(action)
        score += reward

    print(f'Episode: {episode},  Score: {score:.3f}')

方法1：用环境工厂函数（推荐，无需注册） #

Ray Tune允许通过env_creator参数传入自定义的环境创建函数，在这个函数里完成环境的包装逻辑，不用注册环境，Ray也能正确初始化带Wrapper的环境。

修改训练代码如下：

import ray
from ray import tune
from my_file import CustomEnv
from wrappers import ObservationWrapper, RewardWrapper, ActionWrapper  # 导入你的Wrapper

def env_creator(env_config):
    # 创建原始环境并依次包装
    env = CustomEnv(**env_config)
    wrapped_env = ObservationWrapper(RewardWrapper(ActionWrapper(env)))
    return wrapped_env

tune.run(
        "PPO",
        config={
            # 这里不再直接传env，而是通过env_creator指定
            "framework": "tf2",
            "evaluation_interval": 100, 
            "evaluation_duration": 100,
        },
        # 传入环境创建函数
        env=env_creator,
        checkpoint_freq=100,
        local_dir=checkpoint_dir,
        stop={"episode_reward_mean": 250},
        )

如果需要给CustomEnv传参数，可以在config里加"env_config": {"param1": value1}，然后在env_creator里用CustomEnv(**env_config)接收。

方法2：在CustomEnv内部集成Wrapper逻辑 #

如果不想单独维护Wrapper类，可以直接把缩放逻辑整合到CustomEnv的核心方法里，这样直接传CustomEnv给tune.run()的env参数即可。

修改后的CustomEnv示例：

class CustomEnv(gym.Env):
    # 保留原有的__init__、action_space、observation_space定义...

    def reset(self):
        obs = super().reset()
        # 归一化观测值
        obs = (obs - 0.)/(5000. - 0.)
        return obs

    def step(self, action):
        # 先缩放动作值
        action = action / 10
        obs, reward, done, info = super().step(action)
        # 归一化观测值
        obs = (obs - 0.)/(5000. - 0.)
        # 缩放奖励值
        reward = reward/(100 - (-500))
        return obs, reward, done, info

这种方法的好处是无需额外管理Wrapper，环境本身就自带缩放逻辑，Ray可以直接识别。

方法3：临时注册环境（可选，不推荐但可行） #

如果一定要用独立的Wrapper包装环境，也可以临时注册到Gym中，然后用注册的ID传给tune.run()。

代码示例：

import gym
from my_file import CustomEnv
from wrappers import ObservationWrapper, RewardWrapper, ActionWrapper

# 临时注册包装后的环境
gym.envs.register(
    id='WrappedCustomEnv-v0',
    entry_point=lambda env_config: ObservationWrapper(RewardWrapper(ActionWrapper(CustomEnv(**env_config)))),
)

# 训练时传入注册的ID
tune.run(
        "PPO",
        config={
            "env": "WrappedCustomEnv-v0",
            "framework": "tf2",
            "evaluation_interval": 100, 
            "evaluation_duration": 100,
            # 如果需要传参数给CustomEnv
            "env_config": {"param1": value1}
        },
        checkpoint_freq=100,
        local_dir=checkpoint_dir,
        stop={"episode_reward_mean": 250},
        )

注意要避免环境ID冲突，训练结束后如果不需要可以注销，但一般方法1更简洁灵活。

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