Table of Contents
  • Tutorials >
  • 摆：使用 TorchRL 编写环境和变换
Shortcuts
advanced/pendulum

备注

点击 这里 下载完整示例代码

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

Created On: Nov 09, 2023 | Last Updated: Jan 27, 2025 | Last Verified: Nov 05, 2024

作者: Vincent Moens

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

TorchRL 提供一套工具，用于在多种场景中实现这一目标。本教程展示如何使用 PyTorch 和 TorchRL 从头开始编写一个摆动模拟器。它灵感来源于 OpenAI-Gym/Farama-Gymnasium 控制库 中的 Pendulum-v1。

摆

简单摆

关键学习点：

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

  • 转换环境输入和输出，并编写自己定义的变换；

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

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

为了展示使用 TorchRL 环境可以实现的目标，我们将设计一个*无状态*环境。虽然有状态环境跟踪最新的物理状态并依赖于这一点来模拟状态间的转移，但无状态环境期望每一步提供当前状态以及执行的动作。TorchRL 支持两种类型的环境，但无状态环境更为通用，因此涵盖了 TorchRL 环境 API 的更多功能。

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

无状态环境的另一个优势是它们可以实现转移模拟的批量执行。如果后端和实现允许，将可以无缝地对标量、向量或张量执行代数操作。本教程给出了这样的示例。

本教程将结构化地包括以下部分：

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

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

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

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

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 BoundedTensorSpec, CompositeSpec, UnboundedContinuousTensorSpec
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()，实现设定种子机制；

  • 环境规范。

我们首先来描述当前的问题：我们希望对一个简单的摆模型进行建模，其中可以控制其固定点施加的扭矩。我们的目标是将摆置于向上位置（按照约定，角位置为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中，:class:`~torchrl.envs.EnvBase`类有一个:meth:`EnvBase.step`方法，该方法接收一个带有“action”条目的:class:`tensordict.TensorDict`实例，指示要执行的动作。

为了便于从该“tensordict”读取和写入，并确保键与库期望的一致，模拟部分已委托给一个私有抽象方法:meth:_step，它从一个“tensordict”读取输入数据，并使用输出数据写入一个*新的*“tensordict”。

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

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

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

  3. 将观测值集合与奖励和完成状态写入到新的:class:`TensorDict`的相应条目中。

接着，:meth:`~torchrl.envs.EnvBase.step`方法将会合并:meth:`~torchrl.envs.EnvBase.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”条目，其中包含新信息。

在摆模型的例子中，我们的:meth:`_step`方法将从输入的“tensordict”读取相关条目，并计算在“action”键代表的力作用后摆的运动和速度。我们将新的摆角位置“new_th”计算为此前位置“th”加上覆盖时间间隔“dt”中的新速度“new_thdot”的结果。

鉴于我们的目标是让摆向上并保持静止，我们的“cost”（负奖励）函数在位置靠近目标且速度较低时更低。确实，我们希望减少远离“向上”的位置和远离0的速度。

在我们的例子中，:meth:`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()

我们需要关注的第二个方法是:meth:`~torchrl.envs.EnvBase._reset`方法。像:meth:`~torchrl.envs.EnvBase._step`一样，它应在输出的“tensordict”中写入观测条目以及可能的完成状态（如果省略完成状态，父方法:meth:`~torchrl.envs.EnvBase.reset`会将其填充为“False”）。在某些情况下，_reset方法需要接收到调用它的函数的命令（例如，在多智能体设置中我们可能需要指示需要重置的智能体）。这就是:meth:`~torchrl.envs.EnvBase._reset`方法还需要一个作为输入的“tensordict”的原因，尽管它完全可以是空的或“None”。

父:meth:`EnvBase.reset`做了一些简单的检查，如:meth:`EnvBase.step`所做的检查，包括确保在输出“tensordict”中返回一个“done”状态，以及形状符合规范的期望。

对我们来说，唯一重要的事情是确保:meth:`EnvBase._reset`包含所有期望的观测值。再次强调，由于我们工作的是一个无状态环境，我们将摆的配置传递到一个嵌套的“tensordict”中，命名为“params”。

在这个例子中，我们没有传递完成状态，因为这对于:meth:`_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：它会是:class:`~torchrl.data.CompositeSpec`实例，其中每个键是一个观测值（:class:`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提供了多个:class:~torchrl.data.TensorSpec `子类 <https://pytorch.org/rl/stable/reference/data.html#tensorspec>`_用于编码环境的输入和输出特征。

规格形状

环境规格的领先维度必须与环境批量大小匹配。这是为了确保环境的每个组件（包括其变换）准确地表示期望的输入和输出形状。这是有状态设置中应准确编码的内容。

对于无批量锁定的环境，例如示例中的环境（见下文），这与之无关，因为环境批量大小很可能是空的。

def _make_spec(self, td_params):
    # Under the hood, this will populate self.output_spec["observation"]
    self.observation_spec = CompositeSpec(
        th=BoundedTensorSpec(
            low=-torch.pi,
            high=torch.pi,
            shape=(),
            dtype=torch.float32,
        ),
        thdot=BoundedTensorSpec(
            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 = BoundedTensorSpec(
        low=-td_params["params", "max_torque"],
        high=td_params["params", "max_torque"],
        shape=(1,),
        dtype=torch.float32,
    )
    self.reward_spec = UnboundedContinuousTensorSpec(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 = CompositeSpec(
        {
            key: make_composite_from_td(tensor)
            if isinstance(tensor, TensorDictBase)
            else UnboundedContinuousTensorSpec(
                dtype=tensor.dtype, device=tensor.device, shape=tensor.shape
            )
            for key, tensor in td.items()
        },
        shape=td.shape,
    )
    return composite

可重复的实验：设置种子

为环境设置种子是初始化实验时的常见操作。:func:`EnvBase._set_seed`的唯一目标是设置包含的模拟器的种子。如果可能，该操作不应调用``reset()``或与环境执行交互。父:func:`EnvBase.set_seed`方法集成了一个机制，可使用不同的伪随机且可重复种子为多个环境设置种子。

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

将内容整合在一起：:class:`~torchrl.envs.EnvBase`类

我们终于可以将这些部分整合在一起并设计我们的环境类了。规格初始化需要在环境构造期间完成，因此我们必须在:func:`PendulumEnv.__init__`中调用:func:`_make_spec`方法。

我们添加一个静态方法:meth: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

我们通过将其“同名”属性设置为``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提供了一个简单的函数:func:~torchrl.envs.utils.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)

我们可以运行:func:env.rand_step`从``action_spec``域随机生成一个动作。由于我们的环境是无状态的，因此**必须**传递一个包含超参数和当前状态的``tensordict`。在有状态上下文中，``env.rand_step()``也能正常运行。

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

变换环境

为无状态模拟器编写环境变换比有状态的稍微复杂一些：变换一个需要在下一次迭代中读取的输出条目要求在接下来的步骤调用:func:`meth.step`之前应用反变换。这是展示TorchRL变换所有特性的理想场景！

例如，在下面的变换环境中，我们对条目``[“th”, “thdot”]``执行``unsqueeze``以便能够在最后一个维度上堆叠它们。我们也将它们作为``in_keys_inv``传递，以便在下一次迭代时将它们恢复到原始形状后再作为输入传递。

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。在这些情况下，应该重写 _call() 和 forward() 方法，并且可以跳过 _apply_transform() 方法。

让我们编写新的变换来计算位置角度的 正弦 和 余弦 值，因为这些值比原始角度值更有助于我们学习策略：

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 BoundedTensorSpec(
            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 BoundedTensorSpec(
            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)

执行回合

执行回合是简单步骤的连续体：

  • 重置环境

  • 在某些条件未满足时：

    • 根据策略计算行动

    • 根据此行动执行一步操作

    • 收集数据

    • 进行 MDP 步骤

  • 收集数据并返回

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

def simple_rollout(steps=100):
    # preallocate:
    data = TensorDict({}, [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 中批量计算的能力。因为我们的环境并不对输入数据的形状作任何假设，所以可以无缝地对数据批次执行操作。更妙的是：对于像我们的 Pendulum 这样不受批处理限制的环境，我们可以随意更改批量大小而无需重新创建环境。为此，我们只需生成具有所需形状的参数。

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 = 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
pbar = tqdm.tqdm(range(20_000 // batch_size))
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optim, 20_000)
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() 测试环境是否编写正确；

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

  • 如何在完全可微的仿真器上训练策略。

**脚本的总运行时间:**（0分钟 0.000秒）

由Sphinx-Gallery生成的图库

文档

访问 PyTorch 的详细开发者文档

查看文档

教程

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

查看教程

资源

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

查看资源

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