TorchRL 环境

作者：Vincent Moens

环境在强化学习（RL）设置中起着至关重要的作用，通常与监督和无监督设置中的数据集有些相似。RL 社区已经非常熟悉 OpenAI gym API，它提供了一种灵活的方式来构建、初始化和交互环境。然而，还存在许多其他库，并且与它们交互的方式可能与 gym 所期望的非常不同。

让我们从描述 TorchRL 如何与 gym 交互开始，这将作为介绍其他框架的引子。

Gym 环境

要运行本教程的这部分，您需要安装一个最新版本的 gym 库以及 atari suite。您可以通过安装以下包来完成安装：

为了统一所有框架，TorchRL 环境在 `__init__` 方法中构建，并使用一个私有方法 `_build_env`，该方法会将参数和关键字参数传递给根库构建器。

使用 gym，这意味着构建一个环境就像：

import torch
from matplotlib import pyplot as plt
from tensordict import TensorDict

from torchrl.envs.libs.gym import GymEnv

env = GymEnv("Pendulum-v1")

可以通过此命令访问可用环境的列表：

list(GymEnv.available_envs)[:10]

环境规范

与其他框架一样，TorchRL 环境具有指示观测、动作、完成和奖励的空间的属性。因为经常会检索到不止一个观测，所以我们期望观测规范的类型为 `CompositeSpec`。奖励和动作没有这个限制。

print("Env observation_spec: \n", env.observation_spec)
print("Env action_spec: \n", env.action_spec)
print("Env reward_spec: \n", env.reward_spec)

这些规范附带一系列有用的工具：可以断言一个样本是否在定义的空间内。我们还可以使用一些启发式方法将样本投影到空间内（如果它超出空间），并在该空间内生成随机数（可能是均匀分布的）。

action = torch.ones(1) * 3
print("action is in bounds?\n", bool(env.action_spec.is_in(action)))
print("projected action: \n", env.action_spec.project(action))

print("random action: \n", env.action_spec.rand())

在这些规范中，`done_spec` 值得特别关注。在 TorchRL 中，所有环境都会写出至少两种类型的轨迹结束信号：`"terminated"`（表示马尔可夫决策过程已达到最终状态——__episode__ 已结束）和`"done"`，表示这是 __trajectory__ 的最后一步（但不一定是任务的结束）。通常，当 `terminal` 为 `False` 时为 `True` 的 `"done"` 条目是由 `"truncated"` 信号引起的。Gym 环境考虑了这三个信号。

print(env.done_spec)

环境还带有一个类型为 `CompositeSpec` 的 `env.state_spec` 属性，其中包含所有作为环境输入的但非动作的规范。对于有状态环境（例如 gym），这在大多数情况下是空的。对于无状态环境（例如 Brax），这也应该包括前一个状态的表示，或者环境的任何其他输入（包括重置时的输入）。

播种、重置和步进

环境的基本操作是（1）`set_seed`、（2）`reset` 和（3）`step`。

让我们看看这些方法在 TorchRL 中是如何工作的。

torch.manual_seed(0)  # make sure that all torch code is also reproductible
env.set_seed(0)
reset_data = env.reset()
print("reset data", reset_data)

我们现在可以在环境中执行一步。由于我们没有策略，我们可以生成一个随机动作。

policy = TensorDictModule(env.action_spec.rand, in_keys=[], out_keys=["action"])


policy(reset_data)
tensordict_out = env.step(reset_data)

默认情况下，`step` 返回的 tensordict 与输入相同……

assert tensordict_out is reset_data

……但带有新键。

tensordict_out

我们刚才所做的（使用 `action_spec.rand()` 进行随机步进）也可以通过简单的快捷方式完成。

env.rand_step()

新键 `("next", "observation")`（如 `"next"` tensordict 下的所有键）在 TorchRL 中有一个特殊的作用：它们表示它们出现在没有前缀的同名键之后。

我们提供了一个函数 `step_mdp`，它在 tensordict 中执行一步：它返回一个更新后的新 tensordict，使得 *t < -t’*。

from torchrl.envs.utils import step_mdp

tensordict_out.set("some other key", torch.randn(1))
tensordict_tprime = step_mdp(tensordict_out)

print(tensordict_tprime)
print(
    (
        tensordict_tprime.get("observation")
        == tensordict_out.get(("next", "observation"))
    ).all()
)

我们可以注意到 `step_mdp` 移除了所有与时间相关的键值对，但保留了 `"some other key"`。此外，新观测与之前的观测匹配。

最后，请注意 `env.reset` 方法也接受一个 tensordict 进行更新。

data = TensorDict()
assert env.reset(data) is data
data

Rollouts

TorchRL 提供的通用环境类允许您轻松运行给定步数的 rollout。

tensordict_rollout = env.rollout(max_steps=20, policy=policy)
print(tensordict_rollout)

生成的 tensordict 的 `batch_size` 为 `[20]`，这是轨迹的长度。我们可以检查观测值是否与其下一个值匹配。

(
    tensordict_rollout.get("observation")[1:]
    == tensordict_rollout.get(("next", "observation"))[:-1]
).all()

`frame_skip`

在某些情况下，使用 `frame_skip` 参数将相同的动作用于几个连续的帧很有用。

生成的 tensordict 将只包含序列中最后观察到的帧，但奖励将是多帧的总和。

如果环境在此过程中达到完成状态，它将停止并返回截断链的结果。

env = GymEnv("Pendulum-v1", frame_skip=4)
env.reset()

渲染

渲染在许多 RL 设置中起着重要作用，这就是为什么 TorchRL 的通用环境类提供了一个 `from_pixels` 关键字参数，允许用户快速请求基于图像的环境。

env = GymEnv("Pendulum-v1", from_pixels=True)

data = env.reset()
env.close()

plt.imshow(data.get("pixels").numpy())

让我们看看 tensordict 中包含什么。

data

我们仍然有一个 `"state"`，它描述了在上一例中 `"observation"` 所描述的内容（命名差异来自于 gym 现在返回一个字典，并且 TorchRL 从字典中获取名称（如果存在），否则会将步骤输出命名为 `"observation"`：简而言之，这是由于 gym 环境的 step 方法返回的对象类型不一致）。

也可以通过仅请求像素来丢弃此附加输出。

env = GymEnv("Pendulum-v1", from_pixels=True, pixels_only=True)
env.reset()
env.close()

有些环境只有基于图像的格式。

env = GymEnv("ALE/Pong-v5")
print("from pixels: ", env.from_pixels)
print("data: ", env.reset())
env.close()

DeepMind 控制环境

要运行本教程的这部分，请确保您已安装 dm_control。

$ pip install dm_control

我们还为 DM Control suite 提供了一个包装器。同样，构建环境也很容易：首先让我们看看可以访问哪些环境。`available_envs` 现在返回一个包含环境和可能任务的字典。

from matplotlib import pyplot as plt

from torchrl.envs.libs.dm_control import DMControlEnv

DMControlEnv.available_envs

env = DMControlEnv("acrobot", "swingup")
data = env.reset()
print("result of reset: ", data)
env.close()

当然，我们也可以使用基于像素的环境。

env = DMControlEnv("acrobot", "swingup", from_pixels=True, pixels_only=True)
data = env.reset()
print("result of reset: ", data)
plt.imshow(data.get("pixels").numpy())
env.close()

转换环境

在策略读取环境输出或将其存储在缓冲区之前对其进行预处理是很常见的。

在许多情况下，RL 社区已经采用了以下类型的包装方案：

$ env_transformed = wrapper1(wrapper2(env))

来转换环境。这有很多优点：它使访问环境规范变得显而易见（外部包装器是外部世界的真实来源），并且它使得与向量化环境交互变得容易。然而，它也使得访问内部环境变得困难：假设我们想从链中删除一个包装器（例如 `wrapper2`），这个操作需要我们收集

$ env0 = env.env.env

$ env_transformed_bis = wrapper1(env0)

TorchRL 采取使用转换序列的立场，就像在其他 PyTorch 领域库（例如 `torchvision`）中那样。这种方法也类似于 `torch.distribution` 中转换分布的方式，其中 `TransformedDistribution` 对象围绕 `base_dist` 分布和（一系列）`transforms` 构建。

from torchrl.envs.transforms import ToTensorImage, TransformedEnv

# ToTensorImage transforms a numpy-like image into a tensor one,
env = DMControlEnv("acrobot", "swingup", from_pixels=True, pixels_only=True)
print("reset before transform: ", env.reset())

env = TransformedEnv(env, ToTensorImage())
print("reset after transform: ", env.reset())
env.close()

要组合转换，只需使用 `Compose` 类。

from torchrl.envs.transforms import Compose, Resize

env = DMControlEnv("acrobot", "swingup", from_pixels=True, pixels_only=True)
env = TransformedEnv(env, Compose(ToTensorImage(), Resize(32, 32)))
env.reset()

也可以一次添加一个转换。

from torchrl.envs.transforms import GrayScale

env.append_transform(GrayScale())
env.reset()

如预期，元数据也会更新。

print("original obs spec: ", env.base_env.observation_spec)
print("current obs spec: ", env.observation_spec)

如果需要，我们也可以连接张量。

from torchrl.envs.transforms import CatTensors

env = DMControlEnv("acrobot", "swingup")
print("keys before concat: ", env.reset())

env = TransformedEnv(
    env,
    CatTensors(in_keys=["orientations", "velocity"], out_key="observation"),
)
print("keys after concat: ", env.reset())

此功能可以轻松修改应用于环境输入和输出的转换集。事实上，转换在执行步骤之前和之后都会运行：对于预步骤传递，`in_keys_inv` 键列表将被传递给 `_inv_apply_transform` 方法。这种转换的一个例子是将浮点动作（来自神经网络的输出）转换为双精度类型（被包装环境需要）。在执行步骤之后，`_apply_transform` 方法将在 `in_keys` 键列表指定的键上执行。

环境转换的另一个有趣之处在于，它们允许用户检索已包装情况下的 `env.env` 的等效项，换句话说，父环境。父环境可以通过调用 `transform.parent` 来检索：返回的环境将由一个 `TransformedEnvironment` 组成，其中包含当前转换之前的所有转换。例如，这可以在 `NoopResetEnv` 的情况下使用，它在重置时执行以下步骤：在执行该环境中的一定数量的随机步骤之前重置父环境。

env = DMControlEnv("acrobot", "swingup")
env = TransformedEnv(env)
env.append_transform(
    CatTensors(in_keys=["orientations", "velocity"], out_key="observation")
)
env.append_transform(GrayScale())

print("env: \n", env)
print("GrayScale transform parent env: \n", env.transform[1].parent)
print("CatTensors transform parent env: \n", env.transform[0].parent)

环境设备

转换可以在设备上工作，当操作计算量适中或高时，这可以带来显著的加速。这些操作包括 `ToTensorImage`、`Resize`、`GrayScale` 等。

有人可能会合理地问这对被包装的环境意味着什么。对于常规环境来说，影响很小：操作仍将在应有的设备上进行。TorchRL 中的环境设备属性指示传入数据应该在哪个设备上，以及传出数据将在哪个设备上。从该设备到该设备的转换由 TorchRL 环境类负责。将数据存储在 GPU 上的主要优点是（1）如上所述的转换加速和（2）在多进程设置中共享数据。

from torchrl.envs.transforms import CatTensors, GrayScale, TransformedEnv

env = DMControlEnv("acrobot", "swingup")
env = TransformedEnv(env)
env.append_transform(
    CatTensors(in_keys=["orientations", "velocity"], out_key="observation")
)

if torch.has_cuda and torch.cuda.device_count():
    env.to("cuda:0")
    env.reset()

并行运行环境

TorchRL 提供了并行运行环境的实用工具。预计各种环境读取并返回形状和数据类型相似的张量（但也可以设计掩码函数，以便在这些张量形状不同时使其可行）。创建这样的环境非常容易。让我们看看最简单的情况。

from torchrl.envs import ParallelEnv


def env_make():
    return GymEnv("Pendulum-v1")


parallel_env = ParallelEnv(3, env_make)  # -> creates 3 envs in parallel
parallel_env = ParallelEnv(
    3, [env_make, env_make, env_make]
)  # similar to the previous command

`SerialEnv` 类与 `ParallelEnv` 类似，只是环境是按顺序运行的。这主要用于调试目的。

`ParallelEnv` 实例以惰性模式创建：环境将在被调用时才开始运行。这使得我们可以将 `ParallelEnv` 对象从一个进程移动到另一个进程，而无需过多担心正在运行的进程。可以通过调用 `start`、`reset` 或简单地调用 `step`（如果不需要先调用 `reset`）来启动 `ParallelEnv`。

parallel_env.reset()

可以检查并行环境是否具有正确的批处理大小。通常，`batch_size` 的第一部分表示批次，第二部分表示时间帧。让我们用 `rollout` 方法来检查这一点。

parallel_env.rollout(max_steps=20)

关闭并行环境

重要提示：在关闭程序之前，关闭并行环境很重要。通常，即使对于常规环境，用 `close` 调用结束函数也是一个好习惯。在某些情况下，如果未这样做，TorchRL 会抛出错误（通常会在程序结束时，当环境超出作用域时！）。

parallel_env.close()

播种

在播种并行环境时，我们面临的困难是我们不希望为所有环境提供相同的种子。TorchRL 使用的启发式方法是，我们给定输入种子，以一种可以说是马尔可夫的方式生成一个确定的种子链，使其可以从任何元素重建。所有 `set_seed` 方法都将返回下一个要使用的种子，这样我们就可以轻松地根据最后一个种子保持链的连续性。当几个收集器都包含 `ParallelEnv` 实例，并且我们希望每个子子环境都有不同的种子时，这很有用。

out_seed = parallel_env.set_seed(10)
print(out_seed)

del parallel_env

访问环境属性

有时会出现包装环境具有感兴趣的属性的情况。首先，请注意，TorchRL 环境包装器限制了访问此属性的工具。这是一个例子。

from time import sleep

from uuid import uuid1


def env_make():
    env = GymEnv("Pendulum-v1")
    env._env.foo = f"bar_{uuid1()}"
    env._env.get_something = lambda r: r + 1
    return env


env = env_make()

# Goes through env._env
env.foo

parallel_env = ParallelEnv(3, env_make)  # -> creates 3 envs in parallel

# env has not been started --> error:
try:
    parallel_env.foo
except RuntimeError:
    print("Aargh what did I do!")
    sleep(2)  # make sure we don't get ahead of ourselves

if parallel_env.is_closed:
    parallel_env.start()
foo_list = parallel_env.foo
foo_list  # needs to be instantiated, for instance using list

list(foo_list)

类似地，方法也可以被访问。

something = parallel_env.get_something(0)
print(something)

parallel_env.close()
del parallel_env

并行环境的 kwargs

您可能希望为各种环境提供 kwargs。这可以在构造时或之后完成。

from torchrl.envs import ParallelEnv


def env_make(env_name):
    env = TransformedEnv(
        GymEnv(env_name, from_pixels=True, pixels_only=True),
        Compose(ToTensorImage(), Resize(64, 64)),
    )
    return env


parallel_env = ParallelEnv(
    2,
    [env_make, env_make],
    create_env_kwargs=[{"env_name": "ALE/AirRaid-v5"}, {"env_name": "ALE/Pong-v5"}],
)
data = parallel_env.reset()

plt.figure()
plt.subplot(121)
plt.imshow(data[0].get("pixels").permute(1, 2, 0).numpy())
plt.subplot(122)
plt.imshow(data[1].get("pixels").permute(1, 2, 0).numpy())
parallel_env.close()
del parallel_env

from matplotlib import pyplot as plt

转换并行环境

有两种等效的方式可以转换并行环境：在每个进程中单独转换，或在主进程中转换。甚至可以同时执行两者。因此，可以仔细考虑转换设计，以利用设备功能（例如，CUDA 设备上的转换）并在可能的情况下在主进程上进行向量化操作。

from torchrl.envs import (
    Compose,
    GrayScale,
    ParallelEnv,
    Resize,
    ToTensorImage,
    TransformedEnv,
)


def env_make(env_name):
    env = TransformedEnv(
        GymEnv(env_name, from_pixels=True, pixels_only=True),
        Compose(ToTensorImage(), Resize(64, 64)),
    )  # transforms on remote processes
    return env


parallel_env = ParallelEnv(
    2,
    [env_make, env_make],
    create_env_kwargs=[{"env_name": "ALE/AirRaid-v5"}, {"env_name": "ALE/Pong-v5"}],
)
parallel_env = TransformedEnv(parallel_env, GrayScale())  # transforms on main process
data = parallel_env.reset()

print("grayscale data: ", data)
plt.figure()
plt.subplot(121)
plt.imshow(data[0].get("pixels").permute(1, 2, 0).numpy())
plt.subplot(122)
plt.imshow(data[1].get("pixels").permute(1, 2, 0).numpy())
parallel_env.close()
del parallel_env

VecNorm

在 RL 中，我们经常面临在将数据输入模型之前对其进行标准化的问​​题。有时，我们可以通过例如随机策略（或演示）从环境中收集的数据来获得标准统计数据的良好近似值。然而，有时建议“即时”标准化数据，逐步更新标准化常数以反映到目前为止已观察到的情况。当期望标准化统计数据随着任务性能的变化而变化，或者环境由于外部因素而演变时，这一点尤其有用。

警告：此功能应谨慎用于离策略学习，因为旧数据将因其使用以前有效的标准化统计数据进行标准化而被“弃用”。在同策略设置中，此功能也会使学习不稳定，并可能产生意外效果。因此，建议用户谨慎使用此功能，并将其与使用固定版本标准化常数的数据标准化进行比较。

在常规设置中，使用 VecNorm 非常简单。

from torchrl.envs.libs.gym import GymEnv
from torchrl.envs.transforms import TransformedEnv, VecNorm

env = TransformedEnv(GymEnv("Pendulum-v1"), VecNorm())
data = env.rollout(max_steps=100)

print("mean: :", data.get("observation").mean(0))  # Approx 0
print("std: :", data.get("observation").std(0))  # Approx 1

在 **并行环境** 中，事情会稍微复杂一些，因为我们需要在进程之间共享运行统计数据。我们创建了一个 `EnvCreator` 类，该类负责查看环境创建方法，检索要在环境类中共享的 tensordicts，并指向每个进程到正确（公共、共享）数据一旦创建。

from torchrl.envs import EnvCreator, ParallelEnv
from torchrl.envs.libs.gym import GymEnv
from torchrl.envs.transforms import TransformedEnv, VecNorm

make_env = EnvCreator(lambda: TransformedEnv(GymEnv("CartPole-v1"), VecNorm(decay=1.0)))
env = ParallelEnv(3, make_env)
print("env state dict:")
sd = TensorDict(make_env.state_dict())
print(sd)
# Zeroes all tensors
sd *= 0

data = env.rollout(max_steps=5)

print("data: ", data)
print("mean: :", data.get("observation").view(-1, 3).mean(0))  # Approx 0
print("std: :", data.get("observation").view(-1, 3).std(0))  # Approx 1

计数略高于步数（因为我们没有使用任何衰减）。两者之间的差异是由于 `ParallelEnv` 创建了一个虚拟环境来初始化共享的 `TensorDict`，该 `TensorDict` 用于从分派的环境中收集数据。这个微小的差异通常会在训练过程中被吸收。

print(
    "update counts: ",
    make_env.state_dict()["_extra_state"]["observation_count"],
)

env.close()
del env