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

注意

跳至末尾下载完整的示例代码。

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

作者：Vincent Moens

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

TorchRL提供了一套工具，可以在多种上下文中完成此操作。本教程将演示如何使用PyTorch和TorchRL从头开始编写一个pendulum模拟器。该教程受到OpenAI-Gym/Farama-Gymnasium控制库中Pendulum-v1实现的启发。

Pendulum

简单摆锤

主要学习内容

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

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

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

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

为了让大家了解使用TorchRL的环境可以实现什么，我们将设计一个无状态环境。无状态环境在每次步进时期望当前状态被提供给它们，以及所执行的操作。TorchRL同时支持这两种环境，但无状态环境更通用，因此涵盖了TorchRL环境API更广泛的特性。

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

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

本教程的结构如下

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

  • 在编写完模拟器后，我们将展示如何在训练过程中使用变换来使用它。

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

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

此环境的内置版本可以在类~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()，它实现了播种机制；

  • 环境规范。

首先，让我们描述一下当前的问题：我们希望模拟一个简单的摆锤，我们可以控制其固定点的力矩。我们的目标是将摆锤放置在向上位置（按约定角位置为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()

步进方法是第一个需要考虑的，因为它将编码我们感兴趣的模拟。在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"条目，其中包含新信息。

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

由于我们的目标是使摆锤向上并且保持静止，我们的cost（负奖励）函数对于接近目标的位置和低速度来说更低。事实上，我们想阻止那些离“向上”位置很远的位置和/或速度很大的情况。

在我们的示例中，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()是否包含所有预期的观测。再次强调，由于我们处理的是无状态环境，我们将摆锤的配置传递到一个名为"params"的嵌套tensordict中。

在此示例中，我们不传递完成状态，因为这对于_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，其中包含步进函数读取的信息的规范（分为action_spec，包含动作，以及state_spec，包含其余所有内容），以及output_spec，它编码步进输出的规范（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]) -> None:
    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

我们将环境定义为非batch_locked，将homonymous属性设置为False。这意味着我们**不会**强制输入tensordict具有与环境batch-size匹配的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变换所有特性的理想场景！

例如，在下面的转换环境中，我们将unsqueeze条目["th", "thdot"]，以便能够沿最后一个维度堆叠它们。我们还将它们作为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()方法。

让我们编写新的变换，计算位置角的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)

执行采样

执行采样是一系列简单的步骤

  • 重置环境

  • 直到某个条件不满足

    • 根据策略计算动作

    • 根据此动作执行步进

    • 收集数据

    • 进行一次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中进行批量计算的能力。因为我们的环境不对输入数据的形状做任何假设，所以我们可以无缝地在数据批次上执行它。更好的是：对于像我们的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)

使用数据批次执行采样需要我们在采样函数之外重置环境，因为我们需要动态定义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 政策。