注意
跳转至页面底部 下载完整示例代码。
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（例如对于行动者网络），target_actor_network_params 属性仍然可以访问，但这将只返回行动者参数的分离（detached）版本。
稍后我们将看到在 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
我们还需要根据用户需求给出关于如何构建价值估计器的说明。根据提供的估计器，我们将构建相应的模块以在训练时使用。
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 将使用其默认估计器查询此方法。
行动者损失方法#
强化学习算法的核心部分是行动者的训练损失。对于 DDPG，此函数非常简单：我们只需要计算与使用策略计算出的动作相关的价值，并优化行动者权重以最大化此价值。
计算此价值时，必须确保将价值参数从计算图中移除，否则行动者损失和价值损失会混在一起。为此，可以使用 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 方法，它将整合价值损失和行动者损失，收集成本值，并将其写入交付给用户的 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" 任务。目标是让半猎豹（half-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 中，变换（transforms）被附加到专门的 torchrl.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），以确保跨实验收集到的总帧数一致。这一点很重要，因为提高 frame-skip 但保持总帧数不变可能看起来像是在作弊：对比来看，一个使用 frame-skip 为 2 收集的 1000 万个元素的数据集，与另一个使用 frame-skip 为 1 的数据集，实际上与环境交互的比例是 2:1！简而言之，在处理跳帧时，应谨慎对待训练脚本的帧计数，因为这可能导致训练策略之间的比较出现偏差。
缩放奖励有助于我们控制信号幅度，从而实现更高效的学习。
reward_scaling = 5.0
我们还定义了何时截断轨迹。对于 cheetah 任务，一千步（如果 frame-skip = 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 需要一个训练用于估计状态-动作对价值的价值网络，以及一个学习如何选择最大化该价值的动作的参数化行动者。
回想一下，构建 TorchRL 模块需要两个步骤：
编写将用作网络的
torch.nn.Module；将网络包装在
tensordict.nn.TensorDictModule中，通过指定输入和输出键来处理数据流。
在更复杂的场景中，还可以使用 tensordict.nn.TensorDictSequential。
Q-Value 网络被包装在 ValueOperator 中，它会自动将 out_keys 设置为 Q-Value 网络的 "state_action_value" 和其他价值网络的 "state_value"。
TorchRL 提供了原始论文中展示的 DDPG 网络的内置版本。这些可以在 DdpgMlpActor 和 DdpgMlpQNet 下找到。
由于我们使用惰性（lazy）模块，因此在能够将策略从一个设备移动到另一个设备并执行其他操作之前，有必要实例化这些惰性模块。因此，用一小段数据运行模块是一种好习惯。为此，我们从环境规格生成模拟数据。
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:7085: DeprecationWarning: The CompositeSpec has been deprecated and will be removed in v0.8. Please use Composite instead.
warnings.warn(
探索#
该策略被传入 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，它以异步方式执行回放（例如，数据收集将在策略优化时进行，从而解耦训练和数据收集）。
需要指定的参数有：
环境工厂或环境；
策略；
在收集器被视为空之前的总帧数；
每个轨迹的最大帧数（对非终止环境很有用，例如
dm_control环境）。注意
传递给收集器的
max_frames_per_traj将起到为推理环境注册一个新的StepCounter变换的作用。我们可以像在此脚本中所做的那样手动实现相同的结果。
还应传递：
每个收集批次中的帧数；
独立于策略执行的随机步数；
用于策略执行的设备；
在数据传递给主进程之前用于存储数据的设备。
我们在训练期间使用的总帧数应在 100 万左右。
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,
)
回放缓冲区#
回放缓冲区有两种类型：优先经验回放（Prioritized Replay，其中使用某种误差信号来赋予某些项目比其他项目更高的采样概率）和常规的循环经验回放。
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
我们还需要定义每收集一批数据进行多少次更新。这被称为更新数据比（update-to-data, UTD ratio）。
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,
)
损失模块构建#
我们使用刚刚创建的行动者和 qnet 构建我们的损失模块。因为我们有需要更新的目标参数，所以_必须_创建一个目标网络更新器。
让我们使用 TD(lambda) 估计器！
loss_module.make_value_estimator(ValueEstimators.TDLambda, gamma=gamma, lmbda=lmbda, device=device)
注意
离线策略（Off-policy）算法通常使用 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, 1487.50it/s]
16%|█▌ | 1600/10000 [00:02<00:15, 535.18it/s]
24%|██▍ | 2400/10000 [00:03<00:09, 812.26it/s]
32%|███▏ | 3200/10000 [00:03<00:06, 1071.46it/s]
40%|████ | 4000/10000 [00:03<00:04, 1300.47it/s]
48%|████▊ | 4800/10000 [00:04<00:03, 1492.71it/s]
56%|█████▌ | 5600/10000 [00:04<00:02, 1577.52it/s]
reward: -2.96 (r0 = -2.07), reward eval: reward: 0.00, reward normalized=-3.03/6.02, grad norm= 298.69, loss_value= 344.81, loss_actor= 14.83, target value: -20.17: 56%|█████▌ | 5600/10000 [00:06<00:02, 1577.52it/s]
reward: -2.96 (r0 = -2.07), reward eval: reward: 0.00, reward normalized=-3.03/6.02, grad norm= 298.69, loss_value= 344.81, loss_actor= 14.83, target value: -20.17: 64%|██████▍ | 6400/10000 [00:07<00:05, 705.70it/s]
reward: -1.90 (r0 = -2.07), reward eval: reward: 0.00, reward normalized=-1.73/5.72, grad norm= 51.33, loss_value= 237.76, loss_actor= 10.71, target value: -11.86: 64%|██████▍ | 6400/10000 [00:08<00:05, 705.70it/s]
reward: -1.90 (r0 = -2.07), reward eval: reward: 0.00, reward normalized=-1.73/5.72, grad norm= 51.33, loss_value= 237.76, loss_actor= 10.71, target value: -11.86: 72%|███████▏ | 7200/10000 [00:09<00:05, 517.92it/s]
reward: -4.37 (r0 = -2.07), reward eval: reward: 0.00, reward normalized=-1.80/5.27, grad norm= 73.07, loss_value= 166.02, loss_actor= 11.57, target value: -11.65: 72%|███████▏ | 7200/10000 [00:11<00:05, 517.92it/s]
reward: -4.37 (r0 = -2.07), reward eval: reward: 0.00, reward normalized=-1.80/5.27, grad norm= 73.07, loss_value= 166.02, loss_actor= 11.57, target value: -11.65: 80%|████████ | 8000/10000 [00:12<00:04, 441.49it/s]
reward: -4.33 (r0 = -2.07), reward eval: reward: 0.00, reward normalized=-2.39/5.37, grad norm= 203.51, loss_value= 197.36, loss_actor= 15.75, target value: -16.16: 80%|████████ | 8000/10000 [00:13<00:04, 441.49it/s]
reward: -4.33 (r0 = -2.07), reward eval: reward: 0.00, reward normalized=-2.39/5.37, grad norm= 203.51, loss_value= 197.36, loss_actor= 15.75, target value: -16.16: 88%|████████▊ | 8800/10000 [00:14<00:02, 402.09it/s]
reward: -4.94 (r0 = -2.07), reward eval: reward: -5.91, reward normalized=-2.81/5.23, grad norm= 119.57, loss_value= 243.72, loss_actor= 21.55, target value: -20.66: 88%|████████▊ | 8800/10000 [00:17<00:02, 402.09it/s]
reward: -4.94 (r0 = -2.07), reward eval: reward: -5.91, reward normalized=-2.81/5.23, grad norm= 119.57, loss_value= 243.72, loss_actor= 21.55, target value: -20.66: 96%|█████████▌| 9600/10000 [00:18<00:01, 304.07it/s]
reward: -5.04 (r0 = -2.07), reward eval: reward: -5.91, reward normalized=-2.24/4.86, grad norm= 249.82, loss_value= 161.07, loss_actor= 19.03, target value: -16.19: 96%|█████████▌| 9600/10000 [00:20<00:01, 304.07it/s]
reward: -5.04 (r0 = -2.07), reward eval: reward: -5.91, reward normalized=-2.24/4.86, grad norm= 249.82, loss_value= 161.07, loss_actor= 19.03, target value: -16.19: : 10400it [00:21, 301.70it/s]
reward: -3.74 (r0 = -2.07), reward eval: reward: -5.91, reward normalized=-3.11/4.37, grad norm= 139.91, loss_value= 139.23, loss_actor= 21.06, target value: -21.19: : 10400it [00:23, 301.70it/s]
实验结果#
我们制作了一个训练期间平均奖励的简单绘图。我们可以观察到我们的策略很好地学会了解决该任务。
注意
正如上面已经提到的，要获得更合理的性能，请使用更大的 total_frames 值，例如 100 万。
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()
结论#
在本教程中，我们学习了如何在 TorchRL 中通过 DDPG 的具体示例来编写损失模块。
关键要点是：
如何使用
LossModule类编写一个新的损失组件；如何使用（或不使用）目标网络，以及如何更新其参数；
如何创建与损失模块关联的优化器。
下一步#
要进一步迭代此损失模块，我们可以考虑：
使用 @dispatch（参见 [Feature] Distpatch IQL loss module。）
允许灵活的 TensorDict 键。
脚本运行总时间：（0 分 29.552 秒）