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 预期所在的设备。环境设备并不意味着实际的 step 操作将在设备上计算(这是后端的责任,TorchRL 对此无能为力)。环境的设备仅表示当数据输入到环境或从环境中检索时,数据预期所在的设备。TorchRL 负责将数据映射到所需的设备。这对于转换(见下文)尤其有用。对于参数化环境(例如,基于模型的环境),设备确实代表将用于计算操作的硬件。env.observation_spec
:一个Composite
对象,包含所有观察键-规格对。env.state_spec
:一个Composite
对象,包含所有输入键-规格对(除了 action)。对于大多数有状态环境,此容器将为空。env.action_spec
:一个TensorSpec
对象,表示 action 规格。env.reward_spec
:一个TensorSpec
对象,表示 reward 规格。env.done_spec
:一个TensorSpec
对象,表示 done 标志规格。请参阅下面关于轨迹终止的部分。env.input_spec
:一个Composite
对象,包含所有输入键("full_action_spec"
和"full_state_spec"
)。它是锁定的,不应直接修改。env.output_spec
:一个Composite
对象,包含所有输出键("full_observation_spec"
、"full_reward_spec"
和"full_done_spec"
)。它是锁定的,不应直接修改。
如果环境携带非张量数据,可以使用 NonTensorSpec
实例。
重要的是,环境规格形状应包含批大小,例如,env.batch_size == torch.Size([4])
的环境应具有形状为 torch.Size([4, action_size])
的 env.action_spec
。这在预分配张量、检查形状一致性等方面很有帮助。
有了这些,以下方法得以实现
env.reset()
:一个 reset 方法,可以(但不一定需要)接受tensordict.TensorDict
输入。它返回 rollout 的第一个 tensordict,通常包含一个"done"
状态和一组观察。如果不存在,将使用 0 和适当的形状实例化 “reward” 键。env.step()
:一个 step 方法,接受一个tensordict.TensorDict
输入,其中包含输入 action 以及其他输入(例如,对于基于模型或无状态环境)。env.step_and_maybe_reset()
:执行一个 step,并在需要时(部分)重置环境。它返回更新后的输入,其中包含一个"next"
键,包含下一步的数据,以及一个包含下一步输入数据的 tensordict(即,reset 或 result 或step_mdp()
)。这是通过读取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=model
)。策略应使用tensordict.nn.TensorDictModule
(或任何其他tensordict.TensorDict
兼容模块)进行编码。生成的tensordict.TensorDict
实例将被标记一个尾随"time"
命名维度,其他模块可以使用该维度来按应有的方式处理此批处理维度。
下图总结了 rollout 在 torchrl 中的执行方式。

使用 TensorDict 的 TorchRL rollouts。¶
简而言之,TensorDict 由 reset()
方法创建,然后由策略填充 action,再传递给 step()
方法,该方法将观察、done 标志和 reward 写入 "next"
条目下。此调用的结果被存储以供交付,"next"
条目由 step_mdp()
函数收集。
注意
总的来说,所有的 TorchRL 环境在其输出 tensordict 中都有一个 "done"
和 "terminated"
条目。如果它们在设计上不存在,EnvBase
元类将确保每个 done 或 terminated 都与其对偶项配对。在 TorchRL 中,"done"
严格指代所有轨迹结束信号的并集,应被解释为“轨迹的最后一步”或等效地“指示需要重置的信号”。如果环境提供截断信息(例如,Gymnasium),截断条目也会写入 EnvBase.step()
输出的 "truncated"
条目下。如果环境携带单个值,默认情况下它将被解释为 "terminated"
信号。默认情况下,TorchRL 的收集器和 rollout 方法将查找 "done"
条目,以评估环境是否应重置。
注意
可以使用 torchrl.collectors.utils.split_trajectories 函数来分割相邻的轨迹。它依赖于输入 tensordict 中的 "traj_ids"
条目,或者在缺少 "traj_ids"
时依赖于 "done"
和 "truncated"
键的交集。
注意
在某些情况下,标记轨迹的第一步可能很有用。TorchRL 通过 InitTracker
转换提供了这种功能。
我们的环境 教程 提供了关于如何从头开始设计自定义环境的更多信息。
|
抽象环境父类。 |
|
类 gym 环境是一个环境。 |
|
用于环境元数据存储和在多进程设置中传递的类。 |
向量化环境¶
向量化(或更好:并行)环境是强化学习中的一个常见特征,其中执行环境步骤可能是 CPU 密集型的。诸如 gym3 或 EnvPool 等一些库提供了同时执行批量环境的接口。虽然它们通常提供非常有竞争力的计算优势,但它们不一定能扩展到 TorchRL 支持的各种环境库。因此,TorchRL 提供了其自己的通用 ParallelEnv
类以并行运行多个环境。由于此类继承自 SerialEnv
,因此它享有与其他环境完全相同的 API。当然,ParallelEnv
将具有与其环境计数相对应的批大小
注意
鉴于库的许多可选依赖项(例如,Gym、Gymnasium 和许多其他库),警告在多进程/分布式设置中可能很快变得非常烦人。默认情况下,TorchRL 会过滤掉子进程中的这些警告。如果仍然希望看到这些警告,可以通过设置 torchrl.filter_warnings_subprocess=False
来显示它们。
重要的是,您的环境规范与它发送和接收的输入和输出相匹配,因为 ParallelEnv
将从这些规范创建缓冲区,以便与 spawn 进程通信。查看 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()
的内部逻辑。然而,在设计该方法时,应牢记以下几点。
在调用
_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 second 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、经典控制等),但不能使用任意 TorchRL 环境,这在使用 ParallelEnv
时是可能的。运行 benchmarks/benchmark_batched_envs.py 以比较并行化批量环境的不同方法的性能。
|
在同一进程中创建一系列环境。批量环境允许用户查询远程运行的环境的任意方法/属性。 |
|
每个进程创建一个环境。 |
|
环境创建器类。 |
自定义原生 TorchRL 环境¶
TorchRL 提供了一系列自定义的内置环境。
|
无状态的 Pendulum 环境。 |
|
井字棋实现。 |
多智能体环境¶
TorchRL 开箱即用地支持多智能体学习。在单智能体学习管道中使用的相同类可以无缝地用于多智能体环境中,无需任何修改或专用的多智能体基础设施。
在这种观点下,环境在多智能体系统中扮演着核心角色。在多智能体环境中,许多决策智能体在一个共享的世界中行动。智能体可以观察到不同的事物,以不同的方式行动,并获得不同的奖励。因此,存在许多范例来建模多智能体环境 (DecPODPs, Markov Games)。这些范例之间的一些主要区别包括
observation 可以是每个智能体独立的,也可以包含一些共享的组件
reward 可以是每个智能体独立的或共享的
done (以及
"truncated"
或"terminated"
) 可以是每个智能体独立的或共享的。
借助 TorchRL 的 tensordict.TensorDict
数据载体,它可以适应所有这些可能的范例。特别是在多智能体环境中,每个智能体独立的键值将放在嵌套的 “agents” TensorDict 中。这个 TensorDict 将具有额外的智能体维度,从而对每个智能体不同的数据进行分组。另一方面,共享的键值将保留在第一层,就像在单智能体情况下一样。
让我们看一个例子来更好地理解这一点。在这个例子中,我们将使用 VMAS,这是一个也基于 PyTorch 的多机器人任务模拟器,它在设备上运行并行批量模拟。
我们可以创建一个 VMAS 环境,并查看随机步骤的输出结果
>>> 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” 条目中。
这是一个关于如何在多智能体环境中创建规格的示例,在该环境中,只有 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 抽象了这些嵌套的规格。这意味着如果访问的规格是 Composite,则访问 env.reward_spec 将始终返回叶规格。因此,如果在上面的示例中,我们在环境创建后运行 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]。
|
Marl Group Map 类型。 |
|
检查 MARL 组映射。 |
自动重置环境¶
自动重置环境是指在 rollout 期间环境达到 "done"
状态时,不需要调用 reset()
的环境,因为重置会自动发生。通常,在这种情况下,与 done 和 reward 一起传递的观察结果(实际上是由于在环境中执行操作而产生的)实际上是新 episode 的第一个观察结果,而不是当前 episode 的最后一个观察结果。
为了处理这些情况,torchrl 提供了一个 AutoResetTransform
,它会将调用 step 产生的观察结果复制到下一个 reset,并在 rollout 期间跳过对 reset 的调用(在 rollout()
和 SyncDataCollector
迭代中)。这个 transform 类还提供了对无效观察结果要采取的行为进行细粒度控制,可以使用 “nan” 或任何其他值进行掩码,或者完全不掩码。
要告诉 torchrl 环境是自动重置的,只需在构造期间提供一个 auto_reset
参数即可。如果提供,auto_reset_replace
参数还可以控制是否应将 episode 的最后一个观察结果的值替换为某些占位符。
>>> 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
。由于数据无法连续堆叠,因此需要使用 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]):
... 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"
关键字参数传递要包含在此逆向转换中的键。
>>> env.append_transform(DoubleToFloat(in_keys_inv=["action"])) # will map the action from float32 to float64 before calling the base_env.step
通过将基础环境视为转换的“内部”部分,可以理解 in_keys
与 in_keys_inv
之间的关系。相反,用户输入和输出到转换以及从转换输出的数据应被视为外部世界。下图显示了这对于 RenameTransform
类在实践中的含义:step
函数的输入 TensorDict
必须在其条目中列出 out_keys_inv
,因为它们是外部世界的一部分。转换会更改这些名称,以使它们与使用 in_keys_inv
的内部基础环境的名称匹配。逆向过程使用输出 tensordict 执行,其中 in_keys
映射到相应的 out_keys
。

重命名转换逻辑¶
克隆转换¶
因为附加到环境的转换通过 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
|
环境转换父类。 |
|
一个 transformed_in 环境。 |
|
一种离散化连续动作空间的转换。 |
|
一个自适应动作掩码器。 |
|
自动重置环境的子类。 |
|
用于自动重置环境的转换。 |
|
一个用于修改环境批大小的转换。 |
|
如果奖励为空则将奖励映射为二元值 0,否则映射为 1。 |
|
用于部分预热数据序列的转换。 |
|
将连续的观测帧连接成一个张量。 |
|
将多个键连接成一个张量。 |
|
裁剪图像的中心区域。 |
|
一个用于裁剪输入(状态、动作)或输出(观测、奖励)值的转换。 |
|
组合一系列转换。 |
|
在指定位置和输出大小裁剪输入图像。 |
|
为选定的键将一个数据类型转换为另一个数据类型。 |
|
将数据从一个设备移动到另一个设备。 |
|
将离散动作从高维空间投影到低维空间。 |
|
为选定的键将一个数据类型转换为另一个数据类型。 |
|
从具有 lives 方法的 Gym 环境中注册生命周期结束信号。 |
|
从数据中排除键。 |
此转换将检查 tensordict 的所有项是否为有限值,如果不是,则引发异常。 |
|
|
展平张量的相邻维度。 |
|
帧跳跃转换。 |
|
将像素观测转换为灰度。 |
|
重置跟踪器。 |
|
一个用于向奖励添加 KL[pi_current||pi_0] 校正项的转换。 |
|
当环境重置时,运行一系列随机动作。 |
|
观测仿射变换层。 |
|
观测转换的抽象类。 |
|
排列转换。 |
在 tensordict 上调用 pin_memory 以方便写入 CUDA 设备。 |
|
|
R3M 转换类。 |
|
用于 ReplayBuffer 和模块的轨迹子采样器。 |
|
从环境中移除空 specs 和内容。 |
|
一个用于重命名输出 tensordict(或通过反向键重命名输入 tensordict)中条目的转换。 |
|
调整像素观测的大小。 |
|
基于 episode 奖励和折扣因子计算 reward-to-go。 |
|
将奖励裁剪在 clamp_min 和 clamp_max 之间。 |
|
奖励的仿射变换。 |
|
跟踪 episode 累积奖励。 |
|
从输入 tensordict 中选择键。 |
|
一个用于计算 TensorDict 值符号的转换。 |
|
在指定位置移除大小为 1 的维度。 |
|
从重置开始计数步数,并可选择在一定步数后将截断状态设置为 |
|
为 agent 设置在环境中要达成的目标回报。 |
|
用于在重置时初始化 TensorDict 的 primer。 |
|
在最后 T 个观测中,取每个位置的最大值。 |
|
将类似 numpy 的图像 (W x H x C) 转换为 pytorch 图像 (C x W x H)。 |
|
在指定位置插入大小为 1 的维度。 |
|
VC1 转换类。 |
|
一个 VIP 转换,用于基于嵌入相似性计算奖励。 |
|
VIP 转换类。 |
|
用于 GymWrapper 子类的转换,以一致的方式处理自动重置。 |
|
torchrl 环境的移动平均归一化层。 |
|
gSDE 噪声初始化器。 |
具有掩码动作的环境¶
在某些具有离散动作的环境中,agent 可用的动作可能会在执行过程中发生变化。在这种情况下,环境将输出一个动作掩码(默认情况下在 "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 迭代时(如果需要跳过某些迭代,则应添加一个内部变量来跟踪 callback
内的调用计数)对收集到的张量字典调用。
要将收集到的张量字典保存到磁盘,可以使用 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。用户有责任在需要时调用 dumpy 以避免 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()
记录器是转换,它们在数据传入时注册数据,用于日志记录目的。
|
张量字典记录器。 |
|
视频记录器转换。 |
|
一个在父环境中调用 render 并在张量字典中注册像素观测的转换。 |
助手函数¶
|
用于数据收集器的随机策略。 |
|
根据短期 rollout 的结果测试环境规格。 |
返回当前的采样类型。 |
|
返回所有受支持的库。 |
|
|
从张量字典创建 Composite 实例,假设所有值都是无界的。 |
|
|
|
创建一个新的张量字典,反映输入张量字典的时间步进。 |
|
读取张量字典中的 done / terminated / truncated 键,并写入一个新的张量,其中聚合了两个信号的值。 |
领域特定¶
|
用于基于模型的强化学习最先进实现的基本环境。 |
|
Dreamer 模拟环境。 |
用于记录 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
|
使用环境名称构建的 Google Brax 环境包装器。 |
|
Google Brax 环境包装器。 |
|
DeepMind Control lab 环境包装器。 |
|
DeepMind Control lab 环境包装器。 |
|
直接通过环境 ID 构建的 OpenAI Gym 环境包装器。 |
|
OpenAI Gym 环境包装器。 |
|
Habitat 环境的包装器。 |
|
用于 IsaacGym 环境的 TorchRL Env 接口。 |
|
用于 IsaacGymEnvs 环境的包装器。 |
|
使用环境名称构建的 Jumanji 环境包装器。 |
|
Jumanji 环境包装器。 |
|
Meltingpot 环境包装器。 |
|
Meltingpot 环境包装器。 |
|
FARAMA MO-Gymnasium 环境包装器。 |
|
FARAMA MO-Gymnasium 环境包装器。 |
|
基于 EnvPool 的环境多线程执行。 |
|
用于基于 envpool 的多线程环境的包装器。 |
|
用于在 bandits 上下文中使用的 OpenML 数据的环境接口。 |
|
Google DeepMind OpenSpiel 环境包装器。 |
|
使用游戏字符串构建的 Google DeepMind OpenSpiel 环境包装器。 |
|
PettingZoo 环境。 |
|
PettingZoo 环境包装器。 |
|
RoboHive gym 环境的包装器。 |
|
SMACv2 (星际争霸多智能体挑战 v2) 环境包装器。 |
|
SMACv2 (星际争霸多智能体挑战 v2) 环境包装器。 |
|
Unity ML-Agents 环境包装器。 |
|
Unity ML-Agents 环境包装器。 |
|
Vmas 环境包装器。 |
|
Vmas 环境包装器。 |
|
返回 gym 后端或其子模块。 |
|
将 gym 后端设置为特定值。 |