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
:一个CompositeSpec
对象,包含所有观察键-规范对。env.state_spec
:一个CompositeSpec
对象,包含所有输入键-规范对(动作除外)。对于大多数有状态环境,此容器将为空。env.action_spec
:一个TensorSpec
对象,表示动作规范。env.reward_spec
:一个表示奖励规格的TensorSpec
对象。env.done_spec
:一个表示 done 标志规格的TensorSpec
对象。请参阅下面的轨迹终止部分。env.input_spec
:一个CompositeSpec
对象,包含所有输入键("full_action_spec"
和"full_state_spec"
)。它是锁定的,不应该直接修改。env.output_spec
:一个CompositeSpec
对象,包含所有输出键("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()
:一个重置方法,可以(但不一定需要)接受tensordict.TensorDict
输入。它返回展开的第一个 tensordict,通常包含一个"done"
状态和一组观测值。如果不存在,则会使用 0 和适当的形状实例化 “reward” 键。env.step()
:一个步进方法,它接受一个tensordict.TensorDict
输入,其中包含输入动作以及其他输入(例如,对于基于模型或无状态的环境)。env.step_and_maybe_reset()
:执行一步,并在需要时(部分)重置环境。它返回更新的输入,其中包含一个"next"
键,其中包含下一步的数据,以及一个包含下一步输入数据的 tensordict(即,重置或结果或step_mdp()
)。这是通过读取done_keys
并为每个 done 状态分配"_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 由 reset()
方法创建,然后由策略填充动作,然后传递给 step()
方法,该方法将观测值、done 标志和奖励写入 "next"
条目下。此调用的结果将被存储以进行传递,并且 "next"
条目由 step_mdp()
函数收集。
注意
通常,所有 TorchRL 环境在其输出张量字典中都包含 "done"
和 "terminated"
项。如果出于设计原因不存在这些项,则 EnvBase
元类将确保每个 done 或 terminated 都有其对应的另一项。在 TorchRL 中,"done"
严格指代所有轨迹结束信号的并集,应理解为“轨迹的最后一步”或等效地“指示需要重置的信号”。如果环境提供(例如,Gymnasium),则截断项也会在 EnvBase.step()
输出中以 "truncated"
项的形式写入。如果环境携带单个值,则默认情况下将其解释为 "terminated"
信号。默认情况下,TorchRL 的收集器和 rollout 方法将查找 "done"
项以评估是否应重置环境。
注意
torchrl.collectors.utils.split_trajectories 函数可用于切分相邻轨迹。它依赖于输入张量字典中的 "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
将根据这些规范创建缓冲区来与生成进程通信。查看 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"
键可能存在也可能不
存在于输入张量字典中。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"
是否对应于嵌套"done"
项的all()
、any()
或自定义调用,无法预先知道,并且明确假设根级别的"_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
(以及数据收集器)将根据环境规范创建数据缓冲区,以将数据从一个进程传递到另一个进程。这意味着规范(输入、观察或奖励)错误指定会导致运行时错误,因为数据无法写入预分配的缓冲区。通常,应使用 check_env_specs()
测试函数测试环境,然后再在 ParallelEnv
中使用。此函数将在预分配缓冲区与收集到的数据不匹配时引发断言错误。
我们还提供 SerialEnv
类,它拥有完全相同的 API,但以串行方式执行。这主要用于测试目的,当想要评估 ParallelEnv
的行为而无需启动子进程时。
除了提供基于进程并行的ParallelEnv
,我们还提供了一种使用MultiThreadedEnv
创建多线程环境的方法。此类在底层使用了EnvPool库,这使得性能更高,但同时也限制了灵活性——只能创建在EnvPool
中实现的环境。这涵盖了许多流行的RL环境类型(Atari、经典控制等),但不能像使用ParallelEnv
那样使用任意的TorchRL环境。运行benchmarks/benchmark_batched_envs.py以比较并行化批量环境的不同方式的性能。
|
在同一个进程中创建一系列环境。批量环境允许用户查询远程运行的环境的任意方法/属性。 |
|
为每个进程创建一个环境。 |
|
环境创建器类。 |
自定义原生TorchRL环境¶
TorchRL提供了一系列自定义的内置环境。
|
一个无状态的摆动环境。 |
|
井字棋实现。 |
多智能体环境¶
TorchRL开箱即用地支持多智能体学习。在单智能体学习管道中使用的相同类可以在多智能体上下文中无缝使用,无需任何修改或专门的多智能体基础设施。
从这个角度来看,环境在多智能体中扮演着核心角色。在多智能体环境中,许多决策智能体在一个共享的世界中行动。智能体可以观察不同的东西,以不同的方式行动,并且还可以获得不同的奖励。因此,存在许多范式来建模多智能体环境(DecPODPs、马尔可夫博弈)。这些范式之间的一些主要区别包括
观察可以是针对每个智能体的,也可以有一些共享的组件
奖励可以是针对每个智能体的,也可以是共享的
完成(以及
"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 = CompositeSpec(
... {
... "agents": CompositeSpec(
... {"action": torch.stack(action_specs)}, shape=(env.n_agents,)
... )
... }
...)
>>> env.reward_spec = CompositeSpec(
... {
... "agents": CompositeSpec(
... {"reward": torch.stack(reward_specs)}, shape=(env.n_agents,)
... )
... }
...)
>>> env.observation_spec = CompositeSpec(
... {
... "agents": CompositeSpec(
... {"observation": torch.stack(observation_specs)}, shape=(env.n_agents,)
... )
... }
...)
>>> env.done_spec = DiscreteTensorSpec(
... 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将始终返回叶规范,如果访问的规范是复合的。因此,如果在上面的示例中,我们在环境创建后运行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组映射类型。 |
|
检查MARL组映射。 |
自动重置环境¶
自动重置环境是指在环境在展开过程中达到"done"
状态时,不需要调用reset()
的环境,因为重置会自动发生。通常,在这种情况下,随done和reward一起提供的观察结果(实际上是执行环境中的动作产生的结果)实际上是新剧集的第一个观察结果,而不是当前剧集的最后一个观察结果。
为了处理这些情况,torchrl 提供了一个 AutoResetTransform
,它会将 step 调用产生的观测结果复制到下一个 reset,并在 rollout 期间跳过对 reset 的调用(在 rollout()
和 SyncDataCollector
迭代中)。此转换类还提供了对无效观测行为的细粒度控制,可以使用 “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 UnboundedContinuousTensorSpec, CompositeSpec, BoundedTensorSpec, BinaryDiscreteTensorSpec
>>> import torch
>>> from tensordict import TensorDict, TensorDictBase
>>>
>>> class EnvWithDynamicSpec(EnvBase):
... def __init__(self, max_count=5):
... super().__init__(batch_size=())
... self.observation_spec = CompositeSpec(
... observation=UnboundedContinuousTensorSpec(shape=(3, -1, 2)),
... )
... self.action_spec = BoundedTensorSpec(low=-1, high=1, shape=(2,))
... self.full_done_spec = CompositeSpec(
... done=BinaryDiscreteTensorSpec(1, shape=(1,), dtype=torch.bool),
... terminated=BinaryDiscreteTensorSpec(1, shape=(1,), dtype=torch.bool),
... truncated=BinaryDiscreteTensorSpec(1, shape=(1,), dtype=torch.bool),
... )
... self.reward_spec = UnboundedContinuousTensorSpec((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()
将通过形状沿某些维度变化的动态规格,但不通过在 step 期间存在而在其他期间不存在的键,或者维度数量变化的情况。
转换¶
在大多数情况下,必须在将环境的原始输出传递给另一个对象(例如策略或值运算符)之前对其进行处理。为此,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
克隆转换¶
由于附加到环境的转换通过 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
|
环境转换父类。 |
|
一个经过转换的环境。 |
|
一个用于离散化连续动作空间的转换。 |
|
一个自适应动作掩码器。 |
|
一个用于自动重置环境的子类。 |
|
一个用于自动重置环境的转换。 |
|
一个用于修改环境批大小的转换。 |
|
如果奖励为空或非空,则将奖励映射到二进制值(0 或 1)。 |
|
用于部分预加载数据序列的转换。 |
|
将连续的观察帧连接成一个张量。 |
|
将多个键连接到一个张量中。 |
|
裁剪图像的中心。 |
|
一个用于裁剪输入(状态、动作)或输出(观察、奖励)值的转换。 |
|
组合一系列转换。 |
|
在指定位置和输出大小处裁剪输入图像。 |
|
将选定键的数据类型从一种转换为另一种。 |
|
将数据从一个设备移动到另一个设备。 |
|
将离散动作从高维空间投影到低维空间。 |
|
将选定键的数据类型从一种转换为另一种。 |
|
使用具有lives方法的Gym环境注册生命终结信号。 |
|
从数据中排除键。 |
此转换将检查tensordict的所有项是否都是有限的,如果不是,则引发异常。 |
|
|
展平张量的相邻维度。 |
|
一个帧跳跃转换。 |
|
将像素观测转换为灰度。 |
|
重置跟踪器。 |
|
一个将 KL[pi_current||pi_0] 校正项添加到奖励的转换。 |
|
在环境重置时运行一系列随机动作。 |
|
观测仿射变换层。 |
|
观测变换的抽象类。 |
|
置换变换。 |
在 tensordict 上调用 pin_memory 以便于写入 CUDA 设备。 |
|
|
R3M 变换类。 |
|
用于 ReplayBuffer 和模块的轨迹子采样器。 |
|
从环境中删除空规范和内容。 |
|
一个用于重命名输出 tensordict 中条目的转换。 |
|
调整像素观测的大小。 |
|
根据剧集奖励和折扣因子计算奖励到结束。 |
|
将奖励裁剪到 clamp_min 和 clamp_max 之间。 |
|
奖励的仿射变换。 |
|
跟踪剧集累积奖励。 |
|
从输入 tensordict 中选择键。 |
|
一个用于计算 TensorDict 值符号的转换。 |
|
在指定位置删除大小为一的维度。 |
|
计算从重置开始的步数,并在一定步数后选择性地将截断状态设置为 |
|
为代理设置一个目标回报,以便在环境中实现。 |
|
TensorDict 在重置时初始化的引导程序。 |
|
在过去 T 个观测值中获取每个位置的最大值。 |
|
将类似 numpy 的图像 (W x H x C) 转换为 pytorch 图像 (C x W x H)。 |
|
在指定位置插入大小为一的维度。 |
|
VC1 变换类。 |
|
一个 VIP 变换,用于根据嵌入相似性计算奖励。 |
|
VIP 变换类。 |
|
一个用于 GymWrapper 子类的转换,以一致的方式处理自动重置。 |
|
用于 torchrl 环境的移动平均归一化层。 |
|
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
中的调用计数)。
要将收集到的张量字典保存到磁盘上,可以使用 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。用户有责任在需要时调用 dumpy 以避免 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()
记录器是变换,它们在数据传入时进行注册,用于日志记录目的。
|
张量字典记录器。 |
|
视频记录器变换。 |
|
一个在父环境上调用渲染并在张量字典中注册像素观察的变换。 |
帮助程序¶
|
用于数据收集器的随机策略。 |
|
根据短展开的结果测试环境规范。 |
已弃用 返回当前采样模式。 |
|
返回当前采样类型。 |
|
返回所有支持的库。 |
|
|
从张量字典创建 CompositeSpec 实例,假设所有值都是无界的。 |
|
|
|
|
|
创建一个新的张量字典,反映输入张量字典的时间步长。 |
|
读取张量字典中的 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 实验室环境包装器。 |
|
DeepMind Control 实验室环境包装器。 |
|
由环境 ID 直接构建的 OpenAI Gym 环境包装器。 |
|
OpenAI Gym 环境包装器。 |
|
Habitat 环境的包装器。 |
|
IsaacGym 环境的 TorchRL Env 接口。 |
|
IsaacGymEnvs 环境的包装器。 |
|
使用环境名称构建的 Jumanji 环境包装器。 |
|
Jumanji 环境包装器。 |
|
Meltingpot 环境包装器。 |
|
Meltingpot 环境包装器。 |
|
FARAMA MO-Gymnasium 环境包装器。 |
|
FARAMA MO-Gymnasium 环境包装器。 |
|
基于 EnvPool 的多线程环境执行。 |
|
基于 envpool 的多线程环境的包装器。 |
|
用于 bandit 上下文中 OpenML 数据的环境接口。 |
|
PettingZoo 环境。 |
|
PettingZoo 环境包装器。 |
|
RoboHive Gym 环境的包装器。 |
|
SMACv2(星际争霸多人代理挑战v2)环境包装器。 |
|
SMACv2(星际争霸多人代理挑战v2)环境包装器。 |
|
Vmas环境包装器。 |
|
Vmas环境包装器。 |
|
返回gym后端,或其子模块。 |
|
将gym后端设置为特定值。 |