torchrl.envs package

TorchRL 提供了一个 API，用于处理不同后端的环境，例如 gym、dm-control、dm-lab、基于模型的环境以及自定义环境。目标是能够在实验中轻松切换环境，即使这些环境是使用不同的库模拟的。TorchRL 在 torchrl.envs.libs 下提供了一些开箱即用的环境包装器，我们希望这些包装器可以轻松地被其他库模仿。父类 EnvBase 是 torch.nn.Module 的子类，它使用 tensordict.TensorDict 作为数据组织者实现了一些典型的环境方法。这使得这个类具有通用性，并且能够处理任意数量的输入和输出，以及嵌套或批量的数据结构。

每个环境都有以下属性

• env.batch_size：表示批量处理的环境数量的 torch.Size。

• env.device：输入和输出 tensordict 预期所在的设备。环境设备不意味着实际的 step 操作将在设备上计算（这是后端的责任，TorchRL 对此能做的很少）。环境的设备仅表示数据在输入环境或从环境检索时预期所在的设备。TorchRL 负责将数据映射到所需的设备。这对于 transforms 尤其有用（见下文）。对于参数化环境（例如基于模型的环境），设备确实代表了将用于计算操作的硬件。

• env.observation_spec：一个 Composite 对象，包含所有观察（observation）的键-spec 对。

• env.state_spec：一个 Composite 对象，包含所有输入的键-spec 对（除了 action）。对于大多数有状态环境，这个容器将为空。

• env.action_spec：一个 TensorSpec 对象，表示 action spec。

• env.reward_spec：一个 TensorSpec 对象，表示 reward spec。

• env.done_spec：一个 TensorSpec 对象，表示 done 标志 spec。请参阅下面的轨迹终止部分。

• env.input_spec：一个 Composite 对象，包含所有输入键（"full_action_spec" 和 "full_state_spec"）。

• env.output_spec：一个 Composite 对象，包含所有输出键（"full_observation_spec"、"full_reward_spec" 和 "full_done_spec"）。

如果环境携带非张量（non-tensor）数据，可以使用 NonTensor 实例。

环境 Specs：锁和批量大小

环境 spec 默认是锁定的（通过传递给环境构造函数的 spec_locked 参数）。锁定 spec 意味着对 spec（或其子项，如果它是 Composite 实例）的任何修改都需要先解锁。这可以通过 set_spec_lock_() 方法完成。默认锁定 spec 的原因是这样可以轻松缓存值，例如 action 或 reset 键等。只有当预期 spec 会经常被修改时才应该解锁环境（原则上应该避免这种情况）。允许进行诸如 env.observation_spec = new_spec 的 spec 修改：在底层，如果环境之前是锁定的，TorchRL 将清除缓存，解锁 spec，进行修改，然后重新锁定 spec。

重要的是，环境 spec 的形状应该包含批量大小，例如，env.batch_size == torch.Size([4]) 的环境的 env.action_spec 应该具有 torch.Size([4, action_size]) 的形状。这在预分配张量、检查形状一致性等方面很有帮助。

环境方法

基于这些，实现了以下方法

• env.reset()：一个 reset 方法，它可能（但不一定需要）接受 tensordict.TensorDict 输入。它返回 rollout 的第一个 tensordict，通常包含一个 "done" 状态和一组 observations。如果不存在，将使用 0 和适当的形状实例化一个 “reward” 键。

• env.step()：一个 step 方法，它接受包含输入 action 以及其他输入（例如，对于基于模型的或无状态环境）的 tensordict.TensorDict 输入。

• env.step_and_maybe_reset()：执行一个 step，并在需要时（部分）重置环境。它返回带有包含下一步数据的 "next" 键的更新输入，以及一个包含下一步输入数据（即 reset 或结果或 step_mdp() ）的 tensordict。这是通过读取 done_keys 并为每个 done 状态分配一个 "_reset" 信号来完成的。这个方法可以轻松编写不间断的 rollout 函数

  >>> data_ = env.reset()
  >>> result = []
  >>> for i in range(N):
  ...     data, data_ = env.step_and_maybe_reset(data_)
  ...     result.append(data)
  ...
  >>> result = torch.stack(result)
  

• env.set_seed()：一个 seeding 方法，它将返回在多环境设置中使用的下一个种子。这个下一个种子是根据前一个种子确定性地计算出来的，这样就可以用不同的种子为多个环境播种，而不会在连续的实验中出现种子重叠的风险，同时仍然获得可复现的结果。

• env.rollout()：在环境中执行一个 rollout，最大步数（max_steps=N），并使用 policy（policy=model）。policy 应该使用 tensordict.nn.TensorDictModule （或任何其他与 tensordict.TensorDict 兼容的模块）编写。生成的 tensordict.TensorDict 实例将标记一个尾部的 "time" 命名维度，其他模块可以使用这个维度来适当地处理这个批量维度。

下图总结了 TorchRL 中如何执行 rollout。

使用 TensorDict 的 TorchRL Rollouts。

简而言之，TensorDict 由 reset() 方法创建，然后由 policy 用 action 填充，再传递给 step() 方法，后者将 observations、done 标志和 reward 写入 "next" 条目下。此调用的结果被存储以供后续使用，并且 "next" 条目由 step_mdp() 函数收集。

注意

通常，所有 TorchRL 环境在其输出 tensordict 中都有 "done" 和 "terminated" 条目。如果设计中未包含它们，EnvBase 元类将确保每个 done 或 terminated 都伴随其对应的条目。在 TorchRL 中，"done" 严格指代所有轨迹结束信号的联合，应解释为“轨迹的最后一步”或等同于“指示需要重置的信号”。如果环境提供了截断信息（例如 Gymnasium），截断条目也会在 EnvBase.step() 输出中写入一个 "truncated" 条目下。如果环境携带一个单独的值，默认情况下它将被解释为 "terminated" 信号。默认情况下，TorchRL 的 collectors 和 rollout 方法将查找 "done" 条目来判断环境是否应该重置。

注意

torchrl.collectors.utils.split_trajectories 函数可用于分割相邻的轨迹。它依赖于输入 tensordict 中的 "traj_ids" 条目，如果 "traj_ids" 缺失，则依赖于 "done" 和 "truncated" 键的组合。

注意

在某些上下文中，标记轨迹的第一步会很有用。TorchRL 通过 InitTracker transform 提供了此功能。

我们的环境 教程 提供了更多关于如何从头设计自定义环境的信息。

EnvBase(*args, **kwargs)	抽象环境父类。
GymLikeEnv(*args, **kwargs)	类似 gym 的环境是一个环境。
EnvMetaData(*, tensordict, specs, ...)	一个用于在多进程设置中存储和传递环境元数据的类。

向量化环境

向量化（或者更确切地说：并行）环境是强化学习中的一个常见特性，其中执行环境 step 可能占用大量 CPU 资源。一些库，例如 gym3 或 EnvPool，提供了接口来同时执行批量环境。虽然它们通常提供非常有竞争力的计算优势，但它们不一定能扩展到 TorchRL 支持的各种环境库。因此，TorchRL 提供了自己的通用 ParallelEnv 类，用于并行运行多个环境。由于此类继承自 SerialEnv，它享有与其他环境完全相同的 API。当然，ParallelEnv 将具有与其环境数量相对应的批量大小

注意

考虑到该库的许多可选依赖项（例如 Gym、Gymnasium 等许多库），在多进程/分布式设置中，警告可能很快变得相当烦人。默认情况下，TorchRL 在子进程中过滤掉这些警告。如果仍然希望看到这些警告，可以通过设置 torchrl.filter_warnings_subprocess=False 来显示它们。

重要的是，你的环境 spec 必须与它发送和接收的输入和输出相匹配，因为 ParallelEnv 将根据这些 spec 创建缓冲区以便与 spawned 进程通信。使用 check_env_specs() 方法进行健全性检查。

并行环境

     >>> def make_env():
     ...     return GymEnv("Pendulum-v1", from_pixels=True, g=9.81, device="cuda:0")
     >>> check_env_specs(env)  # this must pass for ParallelEnv to work
     >>> env = ParallelEnv(4, make_env)
     >>> print(env.batch_size)
     torch.Size([4])

ParallelEnv 允许检索其所包含环境的属性：可以直接调用

并行环境属性

     >>> a, b, c, d = env.g  # gets the g-force of the various envs, which we set to 9.81 before
     >>> print(a)
     9.81

TorchRL 使用私有的 "_reset" 键来指示环境哪些组件（子环境或智能体）应该被重置。

这允许只重置部分组件。

1. 在调用 _reset() 期间，输入 tensordict 中可能存在或不存在 "_reset" 键。TorchRL 的约定是，在给定 "done" 级别上缺少 "_reset" 键表示该级别的完全重置（除非在更高级别上找到 "_reset" 键，详见下文）。如果存在，则预期仅那些 "_reset" 条目为 True 的条目和组件（沿键和形状维度）将被重置。

   环境在其 _reset() 方法中处理 "_reset" 键的方式特定于其类。根据 "_reset" 输入设计环境行为是开发者的责任，因为 TorchRL 无法控制 _reset() 的内部逻辑。

2. 尽管如此，在设计该方法时应牢记以下几点。调用 _reset() 后，输出将用 "_reset" 条目进行掩码，并且将前一个 step() 的输出写入 "_reset" 为 False 的位置。实际上，这意味着如果 "_reset" 修改了未被它暴露的数据，这个修改将会丢失。进行此掩码操作后，"_reset" 条目将从 reset() 输出中删除。

必须指出，"_reset" 是一个私有键，它只能在编写内部使用的特定环境功能时使用。换句话说，这不应该在库外部使用，并且只要不影响 TorchRL 的内部测试，开发者将保留修改通过设置 "_reset" 实现部分重置逻辑的权利，恕不事先保证。

最后，在设计重置功能时，需要记住以下假设：

• 每个 "_reset" 都与一个 "done" 条目配对（+ "terminated"，并且可能还有 "truncated"）。这意味着以下结构是不允许的：TensorDict({"done": done, "nested": {"_reset": reset}}, [])，因为 "_reset" 所在的嵌套级别与 "done" 不同。

• 在一个级别上的重置并不排除在较低级别上存在 "_reset"，但它会使其效果失效。原因很简单，根级别的 "_reset" 是否对应于对嵌套 "done" 条目的 all()、any() 或自定义调用无法提前得知，并且明确假定根级别的 "_reset" 被放置在那里是为了取代嵌套值（例如，请查看 PettingZooWrapper 的实现，其中每个组关联有一个或多个 "done" 条目，这些条目根据任务以 any 或 all 逻辑在根级别进行聚合）。

• 当调用 env.reset(tensordict)() 并带有将重置部分而非全部已完成子环境的部分 "_reset" 条目时，输入数据应包含__未__被重置的子环境的数据。这个限制的原因在于，env._reset(data) 的输出只能针对被重置的条目进行预测。

对于其他条目，TorchRL 无法提前知道它们是否有意义。例如，可以完全用填充值填充未重置组件的值，在这种情况下，未重置的数据将是无意义的，应该被丢弃。

>>> # single reset at the root
>>> data = TensorDict({"val": [1, 1], "_reset": [False, True]}, [])
>>> env.reset(data)
>>> print(data.get("val"))  # only the second value is 0
tensor([1, 0])
>>> # nested resets
>>> data = TensorDict({
...     ("agent0", "val"): [1, 1], ("agent0", "_reset"): [False, True],
...     ("agent1", "val"): [2, 2], ("agent1", "_reset"): [True, False],
... }, [])
>>> env.reset(data)
>>> print(data.get(("agent0", "val")))  # only the second value is 0
tensor([1, 0])
>>> print(data.get(("agent1", "val")))  # only the first value is 0
tensor([0, 2])
>>> # nested resets are overridden by a "_reset" at the root
>>> data = TensorDict({
...     "_reset": [True, True],
...     ("agent0", "val"): [1, 1], ("agent0", "_reset"): [False, True],
...     ("agent1", "val"): [2, 2], ("agent1", "_reset"): [True, False],
... }, [])
>>> env.reset(data)
>>> print(data.get(("agent0", "val")))  # reset at the root overrides nested
tensor([0, 0])
>>> print(data.get(("agent1", "val")))  # reset at the root overrides nested
tensor([0, 0])

下面，我们给出 "_reset" 键对重置后返回零的环境产生预期效果的一些示例

     >>> tensordict = TensorDict({"_reset": [[True], [False], [True], [True]]}, [4])
     >>> env.reset(tensordict)  # eliminates the "_reset" entry
     TensorDict(
         fields={
             terminated: Tensor(torch.Size([4, 1]), dtype=torch.bool),
             done: Tensor(torch.Size([4, 1]), dtype=torch.bool),
             pixels: Tensor(torch.Size([4, 500, 500, 3]), dtype=torch.uint8),
             truncated: Tensor(torch.Size([4, 1]), dtype=torch.bool),
         batch_size=torch.Size([4]),
         device=None,
         is_shared=True)

注意

并行环境重置

注意

关于性能的注意事项：启动 ParallelEnv 可能需要相当长的时间，因为它需要启动与进程数量相同的 python 实例。由于运行 import torch （以及其他 imports）所需的时间，启动并行环境可能成为瓶颈。例如，这就是 TorchRL 测试如此缓慢的原因。一旦环境启动，应该会观察到显著的加速。

TorchRL 需要精确的 specs：另一个需要考虑的问题是，ParallelEnv （以及数据收集器）将根据环境 spec 创建数据缓冲区，以便在不同进程之间传递数据。这意味着 spec 错误（输入、观察或 reward）将在运行时导致程序中断，因为数据无法写入预分配的缓冲区。通常，在使用 ParallelEnv 之前，应该使用 check_env_specs() 测试函数对环境进行测试。当预分配的缓冲区与收集到的数据不匹配时，此函数将引发一个断言错误。

我们还提供了 SerialEnv 类，它享有完全相同的 API，但以串行方式执行。这主要用于测试目的，当想要评估 ParallelEnv 的行为而不启动子进程时非常有用。

除了提供基于进程的并行性的 ParallelEnv 之外，我们还提供了使用 MultiThreadedEnv 创建多线程环境的方法。此类底层使用 EnvPool 库，这可以实现更高的性能，但同时限制了灵活性 - 只能创建在 EnvPool 中实现的环境。这涵盖了许多流行的强化学习环境类型（Atari、经典控制等），但不能使用任意的 TorchRL 环境，不像 ParallelEnv 那样灵活。运行 benchmarks/benchmark_batched_envs.py 来比较并行化批量环境的不同方法的性能。	SerialEnv(*args, **kwargs)
在同一进程中创建一系列环境。批量环境允许用户查询远程运行环境的任意方法/属性。	ParallelEnv(*args, **kwargs)
为每个进程创建一个环境。	EnvCreator(create_env_fn[, ...])

环境创建器类。

自定义原生 TorchRL 环境

TorchRL 提供了一系列自定义的内置环境。	ChessEnv(*args, **kwargs)
一个遵循 TorchRL API 的国际象棋环境。	PendulumEnv(*args, **kwargs)
一个无状态的摆锤环境。	TicTacToeEnv(*args, **kwargs)
一个井字棋实现。	LLMHashingEnv(*args, **kwargs)

一个文本生成环境，使用哈希模块识别独特的观察。

多智能体环境

TorchRL 开箱即用地支持多智能体学习。在单智能体学习管线中使用的相同类可以无缝地用于多智能体场景，无需任何修改或专用的多智能体基础设施。

• 观测（observation） 可以是每个智能体的，也可以包含一些共享部分

• 奖励（reward） 可以是每个智能体的，也可以是共享的

• 完成（done）（以及 "truncated" 或 "terminated"）可以是每个智能体的，也可以是共享的。

TorchRL 凭借其 tensordict.TensorDict 数据载体，能够适应所有这些可能的范例。特别是在多智能体环境中，按智能体的键将保存在嵌套的“agents” TensorDict 中。这个 TensorDict 将具有额外的智能体维度，从而将每个智能体不同的数据分组。另一方面，共享键将保留在第一层，就像单智能体情况中一样。

让我们看一个例子来更好地理解这一点。在本例中，我们将使用 VMAS，这是一个同样基于 PyTorch 的多机器人任务模拟器，它在设备上运行并行批处理模拟。

我们可以创建一个 VMAS 环境，并查看随机步骤的输出是什么样子

多智能体步骤 tensordict 示例

     >>> from torchrl.envs.libs.vmas import VmasEnv
     >>> env = VmasEnv("balance", num_envs=3, n_agents=5)
     >>> td = env.rand_step()
     >>> td
     TensorDict(
         fields={
             agents: TensorDict(
                 fields={
                     action: Tensor(shape=torch.Size([3, 5, 2]))},
                 batch_size=torch.Size([3, 5])),
             next: TensorDict(
                 fields={
                     agents: TensorDict(
                         fields={
                             info: TensorDict(
                                 fields={
                                     ground_rew: Tensor(shape=torch.Size([3, 5, 1])),
                                     pos_rew: Tensor(shape=torch.Size([3, 5, 1]))},
                                 batch_size=torch.Size([3, 5])),
                             observation: Tensor(shape=torch.Size([3, 5, 16])),
                             reward: Tensor(shape=torch.Size([3, 5, 1]))},
                         batch_size=torch.Size([3, 5])),
                     done: Tensor(shape=torch.Size([3, 1]))},
                 batch_size=torch.Size([3]))},
         batch_size=torch.Size([3]))

我们可以观察到，所有智能体共享的键，例如 done，存在于根 tensordict 中，其批大小为 (num_envs,)，这代表了模拟的环境数量。

另一方面，智能体之间不同的键，例如 action、reward、observation 和 info，存在于嵌套的“agents” tensordict 中，其批大小为 (num_envs, n_agents)，这代表了额外的智能体维度。

多智能体张量规范将遵循与 tensordict 相同的风格。与智能体之间变化的值相关的规范需要嵌套在“agents”条目中。

这里有一个示例，说明如何在只有完成标志在智能体之间共享的多智能体环境中创建规范（如在 VMAS 中）

多智能体规范创建示例

     >>> action_specs = []
     >>> observation_specs = []
     >>> reward_specs = []
     >>> info_specs = []
     >>> for i in range(env.n_agents):
     ...    action_specs.append(agent_i_action_spec)
     ...    reward_specs.append(agent_i_reward_spec)
     ...    observation_specs.append(agent_i_observation_spec)
     >>> env.action_spec = Composite(
     ...    {
     ...        "agents": Composite(
     ...            {"action": torch.stack(action_specs)}, shape=(env.n_agents,)
     ...        )
     ...    }
     ...)
     >>> env.reward_spec = Composite(
     ...    {
     ...        "agents": Composite(
     ...            {"reward": torch.stack(reward_specs)}, shape=(env.n_agents,)
     ...        )
     ...    }
     ...)
     >>> env.observation_spec = Composite(
     ...    {
     ...        "agents": Composite(
     ...            {"observation": torch.stack(observation_specs)}, shape=(env.n_agents,)
     ...        )
     ...    }
     ...)
     >>> env.done_spec = Categorical(
     ...    n=2,
     ...    shape=torch.Size((1,)),
     ...    dtype=torch.bool,
     ... )

如你所见，这非常简单！按智能体的键将具有嵌套的复合规范，而共享键将遵循单智能体标准。

注意

由于 reward、done 和 action 键可能带有额外的“agent”前缀（例如，(“agents”,”action”)），因此在其他 TorchRL 组件的参数中使用的默认键（例如，“action”）将无法完全匹配。因此，TorchRL 提供了 env.action_key、env.reward_key 和 env.done_key 属性，它们会自动指向要使用的正确键。确保将这些属性传递给 TorchRL 中的各种组件，以告知它们正确的键（例如，loss.set_keys() 函数）。

注意

TorchRL 抽象化了这些嵌套规范，以便于使用。这意味着，如果访问的规范是 Composite，则访问 env.reward_spec 将始终返回叶规范。因此，如果在上面的示例中，我们在创建环境后运行 env.reward_spec，我们将获得与 torch.stack(reward_specs)} 相同的输出。要获取带有“agents”键的完整复合规范，可以运行 env.output_spec[“full_reward_spec”]。动作和完成规范也是如此。请注意，env.reward_spec == env.output_spec[“full_reward_spec”][env.reward_key]。

MarlGroupMapType(value)	Marl 分组映射类型。
check_marl_grouping(group_map, agent_names)	检查 MARL 分组映射。

自动重置环境

自动重置环境是指，当环境在 rollout 期间达到 "done" 状态时，不需要调用 reset()，因为重置会自动发生。通常，在这种情况下，随同 done 和 reward（这些实际上是在环境中执行动作的结果）传递的观测是新回合的第一个观测，而不是当前回合的最后一个观测。

为了处理这些情况，torchrl 提供了一个 AutoResetTransform，它会将调用 step 产生的观测复制到下一次 reset，并在 rollouts 期间（在 rollout() 和 SyncDataCollector 迭代中）跳过 reset 的调用。这个变换类还对无效观测采取的行为提供了细粒度控制，这些观测可以用 “nan” 或任何其他值进行屏蔽，或者完全不屏蔽。

要告知 torchrl 某个环境是自动重置的，只需在构造时提供一个 auto_reset 参数即可。如果提供，auto_reset_replace 参数还可以控制回合中最后一个观测的值是否应被某些占位符替换。

>>> from torchrl.envs import GymEnv
>>> from torchrl.envs import set_gym_backend
>>> import torch
>>> torch.manual_seed(0)
>>>
>>> class AutoResettingGymEnv(GymEnv):
...     def _step(self, tensordict):
...         tensordict = super()._step(tensordict)
...         if tensordict["done"].any():
...             td_reset = super().reset()
...             tensordict.update(td_reset.exclude(*self.done_keys))
...         return tensordict
...
...     def _reset(self, tensordict=None):
...         if tensordict is not None and "_reset" in tensordict:
...             return tensordict.copy()
...         return super()._reset(tensordict)
>>>
>>> with set_gym_backend("gym"):
...     env = AutoResettingGymEnv("CartPole-v1", auto_reset=True, auto_reset_replace=True)
...     env.set_seed(0)
...     r = env.rollout(30, break_when_any_done=False)
>>> print(r["next", "done"].squeeze())
tensor([False, False, False, False, False, False, False, False, False, False,
        False, False, False,  True, False, False, False, False, False, False,
        False, False, False, False, False,  True, False, False, False, False])
>>> print("observation after reset are set as nan", r["next", "observation"])
observation after reset are set as nan tensor([[-4.3633e-02, -1.4877e-01,  1.2849e-02,  2.7584e-01],
        [-4.6609e-02,  4.6166e-02,  1.8366e-02, -1.2761e-02],
        [-4.5685e-02,  2.4102e-01,  1.8111e-02, -2.9959e-01],
        [-4.0865e-02,  4.5644e-02,  1.2119e-02, -1.2542e-03],
        [-3.9952e-02,  2.4059e-01,  1.2094e-02, -2.9009e-01],
        [-3.5140e-02,  4.3554e-01,  6.2920e-03, -5.7893e-01],
        [-2.6429e-02,  6.3057e-01, -5.2867e-03, -8.6963e-01],
        [-1.3818e-02,  8.2576e-01, -2.2679e-02, -1.1640e+00],
        [ 2.6972e-03,  1.0212e+00, -4.5959e-02, -1.4637e+00],
        [ 2.3121e-02,  1.2168e+00, -7.5232e-02, -1.7704e+00],
        [ 4.7457e-02,  1.4127e+00, -1.1064e-01, -2.0854e+00],
        [ 7.5712e-02,  1.2189e+00, -1.5235e-01, -1.8289e+00],
        [ 1.0009e-01,  1.0257e+00, -1.8893e-01, -1.5872e+00],
        [        nan,         nan,         nan,         nan],
        [-3.9405e-02, -1.7766e-01, -1.0403e-02,  3.0626e-01],
        [-4.2959e-02, -3.7263e-01, -4.2775e-03,  5.9564e-01],
        [-5.0411e-02, -5.6769e-01,  7.6354e-03,  8.8698e-01],
        [-6.1765e-02, -7.6292e-01,  2.5375e-02,  1.1820e+00],
        [-7.7023e-02, -9.5836e-01,  4.9016e-02,  1.4826e+00],
        [-9.6191e-02, -7.6387e-01,  7.8667e-02,  1.2056e+00],
        [-1.1147e-01, -9.5991e-01,  1.0278e-01,  1.5219e+00],
        [-1.3067e-01, -7.6617e-01,  1.3322e-01,  1.2629e+00],
        [-1.4599e-01, -5.7298e-01,  1.5848e-01,  1.0148e+00],
        [-1.5745e-01, -7.6982e-01,  1.7877e-01,  1.3527e+00],
        [-1.7285e-01, -9.6668e-01,  2.0583e-01,  1.6956e+00],
        [        nan,         nan,         nan,         nan],
        [-4.3962e-02,  1.9845e-01, -4.5015e-02, -2.5903e-01],
        [-3.9993e-02,  3.9418e-01, -5.0196e-02, -5.6557e-01],
        [-3.2109e-02,  5.8997e-01, -6.1507e-02, -8.7363e-01],
        [-2.0310e-02,  3.9574e-01, -7.8980e-02, -6.0090e-01]])

动态规范

通常通过创建用于在进程之间传递信息的内存缓冲区来并行运行环境。在某些情况下，可能无法预测环境在 rollout 期间是否会有一致的输入或输出，因为它们的形状可能是可变的。我们将此称为动态规范。

TorchRL 能够处理动态规范，但需要让批处理环境和收集器了解此特性。请注意，在实践中，这是自动检测的。

要指示张量沿某个维度具有可变大小，可以将所需维度的尺寸值设置为 -1。由于数据无法连续堆叠，调用 env.rollout 时需要使用 return_contiguous=False 参数。这里是一个工作示例

>>> from torchrl.envs import EnvBase
>>> from torchrl.data import Unbounded, Composite, Bounded, Binary
>>> import torch
>>> from tensordict import TensorDict, TensorDictBase
>>>
>>> class EnvWithDynamicSpec(EnvBase):
...     def __init__(self, max_count=5):
...         super().__init__(batch_size=())
...         self.observation_spec = Composite(
...             observation=Unbounded(shape=(3, -1, 2)),
...         )
...         self.action_spec = Bounded(low=-1, high=1, shape=(2,))
...         self.full_done_spec = Composite(
...             done=Binary(1, shape=(1,), dtype=torch.bool),
...             terminated=Binary(1, shape=(1,), dtype=torch.bool),
...             truncated=Binary(1, shape=(1,), dtype=torch.bool),
...         )
...         self.reward_spec = Unbounded((1,), dtype=torch.float)
...         self.count = 0
...         self.max_count = max_count
...
...     def _reset(self, tensordict=None):
...         self.count = 0
...         data = TensorDict(
...             {
...                 "observation": torch.full(
...                     (3, self.count + 1, 2),
...                     self.count,
...                     dtype=self.observation_spec["observation"].dtype,
...                 )
...             }
...         )
...         data.update(self.done_spec.zero())
...         return data
...
...     def _step(
...         self,
...         tensordict: TensorDictBase,
...     ) -> TensorDictBase:
...         self.count += 1
...         done = self.count >= self.max_count
...         observation = TensorDict(
...             {
...                 "observation": torch.full(
...                     (3, self.count + 1, 2),
...                     self.count,
...                     dtype=self.observation_spec["observation"].dtype,
...                 )
...             }
...         )
...         done = self.full_done_spec.zero() | done
...         reward = self.full_reward_spec.zero()
...         return observation.update(done).update(reward)
...
...     def _set_seed(self, seed: Optional[int]):
...         self.manual_seed = seed
...         return seed
>>> env = EnvWithDynamicSpec()
>>> print(env.rollout(5, return_contiguous=False))
LazyStackedTensorDict(
    fields={
        action: Tensor(shape=torch.Size([5, 2]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: LazyStackedTensorDict(
            fields={
                done: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([5, 3, -1, 2]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            exclusive_fields={
            },
            batch_size=torch.Size([5]),
            device=None,
            is_shared=False,
            stack_dim=0),
        observation: Tensor(shape=torch.Size([5, 3, -1, 2]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    exclusive_fields={
    },
    batch_size=torch.Size([5]),
    device=None,
    is_shared=False,
    stack_dim=0)

警告

在 ParallelEnv 和数据收集器中缺少内存缓冲区会极大地影响这些类的性能。任何此类用法都应与在单个进程上的普通执行仔细进行基准测试，因为序列化和反序列化大量张量可能非常昂贵。

目前，check_env_specs() 对于沿某些维度形状变化的动态规范会通过，但对于在某个步骤中存在而在其他步骤中缺失的键，或维度数量变化的情况，则不会通过。

变换

在大多数情况下，环境的原始输出在传递给其他对象（如策略或值运算符）之前必须经过处理。为此，TorchRL 提供了一系列变换，旨在复现 torch.distributions.Transform 和 torchvision.transforms 的变换逻辑。我们的环境 教程 提供了关于如何设计自定义变换的更多信息。

变换环境使用 TransformedEnv 基本单元构建。复合变换使用 Compose 类构建

变换环境

     >>> base_env = GymEnv("Pendulum-v1", from_pixels=True, device="cuda:0")
     >>> transform = Compose(ToTensorImage(in_keys=["pixels"]), Resize(64, 64, in_keys=["pixels"]))
     >>> env = TransformedEnv(base_env, transform)

变换通常是 Transform 的子类，但也可以是任何 Callable[[TensorDictBase], TensorDictBase]。

默认情况下，变换环境将继承传递给它的 base_env 的设备。然后变换将在该设备上执行。现在很明显，根据要计算的操作类型，这可以带来显著的加速。

环境包装器的一个巨大优势是，可以查看直到该包装器的环境。使用 TorchRL 变换环境也可以实现同样的效果：parent 属性将返回一个新的 TransformedEnv，其中包含直到目标变换的所有变换。重用上面的示例

变换父级

     >>> resize_parent = env.transform[-1].parent  # returns the same as TransformedEnv(base_env, transform[:-1])

变换环境可以与向量化环境一起使用。由于每个变换都使用 "in_keys"/"out_keys" 关键字参数集，因此也很容易将变换图连接到观测数据的每个组件（例如像素或状态等）。

正向和逆向变换

变换还具有一个 inv() 方法，该方法在动作以逆序应用于复合变换链之前被调用。这允许在环境中执行动作之前对环境中的数据应用变换。要包含在此逆向变换中的键通过 “in_keys_inv” 关键字参数传递，并且在大多数情况下，out-keys 默认使用这些值

逆向变换

     >>> env.append_transform(DoubleToFloat(in_keys_inv=["action"]))  # will map the action from float32 to float64 before calling the base_env.step

以下段落详细阐述了如何考虑哪些应被视为 in_ 或 out_ 特性。

理解变换键

在变换中，in_keys 和 out_keys 定义了基础环境与外部世界（例如你的策略）之间的交互

• in_keys 指的是基础环境的视角（内部 = TransformedEnv 的 base_env）。

• out_keys 指的是外部世界（外部 = policy、agent 等）。

例如，使用 in_keys=[“obs”] 和 out_keys=[“obs_standardized”] 时，策略将“看到”一个标准化观测，而基础环境输出的是常规观测。

类似地，对于逆向键

• in_keys_inv 指基础环境看到的条目。

• out_keys_inv 指策略看到或产生的条目。

下图说明了 RenameTransform 类的这个概念：step 函数的输入 TensorDict 必须包含 out_keys_inv，因为它们是外部世界的一部分。变换使用 in_keys_inv 将这些名称更改为与内部基础环境的名称匹配。逆向过程使用输出 tensordict 执行，其中 in_keys 映射到对应的 out_keys。

Rename 变换逻辑

变换 Tensor 和规范

变换实际张量（来自策略）时，该过程示意性地表示为

>>> for t in reversed(self.transform):
...     td = t.inv(td)

这从最外层变换开始到最内层变换结束，确保暴露给策略的动作值被正确变换。

对于变换动作规范，过程应从最内层到最外层（类似于观测规范）

>>> def transform_action_spec(self, action_spec):
...     for t in self.transform:
...         action_spec = t.transform_action_spec(action_spec)
...     return action_spec

单个 transform_action_spec 的伪代码可以是

>>> def transform_action_spec(self, action_spec):
...    return spec_from_random_values(self._apply_transform(action_spec.rand()))

这种方法确保“外部”规范是从“内部”规范推断出来的。请注意，我们有意调用的是 _apply_transform 而不是 _inv_apply_transform！

向外部世界暴露规范

TransformedEnv 将暴露对应于动作和状态的 out_keys_inv 的规范。例如，使用 ActionDiscretizer 时，环境的动作（例如 “action”）是一个浮点值张量，在使用变换环境调用 rand_action() 时不应生成该张量。相反，应该生成 “action_discrete”，并通过变换获取其连续对应项。因此，用户应看到 “action_discrete” 条目被暴露，而不是 “action”。

克隆变换

由于附加到环境的变换通过 transform.parent 属性“注册”到该环境，因此在操作变换时，我们应记住父级可能根据对变换进行的操作出现和消失。这里有一些示例：如果我们从一个 Compose 对象获取单个变换，这个变换将保留其父级

>>> third_transform = env.transform[2]
>>> assert third_transform.parent is not None

这意味着禁止将此变换用于另一个环境，因为另一个环境会替换父级，这可能导致意外行为。幸运的是，Transform 类带有一个 clone() 方法，该方法将擦除父级，同时保留所有注册缓冲区的标识。

>>> TransformedEnv(base_env, third_transform)  # raises an Exception as third_transform already has a parent
>>> TransformedEnv(base_env, third_transform.clone())  # works

在单个进程上或如果缓冲区放置在共享内存中，这将导致所有克隆的变换保持相同的行为，即使缓冲区被原地修改（例如，CatFrames 变换就是如此）。在分布式设置中，这可能不适用，在此上下文中应小心处理克隆变换的预期行为。最后，请注意，从 Compose 变换中索引多个变换也可能导致这些变换丢失父级：原因是索引 Compose 变换会产生另一个 Compose 变换，该变换没有父环境。因此，我们必须克隆子变换才能创建这个其他组合

>>> env = TransformedEnv(base_env, Compose(transform1, transform2, transform3))
>>> last_two = env.transform[-2:]
>>> assert isinstance(last_two, Compose)
>>> assert last_two.parent is None
>>> assert last_two[0] is not transform2
>>> assert isinstance(last_two[0], type(transform2))  # and the buffers will match
>>> assert last_two[1] is not transform3
>>> assert isinstance(last_two[1], type(transform3))  # and the buffers will match

Transform([in_keys, out_keys, in_keys_inv, ...])	环境变换父类。
TransformedEnv(*args, **kwargs)	一个 transformed_in 环境。
ActionDiscretizer(num_intervals[, ...])	将连续动作空间离散化的变换。
ActionMask([action_key, mask_key])	自适应动作掩码器。
AutoResetEnv(*args, **kwargs)	用于自动重置环境的子类。
AutoResetTransform(*[, replace, fill_float, ...])	用于自动重置环境的变换。
BatchSizeTransform(*[, batch_size, ...])	修改环境批大小的变换。
BinarizeReward([in_keys, out_keys])	如果奖励为 null 或非 null，则将奖励映射到二进制值（分别为 0 或 1）。
BurnInTransform(modules, burn_in[, out_keys])	用于部分烧入数据序列的变换。
CatFrames(N, dim[, in_keys, out_keys, ...])	将连续观测帧连接到单个张量中。
CatTensors([in_keys, out_key, dim, ...])	将多个键连接到单个张量中。
CenterCrop(w[, h, in_keys, out_keys])	裁剪图像中心。
ClipTransform([in_keys, out_keys, ...])	用于裁剪输入（状态、动作）或输出（观测、奖励）值的变换。
Compose(*transforms)	组合一系列变换。
Crop(w[, h, top, left, in_keys, out_keys])	在指定位置和输出尺寸裁剪输入图像。
DTypeCastTransform(dtype_in, dtype_out[, ...])	将选定键的一种数据类型转换为另一种。
DeviceCastTransform(device[, orig_device, ...])	将数据从一个设备移动到另一个设备。
DiscreteActionProjection(...[, action_key, ...])	将离散动作从高维空间投射到低维空间。
DoubleToFloat([in_keys, out_keys, ...])	将选定键的一种数据类型转换为另一种。
EndOfLifeTransform([eol_key, lives_key, ...])	使用 lives 方法注册来自具有 lives 方法的 Gym 环境的生命周期结束信号。
ExcludeTransform(*excluded_keys[, inverse])	从数据中排除键。
FiniteTensorDictCheck()	此变换将检查 tensordict 的所有项是否为有限值，如果不是则引发异常。
FlattenObservation(first_dim, last_dim[, ...])	展平张量的相邻维度。
FrameSkipTransform([frame_skip])	帧跳跃变换。
GrayScale([in_keys, out_keys])	将像素观测转换为灰度图像。
Hash(in_keys, out_keys, *[, hash_fn, seed, ...])	向 tensordict 添加哈希值。
InitTracker	重置跟踪器。
KLRewardTransform(actor[, coef, in_keys, ...])	向奖励添加 KL[pi_current||pi_0] 校正项的变换。
LineariseRewards(in_keys[, out_keys, weights])	通过加权求和将多目标奖励信号转换为单目标信号。
NoopResetEnv([noops, random])	环境重置时运行一系列随机动作。
ObservationNorm([loc, scale, in_keys, ...])	观测仿射变换层。
ObservationTransform([in_keys, out_keys, ...])	观测变换的抽象类。
PermuteTransform(dims[, in_keys, out_keys, ...])	置换变换。
PinMemoryTransform()	对 tensordict 调用 pin_memory 以方便在 CUDA 设备上写入。
R3MTransform(*args, **kwargs)	R3M 变换类。
RandomCropTensorDict(sub_seq_len[, ...])	ReplayBuffer 和模块的轨迹子采样器。
RemoveEmptySpecs([in_keys, out_keys, ...])	从环境中移除空的规范和内容。
RenameTransform(in_keys, out_keys[, ...])	用于重命名输出 tensordict（或通过逆向键重命名输入 tensordict）中条目的变换。
Resize(w[, h, interpolation, in_keys, out_keys])	调整像素观测的大小。
Reward2GoTransform([gamma, in_keys, ...])	根据回合奖励和折扣因子计算剩余奖励 (reward to go)。
RewardClipping([clamp_min, clamp_max, ...])	将奖励裁剪在 clamp_min 和 clamp_max 之间。
RewardScaling(loc, scale[, in_keys, ...])	奖励的仿射变换。
RewardSum([in_keys, out_keys, reset_keys, ...])	跟踪回合累积奖励。
SelectTransform(*selected_keys[, ...])	从输入 tensordict 中选择键。
SignTransform([in_keys, out_keys, ...])	计算 TensorDict 值符号的变换。
SqueezeTransform(*args, **kwargs)	移除指定位置上大小为 1 的维度。
Stack(in_keys, out_key[, in_key_inv, ...])	堆叠张量和 tensordict。
StepCounter([max_steps, truncated_key, ...])	计算从重置以来的步数，并在达到一定步数后可选地将截断状态设置为 True。
TargetReturn(target_return[, mode, in_keys, ...])	为智能体设置一个在环境中达成的目标回报。
TensorDictPrimer([primers, random, ...])	在重置时用于 TensorDict 初始化的引导器。
TimeMaxPool([in_keys, out_keys, T, reset_key])	在过去 T 个观测中，取每个位置的最大值。
Tokenizer(in_keys, out_keys[, in_keys_inv, ...])	对指定输入应用分词操作。
ToTensorImage([from_int, unsqueeze, dtype, ...])	将类似 numpy 的图像 (W x H x C) 转换为 pytorch 图像 (C x W x H)。
TrajCounter([out_key])	全局轨迹计数变换。
UnaryTransform(in_keys, out_keys[, ...])	对指定输入应用一元操作。
UnsqueezeTransform(*args, **kwargs)	在指定位置插入一个大小为 1 的维度。
VC1Transform(in_keys, out_keys, model_name)	VC1 变换类。
VIPRewardTransform(*args, **kwargs)	基于嵌入相似性计算奖励的 VIP 变换。
VIPTransform(*args, **kwargs)	VIP 变换类。
VecGymEnvTransform([final_name])	用于 GymWrapper 子类的变换，以一致的方式处理自动重置。
VecNorm([in_keys, out_keys, shared_td, ...])	torchrl 环境的移动平均归一化层。
gSDENoise([state_dim, action_dim, shape])	一个 gSDE 噪声初始化器。

带有遮罩动作（masked actions）的环境

在一些具有离散动作的环境中，智能体可用的动作可能会在执行过程中发生变化。在这种情况下，环境会输出一个动作遮罩（默认键为 "action_mask"）。需要使用此遮罩来过滤掉该步骤中不可用的动作。

如果您正在使用自定义策略，可以将此遮罩传递给您的概率分布，如下所示

带有动作遮罩的类别策略

     >>> from tensordict.nn import TensorDictModule, ProbabilisticTensorDictModule, TensorDictSequential
     >>> import torch.nn as nn
     >>> from torchrl.modules import MaskedCategorical
     >>> module = TensorDictModule(
     >>>     nn.Linear(in_feats, out_feats),
     >>>     in_keys=["observation"],
     >>>     out_keys=["logits"],
     >>> )
     >>> dist = ProbabilisticTensorDictModule(
     >>>     in_keys={"logits": "logits", "mask": "action_mask"},
     >>>     out_keys=["action"],
     >>>     distribution_class=MaskedCategorical,
     >>> )
     >>> actor = TensorDictSequential(module, dist)

如果您想使用默认策略，则需要使用 ActionMask 转换来包装您的环境。此转换可以负责更新动作规范中的动作遮罩，以便默认策略始终知道最新的可用动作。您可以按如下方式进行操作

如何使用动作遮罩转换

     >>> from tensordict.nn import TensorDictModule, ProbabilisticTensorDictModule, TensorDictSequential
     >>> import torch.nn as nn
     >>> from torchrl.envs.transforms import TransformedEnv, ActionMask
     >>> env = TransformedEnv(
     >>>     your_base_env
     >>>     ActionMask(action_key="action", mask_key="action_mask"),
     >>> )

注意

如果您使用并行环境，将转换添加到并行环境本身而不是其子环境非常重要。

记录器

在环境 rollout 执行期间记录数据对于关注算法性能以及训练后报告结果至关重要。

TorchRL 提供了几个工具来与环境输出交互：首先，可调用的回调函数 callback 可以传递给 rollout() 方法。此函数将在每次 rollout 迭代时对收集到的 tensordict 进行调用（如果某些迭代需要跳过，应在 callback 内部添加一个内部变量来跟踪调用次数）。

要将收集到的 tensordicts 保存到磁盘，可以使用 TensorDictRecorder。

录制视频

几个后端提供了录制环境渲染图像的可能性。如果像素已经是环境输出的一部分（例如 Atari 或其他游戏模拟器），可以将 VideoRecorder 添加到环境中。此环境转换将能够录制视频的日志记录器（例如 CSVLogger、WandbLogger 或 TensorBoardLogger）以及指示视频应保存位置的标签作为输入。例如，要将 mp4 视频保存到磁盘，可以使用带有 video_format=”mp4” 参数的 CSVLogger。

VideoRecorder 转换可以处理批量图像，并自动检测 numpy 或 PyTorch 格式的图像（WHC 或 CWH）。

>>> logger = CSVLogger("dummy-exp", video_format="mp4")
>>> env = GymEnv("ALE/Pong-v5")
>>> env = env.append_transform(VideoRecorder(logger, tag="rendered", in_keys=["pixels"]))
>>> env.rollout(10)
>>> env.transform.dump()  # Save the video and clear cache

请注意，转换的缓存会一直增长，直到调用 dump。用户有责任根据需要在必要时调用 dump 以避免 OOM（内存不足）问题。

在某些情况下，创建一个可以收集图像的测试环境既繁琐又昂贵，或者根本不可能（某些库只允许每个工作区一个环境实例）。在这种情况下，假设环境中有 render 方法可用，可以使用 PixelRenderTransform 在父环境上调用 render 并将图像保存在 rollout 数据流中。此类别同样适用于单个环境和批量环境。

>>> from torchrl.envs import GymEnv, check_env_specs, ParallelEnv, EnvCreator
>>> from torchrl.record.loggers import CSVLogger
>>> from torchrl.record.recorder import PixelRenderTransform, VideoRecorder
>>>
>>> def make_env():
>>>     env = GymEnv("CartPole-v1", render_mode="rgb_array")
>>>     # Uncomment this line to execute per-env
>>>     # env = env.append_transform(PixelRenderTransform())
>>>     return env
>>>
>>> if __name__ == "__main__":
...     logger = CSVLogger("dummy", video_format="mp4")
...
...     env = ParallelEnv(16, EnvCreator(make_env))
...     env.start()
...     # Comment this line to execute per-env
...     env = env.append_transform(PixelRenderTransform())
...
...     env = env.append_transform(VideoRecorder(logger=logger, tag="pixels_record"))
...     env.rollout(3)
...
...     check_env_specs(env)
...
...     r = env.rollout(30)
...     env.transform.dump()
...     env.close()

记录器是注册接收到的数据以便进行日志记录的转换。

TensorDictRecorder(out_file_base[, ...])	TensorDict 记录器。
VideoRecorder(logger, tag[, in_keys, skip, ...])	视频记录器转换。
PixelRenderTransform([out_keys, preproc, ...])	用于在父环境上调用 render 并在 tensordict 中注册像素观察结果的转换。

助手/工具函数

RandomPolicy(action_spec[, action_key])	用于数据收集器的随机策略。
check_env_specs(env[, return_contiguous, ...])	根据短期 rollout 的结果测试环境规范。
exploration_type()	返回当前的采样类型。
get_available_libraries()	返回所有支持的库。
make_composite_from_td(data[, ...])	从 tensordict 创建一个 Composite 实例，假设所有值都是无界的。
set_exploration_type	`set_interaction_type` 的别名
step_mdp(tensordict[, next_tensordict, ...])	创建一个反映输入 tensordict 时间步长的新 tensordict。
terminated_or_truncated(data[, ...])	读取 tensordict 中的 done / terminated / truncated 键，并写入一个新的 tensor，其中聚合了两个信号的值。

特定领域

ModelBasedEnvBase(*args, **kwargs)	用于基于模型的强化学习 sota 实现的基本环境。
model_based.dreamer.DreamerEnv(*args, **kwargs)	Dreamer 模拟环境。
model_based.dreamer.DreamerDecoder([...])	用于记录 Dreamer 中解码观察结果的转换。

库

TorchRL 的使命是让控制和决策算法的训练尽可能简单，无论使用何种模拟器（如果使用）。针对 DMControl、Habitat、Jumanji 当然还有 Gym 提供了多个包装器。

后一个库在强化学习社区中具有特殊地位，是编写模拟器最常用的框架。其成功的 API 奠定了基础并启发了许多其他框架，其中包括 TorchRL。然而，Gym 经历了多次设计变更，作为外部采用库有时难以适应这些变化：用户通常有他们“偏好”的库版本。此外，Gym 现在由另一个团队以“gymnasium”的名义维护，这不利于代码兼容性。实际上，我们必须考虑到用户可能在同一个虚拟环境中安装了 Gym 和 Gymnasium 的版本，并且必须允许两者同时工作。幸运的是，TorchRL 为此问题提供了解决方案：特殊的装饰器 set_gym_backend 允许控制相关函数中使用哪个库。

>>> from torchrl.envs.libs.gym import GymEnv, set_gym_backend, gym_backend
>>> import gymnasium, gym
>>> with set_gym_backend(gymnasium):
...     print(gym_backend())
...     env1 = GymEnv("Pendulum-v1")
<module 'gymnasium' from '/path/to/venv/python3.9/site-packages/gymnasium/__init__.py'>
>>> with set_gym_backend(gym):
...     print(gym_backend())
...     env2 = GymEnv("Pendulum-v1")
<module 'gym' from '/path/to/venv/python3.9/site-packages/gym/__init__.py'>
>>> print(env1._env.env.env)
<gymnasium.envs.classic_control.pendulum.PendulumEnv at 0x15147e190>
>>> print(env2._env.env.env)
<gym.envs.classic_control.pendulum.PendulumEnv at 0x1629916a0>

我们可以看到这两个库修改了 gym_backend() 返回的值，该值可以进一步用于指示当前计算需要使用哪个库。set_gym_backend 也是一个装饰器：我们可以用它来告诉特定函数在其执行期间需要使用哪个 Gym 后端。torchrl.envs.libs.gym.gym_backend() 函数允许您获取当前的 Gym 后端或其任何模块。

>>> import mo_gymnasium
>>> with set_gym_backend("gym"):
...     wrappers = gym_backend('wrappers')
...     print(wrappers)
<module 'gym.wrappers' from '/path/to/venv/python3.9/site-packages/gym/wrappers/__init__.py'>
>>> with set_gym_backend("gymnasium"):
...     wrappers = gym_backend('wrappers')
...     print(wrappers)
<module 'gymnasium.wrappers' from '/path/to/venv/python3.9/site-packages/gymnasium/wrappers/__init__.py'>

另一个与 Gym 和其他外部依赖项相关的有用工具是 torchrl._utils.implement_for 类。使用 @implement_for 装饰函数会告诉 torchrl，根据指定的版本，预期会有特定的行为。这使得我们能够轻松支持多个版本的 Gym，而无需用户付出任何努力。例如，假设我们的虚拟环境安装了 v0.26.2，那么查询时以下函数将返回 1。

>>> from torchrl._utils import implement_for
>>> @implement_for("gym", None, "0.26.0")
... def fun():
...     return 0
>>> @implement_for("gym", "0.26.0", None)
... def fun():
...     return 1
>>> fun()
1

BraxEnv(*args, **kwargs)	Google Brax 环境包装器，使用环境名称构建。
BraxWrapper(*args, **kwargs)	Google Brax 环境包装器。
DMControlEnv(*args, **kwargs)	DeepMind Control lab 环境包装器。
DMControlWrapper(*args, **kwargs)	DeepMind Control lab 环境包装器。
GymEnv(*args, **kwargs)	OpenAI Gym 环境包装器，直接由环境 ID 构建。
GymWrapper(*args, **kwargs)	OpenAI Gym 环境包装器。
HabitatEnv(*args, **kwargs)	Habitat 环境的包装器。
IsaacGymEnv(*args, **kwargs)	IsaacGym 环境的 TorchRL 环境接口。
IsaacGymWrapper(*args, **kwargs)	IsaacGymEnvs 环境的包装器。
JumanjiEnv(*args, **kwargs)	Jumanji 环境包装器，使用环境名称构建。
JumanjiWrapper(*args, **kwargs)	Jumanji 的环境包装器。
MeltingpotEnv(*args, **kwargs)	Meltingpot 环境包装器。
MeltingpotWrapper(*args, **kwargs)	Meltingpot 环境包装器。
MOGymEnv(*args, **kwargs)	FARAMA MO-Gymnasium 环境包装器。
MOGymWrapper(*args, **kwargs)	FARAMA MO-Gymnasium 环境包装器。
MultiThreadedEnv(*args, **kwargs)	基于 EnvPool 的环境的多线程执行。
MultiThreadedEnvWrapper(*args, **kwargs)	基于 envpool 的多线程环境的包装器。
OpenMLEnv(*args, **kwargs)	用于 bandit 上下文中使用的 OpenML 数据的环境接口。
OpenSpielWrapper(*args, **kwargs)	Google DeepMind OpenSpiel 环境包装器。
OpenSpielEnv(*args, **kwargs)	Google DeepMind OpenSpiel 环境包装器，使用游戏字符串构建。
PettingZooEnv(*args, **kwargs)	PettingZoo 环境。
PettingZooWrapper(*args, **kwargs)	PettingZoo 环境包装器。
RoboHiveEnv(*args, **kwargs)	用于 RoboHive gym 环境的包装器。
SMACv2Env(*args, **kwargs)	SMACv2 (StarCraft Multi-Agent Challenge v2) 环境包装器。
SMACv2Wrapper(*args, **kwargs)	SMACv2 (StarCraft Multi-Agent Challenge v2) 环境包装器。
UnityMLAgentsEnv(*args, **kwargs)	Unity ML-Agents 环境包装器。
UnityMLAgentsWrapper(*args, **kwargs)	Unity ML-Agents 环境包装器。
VmasEnv(*args, **kwargs)	Vmas 环境包装器。
VmasWrapper(*args, **kwargs)	Vmas 环境包装器。
gym_backend([submodule])	返回 gym 后端或其子模块。
set_gym_backend(backend)	将 gym 后端设置为某个值。
register_gym_spec_conversion(spec_type)	注册特定 spec 类型转换函数的装饰器。