目录
  • 文档 >
  • 开始使用环境、TED 和变换
快捷方式
tutorials/getting-started-0

注意

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

开始使用环境、TED 和变换

作者: Vincent Moens

注意

要在 notebook 中运行本教程，请在开头添加一个安装单元格，内容如下：

!pip install tensordict
!pip install torchrl

欢迎来到入门教程！

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

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

RL 中的环境

标准的 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()，它启动一个 episode，以及 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() 获得的结构相同，只是增加了一个 "action" 条目。您可以像使用常规字典一样轻松访问动作

print(reset_with_action["action"])

tensor([1.5947])

我们现在需要将此动作传递给环境。我们将把整个 tensordict 传递给 step 方法，因为在更高级的案例中，例如多智能体 RL 或无状态环境，可能需要读取多个张量

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 与之前的 tensordict 相同，除了它有一个 "next" 条目（它本身是一个 tensordict！）包含由我们的动作产生的观测、奖励和完成状态。

我们将此格式称为 TED，即 TorchRL Episode Data 格式。它是库中表示数据的通用方式，无论是像这里这样动态的，还是像离线数据集那样静态的。

您需要在环境中运行 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)

环境 rollout

写下这三个步骤（计算动作、执行步骤、在 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 步的完整 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，它允许您在多个进程上运行同一（或不同！）环境的多个副本。

  • 使用 Pendulum 教程 设计您自己的环境，并了解规范和无状态环境。

  • 在 专门的教程中 查看关于环境的更深入的教程；

  • 如果您对 MARL 感兴趣，请查看 多智能体环境 API；

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

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

估计内存使用量： 320 MB

由 Sphinx-Gallery 生成的图库

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

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

查看教程

资源

查找开发资源并获得问题解答

查看资源

为了分析流量并优化您的体验，我们在本网站上使用 Cookie。点击或导航即表示您同意允许我们使用 Cookie。作为本网站的当前维护者，Facebook 的 Cookie 政策适用。了解更多信息，包括关于可用控件：Cookie 政策。