torchrl.envs 包

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 期望存在的设备。环境设备并不意味着实际的步进操作将在设备上计算（这是后端的责任，TorchRL 在这方面能做的很少）。环境的设备仅表示数据输入或从环境中检索时期望的设备。TorchRL 负责将数据映射到所需的设备。这对于转换（见下文）特别有用。对于参数化环境（例如，基于模型的环境），设备确实代表了将用于计算操作的硬件。

• env.observation_spec：一个 Composite 对象，包含所有观察键-规格对。

• env.state_spec：一个 Composite 对象，包含所有输入键-规格对（动作除外）。对于大多数有状态环境，此容器将为空。

• env.action_spec：一个 TensorSpec 对象，表示动作规格。

• env.reward_spec：一个 TensorSpec 对象，表示奖励规格。

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

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

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

如果环境包含非张量数据，则可以使用 NonTensor 实例。

环境规格：锁定和批量大小

环境规格默认是锁定的（通过传递给环境构造函数的 spec_locked 参数）。锁定规格意味着任何对规格（或其子项，如果它是 Composite 实例）的修改都需要先解锁。这可以通过 set_spec_lock_() 来完成。规格默认被锁定的原因是，这使得缓存动作或重置键等值变得容易。解锁环境仅应在预期规格会经常被修改时进行（原则上，这应该避免）。规格的修改，如 env.observation_spec = new_spec 是允许的：在幕后，TorchRL 会清除缓存，解锁规格，进行修改，并在环境先前被锁定时重新锁定规格。

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

环境方法

通过这些，实现了以下方法：

• env.reset()：一个重置方法，可能（但不一定需要）接受一个 tensordict.TensorDict 输入。它返回一个轨迹的第一个 tensordict，通常包含一个 "done" 状态和一组观察。如果不存在，“reward”键将被初始化为 0 和适当的形状。

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

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

  >>> 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()：一个播种方法，将在多环境设置中返回下一个要使用的种子。此下一个种子是从前一个种子确定性计算的，这样就可以为多个环境播种不同的种子，而不会在连续的实验中冒种子重叠的风险，同时仍然具有可复现的结果。

• env.rollout()：在环境中执行一个轨迹，最多 max_steps=N 步，并使用策略（policy=model）。策略应使用 tensordict.nn.TensorDictModule（或其他 tensordict.TensorDict 兼容模块）进行编码。生成的 tensordict.TensorDict 实例将带有尾随的 "time" 命名维度，该维度可用于其他模块正确处理此批处理维度。

下图总结了 torchrl 中轨迹的执行方式。

使用 TensorDict 的 TorchRL 轨迹。

简而言之，一个 TensorDict 由 reset() 方法创建，然后由策略填充动作，再传递给 step() 方法，该方法在 "next" 条目下写入观察、完成标志和奖励。此调用的结果被存储以供传递，并且 "next" 条目由 step_mdp() 函数收集。

注意

通常，所有 TorchRL 环境在其输出 tensordict 中都有 "done" 和 "terminated" 条目。如果它们不是按设计存在的，EnvBase 元类将确保每个完成或终止都与其对偶配对。在 TorchRL 中，"done" 严格指所有轨迹结束信号的并集，应解释为“轨迹的最后一步”或等效地“需要重置的信号”。如果环境提供它（例如，Gymnasium），截断条目也会写入 EnvBase.step() 输出，并显示在 "truncated" 条目下。如果环境包含单个值，则默认将其解释为 "terminated" 信号。默认情况下，TorchRL 的收集器和轨迹函数将查找 "done" 条目来评估环境是否应重置。

注意

可以使用 torchrl.collectors.utils.split_trajectories 函数来切片相邻的轨迹。它依赖于输入 tensordict 中的 "traj_ids" 条目，或者在 "traj_ids" 缺失时依赖于 "done" 和 "truncated" 键的连接。

注意

在某些情况下，标记轨迹的第一步可能很有用。TorchRL 通过 InitTracker 转换提供了此功能。

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

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

部分步进和部分重置

TorchRL 允许环境重置部分而非全部环境，或者在一个环境但非所有环境中执行一个步进。如果批次中只有一个环境，则也允许进行部分重置/步进，其行为如下所述。

批处理环境和锁定批处理

在详细说明部分重置和部分步进的作用之前，我们必须区分环境自身具有批量大小（主要是状态化环境）的情况，或者环境只是一个简单的模块，给定任意大小的输入，它会对所有元素进行批处理操作（主要是无状态环境）的情况。

这由 batch_locked 属性控制：批处理锁定的环境要求所有输入 tensordict 具有与环境相同的批量大小。这些环境的典型示例是 GymEnv 及其相关类。相反，批处理未锁定的环境可以处理任何输入大小。值得注意的例子有 BraxEnv 或 JumanjiEnv。

在批处理未锁定的环境中执行部分步进很简单：只需屏蔽不需要执行的 tensordict 部分，将另一部分传递给 step，然后将结果与先前的输入合并。

批处理环境（ParallelEnv 和 SerialEnv）也可以轻松处理部分步进，它们只是将动作传递给需要执行的子环境。

在所有其他情况下，TorchRL 假定环境能够正确处理部分步进。

警告

这意味着自定义环境可能会默默地运行非必需的步进，因为 torchrl 无法控制 _step 方法内部发生的事情！

部分步进

部分步进通过临时键 “_step” 进行控制，该键指向一个布尔掩码，其大小与持有它的 tensordict 相同。能够处理此情况的类是：

• 批处理环境：ParallelEnv 和 SerialEnv 将动作分派给且仅分派给 “_step” 为 True 的环境；

• 批处理未解锁的环境；

• 未批处理的环境（即，没有批量大小的环境）。在这些环境中，step() 方法将首先查找 “_step” 条目，如果存在，则相应操作。如果 Transform 实例将 “_step” 条目传递给 tensordict，它也会被 TransformedEnv 自己的 _step 方法捕获，该方法将跳过 base_env.step 以及任何进一步的转换。

在处理部分步进时，策略始终是使用步进输出，并使用输入 tensordict 的先前内容（如果存在）或 0 值张量（如果找不到张量）来屏蔽缺失值。这意味着如果输入 tensordict 不包含所有先前的观察，那么对于所有未执行步进的元素，输出 tensordict 将为 0 值。在批处理环境、数据收集器和轨迹实用程序中，不会出现此问题，因为这些类会正确处理数据传递。

当 break_when_all_done 为 True 时，部分步进是 rollout() 的一个基本特性，因为在调用 _step 时需要跳过具有 True 完成状态的环境。

ConditionalSkip 转换允许您以编程方式请求（部分）步进跳过。

部分重置

部分重置的工作方式与部分步进几乎相同，但使用 “_reset” 条目。

部分步进的相同限制也适用于部分重置。

同样，当 break_when_any_done 为 True 时，部分重置是 rollout() 的一个基本特性，因为需要重置具有 True 完成状态的环境，但不是其他环境。

有关批处理和向量化环境中部分重置的深入分析，请参阅以下段落。

向量化环境

向量化（或更准确地说，并行）环境是强化学习中的一个常见特性，因为执行环境步进可能非常耗费 CPU。一些库，如 gym3 或 EnvPool，提供了同时执行环境批次的接口。虽然它们通常提供非常有竞争力的计算优势，但它们不一定能支持 TorchRL 所支持的广泛的环境库。因此，TorchRL 提供了自己的通用 ParallelEnv 类来并行运行多个环境。由于此类继承自 SerialEnv，因此它享有与其他环境完全相同的 API。当然，ParallelEnv 的批量大小将对应于其环境计数。

注意

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

重要的是，您的环境规格应与它发送和接收的输入和输出相匹配，因为 ParallelEnv 将根据这些规格创建缓冲区，以与生成的进程进行通信。请检查 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" 来指示环境的哪个组件（子环境或代理）应该被重置。这允许重置部分但非全部组件。

"_reset" 键有两个不同的功能：

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

   环境在其 _reset() 方法中处理 "_reset" 键的方式是其类特有的。设计一个根据 "_reset" 输入进行行为的环境是开发者的责任，因为 TorchRL 无法控制 _reset() 的内部逻辑。但是，在设计该方法时应牢记以下几点。

2. 在调用 _reset() 之后，输出将使用 "_reset" 条目进行掩码，并且先前 step() 的输出将写入 "_reset" 为 False 的地方。实际上，这意味着如果 "_reset" 修改了未被其暴露的数据，这些修改将丢失。在此掩码操作之后，"_reset" 条目将从 reset() 输出中删除。

必须指出的是，"_reset" 是一个私有键，仅在编写面向内部的特定环境功能时使用。换句话说，它不应在库外使用，并且开发人员保留在不事先保证的情况下修改通过 "_reset" 设置进行部分重置的逻辑的权利，只要它们不影响 TorchRL 的内部测试。

最后，在设计重置功能时，应牢记以下假设：

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

• 一个级别的重置并不排除较低级别上 "_reset" 的存在，但它会消除其影响。原因很简单，根级别的 "_reset" 对应于 all()、any() 还是对嵌套 "done" 条目的自定义调用是无法预先知道的，并且明确假设根级别的 "_reset" 是为了覆盖嵌套值（例如，请参阅 PettingZooWrapper 实现，其中每个组都有一个或多个与 "done" 条目相关联，该条目在根级别使用 any 或 all 逻辑进行聚合，具体取决于任务）。

• 在调用 env.reset(tensordict)() 时，如果有一个部分 "_reset" 条目将重置部分而非全部完成的子环境，则输入数据应包含 __未__ 被重置的子环境的数据。此约束的原因在于 env._reset(data) 的输出仅对已重置的条目有意义。对于其他条目，TorchRL 无法提前知道它们是否有意义。例如，可以完美地用 0 填充未重置组件的值，在这种情况下，未重置的数据将无意义，应被丢弃。

下面，我们提供了一些 "_reset" 键对重置后返回零的环境的预期效果的示例：

>>> # 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])

并行环境重置

     >>> 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（和其他导入）所需的时间，启动并行环境可能成为瓶颈。这就是为什么 TorchRL 测试如此缓慢的原因。环境启动后，应观察到速度大幅提升。

注意

TorchRL 要求精确的规格：另一个需要考虑的事情是，ParallelEnv（以及数据收集器）将根据环境规格创建数据缓冲区，以在不同进程之间传递信息。这意味着规格错误（输入、观察或奖励）将在运行时导致中断，因为数据无法写入预先分配的缓冲区。通常，在使用 ParallelEnv 之前，应使用 check_env_specs() 测试函数来测试环境。该函数将在预分配的缓冲区和收集的数据不匹配时引发断言错误。

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

除了提供基于进程的并行性的 ParallelEnv 之外，我们还提供了一种使用 MultiThreadedEnv 创建多线程环境的方法。该类底层使用 EnvPool 库，可以实现更高的性能，但同时限制了灵活性——只能创建 EnvPool 中实现的环境。这涵盖了许多流行的 RL 环境类型（Atari、Classic Control 等），但不能使用任意 TorchRL 环境，而使用 ParallelEnv 则可以。运行 benchmarks/benchmark_batched_envs.py 以比较不同并行化批处理环境的方法。

SerialEnv(*args, **kwargs)	在同一进程中创建一系列环境。批处理环境允许用户查询远程运行的环境的任意方法/属性。
ParallelEnv(*args, **kwargs)	每个进程创建一个环境。
EnvCreator(create_env_fn[, ...])	环境创建者类。

异步环境

异步环境允许并行执行多个环境，这可以显著加快强化学习中的数据收集过程。

AsyncEnvPool 类及其子类提供了使用不同后端（如线程和多进程）管理这些环境的灵活接口。

AsyncEnvPool 类作为异步环境池的基类，提供了管理多个环境并发执行的通用接口。它支持并行执行的不同后端（如线程和多进程），并提供了异步步进和重置环境的方法。

与 ParallelEnv 相反，AsyncEnvPool 及其子类允许在执行某个任务的同时执行给定的一组子环境，从而允许同时运行复杂的异步作业。例如，可以在策略根据其他环境的输出来运行时执行某些环境。

当处理具有高（和/或可变）延迟的环境时，此类及其子类尤其有用。

注意

此类及其子类应与 TransformedEnv 和批处理环境一起嵌套使用，但用户目前无法在使用这些类时使用基础环境的异步功能。应优先在 AsyncEnvPool 内嵌套转换后的环境。如果这不可行，请提交 issue。

类

• AsyncEnvPool：异步环境池的基类。它根据提供的参数确定要使用的后端实现，并管理环境的生命周期。

• ProcessorAsyncEnvPool：使用多进程并行执行环境的 AsyncEnvPool 的实现。此类管理一个环境池，每个环境运行在自己的进程中，并提供使用进程间通信进行异步步进和重置环境的方法。在实例化 AsyncEnvPool 时，如果将 “multiprocessing” 作为后端传递，它会被自动实例化。

• ThreadingAsyncEnvPool：使用线程并行执行环境的 AsyncEnvPool 的实现。此类管理一个环境池，每个环境运行在自己的线程中，并提供使用线程池执行器进行异步步进和重置环境的方法。在实例化 AsyncEnvPool 时，如果将 “threading” 作为后端传递，它会被自动实例化。

示例

>>> from functools import partial
>>> from torchrl.envs import AsyncEnvPool, GymEnv
>>> import torch
>>> # Choose backend
>>> backend = "threading"
>>> env = AsyncEnvPool(
>>>     [partial(GymEnv, "Pendulum-v1"), partial(GymEnv, "CartPole-v1")],
>>>     stack="lazy",
>>>     backend=backend
>>> )
>>> # Execute a synchronous reset
>>> reset = env.reset()
>>> print(reset)
>>> # Execute a synchronous step
>>> s = env.rand_step(reset)
>>> print(s)
>>> # Execute an asynchronous step in env 0
>>> s0 = s[0]
>>> s0["action"] = torch.randn(1).clamp(-1, 1)
>>> s0["env_index"] = 0
>>> env.async_step_send(s0)
>>> # Receive data
>>> s0_result = env.async_step_recv()
>>> print('result', s0_result)
>>> # Close env
>>> env.close()

AsyncEnvPool(*args, **kwargs)	异步环境池的基类，提供了管理多个环境并发执行的通用接口。
ProcessorAsyncEnvPool(*args, **kwargs)	使用多进程并行执行环境的 AsyncEnvPool 的实现。
ThreadingAsyncEnvPool(*args, **kwargs)	使用线程并行执行环境的 AsyncEnvPool 的实现。

自定义原生 TorchRL 环境

TorchRL 提供了一系列自定义内置环境。

ChessEnv(*args, **kwargs)	遵循 TorchRL API 的国际象棋环境。
PendulumEnv(*args, **kwargs)	一个无状态的 Pendulum 环境。
TicTacToeEnv(*args, **kwargs)	井字棋实现。
LLMHashingEnv(*args, **kwargs)	一个使用哈希模块来识别唯一观测值的文本生成环境。

多代理环境

TorchRL 支持开箱即用的多代理学习。在单代理学习管道中使用的相同类可以无缝地用于多代理上下文，而无需任何修改或专用的多代理基础设施。

从这个角度来看，环境在多代理学习中起着核心作用。在多代理环境中，许多决策代理在共享的世界中进行操作。代理可以观察到不同的事物，以不同的方式行动，并获得不同的奖励。因此，存在许多范式来模拟多代理环境（DecPODPs、马尔可夫博弈）。这些范式之间的一些主要区别包括：

• 观察可以是每个代理的，也可以包含一些共享组件。

• 奖励可以是每个代理的或共享的。

• 完成（以及 "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)，这表示额外的代理维度。

多代理张量规格将遵循与 tensordicts 相同的样式。与代理相关的值的规格需要嵌套在“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,
     ... )

正如您所见，这非常简单！每个代理的键将具有嵌套的复合规格，而共享键将遵循单代理标准。

注意

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

注意

TorchRL 会抽象这些嵌套的规格以方便使用。这意味着访问 env.reward_spec 将始终返回叶子规格，如果访问的规格是 Composite。因此，如果在上面的示例中，我们在环境创建后运行 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 Group Map 类型。
check_marl_grouping(group_map, agent_names)	检查 MARL 组映射。

自动重置环境

自动重置环境是指在轨迹收集过程中，当环境达到 "done" 状态时，不期望调用 reset()，因为重置是自动发生的。通常，在这种情况下，与完成和奖励一起提供的观察（实际上是执行环境中的动作的结果）是新剧集的第一个观察，而不是当前剧集的最后一个观察。

为了处理这些情况，torchrl 提供了一个 AutoResetTransform，它会将调用 step 产生的结果观察复制到下一个 reset，并在轨迹收集过程中（在 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]])

动态规格

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

TorchRL 能够处理动态规格，但批处理环境和收集器需要了解此功能。请注意，实际上，这是自动检测的。

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

>>> 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]) -> None:
...         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() 对于沿某些维度形状可变的动态规格会通过，但对于在一次步进中存在而在其他步进中不存在的键，或者维度数量可变的键，则不会。

变换 (Transforms)

在大多数情况下，环境的原始输出必须在传递给另一个对象（如策略或值运算符）之前进行处理。为此，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。

重命名转换逻辑

注意

在调用 inv 时，转换按相反的顺序执行（与正向/步进模式相比）。

转换张量和规格

转换实际张量（来自策略）时，过程如图所示：

>>> 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()))

此方法确保“外部”规格是从“内部”规格推断出来的。请注意，我们特意未调用 _inv_apply_transform 而是 _apply_transform！

将规格暴露给外部世界

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

设计自己的转换

要创建基本的自定义转换，您需要继承 Transform 类并实现 _apply_transform() 方法。这是一个向观察张量加 1 的简单转换的示例：

>>> class AddOneToObs(Transform):
...     """A transform that adds 1 to the observation tensor."""
...
...     def __init__(self):
...         super().__init__(in_keys=["observation"], out_keys=["observation"])
...
...     def _apply_transform(self, obs: torch.Tensor) -> torch.Tensor:
...         return obs + 1

继承 Transform 的技巧

有多种方法可以继承转换。需要考虑的事项是：

• 转换对于每个被转换的张量/项是否相同？使用 _apply_transform() 和 _inv_apply_transform()。

• 转换是否需要访问输入数据到 env.step 以及输出？重写 _step()。否则，重写 _call()（或 _inv_call()）。

• 转换是否要在回放缓冲区中使用？重写 forward()、inv()、_apply_transform() 或 _inv_apply_transform()。

• 在转换内部，您可以使用 parent（基础 env + 直到此处的*所有*转换）或 container()（封装转换的对象）来访问（并调用）父环境。

• 请记住，如果需要，请编辑规范：顶级：transform_output_spec()、transform_input_spec()。叶级：transform_observation_spec()、transform_action_spec()、transform_state_spec()、transform_reward_spec() 和 transform_reward_spec()。

有关实际示例，请参阅上述方法。

您可以通过将转换传递给 TransformedEnv 构造函数来在环境中将其用作转换。

>>> env = TransformedEnv(GymEnv("Pendulum-v1"), AddOneToObs())

您可以使用 Compose 类将多个转换组合在一起。

>>> transform = Compose(AddOneToObs(), RewardSum())
>>> env = TransformedEnv(GymEnv("Pendulum-v1"), transform)

反向转换

某些转换具有可以用来撤消转换的反向转换。例如，AddOneToAction 转换有一个反向转换，它从动作张量中减去 1。

>>> class AddOneToAction(Transform):
...     """A transform that adds 1 to the action tensor."""
...     def __init__(self):
...         super().__init__(in_keys=[], out_keys=[], in_keys_inv=["action"], out_keys_inv=["action"])
...     def _inv_apply_transform(self, action: torch.Tensor) -> torch.Tensor:
...         return action + 1

将转换与回放缓冲区结合使用

您可以通过将转换传递给 ReplayBuffer 构造函数来将其与回放缓冲区结合使用。

克隆转换

由于附加到环境的转换是通过 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, ...])	环境转换的基类，用于修改或创建 tensordict 中的新数据。
TransformedEnv(*args, **kwargs)	已转换的环境。
ActionDiscretizer(num_intervals[, ...])	用于离散化连续动作空间的转换。
ActionMask([action_key, mask_key])	自适应动作掩码器。
AutoResetEnv(*args, **kwargs)	用于自动重置环境的子类。
AutoResetTransform(*[, replace, fill_float, ...])	用于自动重置环境的转换。
BatchSizeTransform(*[, batch_size, ...])	用于修改环境批处理大小的转换。
BinarizeReward([in_keys, out_keys])	将奖励映射为二进制值（分别为 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)	组合一系列转换。
ConditionalPolicySwitch(policy, condition)	一种根据指定条件有条件地在策略之间切换的转换。
ConditionalSkip(cond)	一种在满足特定条件时跳过环境中步骤的转换。
Crop(w[, h, top, left, in_keys, out_keys])	在指定位置和输出大小处裁剪输入图像。
DTypeCastTransform(dtype_in, dtype_out[, ...])	将一个 dtype 转换为另一个 dtype，针对选定的键。
DeviceCastTransform(device[, orig_device, ...])	将数据从一个设备移动到另一个设备。
DiscreteActionProjection(...[, action_key, ...])	将离散动作从高维空间投影到低维空间。
DoubleToFloat([in_keys, out_keys, ...])	将一个 dtype 转换为另一个 dtype，针对选定的键。
EndOfLifeTransform([eol_key, lives_key, ...])	使用 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[, in_keys_inv, ...])	向 tensordict 添加哈希值。
InitTracker([init_key])	重置跟踪器。
KLRewardTransform(actor[, coef, in_keys, ...])	一种用于向奖励添加 KL[pi_current||pi_0] 校正项的转换。
LineariseRewards(in_keys[, out_keys, weights])	通过加权求和将多目标奖励信号转换为单目标奖励信号。
MultiAction(*[, dim, stack_rewards, ...])	用于在父环境中执行多个动作的转换。
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[, ...])	用于回放缓冲区和模块的轨迹子采样器。
RemoveEmptySpecs([in_keys, out_keys, ...])	从环境中删除空的规范和内容。
RenameTransform(in_keys, out_keys[, ...])	用于重命名输出 tensordict 中的条目（或通过反向键重命名输入 tensordict）的转换。
Resize(w[, h, interpolation, in_keys, out_keys])	调整像素观测的大小。
Reward2GoTransform([gamma, in_keys, ...])	根据剧集奖励和折扣因子计算目标奖励。
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)	在指定位置删除大小为一的维度。
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 的 primer。
TimeMaxPool([in_keys, out_keys, T, reset_key])	在最后 T 个观测值中，取每个位置的最大值。
Timer([out_keys, time_key])	一种用于测量环境中 inv 和 call 操作之间时间间隔的转换。
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, repeats])	全局轨迹计数器转换。
UnaryTransform(in_keys, out_keys[, ...])	对指定的输入应用一元操作。
UnsqueezeTransform(*args, **kwargs)	在指定位置插入大小为一的维度。
VC1Transform(in_keys, out_keys, model_name)	VC1 转换类。
VIPRewardTransform(*args, **kwargs)	一种 VIP 转换，用于根据嵌入相似性计算奖励。
VIPTransform(*args, **kwargs)	VIP 转换类。
VecGymEnvTransform([final_name, ...])	用于 GymWrapper 子类的转换，可处理一致的自动重置。
VecNorm(*args, **kwargs)	torchrl 环境的移动平均归一化层。
VecNormV2(in_keys[, out_keys, lock, ...])	用于在强化学习环境中对向量化观测和奖励进行归一化的类。
gSDENoise([state_dim, action_dim, shape])	gSDE 噪声初始化器。

具有动作掩码的环境

在某些具有离散动作的环境中，代理可用的动作可能会在执行过程中发生变化。在这种情况下，环境将输出一个动作掩码（默认为 "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"),
     >>> )

注意

如果您使用的是并行环境，则必须将转换添加到并行环境本身，而不是添加到其子环境中。

记录器

在环境滚动执行过程中记录数据对于密切关注算法性能以及在训练后报告结果至关重要。

TorchRL 提供了几种与环境输出交互的工具：最重要的是，可以将 callback 可调用对象传递给 rollout() 方法。此函数将在滚动每次迭代时被调用（如果必须跳过某些迭代，则应添加一个内部变量来跟踪 callback 中的调用计数）。

要将收集的 tensordict 保存到磁盘，可以使用 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 并将图像保存在滚动数据流中。此类同样适用于单环境和批处理环境。

>>> 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 键，并写入一个新张量，其中两个信号的值都被聚合。

特定领域

ModelBasedEnvBase(*args, **kwargs)	基于模型的 RL sota 实现的基础环境。
model_based.dreamer.DreamerEnv(*args, **kwargs)	Dreamer 模拟环境。
model_based.dreamer.DreamerDecoder([...])	用于在 Dreamer 中记录解码观测值的转换。

库

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

最后一个库在 RL 社区中具有特殊地位，因为它是最常用的编码模拟器框架。其成功的 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.10/site-packages/gymnasium/__init__.py'>
>>> with set_gym_backend(gym):
...     print(gym_backend())
...     env2 = GymEnv("Pendulum-v1")
<module 'gym' from '/path/to/venv/python3.10/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.10/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.10/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 实验室环境封装器。
DMControlWrapper(*args, **kwargs)	DeepMind Control 实验室环境封装器。
GymEnv(*args, **kwargs)	直接通过环境 ID 构建的 OpenAI Gym 环境包装器。
GymWrapper(*args, **kwargs)	OpenAI Gym 环境包装器。
HabitatEnv(*args, **kwargs)	Habitat 环境的包装器。
IsaacGymEnv(*args, **kwargs)	IsaacGym 环境的 TorchRL Env 接口。
IsaacGymWrapper(*args, **kwargs)	IsaacGymEnvs 环境的包装器。
IsaacLabWrapper(*args, **kwargs)	IsaacLab 环境的包装器。
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)	用于在 the bandit contexts 中使用的 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)	用于注册特定规范类型的转换函数的装饰器。