目录
目录

注意

转到末尾 以下载完整示例代码。

开始使用环境、TED 和转换

作者: Vincent Moens

注意

如需在笔记本中运行本教程，请在开头添加一个安装单元，内容为：

!pip install tensordict
!pip install torchrl

欢迎来到入门教程！

以下是我们将要涵盖的主题列表。

如果您时间紧迫，可直接跳转到最后一个教程—— 您自己的首个训练循环，再根据需要回溯其他所有“入门”教程，以解决疑问或深入了解特定主题！

RL中的环境

标准的强化学习（Reinforcement Learning，RL）训练循环包含一个模型（也称为策略），该模型被训练以在特定环境中完成某项任务。通常，该环境是一个模拟器，它以动作作为输入，并输出观测结果以及一些元数据。

在本文档中，我们将探索 TorchRL 的环境 API：学习如何创建环境、与环境交互，以及理解其使用的数据格式。

创建环境

本质上，TorchRL 不直接提供环境，而是为其他库提供封装模拟器的包装器。envs 模块可以被视为通用环境 API 的提供者，以及像 gym (GymEnv)、Brax (BraxEnv) 或 DeepMind Control Suite (DMControlEnv) 这样的仿真后端的中心枢纽。

创建环境通常与底层后端 API 所允许的一样简单。以下是使用 gym 的示例：

from torchrl.envs import GymEnv

env = GymEnv("Pendulum-v1")

运行环境

TorchRL中的环境有两个关键方法: reset()，它启动 一个回合，和 step()，它执行由演员选择的 动作。 在TorchRL中，环境方法读取和写入 TensorDict 实例。 基本上，TensorDict 是一个通用的关键数据 载体用于张量。 使用TensorDict而不是普通张量的好处是它使我们能够 交替处理简单和复杂的数据结构。由于我们的函数 签名非常通用，它消除了适应不同数据格式的挑战。更简单地说，在这个简短的教程之后， 你将能够操作简单和非常复杂的 环境，因为它们面向用户的API是相同的并且简单！

让我们将环境投入运行，看看 tensordict 实例是什么样子：

reset = env.reset()
print(reset)

TensorDict(
    fields={
        done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        observation: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

现在让我们在动作空间中采取一个随机动作。首先，采样动作：

reset_with_action = env.rand_action(reset)
print(reset_with_action)

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        observation: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

这个张量字典具有与从 EnvBase() 获得的相同结构，并且有一个额外的 "action" 条目。 您可以轻松访问操作，就像使用常规字典一样：

print(reset_with_action["action"])

tensor([0.5004])

我们现在需要将此动作传递给环境。 我们将把整个张量字典传递给step方法，因为在更复杂的情况下（如多智能体强化学习或无状态环境）可能需要读取多个张量：

stepped_data = env.step(reset_with_action)
print(stepped_data)

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False),
        observation: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

再次，这个新的张量字典与之前的相同，除了它有一个"next"条目（它本身也是一个张量字典！）包含由我们的动作产生的观察、奖励和完成状态。

我们将这种格式称为 TED（ TorchRL 情节数据格式）。它是该库中表示数据的通用方式，既可动态使用（如此处所示），也可静态使用（例如离线数据集）。

您需要在环境中运行一个 rollout 的最后一点信息是 如何将根目录中的 "next" 条目用于执行下一步。 TorchRL 提供了一个专门的 step_mdp() 函数 来完成这一任务：它过滤掉您不需要的信息，并 提供一个数据结构，对应于您在马尔可夫决策过程（MDP）中一步之后的观察。

from torchrl.envs import step_mdp

data = step_mdp(stepped_data)
print(data)

TensorDict(
    fields={
        done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        observation: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

环境滚动

写下这三个步骤（计算一个动作、迈出一步、在MDP中移动）可能会有点繁琐和重复。幸运的是，TorchRL提供了一个不错的rollout()函数，允许您随意在闭环中运行它们：

rollout = env.rollout(max_steps=10)
print(rollout)

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([10]),
            device=None,
            is_shared=False),
        observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([10]),
    device=None,
    is_shared=False)

这些数据看起来很像上面的 stepped_data，除了它的批量大小，现在等于我们通过 max_steps 参数提供的步骤数。tensordict 的魔力还不止于此：如果你对这个环境的单个转换感兴趣，你可以像索引张量一样索引 tensordict：

transition = rollout[3]
print(transition)

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False),
        observation: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

TensorDict 将自动检查您提供的索引是否为键（在这种情况下，我们沿键维度进行索引）或类似此处的空间索引。

以这种方式执行（没有策略），rollout 方法可能看起来相当无用：它只是运行随机动作。如果有可用的策略，可以将其传递给该方法并用于收集数据。

尽管如此，首先运行一次简单的、无策略的 rollout 仍可能很有用，以便快速了解环境的预期表现。

要欣赏TorchRL API的多功能性，请考虑这样一个事实：rollout方法具有普遍适用性。它适用于所有使用场景，无论您是在处理像这样的单一环境、多个进程中的多个副本、多智能体环境，甚至是无状态版本！

转换环境

大多数情况下，您需要修改环境的输出，以更好地满足您的需求。例如，您可能希望监控自上次重置以来执行的步数、调整图像大小，或堆叠连续的观测值。

在本节中，我们将介绍一种简单的变换，即 StepCounter 变换。 所有可用变换的完整列表可在此处查看： 此处。

该变换通过一个 TransformedEnv:

from torchrl.envs import StepCounter, TransformedEnv

transformed_env = TransformedEnv(env, StepCounter(max_steps=10))
rollout = transformed_env.rollout(max_steps=100)
print(rollout)

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                step_count: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.int64, is_shared=False),
                terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([10]),
            device=None,
            is_shared=False),
        observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        step_count: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.int64, is_shared=False),
        terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([10]),
    device=None,
    is_shared=False)

如您所见，我们的环境现在多了一个条目 "step_count"，用于跟踪自上次重置以来的步数。 由于我们在转换构造函数中传递了可选参数 max_steps=10，我们还在10步后截断了轨迹（没有完成我们通过 rollout 调用请求的100步完整滚动）。我们可以通过查看截断条目来确认轨迹已被截断：

print(rollout["next", "truncated"])

tensor([[False],
        [False],
        [False],
        [False],
        [False],
        [False],
        [False],
        [False],
        [False],
        [ True]])

这就是关于TorchRL环境API的简短介绍的全部内容！

下一步

要进一步了解 TorchRL 的环境可以做什么，请查看：

  • The step_and_maybe_reset() 方法将 step(), step_mdp() 和 reset() 包装在一起。

  • 某些环境如 GymEnv 支持通过 from_pixels 参数进行渲染。请查看类的文档字符串以了解更多！

  • 批处理环境，特别是 ParallelEnv 它允许你在多个进程中运行多个相同（或不同！）的环境。

  • 使用倒立摆教程设计您自己的环境，并了解规格说明和无状态环境。

  • 查看关于环境的更深入教程 详见专门教程;

  • 如果您对多智能体强化学习（MARL）感兴趣，请查看 多智能体环境 API ；

  • TorchRL 有许多工具可以与 Gym API 进行交互，例如 一种通过 register_gym() 在 Gym 注册表中注册 TorchRL 环境的方法，一个通过 set_info_dict_reader() 读取 info 字典的 API 或者一种 通过 set_gym_backend() 控制 gym 后端的方法。

脚本总运行时间： (0 分钟 49.147 秒)

估计内存使用量： 316 MB

通过 Sphinx-Gallery 生成的画廊

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源