注意

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

TorchRL 目标：编码 DDPG 损失¶

概述¶

TorchRL 将 RL 算法的训练分成几个部分，这些部分将在您的训练脚本中组装起来：环境、数据收集和存储、模型，最后是损失函数。

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_"))
```

The `init` method¶

所有损失的父类是 LossModule。与库中的许多其他组件一样，它的 forward() 方法期望作为输入接收一个从经验回放缓冲区采样的 tensordict.TensorDict 实例，或任何类似的数据结构。使用这种格式可以在跨模式或模型需要读取多个条目（例如）的复杂设置中重复使用模块。换句话说，它使我们能够编写一个对给定数据类型不敏感的损失模块，该模块专注于运行损失函数的基本步骤，仅此而已。

为了使教程尽可能具有指导意义，我们将独立显示类的每个方法，并在以后阶段填充类。

让我们从 __init__() 方法开始。DDPG 旨在使用一种简单的策略解决控制任务：训练一个策略来输出由价值网络预测的价值最大化的动作。因此，我们的损失模块需要在其构造函数中接收两个网络：一个演员网络和一个价值网络。我们期望这两个网络都是与 TensorDict 兼容的对象，例如 tensordict.nn.TensorDictModule。我们的损失函数需要计算目标值并使价值网络拟合该值，以及生成动作并使策略拟合，以便其价值估计值最大化。

The LossModule.__init__() method 的关键步骤是调用 convert_to_functional()。此方法将从模块中提取参数并将其转换为功能模块。严格来说，这并非必需，即使没有它，也可以完美地编写所有损失。但是，我们鼓励使用它，原因如下。

TorchRL 这样做是因为 RL 算法通常使用不同的参数集执行相同的模型，称为“可训练”参数和“目标”参数。“可训练”参数是优化器需要拟合的参数。“目标”参数通常是前者的副本，但存在一定的时滞（绝对或通过移动平均值稀释）。这些目标参数用于计算与下一个观察相关的价值。使用与当前配置不完全匹配的价值模型目标参数集的优点之一是，它们为计算的价值函数提供了悲观的界限。请注意以下的 create_target_params 关键字参数：此参数告诉 convert_to_functional() 方法在损失模块中创建一组目标参数，用于目标值计算。如果将其设置为 False（例如，参见演员网络），则 target_actor_network_params 属性仍然可以访问，但它将只返回演员参数的分离版本。

稍后，我们将看到如何在 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"

The value estimator loss method¶

在许多 RL 算法中，价值网络（或 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)

The make_value_estimator method 可以调用，也可以不调用：如果不调用，则 LossModule 将使用其默认估计器查询此方法。

The actor loss method¶

RL 算法的核心部分是演员的训练损失。在 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")

The value loss method¶

现在我们需要优化我们的价值网络参数。为此，我们将依靠我们类的价值估计器

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

Putting things together in a forward call¶

唯一缺少的部分是 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" 任务。目标是使一只半猎豹尽可能快地奔跑。

在 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() 时，我们的动作需要以双精度数表示，并且输出需要转换为单精度数。The DoubleToFloat transform 正是执行此操作：in_keys 列表指的是需要从双精度数转换为单精度数的键，而 in_keys_inv 指的是需要在传递给环境之前转换为双精度数的键。
我们使用 CatTensors transform 将状态键连接在一起。
最后，我们还提供了对状态进行归一化的可能性：我们将在稍后处理归一化常数的计算。

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）需要进行调整，以确保在所有实验中收集的总帧数一致。这一点很重要，因为提高帧跳跃但保持总帧数不变可能看起来像是作弊：在所有条件相同的情况下，使用帧跳跃为 2 收集的 1000 万个元素的数据集与帧跳跃为 1 的数据集的实际环境交互比率为 2:1！简而言之，在处理帧跳跃时，应该注意训练脚本的帧计数，因为这可能导致训练策略之间存在偏差比较。

缩放奖励有助于我们控制信号幅度，以实现更有效的学习。

reward_scaling = 5.0

我们还定义了轨迹何时会被截断。对于猎豹任务，1000 个步骤（如果帧跳跃 = 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()

每个数据收集器中的环境数量

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-值网络被封装在一个 ValueOperator 中，该 ValueOperator 会自动将 out_keys 设置为 Q 值网络的 "state_action_value" 和其他价值网络的 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,
)

探索¶

策略被传递给一个 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 提供了专门的类，可以帮助您通过在环境中执行策略来收集数据。这些“数据收集器”会迭代计算在给定时间要执行的动作，然后在环境中执行一步，并在需要时重置它。数据收集器旨在帮助开发人员严格控制每批数据的帧数，以及此收集的（a）同步性质以及分配给数据收集的资源（例如 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.MEAN,
        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 个时间步长的子轨迹。

Storing trajectories in the replay buffer

让我们从存储在缓冲区中的帧数开始

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,
)

损失模块构建¶

我们使用刚刚创建的演员和 qnet 构建损失模块。由于我们有目标参数要更新，因此_必须_创建一个目标网络更新器。

gamma = 0.99
lmbda = 0.9
tau = 0.001  # Decay factor for the target network

loss_module = DDPGLoss(actor, qnet)

让我们使用 TD(lambda) 估计器！

loss_module.make_value_estimator(ValueEstimators.TDLambda, gamma=gamma, lmbda=lmbda)

注意

离策略通常决定 TD(0) 估计器。在这里，我们使用 TD(\(\lambda\)) 估计器，它将引入一些偏差，因为遵循特定状态的轨迹是使用过时的策略收集的。这种技巧，以及可以在数据收集过程中使用的多步技巧，是我们通常发现有效的“黑客”的替代版本，尽管它们在收益估计中引入了一些偏差。

目标网络更新器¶

目标网络是离策略 RL 算法的关键部分。使用 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%|8         | 800/10000 [00:00<00:05, 1681.77it/s]
 16%|#6        | 1600/10000 [00:02<00:16, 496.76it/s]
 24%|##4       | 2400/10000 [00:03<00:10, 705.56it/s]
 32%|###2      | 3200/10000 [00:03<00:07, 895.76it/s]
 40%|####      | 4000/10000 [00:04<00:05, 1049.99it/s]
 48%|####8     | 4800/10000 [00:05<00:04, 1175.99it/s]
 56%|#####6    | 5600/10000 [00:05<00:03, 1272.51it/s]
reward: -2.24 (r0 = -3.03), reward eval: reward: -0.01, reward normalized=-1.51/6.64, grad norm= 103.61, loss_value= 520.37, loss_actor= 11.98, target value: -10.00:  56%|#####6    | 5600/10000 [00:06<00:03, 1272.51it/s]
reward: -2.24 (r0 = -3.03), reward eval: reward: -0.01, reward normalized=-1.51/6.64, grad norm= 103.61, loss_value= 520.37, loss_actor= 11.98, target value: -10.00:  64%|######4   | 6400/10000 [00:07<00:04, 882.34it/s]
reward: -0.14 (r0 = -3.03), reward eval: reward: -0.01, reward normalized=-2.27/5.94, grad norm= 124.42, loss_value= 269.42, loss_actor= 10.56, target value: -14.61:  64%|######4   | 6400/10000 [00:07<00:04, 882.34it/s]
reward: -0.14 (r0 = -3.03), reward eval: reward: -0.01, reward normalized=-2.27/5.94, grad norm= 124.42, loss_value= 269.42, loss_actor= 10.56, target value: -14.61:  72%|#######2  | 7200/10000 [00:08<00:03, 716.17it/s]
reward: -3.72 (r0 = -3.03), reward eval: reward: -0.01, reward normalized=-1.87/5.92, grad norm= 77.73, loss_value= 290.63, loss_actor= 12.12, target value: -12.06:  72%|#######2  | 7200/10000 [00:09<00:03, 716.17it/s]
reward: -3.72 (r0 = -3.03), reward eval: reward: -0.01, reward normalized=-1.87/5.92, grad norm= 77.73, loss_value= 290.63, loss_actor= 12.12, target value: -12.06:  80%|########  | 8000/10000 [00:10<00:03, 635.41it/s]
reward: -5.02 (r0 = -3.03), reward eval: reward: -0.01, reward normalized=-2.53/5.25, grad norm= 54.73, loss_value= 260.91, loss_actor= 19.76, target value: -16.34:  80%|########  | 8000/10000 [00:11<00:03, 635.41it/s]
reward: -5.02 (r0 = -3.03), reward eval: reward: -0.01, reward normalized=-2.53/5.25, grad norm= 54.73, loss_value= 260.91, loss_actor= 19.76, target value: -16.34:  88%|########8 | 8800/10000 [00:11<00:02, 580.84it/s]
reward: -5.23 (r0 = -3.03), reward eval: reward: -4.49, reward normalized=-3.37/5.13, grad norm= 370.23, loss_value= 317.04, loss_actor= 16.90, target value: -22.72:  88%|########8 | 8800/10000 [00:14<00:02, 580.84it/s]
reward: -5.23 (r0 = -3.03), reward eval: reward: -4.49, reward normalized=-3.37/5.13, grad norm= 370.23, loss_value= 317.04, loss_actor= 16.90, target value: -22.72:  96%|#########6| 9600/10000 [00:15<00:01, 398.60it/s]
reward: -4.74 (r0 = -3.03), reward eval: reward: -4.49, reward normalized=-2.89/5.05, grad norm= 61.31, loss_value= 202.66, loss_actor= 16.60, target value: -20.96:  96%|#########6| 9600/10000 [00:16<00:01, 398.60it/s]
reward: -4.74 (r0 = -3.03), reward eval: reward: -4.49, reward normalized=-2.89/5.05, grad norm= 61.31, loss_value= 202.66, loss_actor= 16.60, target value: -20.96: : 10400it [00:17, 363.66it/s]
reward:  3.56 (r0 = -3.03), reward eval: reward: -4.49, reward normalized=-1.99/6.94, grad norm= 215.45, loss_value= 493.50, loss_actor= 9.32, target value: -14.82: : 10400it [00:18, 363.66it/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（参见 [Feature] Distpatch IQL loss module。）
允许灵活的 TensorDict 键。

脚本的总运行时间：（0 分钟 28.749 秒）

由 Sphinx-Gallery 生成的画廊