注意
转到末尾 下载完整的示例代码。
TorchRL 目标：编码 DDPG 损失#
创建日期：2023 年 8 月 14 日 | 最后更新：2025 年 3 月 20 日 | 最后验证：未验证
概述#
TorchRL 将强化学习算法的训练分为多个部分，这些部分将在您的训练脚本中进行组装：环境、数据收集和存储、模型以及最终的损失函数。
TorchRL 损失（或“目标”）是包含可训练参数（策略和值模型）的有状态对象。本教程将指导您从头开始使用 TorchRL 编码损失。
为此，我们将重点关注 DDPG，这是一个相对简单的编码算法。深度确定性策略梯度 (DDPG) 是一种简单的连续控制算法。它包括学习动作-观察对的参数值函数，然后学习一个策略，该策略输出在给定特定观察的情况下最大化此值函数的动作。
您将学到什么
如何编写损失模块并自定义其值估计器；
如何在 TorchRL 中构建环境，包括转换（例如，数据归一化）和并行执行；
如何设计策略和值网络；
如何从环境中高效收集数据并将其存储在回放缓冲区中；
如何将轨迹（而非转换）存储在回放缓冲区中；
如何评估您的模型。
先决条件#
本教程假定您已完成PPO 教程，该教程概述了 TorchRL 组件和依赖项，例如 tensordict.TensorDict 和 tensordict.nn.TensorDictModules，尽管它应该足够透明，无需深入了解这些类即可理解。
注意
我们不旨在提供该算法的 SOTA 实现，而是提供 TorchRL 损失实现以及在该算法上下文中使用的库功能的高级说明。
导入和设置#
%%bash pip3 install torchrl mujoco glfw
import torch
import tqdm
如果可用，我们将在 CUDA 上执行策略
is_fork = multiprocessing.get_start_method() == "fork"
device = (
torch.device(0)
if torch.cuda.is_available() and not is_fork
else torch.device("cpu")
)
collector_device = torch.device("cpu") # Change the device to ``cuda`` to use CUDA
TorchRL LossModule#
TorchRL 提供了一系列可在训练脚本中使用的损失。目的是让损失易于重用/互换，并且具有简单的签名。
TorchRL 损失的主要特点是
它们是有状态对象：它们包含可训练参数的副本，因此
loss_module.parameters()提供训练算法所需的一切。它们遵循
TensorDict约定：torch.nn.Module.forward()方法将接收一个 TensorDict 作为输入，其中包含返回损失值所需的所有信息。>>> data = replay_buffer.sample() >>> loss_dict = loss_module(data)
它们输出一个
tensordict.TensorDict实例，其中损失值以"loss_<smth>"的形式写入，其中smth是描述损失的字符串。TensorDict中的附加键可能是在训练期间记录的有用指标。注意
我们返回独立损失的原因是让用户可以为不同的参数集使用不同的优化器。损失的求和可以通过以下方式简单完成
>>> loss_val = sum(loss for key, loss in loss_dict.items() if key.startswith("loss_"))
__init__ 方法#
所有损失的父类是 LossModule。与库中的许多其他组件一样，它的 forward() 方法期望从经验回放缓冲区或任何类似数据结构中采样的 tensordict.TensorDict 实例作为输入。使用此格式可以跨模态重用模块，或在模型需要读取多个条目的复杂设置中重用。换句话说，它允许我们编写一个损失模块，该模块不关心提供给它的数据类型，并且只专注于运行损失函数的基本步骤。
为了使教程尽可能具有教学性，我们将独立显示类的每个方法，稍后我们将填充该类。
让我们从 __init__() 方法开始。DDPG 旨在通过一个简单的策略来解决控制任务：训练一个策略来输出使值网络预测的值最大化的动作。因此，我们的损失模块需要在其构造函数中接收两个网络：一个 actor 网络和一个值网络。我们期望这两个都是 TensorDict 兼容对象，例如 tensordict.nn.TensorDictModule。我们的损失函数需要计算目标值并使值网络适应此值，并生成一个动作并使策略适应，以使其值估计最大化。
LossModule.__init__() 方法的关键步骤是调用 convert_to_functional()。此方法将从模块中提取参数并将其转换为函数模块。严格来说，这并非必要，完全可以在没有它的情况下编码所有损失。但是，我们鼓励使用它，原因如下。
TorchRL 之所以这样做，是因为强化学习算法通常使用不同的参数集（称为“可训练”参数和“目标”参数）执行相同的模型。“可训练”参数是优化器需要拟合的参数。“目标”参数通常是前者的副本，带有一些时间滞后（绝对或通过移动平均稀释）。这些目标参数用于计算与下一个观察相关联的值。使用一组与当前配置不完全匹配的值模型目标参数的优点之一是，它们对正在计算的值函数提供了悲观的界限。请注意下面的 create_target_params 关键字参数：此参数告诉 convert_to_functional() 方法在损失模块中创建一组目标参数，用于目标值计算。如果将其设置为 False（例如，参见 actor 网络），target_actor_network_params 属性仍将可访问，但这只会返回 actor 参数的分离版本。
稍后，我们将看到如何在 TorchRL 中更新目标参数。
from tensordict.nn import TensorDictModule, TensorDictSequential
def _init(
self,
actor_network: TensorDictModule,
value_network: TensorDictModule,
) -> None:
super(type(self), self).__init__()
self.convert_to_functional(
actor_network,
"actor_network",
create_target_params=True,
)
self.convert_to_functional(
value_network,
"value_network",
create_target_params=True,
compare_against=list(actor_network.parameters()),
)
self.actor_in_keys = actor_network.in_keys
# Since the value we'll be using is based on the actor and value network,
# we put them together in a single actor-critic container.
actor_critic = ActorCriticWrapper(actor_network, value_network)
self.actor_critic = actor_critic
self.loss_function = "l2"
值估计器损失方法#
在许多强化学习算法中，值网络（或 Q 值网络）是根据经验值估计进行训练的。这可以是自举（TD(0)，低方差，高偏差），这意味着目标值是使用下一个奖励获得，没有其他，或者可以获得蒙特卡洛估计（TD(1)），在这种情况下，将使用即将到来的整个奖励序列（高方差，低偏差）。也可以使用中间估计器（TD(\(\lambda\))）来折衷偏差和方差。TorchRL 通过 ValueEstimators 枚举类轻松使用其中一个或另一个估计器，该类包含指向所有已实现值估计器的指针。让我们在这里定义默认值函数。我们将采用最简单的版本（TD(0)），稍后将展示如何更改它。
from torchrl.objectives.utils import ValueEstimators
default_value_estimator = ValueEstimators.TD0
我们还需要根据用户查询向 DDPG 提供一些关于如何构建值估计器的说明。根据提供的估计器，我们将构建相应的模块以在训练时使用
from torchrl.objectives.utils import default_value_kwargs
from torchrl.objectives.value import TD0Estimator, TD1Estimator, TDLambdaEstimator
def make_value_estimator(self, value_type: ValueEstimators, **hyperparams):
hp = dict(default_value_kwargs(value_type))
if hasattr(self, "gamma"):
hp["gamma"] = self.gamma
hp.update(hyperparams)
value_key = "state_action_value"
if value_type == ValueEstimators.TD1:
self._value_estimator = TD1Estimator(value_network=self.actor_critic, **hp)
elif value_type == ValueEstimators.TD0:
self._value_estimator = TD0Estimator(value_network=self.actor_critic, **hp)
elif value_type == ValueEstimators.GAE:
raise NotImplementedError(
f"Value type {value_type} it not implemented for loss {type(self)}."
)
elif value_type == ValueEstimators.TDLambda:
self._value_estimator = TDLambdaEstimator(value_network=self.actor_critic, **hp)
else:
raise NotImplementedError(f"Unknown value type {value_type}")
self._value_estimator.set_keys(value=value_key)
make_value_estimator 方法可以但不需要调用：如果未调用，LossModule 将使用其默认估计器查询此方法。
Actor 损失方法#
强化学习算法的核心部分是 actor 的训练损失。在 DDPG 的情况下，这个函数非常简单：我们只需要计算使用策略计算的动作相关联的值，并优化 actor 权重以最大化这个值。
在计算此值时，我们必须确保将值参数从图中取出，否则 actor 和值损失将混淆。为此，可以使用 hold_out_params() 函数。
def _loss_actor(
self,
tensordict,
) -> torch.Tensor:
td_copy = tensordict.select(*self.actor_in_keys)
# Get an action from the actor network: since we made it functional, we need to pass the params
with self.actor_network_params.to_module(self.actor_network):
td_copy = self.actor_network(td_copy)
# get the value associated with that action
with self.value_network_params.detach().to_module(self.value_network):
td_copy = self.value_network(td_copy)
return -td_copy.get("state_action_value")
值损失方法#
现在我们需要优化值网络参数。为此，我们将依赖于类的值估计器
from torchrl.objectives.utils import distance_loss
def _loss_value(
self,
tensordict,
):
td_copy = tensordict.clone()
# V(s, a)
with self.value_network_params.to_module(self.value_network):
self.value_network(td_copy)
pred_val = td_copy.get("state_action_value").squeeze(-1)
# we manually reconstruct the parameters of the actor-critic, where the first
# set of parameters belongs to the actor and the second to the value function.
target_params = TensorDict(
{
"module": {
"0": self.target_actor_network_params,
"1": self.target_value_network_params,
}
},
batch_size=self.target_actor_network_params.batch_size,
device=self.target_actor_network_params.device,
)
with target_params.to_module(self.actor_critic):
target_value = self.value_estimator.value_estimate(tensordict).squeeze(-1)
# Computes the value loss: L2, L1 or smooth L1 depending on `self.loss_function`
loss_value = distance_loss(pred_val, target_value, loss_function=self.loss_function)
td_error = (pred_val - target_value).pow(2)
return loss_value, td_error, pred_val, target_value
在 forward 调用中整合所有内容#
唯一缺少的部分是 forward 方法，它将值损失和 actor 损失粘合在一起，收集成本值并将它们写入交付给用户的 TensorDict 中。
from tensordict import TensorDict, TensorDictBase
def _forward(self, input_tensordict: TensorDictBase) -> TensorDict:
loss_value, td_error, pred_val, target_value = self.loss_value(
input_tensordict,
)
td_error = td_error.detach()
td_error = td_error.unsqueeze(input_tensordict.ndimension())
if input_tensordict.device is not None:
td_error = td_error.to(input_tensordict.device)
input_tensordict.set(
"td_error",
td_error,
inplace=True,
)
loss_actor = self.loss_actor(input_tensordict)
return TensorDict(
source={
"loss_actor": loss_actor.mean(),
"loss_value": loss_value.mean(),
"pred_value": pred_val.mean().detach(),
"target_value": target_value.mean().detach(),
"pred_value_max": pred_val.max().detach(),
"target_value_max": target_value.max().detach(),
},
batch_size=[],
)
from torchrl.objectives import LossModule
class DDPGLoss(LossModule):
default_value_estimator = default_value_estimator
make_value_estimator = make_value_estimator
__init__ = _init
forward = _forward
loss_value = _loss_value
loss_actor = _loss_actor
现在我们有了损失，我们可以用它来训练一个策略来解决控制任务。
环境#
在大多数算法中，首先需要处理的是环境的构建，因为它决定了训练脚本的其余部分。
在这个例子中，我们将使用 "cheetah" 任务。目标是让半猎豹尽可能快地奔跑。
在 TorchRL 中，可以通过依赖 dm_control 或 gym 来创建这样的任务
env = GymEnv("HalfCheetah-v4")
或
env = DMControlEnv("cheetah", "run")
默认情况下，这些环境禁用渲染。从状态进行训练通常比从图像进行训练更容易。为了简单起见，我们只关注从状态学习。要将像素传递给由 env.step() 收集的 tensordicts，只需将 from_pixels=True 参数传递给构造函数
env = GymEnv("HalfCheetah-v4", from_pixels=True, pixels_only=True)
我们编写一个 make_env() 辅助函数，它将使用上述两个后端之一（dm-control 或 gym）创建环境。
from torchrl.envs.libs.dm_control import DMControlEnv
from torchrl.envs.libs.gym import GymEnv
env_library = None
env_name = None
def make_env(from_pixels=False):
"""Create a base ``env``."""
global env_library
global env_name
if backend == "dm_control":
env_name = "cheetah"
env_task = "run"
env_args = (env_name, env_task)
env_library = DMControlEnv
elif backend == "gym":
env_name = "HalfCheetah-v4"
env_args = (env_name,)
env_library = GymEnv
else:
raise NotImplementedError
env_kwargs = {
"device": device,
"from_pixels": from_pixels,
"pixels_only": from_pixels,
"frame_skip": 2,
}
env = env_library(*env_args, **env_kwargs)
return env
转换#
现在我们有了一个基本环境，我们可能希望修改其表示以使其更易于策略友好。在 TorchRL 中，转换被附加到基本环境的特殊 torchr.envs.TransformedEnv 类中。
在 DDPG 中，通常使用某种启发式值来重新缩放奖励。在此示例中，我们将奖励乘以 5。
如果我们使用
dm_control，那么在与使用双精度数字的模拟器和可能使用单精度数字的脚本之间构建一个接口也很重要。这种转换是双向的：当调用env.step()时，我们的动作需要以双精度表示，并且输出需要转换为单精度。DoubleToFloat转换正是这样做的：in_keys列表指的是需要从双精度转换为浮点数的键，而in_keys_inv指的是在传递给环境之前需要转换为双精度的键。我们使用
CatTensors转换将状态键连接在一起。最后，我们还保留了状态归一化的可能性：我们将稍后计算归一化常数。
from torchrl.envs import (
CatTensors,
DoubleToFloat,
EnvCreator,
InitTracker,
ObservationNorm,
ParallelEnv,
RewardScaling,
StepCounter,
TransformedEnv,
)
def make_transformed_env(
env,
):
"""Apply transforms to the ``env`` (such as reward scaling and state normalization)."""
env = TransformedEnv(env)
# we append transforms one by one, although we might as well create the
# transformed environment using the `env = TransformedEnv(base_env, transforms)`
# syntax.
env.append_transform(RewardScaling(loc=0.0, scale=reward_scaling))
# We concatenate all states into a single "observation_vector"
# even if there is a single tensor, it'll be renamed in "observation_vector".
# This facilitates the downstream operations as we know the name of the
# output tensor.
# In some environments (not half-cheetah), there may be more than one
# observation vector: in this case this code snippet will concatenate them
# all.
selected_keys = list(env.observation_spec.keys())
out_key = "observation_vector"
env.append_transform(CatTensors(in_keys=selected_keys, out_key=out_key))
# we normalize the states, but for now let's just instantiate a stateless
# version of the transform
env.append_transform(ObservationNorm(in_keys=[out_key], standard_normal=True))
env.append_transform(DoubleToFloat())
env.append_transform(StepCounter(max_frames_per_traj))
# We need a marker for the start of trajectories for our Ornstein-Uhlenbeck (OU)
# exploration:
env.append_transform(InitTracker())
return env
并行执行#
以下辅助函数允许我们并行运行环境。并行运行环境可以显著加快收集吞吐量。使用转换后的环境时，我们需要选择是为每个环境单独执行转换，还是集中数据并批量转换。这两种方法都很容易编码
env = ParallelEnv(
lambda: TransformedEnv(GymEnv("HalfCheetah-v4"), transforms),
num_workers=4
)
env = TransformedEnv(
ParallelEnv(lambda: GymEnv("HalfCheetah-v4"), num_workers=4),
transforms
)
为了利用 PyTorch 的向量化能力，我们采用第一种方法
def parallel_env_constructor(
env_per_collector,
transform_state_dict,
):
if env_per_collector == 1:
def make_t_env():
env = make_transformed_env(make_env())
env.transform[2].init_stats(3)
env.transform[2].loc.copy_(transform_state_dict["loc"])
env.transform[2].scale.copy_(transform_state_dict["scale"])
return env
env_creator = EnvCreator(make_t_env)
return env_creator
parallel_env = ParallelEnv(
num_workers=env_per_collector,
create_env_fn=EnvCreator(lambda: make_env()),
create_env_kwargs=None,
pin_memory=False,
)
env = make_transformed_env(parallel_env)
# we call `init_stats` for a limited number of steps, just to instantiate
# the lazy buffers.
env.transform[2].init_stats(3, cat_dim=1, reduce_dim=[0, 1])
env.transform[2].load_state_dict(transform_state_dict)
return env
# The backend can be ``gym`` or ``dm_control``
backend = "gym"
注意
frame_skip 将多个步骤与单个动作批处理在一起。如果 > 1，则其他帧计数（例如，frames_per_batch、total_frames）需要调整，以使实验收集的帧总数保持一致。这很重要，因为提高帧跳数但保持总帧数不变可能看起来像作弊：总而言之，一个包含 10M 元素的数据集，其帧跳数为 2，另一个帧跳数为 1，实际上与环境的交互比率为 2:1！简而言之，在处理帧跳数时，应谨慎对待训练脚本的帧计数，因为这可能导致训练策略之间的比较存在偏差。
缩放奖励有助于我们控制信号幅度，以实现更有效的学习。
reward_scaling = 5.0
我们还定义了轨迹何时截断。一千步（如果帧跳数 = 2，则为 500）是用于猎豹任务的不错数字
max_frames_per_traj = 500
观察值的归一化#
为了计算归一化统计量，我们在环境中运行任意数量的随机步骤，并计算所收集观察量的均值和标准差。ObservationNorm.init_stats() 方法可用于此目的。为了获取摘要统计量，我们创建一个虚拟环境，运行一定数量的步骤，收集一定数量的步骤的数据并计算其摘要统计量。
def get_env_stats():
"""Gets the stats of an environment."""
proof_env = make_transformed_env(make_env())
t = proof_env.transform[2]
t.init_stats(init_env_steps)
transform_state_dict = t.state_dict()
proof_env.close()
return transform_state_dict
归一化统计#
使用 ObservationNorm 计算统计数据所用的随机步数
init_env_steps = 5000
transform_state_dict = get_env_stats()
Gym has been unmaintained since 2022 and does not support NumPy 2.0 amongst other critical functionality.
Please upgrade to Gymnasium, the maintained drop-in replacement of Gym, or contact the authors of your software and request that they upgrade.
Users of this version of Gym should be able to simply replace 'import gym' with 'import gymnasium as gym' in the vast majority of cases.
See the migration guide at https://gymnasium.org.cn/introduction/migration_guide/ for additional information.
每个数据收集器中的环境数量
env_per_collector = 4
我们将之前计算的统计数据传递给环境输出进行归一化
parallel_env = parallel_env_constructor(
env_per_collector=env_per_collector,
transform_state_dict=transform_state_dict,
)
from torchrl.data import CompositeSpec
构建模型#
现在我们转向模型的设置。正如我们所见，DDPG 需要一个值网络（训练用于估计状态-动作对的值）和一个参数化 actor（学习如何选择最大化该值的动作）。
回想一下，构建 TorchRL 模块需要两个步骤
编写将用作网络的
torch.nn.Module，通过指定输入和输出键，将网络包装在
tensordict.nn.TensorDictModule中，以处理数据流。
在更复杂的场景中，也可以使用 tensordict.nn.TensorDictSequential。
Q 值网络包装在 ValueOperator 中，该操作器自动将 Q 值网络的 out_keys 设置为 "state_action_value"，将其他值网络的 out_keys 设置为 state_value。
TorchRL 提供了原始论文中提出的 DDPG 网络的内置版本。这些可以在 DdpgMlpActor 和 DdpgMlpQNet 下找到。
由于我们使用惰性模块，因此在能够将策略从一个设备移动到另一个设备并执行其他操作之前，有必要实现惰性模块。因此，使用少量数据运行模块是一个很好的做法。为此，我们从环境规范生成假数据。
from torchrl.modules import (
ActorCriticWrapper,
DdpgMlpActor,
DdpgMlpQNet,
OrnsteinUhlenbeckProcessModule,
ProbabilisticActor,
TanhDelta,
ValueOperator,
)
def make_ddpg_actor(
transform_state_dict,
device="cpu",
):
proof_environment = make_transformed_env(make_env())
proof_environment.transform[2].init_stats(3)
proof_environment.transform[2].load_state_dict(transform_state_dict)
out_features = proof_environment.action_spec.shape[-1]
actor_net = DdpgMlpActor(
action_dim=out_features,
)
in_keys = ["observation_vector"]
out_keys = ["param"]
actor = TensorDictModule(
actor_net,
in_keys=in_keys,
out_keys=out_keys,
)
actor = ProbabilisticActor(
actor,
distribution_class=TanhDelta,
in_keys=["param"],
spec=CompositeSpec(action=proof_environment.action_spec),
).to(device)
q_net = DdpgMlpQNet()
in_keys = in_keys + ["action"]
qnet = ValueOperator(
in_keys=in_keys,
module=q_net,
).to(device)
# initialize lazy modules
qnet(actor(proof_environment.reset().to(device)))
return actor, qnet
actor, qnet = make_ddpg_actor(
transform_state_dict=transform_state_dict,
device=device,
)
/usr/local/lib/python3.10/dist-packages/torchrl/data/tensor_specs.py:6911: DeprecationWarning:
The CompositeSpec has been deprecated and will be removed in v0.8. Please use Composite instead.
探索#
根据原始论文的建议，策略被传递到 OrnsteinUhlenbeckProcessModule 探索模块中。让我们定义 OU 噪声达到最小值的帧数
annealing_frames = 1_000_000
actor_model_explore = TensorDictSequential(
actor,
OrnsteinUhlenbeckProcessModule(
spec=actor.spec.clone(),
annealing_num_steps=annealing_frames,
).to(device),
)
if device == torch.device("cpu"):
actor_model_explore.share_memory()
数据收集器#
TorchRL 提供了专门的类来帮助您通过在环境中执行策略来收集数据。这些“数据收集器”迭代地计算在给定时间执行的动作，然后执行环境中的一步并在需要时重置。数据收集器旨在帮助开发人员严格控制每批数据的帧数、此收集的（异步）性质以及分配给数据收集的资源（例如 GPU、工作器数量等）。
在这里，我们将使用 SyncDataCollector，一个简单的单进程数据收集器。TorchRL 提供其他收集器，例如 MultiaSyncDataCollector，它以异步方式执行 rollout（例如，在策略优化时将收集数据，从而解耦训练和数据收集）。
需要指定的参数是
环境工厂或环境，
策略，
收集器被视为空之前的总帧数，
每个轨迹的最大帧数（对于非终止环境很有用，例如
dm_control环境）。注意
传递给收集器的
max_frames_per_traj将会注册一个新的StepCounter转换到用于推理的环境中。我们可以在此脚本中手动实现相同的效果。
还应该传递
每批收集的帧数，
与策略无关执行的随机步数，
用于策略执行的设备
用于在数据传递给主进程之前存储数据的设备。
训练期间将使用的总帧数应在 1M 左右。
total_frames = 10_000 # 1_000_000
在外部循环的每次迭代中，收集器返回的帧数等于每个子轨迹的长度乘以每个收集器中并行运行的环境数量。
换句话说，我们期望收集器中的批次具有形状 [env_per_collector, traj_len]，其中 traj_len=frames_per_batch/env_per_collector
traj_len = 200
frames_per_batch = env_per_collector * traj_len
init_random_frames = 5000
num_collectors = 2
from torchrl.collectors import SyncDataCollector
from torchrl.envs import ExplorationType
collector = SyncDataCollector(
parallel_env,
policy=actor_model_explore,
total_frames=total_frames,
frames_per_batch=frames_per_batch,
init_random_frames=init_random_frames,
reset_at_each_iter=False,
split_trajs=False,
device=collector_device,
exploration_type=ExplorationType.RANDOM,
)
评估器：构建您的记录器对象#
由于训练数据是通过某种探索策略获得的，因此我们算法的真实性能需要在确定性模式下进行评估。我们通过专用类 Recorder 来实现这一点，该类以给定频率在环境中执行策略并返回从这些模拟中获得的一些统计数据。
以下辅助函数构建此对象
from torchrl.trainers import Recorder
def make_recorder(actor_model_explore, transform_state_dict, record_interval):
base_env = make_env()
environment = make_transformed_env(base_env)
environment.transform[2].init_stats(
3
) # must be instantiated to load the state dict
environment.transform[2].load_state_dict(transform_state_dict)
recorder_obj = Recorder(
record_frames=1000,
policy_exploration=actor_model_explore,
environment=environment,
exploration_type=ExplorationType.DETERMINISTIC,
record_interval=record_interval,
)
return recorder_obj
我们将每收集 10 个批次记录一次性能
record_interval = 10
recorder = make_recorder(
actor_model_explore, transform_state_dict, record_interval=record_interval
)
from torchrl.data.replay_buffers import (
LazyMemmapStorage,
PrioritizedSampler,
RandomSampler,
TensorDictReplayBuffer,
)
回放缓冲区#
回放缓冲区有两种类型：优先级回放（其中使用一些错误信号来提高某些项目的采样可能性）和常规的循环经验回放。
TorchRL 回放缓冲区是可组合的：可以选择存储、采样和写入策略。也可以使用内存映射数组将张量存储在物理内存中。以下函数负责使用所需的超参数创建回放缓冲区
from torchrl.envs import RandomCropTensorDict
def make_replay_buffer(buffer_size, batch_size, random_crop_len, prefetch=3, prb=False):
if prb:
sampler = PrioritizedSampler(
max_capacity=buffer_size,
alpha=0.7,
beta=0.5,
)
else:
sampler = RandomSampler()
replay_buffer = TensorDictReplayBuffer(
storage=LazyMemmapStorage(
buffer_size,
scratch_dir=buffer_scratch_dir,
),
batch_size=batch_size,
sampler=sampler,
pin_memory=False,
prefetch=prefetch,
transform=RandomCropTensorDict(random_crop_len, sample_dim=1),
)
return replay_buffer
我们将把回放缓冲区存储在磁盘上的临时目录中
import tempfile
tmpdir = tempfile.TemporaryDirectory()
buffer_scratch_dir = tmpdir.name
回放缓冲区存储和批大小#
TorchRL 回放缓冲区计算沿第一个维度的元素数量。由于我们将向缓冲区提供轨迹，因此我们需要通过将其除以数据收集器生成的子轨迹长度来调整缓冲区大小。关于批大小，我们的采样策略将包括采样长度为 traj_len=200 的轨迹，然后选择长度为 random_crop_len=25 的子轨迹来计算损失。此策略平衡了存储一定长度的完整轨迹的选择与为我们的损失提供足够异构样本的需求。下图显示了数据流，从一个收集器开始，该收集器在每个批次中获取 8 帧，并行运行 2 个环境，将它们馈送到包含 1000 个轨迹的回放缓冲区，并采样每个轨迹 2 个时间步长的子轨迹。
让我们从缓冲区中存储的帧数开始
def ceil_div(x, y):
return -x // (-y)
buffer_size = 1_000_000
buffer_size = ceil_div(buffer_size, traj_len)
优先级回放缓冲区默认禁用
prb = False
我们还需要定义每批收集的数据将执行多少次更新。这被称为更新-数据比或 UTD 比率
update_to_data = 64
我们将以 25 的长度将轨迹馈送给损失函数
random_crop_len = 25
在原始论文中，作者对每收集一帧数据，都用一个包含 64 个元素的批次进行一次更新。在这里，我们重现相同的比例，但在每次批次收集时实现多次更新。我们调整批次大小以达到相同的每帧更新次数比率
batch_size = ceil_div(64 * frames_per_batch, update_to_data * random_crop_len)
replay_buffer = make_replay_buffer(
buffer_size=buffer_size,
batch_size=batch_size,
random_crop_len=random_crop_len,
prefetch=3,
prb=prb,
)
损失模块构建#
我们使用刚刚创建的 actor 和 qnet 来构建我们的损失模块。因为我们有目标参数需要更新，所以我们_必须_创建一个目标网络更新器。
让我们使用 TD(lambda) 估计器！
loss_module.make_value_estimator(ValueEstimators.TDLambda, gamma=gamma, lmbda=lmbda, device=device)
注意
离策略通常要求使用 TD(0) 估计器。在这里，我们使用 TD(\(\lambda\)) 估计器，这会引入一些偏差，因为某个状态之后的轨迹是使用过时的策略收集的。这个技巧，以及数据收集期间可以使用的多步技巧，都是“黑客”的替代版本，尽管它们在回报估计中引入了一些偏差，但我们通常发现它们在实践中效果很好。
目标网络更新器#
目标网络是离策略强化学习算法的关键部分。借助 HardUpdate 和 SoftUpdate 类，更新目标网络参数变得很容易。它们以损失模块作为参数构建，并通过在训练循环中的适当位置调用 updater.step() 来实现更新。
from torchrl.objectives.utils import SoftUpdate
target_net_updater = SoftUpdate(loss_module, eps=1 - tau)
优化器#
最后，我们将为策略和值网络使用 Adam 优化器
from torch import optim
optimizer_actor = optim.Adam(
loss_module.actor_network_params.values(True, True), lr=1e-4, weight_decay=0.0
)
optimizer_value = optim.Adam(
loss_module.value_network_params.values(True, True), lr=1e-3, weight_decay=1e-2
)
total_collection_steps = total_frames // frames_per_batch
训练策略的时间#
现在我们已经构建了所有需要的模块，训练循环非常简单。
rewards = []
rewards_eval = []
# Main loop
collected_frames = 0
pbar = tqdm.tqdm(total=total_frames)
r0 = None
for i, tensordict in enumerate(collector):
# update weights of the inference policy
collector.update_policy_weights_()
if r0 is None:
r0 = tensordict["next", "reward"].mean().item()
pbar.update(tensordict.numel())
# extend the replay buffer with the new data
current_frames = tensordict.numel()
collected_frames += current_frames
replay_buffer.extend(tensordict.cpu())
# optimization steps
if collected_frames >= init_random_frames:
for _ in range(update_to_data):
# sample from replay buffer
sampled_tensordict = replay_buffer.sample().to(device)
# Compute loss
loss_dict = loss_module(sampled_tensordict)
# optimize
loss_dict["loss_actor"].backward()
gn1 = torch.nn.utils.clip_grad_norm_(
loss_module.actor_network_params.values(True, True), 10.0
)
optimizer_actor.step()
optimizer_actor.zero_grad()
loss_dict["loss_value"].backward()
gn2 = torch.nn.utils.clip_grad_norm_(
loss_module.value_network_params.values(True, True), 10.0
)
optimizer_value.step()
optimizer_value.zero_grad()
gn = (gn1**2 + gn2**2) ** 0.5
# update priority
if prb:
replay_buffer.update_tensordict_priority(sampled_tensordict)
# update target network
target_net_updater.step()
rewards.append(
(
i,
tensordict["next", "reward"].mean().item(),
)
)
td_record = recorder(None)
if td_record is not None:
rewards_eval.append((i, td_record["r_evaluation"].item()))
if len(rewards_eval) and collected_frames >= init_random_frames:
target_value = loss_dict["target_value"].item()
loss_value = loss_dict["loss_value"].item()
loss_actor = loss_dict["loss_actor"].item()
rn = sampled_tensordict["next", "reward"].mean().item()
rs = sampled_tensordict["next", "reward"].std().item()
pbar.set_description(
f"reward: {rewards[-1][1]: 4.2f} (r0 = {r0: 4.2f}), "
f"reward eval: reward: {rewards_eval[-1][1]: 4.2f}, "
f"reward normalized={rn :4.2f}/{rs :4.2f}, "
f"grad norm={gn: 4.2f}, "
f"loss_value={loss_value: 4.2f}, "
f"loss_actor={loss_actor: 4.2f}, "
f"target value: {target_value: 4.2f}"
)
# update the exploration strategy
actor_model_explore[1].step(current_frames)
collector.shutdown()
del collector
0%| | 0/10000 [00:00<?, ?it/s]
8%|▊ | 800/10000 [00:00<00:06, 1375.69it/s]
16%|█▌ | 1600/10000 [00:02<00:15, 533.87it/s]
24%|██▍ | 2400/10000 [00:03<00:09, 789.72it/s]
32%|███▏ | 3200/10000 [00:03<00:06, 1022.64it/s]
40%|████ | 4000/10000 [00:04<00:04, 1222.45it/s]
48%|████▊ | 4800/10000 [00:04<00:03, 1387.16it/s]
56%|█████▌ | 5600/10000 [00:04<00:02, 1511.49it/s]
reward: -2.15 (r0 = -1.90), reward eval: reward: -0.01, reward normalized=-2.38/5.99, grad norm= 131.77, loss_value= 305.84, loss_actor= 13.20, target value: -15.17: 56%|█████▌ | 5600/10000 [00:06<00:02, 1511.49it/s]
reward: -2.15 (r0 = -1.90), reward eval: reward: -0.01, reward normalized=-2.38/5.99, grad norm= 131.77, loss_value= 305.84, loss_actor= 13.20, target value: -15.17: 64%|██████▍ | 6400/10000 [00:07<00:05, 703.62it/s]
reward: -3.00 (r0 = -1.90), reward eval: reward: -0.01, reward normalized=-2.58/5.08, grad norm= 93.33, loss_value= 153.55, loss_actor= 15.74, target value: -17.39: 64%|██████▍ | 6400/10000 [00:09<00:05, 703.62it/s]
reward: -3.00 (r0 = -1.90), reward eval: reward: -0.01, reward normalized=-2.58/5.08, grad norm= 93.33, loss_value= 153.55, loss_actor= 15.74, target value: -17.39: 72%|███████▏ | 7200/10000 [00:09<00:05, 521.67it/s]
reward: -1.65 (r0 = -1.90), reward eval: reward: -0.01, reward normalized=-2.10/5.24, grad norm= 77.63, loss_value= 146.87, loss_actor= 14.82, target value: -14.25: 72%|███████▏ | 7200/10000 [00:11<00:05, 521.67it/s]
reward: -1.65 (r0 = -1.90), reward eval: reward: -0.01, reward normalized=-2.10/5.24, grad norm= 77.63, loss_value= 146.87, loss_actor= 14.82, target value: -14.25: 80%|████████ | 8000/10000 [00:12<00:04, 444.45it/s]
reward: -3.42 (r0 = -1.90), reward eval: reward: -0.01, reward normalized=-2.62/4.53, grad norm= 36.96, loss_value= 130.61, loss_actor= 18.17, target value: -17.53: 80%|████████ | 8000/10000 [00:13<00:04, 444.45it/s]
reward: -3.42 (r0 = -1.90), reward eval: reward: -0.01, reward normalized=-2.62/4.53, grad norm= 36.96, loss_value= 130.61, loss_actor= 18.17, target value: -17.53: 88%|████████▊ | 8800/10000 [00:14<00:02, 403.19it/s]
reward: -3.93 (r0 = -1.90), reward eval: reward: -4.86, reward normalized=-2.38/4.69, grad norm= 81.39, loss_value= 112.50, loss_actor= 16.11, target value: -16.39: 88%|████████▊ | 8800/10000 [00:17<00:02, 403.19it/s]
reward: -3.93 (r0 = -1.90), reward eval: reward: -4.86, reward normalized=-2.38/4.69, grad norm= 81.39, loss_value= 112.50, loss_actor= 16.11, target value: -16.39: 96%|█████████▌| 9600/10000 [00:18<00:01, 306.87it/s]
reward: -4.76 (r0 = -1.90), reward eval: reward: -4.86, reward normalized=-2.48/5.09, grad norm= 83.59, loss_value= 169.57, loss_actor= 16.78, target value: -17.10: 96%|█████████▌| 9600/10000 [00:20<00:01, 306.87it/s]
reward: -4.76 (r0 = -1.90), reward eval: reward: -4.86, reward normalized=-2.48/5.09, grad norm= 83.59, loss_value= 169.57, loss_actor= 16.78, target value: -17.10: : 10400it [00:21, 306.34it/s]
reward: -3.95 (r0 = -1.90), reward eval: reward: -4.86, reward normalized=-2.07/4.72, grad norm= 139.60, loss_value= 159.43, loss_actor= 15.56, target value: -16.38: : 10400it [00:23, 306.34it/s]
实验结果#
我们绘制了一个简单的训练期间平均奖励图。我们可以观察到我们的策略很好地学会了解决任务。
注意
如上所述，为了获得更合理的性能，请将 total_frames 设置为更大的值，例如 1M。
from matplotlib import pyplot as plt
plt.figure()
plt.plot(*zip(*rewards), label="training")
plt.plot(*zip(*rewards_eval), label="eval")
plt.legend()
plt.xlabel("iter")
plt.ylabel("reward")
plt.tight_layout()
结论#
在本教程中，我们学习了如何以 DDPG 的具体示例为例，在 TorchRL 中编码损失模块。
主要收获是
如何使用
LossModule类编码新的损失组件；如何使用（或不使用）目标网络，以及如何更新其参数；
如何创建与损失模块关联的优化器。
后续步骤#
为了进一步迭代这个损失模块，我们可以考虑
使用 @dispatch （参见 [功能] 分发 IQL 损失模块。）
允许灵活的 TensorDict 键。
脚本总运行时间：（0 分 29.511 秒）
