实现自定义包装器

在本教程中,我们将介绍如何实现您自己的自定义包装器。包装器是一种以模块化方式为您的环境添加功能的绝佳方式。这将为您节省大量样板代码。

我们将展示如何创建一个包装器。

在开始本教程之前,请确保查阅了 gymnasium.wrappers 模块的文档。

继承自 gymnasium.ObservationWrapper

观察包装器在你想对环境返回的观察结果应用某些函数时非常有用。如果你实现一个观察包装器,你只需要通过实现 gymnasium.ObservationWrapper.observation() 方法来定义这种转换。此外,如果转换改变了观察的形状(例如,通过将字典转换为numpy数组,如下例所示),你应该记得更新观察空间。

想象你有一个2D导航任务,其中环境返回带有键 "agent_position""target_position" 的字典作为观察结果。一个常见的做法可能是抛弃一些自由度,只考虑目标相对于代理的位置,即 observation["target_position"] - observation["agent_position"]。为此,你可以实现一个观察包装器,如下所示:

import numpy as np
from gym 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-v2")
    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)

继承自 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
        self.reward_range = (min_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