目录
  • 文档 >
  • 摆：使用 TorchRL 编写环境和变换
快捷方式
tutorials/pendulum

注意

跳到最后 下载完整示例代码。

摆：使用 TorchRL 编写环境和变换

作者: Vincent Moens

创建环境（模拟器或物理控制系统的接口）是强化学习和控制工程的一个组成部分。

TorchRL 提供了一套工具来在多种上下文中完成这项工作。本教程演示了如何使用 PyTorch 和 TorchRL 从头开始编写一个摆模拟器。它自由地借鉴了 OpenAI-Gym/Farama-Gymnasium 控制库 中的 Pendulum-v1 实现。

Pendulum

简单摆

主要学习内容

  • 如何在 TorchRL 中设计环境： - 编写规范（输入、观测和奖励）； - 实现行为：设置种子、重置和步进。

  • 变换环境的输入和输出，以及编写自己的变换；

  • 如何使用 TensorDict 在 codebase 中传递任意数据结构。

    在此过程中，我们将涉及 TorchRL 的三个关键组件

为了了解 TorchRL 环境可以实现的功能，我们将设计一个无状态环境。有状态环境跟踪遇到的最新物理状态并依赖此来模拟状态到状态的转换，而无状态环境期望在每一步提供当前状态以及采取的动作。TorchRL 支持这两种类型的环境，但无状态环境更通用，因此涵盖了 TorchRL 中环境 API 更广泛的功能。

对无状态环境进行建模使用户可以完全控制模拟器的输入和输出：可以在任何阶段重置实验或从外部主动修改动力学。然而，它假设我们可以控制任务，这可能并非总是如此：解决一个我们无法控制当前状态的问题更具挑战性，但也具有更广泛的应用范围。

无状态环境的另一个优点是它们可以实现转换模拟的批处理执行。如果后端和实现允许，代数运算可以在标量、向量或张量上无缝执行。本教程提供了此类示例。

本教程将结构化如下

  • 我们将首先熟悉环境属性：其形状（batch_size）、其方法（主要是 step()、reset() 和 set_seed()），最后是其规范（specs）。

  • 编写完模拟器后，我们将演示如何在训练期间使用变换来使用它。

  • 我们将探索 TorchRL API 引出的新途径，包括：转换输入的可能性、模拟的向量化执行以及通过模拟图进行反向传播的可能性。

  • 最后，我们将训练一个简单的策略来解决我们实现的系统。

此环境的内置版本可在 class:~torchrl.envs.PendulumEnv 中找到。

from collections import defaultdict
from typing import Optional

import numpy as np
import torch
import tqdm
from tensordict import TensorDict, TensorDictBase
from tensordict.nn import TensorDictModule
from torch import nn

from torchrl.data import Bounded, Composite, Unbounded
from torchrl.envs import (
    CatTensors,
    EnvBase,
    Transform,
    TransformedEnv,
    UnsqueezeTransform,
)
from torchrl.envs.transforms.transforms import _apply_to_composite
from torchrl.envs.utils import check_env_specs, step_mdp

DEFAULT_X = np.pi
DEFAULT_Y = 1.0

在设计新的环境类时，必须注意四件事

  • EnvBase._reset()，它编码了模拟器在（可能随机的）初始状态下重置的代码；

  • EnvBase._step()，它编码了状态转换动态；

  • EnvBase._set_seed()，它实现了种子设置机制；

  • 环境规范（specs）。

让我们首先描述手头的问题：我们希望模拟一个简单摆，我们可以控制其固定点施加的扭矩。我们的目标是将摆置于向上位置（约定角度位置为 0），并使其在该位置静止。为了设计我们的动力系统，我们需要定义两个方程：跟随动作（施加的扭矩）的运动方程，以及构成我们目标函数的奖励方程。

对于运动方程，我们将根据以下公式更新角速度

\[\dot{\theta}_{t+1} = \dot{\theta}_t + (3 * g / (2 * L) * \sin(\theta_t) + 3 / (m * L^2) * u) * dt\]

其中 \(\dot{\theta}\) 是角速度（弧度/秒），\(g\) 是重力加速度，\(L\) 是摆长，\(m\) 是其质量，\(\theta\) 是其角度位置，\(u\) 是扭矩。然后根据以下公式更新角度位置

\[\theta_{t+1} = \theta_{t} + \dot{\theta}_{t+1} dt\]

我们将奖励定义为

\[r = -(\theta^2 + 0.1 * \dot{\theta}^2 + 0.001 * u^2)\]

当角度接近 0（摆处于向上位置）、角速度接近 0（无运动）且扭矩也为 0 时，该奖励将被最大化。

编写动作的影响：_step()

首先要考虑的是步进（step）方法，因为它将编码我们感兴趣的模拟。在 TorchRL 中，EnvBase 类有一个 EnvBase.step() 方法，它接收一个 tensordict.TensorDict 实例，其中包含一个 "action" 条目，指示要采取的动作。

为了方便从该 tensordict 中读写数据，并确保键与库期望的一致，模拟部分已委托给一个私有抽象方法 _step()，该方法从输入 tensordict 中读取输入数据，并使用输出数据写入一个新的 TensorDict。

_step() 方法应执行以下操作

  1. 读取输入键（例如 "action"），并根据这些键执行模拟；

  2. 检索观测、完成状态和奖励；

  3. 将一组观测值以及奖励和完成状态写入新的 TensorDict 中对应的条目。

接下来，step() 方法会将 step() 的输出合并到输入 tensordict 中，以强制执行输入/输出一致性。

通常，对于有状态环境，这看起来像这样

>>> policy(env.reset())
>>> print(tensordict)
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([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)
>>> env.step(tensordict)
>>> print(tensordict)
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([]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)

请注意，根 tensordict 没有改变，唯一的修改是出现了一个包含新信息的 "next" 条目。

在摆的例子中，我们的 _step() 方法将从输入 tensordict 中读取相关条目，并计算应用了由 "action" 键编码的力后摆的位置和速度。我们将摆的新角度位置 "new_th" 计算为先前位置 "th" 加上新速度 "new_thdot" 在时间间隔 dt 内的结果。

由于我们的目标是将摆向上翻转并使其在该位置保持静止，因此我们的 cost（负奖励）函数对于接近目标的位置和低速度的值较低。实际上，我们希望阻止远离“向上”的位置和/或远离 0 的速度。

在我们的示例中，EnvBase._step() 被编码为一个静态方法，因为我们的环境是无状态的。在有状态设置中，需要 self 参数，因为需要从环境中读取状态。

def _step(tensordict):
    th, thdot = tensordict["th"], tensordict["thdot"]  # th := theta

    g_force = tensordict["params", "g"]
    mass = tensordict["params", "m"]
    length = tensordict["params", "l"]
    dt = tensordict["params", "dt"]
    u = tensordict["action"].squeeze(-1)
    u = u.clamp(-tensordict["params", "max_torque"], tensordict["params", "max_torque"])
    costs = angle_normalize(th) ** 2 + 0.1 * thdot**2 + 0.001 * (u**2)

    new_thdot = (
        thdot
        + (3 * g_force / (2 * length) * th.sin() + 3.0 / (mass * length**2) * u) * dt
    )
    new_thdot = new_thdot.clamp(
        -tensordict["params", "max_speed"], tensordict["params", "max_speed"]
    )
    new_th = th + new_thdot * dt
    reward = -costs.view(*tensordict.shape, 1)
    done = torch.zeros_like(reward, dtype=torch.bool)
    out = TensorDict(
        {
            "th": new_th,
            "thdot": new_thdot,
            "params": tensordict["params"],
            "reward": reward,
            "done": done,
        },
        tensordict.shape,
    )
    return out


def angle_normalize(x):
    return ((x + torch.pi) % (2 * torch.pi)) - torch.pi

重置模拟器：_reset()

我们需要关心的第二个方法是 _reset() 方法。像 _step() 一样，它应该将观测条目和可能的完成状态写入其输出的 tensordict 中（如果省略完成状态，父方法 reset() 会将其填充为 False）。在某些情况下，要求 _reset 方法接收调用它的函数的命令（例如，在多智能体设置中，我们可能需要指示哪些智能体需要重置）。这就是为什么 _reset() 方法也期望 tensordict 作为输入，尽管它可以完全为空或为 None。

父方法 EnvBase.reset() 会像 EnvBase.step() 一样进行一些简单的检查，例如确保输出 tensordict 中返回了 "done" 状态，并且形状与规范期望的一致。

对于我们来说，唯一需要考虑的重要事情是 EnvBase._reset() 是否包含所有期望的观测。同样，由于我们正在使用无状态环境，我们将摆的配置作为嵌套的 tensordict 通过，命名为 "params"。

在本例中，我们不传递完成状态，因为对于 _reset() 而言这不是强制性的，而且我们的环境是非终止的，因此我们始终期望它是 False。

def _reset(self, tensordict):
    if tensordict is None or tensordict.is_empty():
        # if no ``tensordict`` is passed, we generate a single set of hyperparameters
        # Otherwise, we assume that the input ``tensordict`` contains all the relevant
        # parameters to get started.
        tensordict = self.gen_params(batch_size=self.batch_size)

    high_th = torch.tensor(DEFAULT_X, device=self.device)
    high_thdot = torch.tensor(DEFAULT_Y, device=self.device)
    low_th = -high_th
    low_thdot = -high_thdot

    # for non batch-locked environments, the input ``tensordict`` shape dictates the number
    # of simulators run simultaneously. In other contexts, the initial
    # random state's shape will depend upon the environment batch-size instead.
    th = (
        torch.rand(tensordict.shape, generator=self.rng, device=self.device)
        * (high_th - low_th)
        + low_th
    )
    thdot = (
        torch.rand(tensordict.shape, generator=self.rng, device=self.device)
        * (high_thdot - low_thdot)
        + low_thdot
    )
    out = TensorDict(
        {
            "th": th,
            "thdot": thdot,
            "params": tensordict["params"],
        },
        batch_size=tensordict.shape,
    )
    return out

环境元数据：env.*_spec

规范定义了环境的输入和输出域。重要的是规范准确地定义了运行时将接收到的张量，因为它们通常用于在多进程和分布式设置中携带有关环境的信息。它们还可以用于实例化惰性定义的神经网络和测试脚本，而无需实际查询环境（例如，对于真实的物理系统而言，这可能会很昂贵）。

我们需要在环境编码四种规范

  • EnvBase.observation_spec: 这将是一个 CompositeSpec 实例，其中每个键都是一个观测（CompositeSpec 可以被视为规范的字典）。

  • EnvBase.action_spec: 它可以是任何类型的规范，但要求它与输入 tensordict 中的 "action" 条目对应；

  • EnvBase.reward_spec: 提供有关奖励空间的信息；

  • EnvBase.done_spec: 提供有关完成标志空间的信息。

TorchRL 规范组织在两个通用容器中：input_spec，包含 step 函数读取的信息的规范（分为 action_spec，包含动作；state_spec，包含其余所有信息）；以及 output_spec，编码 step 输出的规范（observation_spec、reward_spec 和 done_spec）。通常，您不应直接与 output_spec 和 input_spec 交互，而只与其内容交互：observation_spec、reward_spec、done_spec、action_spec 和 state_spec。原因是这些规范在 output_spec 和 input_spec 内以一种非平凡的方式组织，并且两者都不应直接修改。

换句话说，observation_spec 及相关属性是输出和输入规范容器内容的便捷快捷方式。

TorchRL 提供了多种 TensorSpec 子类 来编码环境的输入和输出特性。

规范形状

环境规范的领先维度必须与环境批次大小匹配。这样做是为了确保环境的每个组件（包括其变换）都能准确地表示预期的输入和输出形状。在有状态设置中，这应该被准确编码。

对于非批次锁定环境，例如我们示例中的环境（见下文），这无关紧要，因为环境批次大小很可能是空的。

def _make_spec(self, td_params):
    # Under the hood, this will populate self.output_spec["observation"]
    self.observation_spec = Composite(
        th=Bounded(
            low=-torch.pi,
            high=torch.pi,
            shape=(),
            dtype=torch.float32,
        ),
        thdot=Bounded(
            low=-td_params["params", "max_speed"],
            high=td_params["params", "max_speed"],
            shape=(),
            dtype=torch.float32,
        ),
        # we need to add the ``params`` to the observation specs, as we want
        # to pass it at each step during a rollout
        params=make_composite_from_td(td_params["params"]),
        shape=(),
    )
    # since the environment is stateless, we expect the previous output as input.
    # For this, ``EnvBase`` expects some state_spec to be available
    self.state_spec = self.observation_spec.clone()
    # action-spec will be automatically wrapped in input_spec when
    # `self.action_spec = spec` will be called supported
    self.action_spec = Bounded(
        low=-td_params["params", "max_torque"],
        high=td_params["params", "max_torque"],
        shape=(1,),
        dtype=torch.float32,
    )
    self.reward_spec = Unbounded(shape=(*td_params.shape, 1))


def make_composite_from_td(td):
    # custom function to convert a ``tensordict`` in a similar spec structure
    # of unbounded values.
    composite = Composite(
        {
            key: make_composite_from_td(tensor)
            if isinstance(tensor, TensorDictBase)
            else Unbounded(dtype=tensor.dtype, device=tensor.device, shape=tensor.shape)
            for key, tensor in td.items()
        },
        shape=td.shape,
    )
    return composite

可重现的实验：种子设置

在初始化实验时，设置环境种子是一种常见的操作。EnvBase._set_seed() 的唯一目标是设置所包含模拟器的种子。如果可能，此操作不应调用 reset() 或与环境执行交互。父方法 EnvBase.set_seed() 包含一种机制，允许使用不同的伪随机和可重现种子来设置多个环境的种子。

def _set_seed(self, seed: Optional[int]):
    rng = torch.manual_seed(seed)
    self.rng = rng

将所有内容打包在一起：EnvBase 类

我们终于可以将各个部分组合在一起，设计我们的环境类。规范的初始化需要在环境构建期间执行，因此我们必须注意在 PendulumEnv.__init__() 中调用 _make_spec() 方法。

我们添加一个静态方法 PendulumEnv.gen_params()，它确定性地生成一组在执行期间使用的超参数

def gen_params(g=10.0, batch_size=None) -> TensorDictBase:
    """Returns a ``tensordict`` containing the physical parameters such as gravitational force and torque or speed limits."""
    if batch_size is None:
        batch_size = []
    td = TensorDict(
        {
            "params": TensorDict(
                {
                    "max_speed": 8,
                    "max_torque": 2.0,
                    "dt": 0.05,
                    "g": g,
                    "m": 1.0,
                    "l": 1.0,
                },
                [],
            )
        },
        [],
    )
    if batch_size:
        td = td.expand(batch_size).contiguous()
    return td

我们通过将 homonymous 属性设置为 False 来将环境定义为非 batch_locked。这意味着我们将不强制输入 tensordict 具有与环境匹配的 batch-size。

以下代码将把我们上面编写的部分组合在一起。

class PendulumEnv(EnvBase):
    metadata = {
        "render_modes": ["human", "rgb_array"],
        "render_fps": 30,
    }
    batch_locked = False

    def __init__(self, td_params=None, seed=None, device="cpu"):
        if td_params is None:
            td_params = self.gen_params()

        super().__init__(device=device, batch_size=[])
        self._make_spec(td_params)
        if seed is None:
            seed = torch.empty((), dtype=torch.int64).random_().item()
        self.set_seed(seed)

    # Helpers: _make_step and gen_params
    gen_params = staticmethod(gen_params)
    _make_spec = _make_spec

    # Mandatory methods: _step, _reset and _set_seed
    _reset = _reset
    _step = staticmethod(_step)
    _set_seed = _set_seed

测试我们的环境

TorchRL 提供了一个简单的函数 check_env_specs() 来检查（变换后的）环境的输入/输出结构是否与规范规定的结构匹配。让我们试一下

env = PendulumEnv()
check_env_specs(env)

我们可以查看我们的规范以获得环境签名的可视化表示

print("observation_spec:", env.observation_spec)
print("state_spec:", env.state_spec)
print("reward_spec:", env.reward_spec)

我们也可以执行一些命令来检查输出结构是否与预期的一致。

td = env.reset()
print("reset tensordict", td)

我们可以运行 env.rand_step() 从 action_spec 域中随机生成一个动作。由于我们的环境是无状态的，必须传递包含超参数和当前状态的 tensordict。在有状态环境中，env.rand_step() 也能完美工作。

td = env.rand_step(td)
print("random step tensordict", td)

变换环境

为无状态模拟器编写环境变换比为有状态模拟器更复杂一些：转换需要在下一次迭代中读取的输出条目，需要在下一次步进调用 meth.step() 之前应用逆变换。这是一个展示 TorchRL 变换所有功能的理想场景！

例如，在以下变换后的环境中，我们对 ["th", "thdot"] 条目进行 unsqueeze，以便能够沿最后一个维度堆叠它们。我们也将它们作为 in_keys_inv 传递，以便在下一次迭代中作为输入传递时将它们 squeeze 回其原始形状。

env = TransformedEnv(
    env,
    # ``Unsqueeze`` the observations that we will concatenate
    UnsqueezeTransform(
        dim=-1,
        in_keys=["th", "thdot"],
        in_keys_inv=["th", "thdot"],
    ),
)

编写自定义变换

TorchRL 的变换可能无法涵盖环境执行后想要执行的所有操作。编写一个变换并不需要太多精力。就像环境设计一样，编写变换有两个步骤

  • 理清动力学（前向和逆向）；

  • 调整环境规范。

变换可以在两种设置下使用：独立使用时，它可以作为 Module 使用。它也可以附加到 TransformedEnv 使用。类的结构允许在不同的上下文中自定义行为。

一个 Transform 骨架可以概括如下

class Transform(nn.Module):
    def forward(self, tensordict):
        ...
    def _apply_transform(self, tensordict):
        ...
    def _step(self, tensordict):
        ...
    def _call(self, tensordict):
        ...
    def inv(self, tensordict):
        ...
    def _inv_apply_transform(self, tensordict):
        ...

有三个入口点（forward()、_step() 和 inv()），它们都接收 tensordict.TensorDict 实例。前两个最终将通过 in_keys 指示的键，并对每个键调用 _apply_transform()。结果将写入由 Transform.out_keys 指示的条目中（如果提供），如果未提供，则 in_keys 将使用转换后的值更新。如果需要执行逆变换，将执行类似的数据流，但使用 Transform.inv() 和 Transform._inv_apply_transform() 方法，并跨越 in_keys_inv 和 out_keys_inv 键列表。下图总结了环境和回放缓冲区的数据流。

变换 API

在某些情况下，变换不会以单一方式处理 tensordict 中的键子集，而是对父环境执行某些操作或处理整个输入 tensordict。在这些情况下，应重写 _call() 和 forward() 方法，并可以跳过 _apply_transform() 方法。

让我们编写新的变换，计算位置角度的 sine 和 cosine 值，因为这些值对于我们学习策略比原始角度值更有用

class SinTransform(Transform):
    def _apply_transform(self, obs: torch.Tensor) -> None:
        return obs.sin()

    # The transform must also modify the data at reset time
    def _reset(
        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
    ) -> TensorDictBase:
        return self._call(tensordict_reset)

    # _apply_to_composite will execute the observation spec transform across all
    # in_keys/out_keys pairs and write the result in the observation_spec which
    # is of type ``Composite``
    @_apply_to_composite
    def transform_observation_spec(self, observation_spec):
        return Bounded(
            low=-1,
            high=1,
            shape=observation_spec.shape,
            dtype=observation_spec.dtype,
            device=observation_spec.device,
        )


class CosTransform(Transform):
    def _apply_transform(self, obs: torch.Tensor) -> None:
        return obs.cos()

    # The transform must also modify the data at reset time
    def _reset(
        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
    ) -> TensorDictBase:
        return self._call(tensordict_reset)

    # _apply_to_composite will execute the observation spec transform across all
    # in_keys/out_keys pairs and write the result in the observation_spec which
    # is of type ``Composite``
    @_apply_to_composite
    def transform_observation_spec(self, observation_spec):
        return Bounded(
            low=-1,
            high=1,
            shape=observation_spec.shape,
            dtype=observation_spec.dtype,
            device=observation_spec.device,
        )


t_sin = SinTransform(in_keys=["th"], out_keys=["sin"])
t_cos = CosTransform(in_keys=["th"], out_keys=["cos"])
env.append_transform(t_sin)
env.append_transform(t_cos)

将观测值连接到“observation”条目上。del_keys=False 确保我们在下一次迭代中保留这些值。

cat_transform = CatTensors(
    in_keys=["sin", "cos", "thdot"], dim=-1, out_key="observation", del_keys=False
)
env.append_transform(cat_transform)

再次，让我们检查我们的环境规范是否与接收到的匹配

check_env_specs(env)

执行 Rollout

执行 rollout 是一系列简单的步骤

  • 重置环境

  • 当某个条件未满足时

    • 根据策略计算动作

    • 根据此动作执行一步

    • 收集数据

    • 进行 MDP 步进

  • 收集数据并返回

这些操作已便捷地封装在 rollout() 方法中，下面我们提供了一个简化版本。

def simple_rollout(steps=100):
    # preallocate:
    data = TensorDict(batch_size=[steps])
    # reset
    _data = env.reset()
    for i in range(steps):
        _data["action"] = env.action_spec.rand()
        _data = env.step(_data)
        data[i] = _data
        _data = step_mdp(_data, keep_other=True)
    return data


print("data from rollout:", simple_rollout(100))

批处理计算

本教程尚未探索的最后一个方面是我们在 TorchRL 中进行批处理计算的能力。由于我们的环境对输入数据形状没有任何假设，我们可以无缝地在数据批次上执行它。更棒的是：对于非批次锁定环境，例如我们的摆，我们可以动态更改批次大小而无需重新创建环境。为此，我们只需生成所需形状的参数。

batch_size = 10  # number of environments to be executed in batch
td = env.reset(env.gen_params(batch_size=[batch_size]))
print("reset (batch size of 10)", td)
td = env.rand_step(td)
print("rand step (batch size of 10)", td)

使用一批数据执行 rollout 需要我们在 rollout 函数之外重置环境，因为我们需要动态定义 batch_size，而 rollout() 不支持此功能

rollout = env.rollout(
    3,
    auto_reset=False,  # we're executing the reset out of the ``rollout`` call
    tensordict=env.reset(env.gen_params(batch_size=[batch_size])),
)
print("rollout of len 3 (batch size of 10):", rollout)

训练一个简单的策略

在本例中，我们将使用奖励作为可微分目标（例如负损失）来训练一个简单的策略。我们将利用我们的动力系统完全可微分这一事实，通过轨迹回报进行反向传播，并直接调整策略的权重以最大化此值。当然，在许多设置中，我们做出的许多假设都不成立，例如可微分系统和对底层机制的完全访问。

尽管如此，这是一个非常简单的示例，展示了如何在 TorchRL 中使用自定义环境编码训练循环。

让我们首先编写策略网络

torch.manual_seed(0)
env.set_seed(0)

net = nn.Sequential(
    nn.LazyLinear(64),
    nn.Tanh(),
    nn.LazyLinear(64),
    nn.Tanh(),
    nn.LazyLinear(64),
    nn.Tanh(),
    nn.LazyLinear(1),
)
policy = TensorDictModule(
    net,
    in_keys=["observation"],
    out_keys=["action"],
)

和我们的优化器

optim = torch.optim.Adam(policy.parameters(), lr=2e-3)

训练循环

我们将依次

  • 生成轨迹

  • 累加奖励

  • 通过这些操作定义的图进行反向传播

  • 裁剪梯度范数并执行优化步骤

  • 重复

在训练循环结束时，我们应该获得一个接近 0 的最终奖励，这表明摆锤如预期般向上且静止。

batch_size = 32
n_iter = 1000  # set to 20_000 for a proper training
pbar = tqdm.tqdm(range(n_iter // batch_size))
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optim, n_iter)
logs = defaultdict(list)

for _ in pbar:
    init_td = env.reset(env.gen_params(batch_size=[batch_size]))
    rollout = env.rollout(100, policy, tensordict=init_td, auto_reset=False)
    traj_return = rollout["next", "reward"].mean()
    (-traj_return).backward()
    gn = torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
    optim.step()
    optim.zero_grad()
    pbar.set_description(
        f"reward: {traj_return: 4.4f}, "
        f"last reward: {rollout[..., -1]['next', 'reward'].mean(): 4.4f}, gradient norm: {gn: 4.4}"
    )
    logs["return"].append(traj_return.item())
    logs["last_reward"].append(rollout[..., -1]["next", "reward"].mean().item())
    scheduler.step()


def plot():
    import matplotlib
    from matplotlib import pyplot as plt

    is_ipython = "inline" in matplotlib.get_backend()
    if is_ipython:
        from IPython import display

    with plt.ion():
        plt.figure(figsize=(10, 5))
        plt.subplot(1, 2, 1)
        plt.plot(logs["return"])
        plt.title("returns")
        plt.xlabel("iteration")
        plt.subplot(1, 2, 2)
        plt.plot(logs["last_reward"])
        plt.title("last reward")
        plt.xlabel("iteration")
        if is_ipython:
            display.display(plt.gcf())
            display.clear_output(wait=True)
        plt.show()


plot()

结论

在本教程中，我们学习了如何从头开始编写一个无状态环境。我们涉及的主题包括

  • 编码环境时需要关注的四个基本组成部分（step、reset、随机种子和构建规范）。我们看到了这些方法和类如何与 TensorDict 类交互；

  • 如何使用 check_env_specs() 测试环境是否正确编码；

  • 如何在无状态环境的上下文中添加变换以及如何编写自定义变换；

  • 如何在一个完全可微的模拟器上训练策略。

由 Sphinx-Gallery 生成的画廊

文档

获取 PyTorch 综合开发者文档

查看文档

教程

获取适合初学者和高级开发者的深入教程

查看教程

资源

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

查看资源

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