备注
点击此处下载完整示例代码
摆锤：使用 TorchRL 编写你的环境和变换
创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日
作者：Vincent Moens
创建环境（模拟器或物理控制系统接口）是强化学习和控制工程的综合部分。
TorchRL 提供了一套工具，在多个环境中实现这一功能。本教程演示了如何使用 PyTorch 和 TorchRL 从头开始编写摆锤模拟器。它自由地受到了 OpenAI-Gym/Farama-Gymnasium 控制库中的 Pendulum-v1 实现的启发。
简单摆锤 ¶
关键学习点：
在 TorchRL 中设计环境的方法：- 编写规范（输入、观察和奖励）；- 实现行为：播种、重置和步骤。
转换你的环境输入和输出，并编写你自己的转换；
如何使用
TensorDict通过codebase传递任意数据结构。在此过程中，我们将触及 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}\) 是角速度（单位：rad/s），\(g\) 是重力，\(L\) 是摆长，\(m\) 是质量，\(\theta\) 是角位置，\(u\) 是扭矩。然后根据以下公式更新角位置：
我们将奖励定义为
当角度接近 0（摆锤向上位置），角速度接近 0（无运动）以及扭矩也为 0 时，该奖励将被最大化。
编码动作的影响： _step() ¶
步骤方法是首先要考虑的，因为它将编码我们感兴趣的模拟。在 TorchRL 中， EnvBase 类有一个 EnvBase.step() 方法，该方法接收一个 tensordict.TensorDict 实例，其中包含一个 "action" 条目，指示要采取的操作。
为了方便从 tensordict 读取和写入，并确保键与库期望的一致，模拟部分已被委托给一个私有抽象方法 _step() ，该方法从 tensordict 读取输入数据，并将新的 tensordict 输出数据写入。
_step() 方法应执行以下操作：
读取输入键（如
"action"），并根据这些键执行模拟。获取观察结果、完成状态和奖励；
将观察值、奖励和完成状态写入新的
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() 是否包含所有预期的观察结果。再次强调，由于我们正在处理无状态环境，我们通过一个名为 "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 = 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
可复现实验：播种 ¶
在初始化实验时，对环境进行播种是一个常见操作。 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
我们将环境定义为非 batch_locked ，通过将 homonymous 属性转换为 False 来实现。这意味着我们不会强制输入 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 传递，以便在它们作为输入传递到下一个迭代时将其压缩回原始形状。
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 键列表中进行。以下图总结了环境和重放缓冲区中的此流程。
Transform 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 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)
将观察结果连接到“观察”条目。 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 秒）
