实现自定义包装器

实现自定义包装器#

在本教程中，将描述如何实现自定义包装器。包装器以模块化方式向环境添加功能。这将为您节省大量样板代码。

我们将展示如何通过以下方式创建包装器：

在遵循此教程之前，请确保查阅 gymnasium.wrappers 模块的文档。

继承自 gymnasium.ObservationWrapper#

观测包装器对环境返回的观测应用某种函数时非常有用。如果您实现了观测包装器，只需通过实现 gymnasium.ObservationWrapper.observation() 方法来定义这种转换即可。此外，如果转换改变了观测的形状（例如，通过将字典转换为 numpy 数组，如下例所示），您应该记得更新观测空间。

想象一下您有 2D 导航任务，其中环境以包含键 "agent_position" 和 "target_position" 的字典作为观测返回。常见的做法可能是丢弃一些自由度，只考虑目标相对于智能体的位置，即 observation["target_position"] - observation["agent_position"]。为此，您可以像这样实现观测包装器：

import numpy as np
from gymnasium import ActionWrapper, ObservationWrapper, RewardWrapper, Wrapper

import gymnasium as gym
from gymnasium.spaces import Box, Discrete


class RelativePosition(ObservationWrapper):
    def __init__(self, env):
        super().__init__(env)
        self.observation_space = Box(shape=(2,), low=-np.inf, high=np.inf)

    def observation(self, obs):
        return obs["target"] - obs["agent"]

继承自 gymnasium.ActionWrapper#

动作包装器可以用来在将动作应用到环境之前对动作进行转换。 如果您实现了一个动作包装器，您需要通过实现 gymnasium.ActionWrapper.action() 来定义该转换。此外，您应该通过更新包装器的动作空间来指定该转换的域。

假设您有一个动作空间为 gymnasium.spaces.Box 类型的环境，但您只想使用有限的动作子集。那么，您可能想实现以下包装器：

class DiscreteActions(ActionWrapper):
    def __init__(self, env, disc_to_cont):
        super().__init__(env)
        self.disc_to_cont = disc_to_cont
        self.action_space = Discrete(len(disc_to_cont))

    def action(self, act):
        return self.disc_to_cont[act]


if __name__ == "__main__":
    env = gym.make("LunarLanderContinuous-v3")
    wrapped_env = DiscreteActions(
        env, [np.array([1, 0]), np.array([-1, 0]), np.array([0, 1]), np.array([0, -1])]
    )
    print(wrapped_env.action_space)  # Discrete(4)

Discrete(4)

继承自 gymnasium.RewardWrapper#

奖励包装器用于转换环境返回的奖励。 与之前的包装器一样，您需要通过实现 gymnasium.RewardWrapper.reward() 方法来指定该转换。

让我们看一个例子：有时（特别是在我们无法控制奖励因为它是内在的），我们希望将奖励裁剪到一定范围内以获得一些数值稳定性。为此，我们可以实现如下包装器：

from typing import SupportsFloat


class ClipReward(RewardWrapper):
    def __init__(self, env, min_reward, max_reward):
        super().__init__(env)
        self.min_reward = min_reward
        self.max_reward = max_reward

    def reward(self, r: SupportsFloat) -> SupportsFloat:
        return np.clip(r, self.min_reward, self.max_reward)

继承自 gymnasium.Wrapper#

有时您可能需要实现一个执行更复杂修改的包装器（例如，根据info中的数据修改奖励或更改渲染行为）。 这样的包装器可以通过继承自 gymnasium.Wrapper 来实现。

如果您这样做，可以通过访问属性 env 来访问传递给您的包装器的原始环境（该环境可能仍然被其他包装器包装）。

让我们也看一下这种情况的一个例子。大多数MuJoCo环境返回的奖励包含不同的项：例如，可能有一项奖励智能体完成任务，另一项惩罚大动作（即能量使用）。通常，您可以在环境初始化期间传递这些项的权重参数。然而，Reacher 不允许这样做！不过，所有单独的奖励项都返回在 info 中，所以让我们为 Reacher 构建一个包装器，允许我们对这些项进行加权：

class ReacherRewardWrapper(Wrapper):
    def __init__(self, env, reward_dist_weight, reward_ctrl_weight):
        super().__init__(env)
        self.reward_dist_weight = reward_dist_weight
        self.reward_ctrl_weight = reward_ctrl_weight

    def step(self, action):
        obs, _, terminated, truncated, info = self.env.step(action)
        reward = (
            self.reward_dist_weight * info["reward_dist"]
            + self.reward_ctrl_weight * info["reward_ctrl"]
        )
        return obs, reward, terminated, truncated, info