torchrl.envs 包

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

每个 env 都将具有以下属性

• env.batch_size：一个 torch.Size，表示一起批处理的 env 数量。

• 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 实例。

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

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

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

环境方法

有了这些，实现了以下方法：

• env.reset()：一个重置方法，可以（但不一定需要）接受 tensordict.TensorDict 作为输入。它返回一个 rollout 的第一个 tensordict，通常包含一个 "done" 状态和一组观测值。如果不存在，则会实例化一个具有合适形状的 “reward” 键，其值为 0。

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

• env.step_and_maybe_reset()：执行一步，并在需要时（部分）重置环境。它返回更新后的输入，其中包含下一步数据的 "next" 键，以及一个包含下一步输入数据（即重置、结果或 step_mdp()）的 tensordict。这是通过读取 done_keys 并为每个完成状态分配一个 "_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()：一个种子设置方法，在多环境设置中返回下一个要使用的种子。这个下一个种子是从前一个种子确定的，这样就可以用不同的种子来设置多个环境，而不会冒着在连续实验中种子重叠的风险，同时仍然保持可复现的结果。

• env.rollout()：在环境中执行一个 rollout，最大步数为 (max_steps=N)，并使用策略 (policy=model)。策略应该使用 tensordict.nn.TensorDictModule（或其他任何 tensordict.TensorDict 兼容模块）进行编码。生成的 tensordict.TensorDict 实例将被标记一个尾随的 "time" 命名维度，其他模块可以使用此维度来正确处理它。

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

使用 TensorDict 的 TorchRL rollout。

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

注意

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

注意

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

注意

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

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

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

部分步进和部分重置

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

环境批处理和锁定批次

在详细说明部分重置和部分步进之前，我们必须区分环境有自己的批量大小（主要是状态化环境）或环境仅仅是一个模块，它接收任意大小的输入，并将操作批量化到所有元素上（主要是无状态环境）的情况。

这由 batch_locked 属性控制：一个批次锁定的环境要求所有输入 tensordict 都具有与 env 相同的批量大小。这些环境的典型例子是 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 值。在批处理环境、数据收集器和 rollout 工具中，这种情况不会出现，因为这些类会正确处理数据的传递。

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

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

部分重置

部分重置的工作方式与部分步进非常相似，但使用的是 “_reset” 条目。

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

同样，部分重置是 rollout() 的一项基本功能，当 break_when_any_done 为 True 时，因为具有 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() 时，输入 tensordict 中可能存在也可能不存在 "_reset" 键。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 逻辑进行聚合）。

• 在调用带有部分 "_reset" 条目的 env.reset(tensordict)() 时，该条目将重置某些但不全部完成的子环境，输入数据应包含__未__重置的子环境的数据。此限制的原因在于，env._reset(data) 的输出只能预测被重置的条目。对于其他条目，TorchRL 无法提前知道它们是否有意义。例如，可以将未重置组件的值填充为任意值，在这种情况下，未重置的数据将无意义，应予以丢弃。

下面，我们给出 "_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)	一个无状态的单摆环境。
TicTacToeEnv(*args, **kwargs)	一个井字棋实现。
LLMHashingEnv(*args, **kwargs)	一个使用哈希模块来识别唯一观测值的文本生成环境。

多智能体环境

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

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

• 观测 可以是每个智能体独有的，也可以包含一些共享的组件。

• 奖励 可以是每个智能体独有的或共享的。

• 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**，存在于具有批次大小 (num_envs,) 的根 tensordict 中，这代表了模拟的环境数量。

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

多智能体张量规范将遵循与 tensordict 相同的样式。与智能体相关的规范需要嵌套在“agents”条目中。

下面是一个如何创建多智能体环境规范的示例，其中只有 done 标志在智能体之间共享（如 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 会抽象掉这些嵌套规范，以方便使用。这意味着访问 env.reward_spec 时，如果访问的规范是 Composite，将始终返回叶子规范。因此，如果我们在创建环境后运行上面示例中的 env.reward_spec，我们将获得与 torch.stack(reward_specs)} 相同的输出。要获取包含“agents”键的完整复合规范，您可以运行 env.output_spec[“full_reward_spec”]。对于 action 和 done 规范也是如此。请注意 env.reward_spec == env.output_spec[“full_reward_spec”][env.reward_key]。

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

自动重置环境

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

为了处理这些情况，torchrl 提供了一个 AutoResetTransform，它会将调用 step 产生的观测复制到下一个 reset，并跳过回滚过程中的 reset 调用（在 rollout() 和 SyncDataCollector 迭代中）。此转换类还提供了对无效观测要采用的行为的细粒度控制，这些行为可以用 “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。由于数据无法连续堆叠，因此调用 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]) -> 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（基础环境 + 直到此处的*;* 所有转换）或 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, ...])	从 Gym 环境向具有 lives 方法的 end-of-life 信号注册。
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初始化的启动器。
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"),
     >>> )

注意

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

录制器

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

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

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

录制视频

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

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

Recorder 是在数据传入时进行注册以供日志记录的转换。

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	alias of set_interaction_type
step_mdp(tensordict[, next_tensordict, ...])	创建一个新的 tensordict，该 tensordict 反映了输入 tensordict 的时间步长。
terminated_or_truncated(data[, ...])	读取 tensordict 中的 done / terminated / truncated 键，并写入一个新张量，其中两个信号的值被聚合。

领域特定

ModelBasedEnvBase(*args, **kwargs)	用于基于模型的强化学习最新实现的基类环境。
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 实验室环境封装器。
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 problem 的 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 类型的转换函数的装饰器。