注意

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

开始使用环境、TED和转换¶

注意

要在笔记本中运行本教程，请在开头添加一个包含以下内容的安装单元格

!pip install tensordict
!pip install torchrl

欢迎来到入门教程！

以下是我们将要介绍的主题列表。

如果您时间紧迫，可以直接跳到最后一个教程，您自己的第一个训练循环，如果您不清楚某些内容或想了解特定主题的更多信息，可以从那里回溯其他每个“入门”教程！

强化学习中的环境¶

标准的强化学习 (Reinforcement Learning) 训练循环涉及一个模型，也称为策略，该模型经过训练以在特定环境中完成任务。通常，此环境是一个模拟器，它接受动作作为输入并生成观察结果以及一些元数据作为输出。

在本文档中，我们将探索 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()，它执行 actor 选择的动作。在 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)

此 tensordict 与从 EnvBase() 获取的 tensordict 结构相同，但增加了一个 "action" 条目。您可以像使用普通字典一样轻松访问动作

print(reset_with_action["action"])

tensor([0.7025])

现在我们需要将此动作传递给环境。我们将整个 tensordict 传递给 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)

同样，这个新的 tensordict 与上一个相同，除了它有一个 "next" 条目（它本身也是一个 tensordict！），其中包含由我们的动作产生的观察结果、奖励和完成状态。

我们将此格式称为 TED，代表 TorchRL 情节数据格式。它是库中表示数据的普遍方式，既可以像这里一样动态地表示，也可以使用离线数据集静态地表示。

在环境中运行回滚所需的最后一部分信息是如何将该 "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 方法可能看起来相当无用：它只是运行随机动作。如果策略可用，则可以将其传递给该方法并用于收集数据。

尽管如此，一开始运行一个简单的、无策略的回滚以快速了解环境的预期结果还是很有用的。

为了更好地理解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步的完整 rollout）。我们可以通过查看truncated条目来确认轨迹已被截断。

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

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

这就是对TorchRL环境API的简短介绍！

后续步骤¶

要进一步探索TorchRL的环境功能，请查看

将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()读取信息字典的 API 或通过set_gym_backend()控制 Gym 后端的方法。

脚本的总运行时间：（0 分钟 43.620 秒）

估计内存使用量：317 MB

由 Sphinx-Gallery 生成的图库