竞争性多智能体强化学习 (DDPG) TorchRL 教程

作者: Matteo Bettini

另请参阅

BenchMARL 库使用 TorchRL 提供了最先进的 MARL 算法实现。

本教程演示了如何使用 PyTorch 和 TorchRL 解决竞争性多智能体强化学习 (MARL) 问题。

为了方便使用，本教程将沿用已有教程 TorchRL 多智能体强化学习 (PPO) 教程 的总体结构。

在本教程中，我们将使用 MADDPG 论文 中的 simple_tag 环境。该环境是随论文引入的 MultiAgentParticleEnvironments (MPE) 环境集的一部分。

当前有多个模拟器提供 MPE 环境。本教程展示了如何使用以下任一方式在 TorchRL 中训练该环境：

• PettingZoo，采用该环境的传统 CPU 版本；

• VMAS，它提供了 PyTorch 中的向量化实现，能够在 GPU 上模拟多个环境以加速计算。

多智能体 simple_tag 场景

主要学习内容

• 如何在 TorchRL 中使用竞争性多智能体环境，了解其规范 (specs) 的工作原理以及如何与库集成；

• 如何在 TorchRL 中使用具有多个智能体组的并行 PettingZoo 和 VMAS 环境；

• 如何在 TorchRL 中创建不同的多智能体网络架构（例如，使用参数共享、中心化评论家）；

• 如何使用 TensorDict 来承载多智能体多组数据；

• 如何在一个离策略多智能体 MADDPG/IDDPG 训练循环中整合所有库组件（收集器、模块、回放缓冲区和损失）。

如果你在 Google Colab 中运行，请确保安装以下依赖项

!pip3 install torchrl
!pip3 install vmas
!pip3 install pettingzoo[mpe]==1.24.3
!pip3 install tqdm

深度确定性策略梯度 (DDPG) 是一种离策略的 actor-critic 算法，其中使用评论家网络的梯度来优化确定性策略。更多信息请参阅 深度确定性策略梯度 论文。这类算法通常进行离策略训练。更多关于离策略学习的信息请参阅 Sutton, Richard S., and Andrew G. Barto. Reinforcement learning: An introduction. MIT press, 2018。

离策略学习

这种方法已在 用于混合合作-竞争环境的多智能体 Actor-Critic 算法 中扩展到多智能体学习，该文引入了多智能体 DDPG (MADDPG) 算法。在多智能体设置中，情况略有不同。我们现在有多个策略 \(\mathbf{\pi}\)，每个智能体一个。策略通常是局部和去中心化的。这意味着单个智能体的策略将仅基于其观测来输出该智能体的动作。在 MARL 文献中，这被称为去中心化执行。另一方面，评论家存在不同的表述形式，主要有

• 在 MADDPG 中，评论家是中心化的，并将系统的全局状态和全局动作作为输入。全局状态可以是全局观测，也可以简单地是智能体观测的拼接。全局动作是智能体动作的拼接。MADDPG 可用于执行中心化训练的场景，因为它需要访问全局信息。

• 在 IDDPG 中，评论家仅将一个智能体的观测和动作作为输入。这使得去中心化训练成为可能，因为评论家和策略都只需要局部信息来计算它们的输出。

中心化评论家有助于克服多个智能体同时学习时的非平稳性问题，但另一方面，它们可能受到其庞大输入空间的影响。在本教程中，我们将能够训练这两种形式，并将讨论参数共享（在智能体之间共享网络参数的做法）如何影响每种形式。

本教程结构如下

1. 首先，我们将建立一组用于使用的超参数。

2. 随后，我们将构建一个多智能体环境，利用 TorchRL 对 PettingZoo 或 VMAS 的封装。

3. 之后，我们将构建策略网络和评论家网络，讨论不同选择对参数共享和评论家中心化的影响。

4. 接下来，我们将创建采样收集器和回放缓冲区。

5. 最后，我们将执行训练循环并检查结果。

如果你在 Colab 或带有 GUI 的机器上运行，你还将有机会在训练过程之前和之后渲染并可视化你训练好的策略。

导入依赖项

import copy
import tempfile

import torch

from matplotlib import pyplot as plt
from tensordict import TensorDictBase

from tensordict.nn import TensorDictModule, TensorDictSequential
from torch import multiprocessing

from torchrl.collectors import SyncDataCollector
from torchrl.data import LazyMemmapStorage, RandomSampler, ReplayBuffer

from torchrl.envs import (
    check_env_specs,
    ExplorationType,
    PettingZooEnv,
    RewardSum,
    set_exploration_type,
    TransformedEnv,
    VmasEnv,
)

from torchrl.modules import (
    AdditiveGaussianModule,
    MultiAgentMLP,
    ProbabilisticActor,
    TanhDelta,
)

from torchrl.objectives import DDPGLoss, SoftUpdate, ValueEstimators

from torchrl.record import CSVLogger, PixelRenderTransform, VideoRecorder

from tqdm import tqdm

# Check if we're building the doc, in which case disable video rendering
try:
    is_sphinx = __sphinx_build__
except NameError:
    is_sphinx = False

定义超参数

我们为教程设置超参数。根据可用资源，可以选择在 GPU 或其他设备上执行策略和模拟器。你可以调整其中一些值来适应计算要求。

# Seed
seed = 0
torch.manual_seed(seed)

# Devices
is_fork = multiprocessing.get_start_method() == "fork"
device = (
    torch.device(0)
    if torch.cuda.is_available() and not is_fork
    else torch.device("cpu")
)

# Sampling
frames_per_batch = 1_000  # Number of team frames collected per sampling iteration
n_iters = 10  # Number of sampling and training iterations
total_frames = frames_per_batch * n_iters

# We will stop training the evaders after this many iterations,
# should be 0 <= iteration_when_stop_training_evaders <= n_iters
iteration_when_stop_training_evaders = n_iters // 2

# Replay buffer
memory_size = 1_000_000  # The replay buffer of each group can store this many frames

# Training
n_optimiser_steps = 100  # Number of optimization steps per training iteration
train_batch_size = 128  # Number of frames trained in each optimiser step
lr = 3e-4  # Learning rate
max_grad_norm = 1.0  # Maximum norm for the gradients

# DDPG
gamma = 0.99  # Discount factor
polyak_tau = 0.005  # Tau for the soft-update of the target network

环境

多智能体环境模拟了多个智能体与世界互动。TorchRL API 允许集成各种类型的多智能体环境形式。本教程将重点介绍多个智能体组并行交互的环境。也就是说：在每一步，所有智能体都会同步获取观测并采取动作。

此外，TorchRL MARL API 允许将智能体分成组。每个组将是 tensordict 中的一个独立条目。组内智能体的数据被堆叠在一起。因此，通过选择如何对智能体进行分组，你可以决定哪些数据被堆叠/哪些数据保持为独立条目。在构建 VMAS 和 PettingZoo 等环境时可以指定分组策略。更多关于分组的信息，请参阅 MarlGroupMapType。

在 simple_tag 环境中，有两组智能体：追击者（或称“对手”）（红色圆圈）和逃避者（或称“智能体”）（绿色圆圈）。追击者因触碰到逃避者而获得奖励 (+10)。一旦接触，追击者团队会集体获得奖励，而逃避者被触碰则受到相同数值的惩罚 (-10)。逃避者比追击者具有更高的速度和加速度。环境中还有障碍物（黑色圆圈）。智能体和障碍物根据均匀随机分布生成。智能体在一个具有阻力和弹性碰撞的二维连续世界中活动。它们的动作是二维连续力，决定了它们的加速度。每个智能体观察其位置、速度、相对于所有其他智能体和障碍物的相对位置以及逃避者的速度。

PettingZoo 和 VMAS 版本在奖励函数上略有不同，因为 PettingZoo 会惩罚超出边界的逃避者，而 VMAS 会物理地阻止这种情况。这就是为什么你会观察到在 VMAS 中，两组的奖励是相同的，只是符号相反，而在 PettingZoo 中，逃避者会有较低的奖励。

现在我们将实例化环境。本教程中，我们将回合限制在 max_steps 步，之后会设置终止标志。PettingZoo 和 VMAS 模拟器已经提供了此功能，但也可以选择使用 TorchRL 的 StepCounter 转换。

max_steps = 100  # Environment steps before done

n_chasers = 2
n_evaders = 1
n_obstacles = 2

use_vmas = True  # Set this to True for a great performance speedup

if not use_vmas:
    base_env = PettingZooEnv(
        task="simple_tag_v3",
        parallel=True,  # Use the Parallel version
        seed=seed,
        # Scenario specific
        continuous_actions=True,
        num_good=n_evaders,
        num_adversaries=n_chasers,
        num_obstacles=n_obstacles,
        max_cycles=max_steps,
    )
else:
    num_vmas_envs = (
        frames_per_batch // max_steps
    )  # Number of vectorized environments. frames_per_batch collection will be divided among these environments
    base_env = VmasEnv(
        scenario="simple_tag",
        num_envs=num_vmas_envs,
        continuous_actions=True,
        max_steps=max_steps,
        device=device,
        seed=seed,
        # Scenario specific
        num_good_agents=n_evaders,
        num_adversaries=n_chasers,
        num_landmarks=n_obstacles,
    )

组映射

PettingZoo 和 VMAS 环境使用 TorchRL MARL 分组 API。我们可以按如下方式访问组映射，将每个组映射到其中的智能体：

print(f"group_map: {base_env.group_map}")

正如我们所见，它包含 2 个组：“agents”（逃避者）和“adversaries”（追击者）。

环境不仅由其模拟器和转换定义，还由一系列元数据定义，这些元数据描述了在执行期间可以期望的内容。出于效率目的，TorchRL 对环境规范 (specs) 要求相当严格，但你可以轻松检查你的环境规范是否足够。在我们的示例中，模拟器封装器会处理好为你的 base_env 设置适当的规范，所以你不必担心这一点。

有四种规范需要查看：

• action_spec 定义了动作空间；

• reward_spec 定义了奖励域；

• done_spec 定义了完成域；

• observation_spec 定义了环境步骤所有其他输出的域；

print("action_spec:", base_env.full_action_spec)
print("reward_spec:", base_env.full_reward_spec)
print("done_spec:", base_env.full_done_spec)
print("observation_spec:", base_env.observation_spec)

使用刚才展示的命令，我们可以访问每个值的域。

我们可以看到所有规范都结构化为一个字典，根始终包含组名。这种结构将用于所有进出环境的 tensordict 数据。此外，每个组的规范具有前导形状 (n_agents_in_that_group)（agents 组为 1，adversaries 组为 2），这意味着该组的张量数据将始终具有该前导形状（组内智能体的数据已堆叠）。

查看 done_spec，我们可以看到有些键位于智能体组之外（"done", "terminated", "truncated"），它们没有前导多智能体维度。这些键由所有智能体共享，表示用于重置的环境全局完成状态。默认情况下，就像本例一样，当任何智能体完成时，并行 PettingZoo 环境就完成，但可以通过在 PettingZoo 环境构造时设置 done_on_any 来覆盖此行为。

要快速访问 tensordicts 中每个这些值的键，我们可以简单地向环境询问相应的键，我们将立即了解哪些是每个智能体的键以及哪些是共享的键。这些信息对于告知所有其他 TorchRL 组件在哪里查找每个值非常有用。

print("action_keys:", base_env.action_keys)
print("reward_keys:", base_env.reward_keys)
print("done_keys:", base_env.done_keys)

转换

我们可以将所需的任何 TorchRL 转换附加到我们的环境中。这些转换将以某种期望的方式修改其输入/输出。我们强调，在多智能体环境中，明确提供要修改的键至关重要。

例如，在本例中，我们将实例化一个 RewardSum 转换，它将在回合中累加奖励。我们将告诉这个转换在哪里找到每个奖励键的重置键。本质上，我们只是说当 "_reset" tensordict 键被设置时，每个组的回合奖励应该被重置，这意味着调用了 env.reset()。转换后的环境将继承被封装环境的设备和元数据，并根据其包含的转换序列进行转换。

env = TransformedEnv(
    base_env,
    RewardSum(
        in_keys=base_env.reward_keys,
        reset_keys=["_reset"] * len(base_env.group_map.keys()),
    ),
)

check_env_specs() 函数会运行一个小规模的 rollout，并将其输出与环境规范 (specs) 进行比较。如果没有引发错误，我们可以确信规范已正确定义。

check_env_specs(env)

Rollout

为了好玩，让我们看看一个简单的随机 rollout 是什么样的。你可以调用 env.rollout(n_steps) 来概览环境输入和输出是什么样的。动作将自动从动作规范域中随机抽取。

n_rollout_steps = 5
rollout = env.rollout(n_rollout_steps)
print(f"rollout of {n_rollout_steps} steps:", rollout)
print("Shape of the rollout TensorDict:", rollout.batch_size)

我们可以看到我们的 rollout 的 batch_size 为 (n_rollout_steps)。这意味着其中的所有张量都将具有这个前导维度。

更深入地看，我们可以看到输出 tensordict 可以按以下方式划分：

• 在根层（通过运行 rollout.exclude("next") 可访问），我们将找到在第一个时间步调用 reset 后所有可用的键。通过索引 n_rollout_steps 维度，我们可以看到它们在 rollout 步骤中的演变。在这些键中，我们将找到在 rollout[group_name] tensordicts 中每个智能体不同的键，它们的批大小将是 (n_rollout_steps, n_agents_in_group)，表示它存储了额外的智能体维度。不在组 tensordicts 中的键将是共享的。

• 在 next 层（通过运行 rollout.get("next") 可访问）。我们将找到与根层相同的结构，但有一些细微差异将在下面突出显示。

在 TorchRL 中，约定是 done 和 observations 将同时存在于 root 和 next 中（因为它们在 reset 时和 step 后都可用）。Action 将仅在 root 中可用（因为 step 不会产生动作），而 reward 将仅在 next 中可用（因为 reset 时没有奖励）。此结构遵循 Reinforcement Learning: An Introduction (Sutton and Barto) 中的结构，其中 root 表示时间 \(t\) 的数据，next 表示世界步时间 \(t+1\) 的数据。

渲染随机 rollout

如果你在 Google Colab 或带有 OpenGL 和 GUI 的机器上，你可以实际渲染一个随机 rollout。这将让你了解随机策略在此任务中的表现，以便与你亲自训练的策略进行比较！

要渲染 rollout，请按照本教程末尾渲染部分中的说明进行操作，并从 env.rollout() 中移除 policy=agents_exploration_policy 这一行即可。

策略

DDPG 利用确定性策略。这意味着我们的神经网络将输出要采取的动作。由于动作是连续的，我们使用 Tanh-Delta 分布来遵守动作空间的边界。该类所做的唯一事情是应用 Tanh 转换，以确保动作在域边界内。

我们需要做出的另一个重要决定是，是否希望团队内的智能体共享策略参数。一方面，共享参数意味着它们都将共享同一策略，这将使它们能够从彼此的经验中受益。这也会加快训练速度。另一方面，由于它们实际上共享相同的模型，这会使它们在行为上变得同质化。对于本示例，我们将启用共享，因为我们不介意同质化并可以受益于计算速度，但在你自己的问题中，考虑这个决定总是很重要的！

我们分三步设计策略。

第一步：定义一个神经网络 n_obs_per_agent -> n_actions_per_agents

为此，我们使用 MultiAgentMLP，这是一个专为多智能体设计的 TorchRL 模块，提供了许多定制选项。

我们将为每个组定义不同的策略并将其存储在一个字典中。

policy_modules = {}
for group, agents in env.group_map.items():
    share_parameters_policy = True  # Can change this based on the group

    policy_net = MultiAgentMLP(
        n_agent_inputs=env.observation_spec[group, "observation"].shape[
            -1
        ],  # n_obs_per_agent
        n_agent_outputs=env.full_action_spec[group, "action"].shape[
            -1
        ],  # n_actions_per_agents
        n_agents=len(agents),  # Number of agents in the group
        centralised=False,  # the policies are decentralised (i.e., each agent will act from its local observation)
        share_params=share_parameters_policy,
        device=device,
        depth=2,
        num_cells=256,
        activation_class=torch.nn.Tanh,
    )

    # Wrap the neural network in a :class:`~tensordict.nn.TensorDictModule`.
    # This is simply a module that will read the ``in_keys`` from a tensordict, feed them to the
    # neural networks, and write the
    # outputs in-place at the ``out_keys``.

    policy_module = TensorDictModule(
        policy_net,
        in_keys=[(group, "observation")],
        out_keys=[(group, "param")],
    )  # We just name the input and output that the network will read and write to the input tensordict
    policy_modules[group] = policy_module

第二步：将 TensorDictModule 封装在 ProbabilisticActor 中。

现在我们需要构建 TanhDelta 分布。我们指示 ProbabilisticActor 类根据策略动作参数构建 TanhDelta。我们还提供了该分布的最小值和最大值，这些值从环境规范 (specs) 中获取。

in_keys 的名称（以及因此上面 TensorDictModule 中 out_keys 的名称）必须以 TanhDelta 分布构造函数的关键字参数 (param) 结尾。

policies = {}
for group, _agents in env.group_map.items():
    policy = ProbabilisticActor(
        module=policy_modules[group],
        spec=env.full_action_spec[group, "action"],
        in_keys=[(group, "param")],
        out_keys=[(group, "action")],
        distribution_class=TanhDelta,
        distribution_kwargs={
            "low": env.full_action_spec_unbatched[group, "action"].space.low,
            "high": env.full_action_spec_unbatched[group, "action"].space.high,
        },
        return_log_prob=False,
    )
    policies[group] = policy

第三步：探索

由于 DDPG 策略是确定性的，我们在收集数据时需要一种进行探索的方法。

为此，我们需要在将策略传递给收集器之前，为其附加一个探索层。在本例中，我们使用 AdditiveGaussianModule，它会给我们的动作添加高斯噪声（如果噪声导致动作超出边界，则会对其进行截断）。

这个探索封装器使用一个 sigma 参数，该参数乘以噪声来确定其幅度。Sigma 可以在整个训练过程中进行退火以减少探索。Sigma 将在 annealing_num_steps 步内从 sigma_init 变化到 sigma_end。

exploration_policies = {}
for group, _agents in env.group_map.items():
    exploration_policy = TensorDictSequential(
        policies[group],
        AdditiveGaussianModule(
            spec=policies[group].spec,
            annealing_num_steps=total_frames
            // 2,  # Number of frames after which sigma is sigma_end
            action_key=(group, "action"),
            sigma_init=0.9,  # Initial value of the sigma
            sigma_end=0.1,  # Final value of the sigma
        ),
    )
    exploration_policies[group] = exploration_policy

评论家网络

评论家网络是 DDPG 算法的关键组成部分，尽管它在采样时并未使用。此模块将读取观测和采取的动作，并返回相应的值估计。

如前所述，应仔细考虑在智能体组内共享评论家参数的决定。一般来说，参数共享会加快训练收敛速度，但有一些重要的注意事项需要考虑

• 当智能体具有不同的奖励函数时，不建议共享，因为评论家需要学习为相同的状态分配不同的值（例如，在混合合作-竞争设置中）。在这种情况下，由于两组已经使用单独的网络，共享决策仅适用于组内的智能体，而我们已经知道这些智能体具有相同的奖励函数。

• 在去中心化训练设置中，如果没有额外的基础设施来同步参数，则无法执行共享。

在所有其他情况下，如果组中所有智能体的奖励函数（区别于奖励本身）是相同的（如当前场景所示），共享可以提供改进的性能。这可能会以智能体策略的同质性为代价。一般来说，了解哪种选择更优的最佳方法是快速实验两种选项。

在这里，我们还需要在 MADDPG 和 IDDPG 之间进行选择。

• 使用 MADDPG，我们将获得一个具有全观测能力的中心评论家（即，它将把所有拼接后的全局智能体观测和动作作为输入）。我们之所以能这样做，是因为我们在模拟器中并且训练是中心化的。

• 使用 IDDPG，我们将拥有一个局部去中心化评论家，就像策略一样。

无论如何，评论家的输出形状将是 (..., n_agents_in_group, 1)。如果评论家是中心化且共享的，那么沿 n_agents_in_group 维度的所有值都将相同。

与策略类似，我们为每个组创建一个评论家网络并将其存储在一个字典中。

critics = {}
for group, agents in env.group_map.items():
    share_parameters_critic = True  # Can change for each group
    MADDPG = True  # IDDPG if False, can change for each group

    # This module applies the lambda function: reading the action and observation entries for the group
    # and concatenating them in a new ``(group, "obs_action")`` entry
    cat_module = TensorDictModule(
        lambda obs, action: torch.cat([obs, action], dim=-1),
        in_keys=[(group, "observation"), (group, "action")],
        out_keys=[(group, "obs_action")],
    )

    critic_module = TensorDictModule(
        module=MultiAgentMLP(
            n_agent_inputs=env.observation_spec[group, "observation"].shape[-1]
            + env.full_action_spec[group, "action"].shape[-1],
            n_agent_outputs=1,  # 1 value per agent
            n_agents=len(agents),
            centralised=MADDPG,
            share_params=share_parameters_critic,
            device=device,
            depth=2,
            num_cells=256,
            activation_class=torch.nn.Tanh,
        ),
        in_keys=[(group, "obs_action")],  # Read ``(group, "obs_action")``
        out_keys=[
            (group, "state_action_value")
        ],  # Write ``(group, "state_action_value")``
    )

    critics[group] = TensorDictSequential(
        cat_module, critic_module
    )  # Run them in sequence

让我们尝试一下我们的策略和评论家模块。如前所述，TensorDictModule 的使用使得可以直接读取环境的输出以运行这些模块，因为它们知道要读取什么信息以及将信息写入何处。

我们可以看到，在每个组的网络运行后，它们的输出键会添加到组条目下的数据中。

从这一点开始，多智能体特定的组件已经实例化完成，我们将像在单智能体学习中一样简单地使用相同的组件。这不是很棒吗？

reset_td = env.reset()
for group, _agents in env.group_map.items():
    print(
        f"Running value and policy for group '{group}':",
        critics[group](policies[group](reset_td)),
    )

数据收集器

TorchRL 提供了一系列数据收集器类。简而言之，这些类执行三个操作：重置环境，使用策略和最新观测计算动作，在环境中执行一步，然后重复后两个步骤，直到环境发出停止信号（或达到完成状态）。

我们将使用最简单的数据收集器，它与环境 rollout 具有相同的输出，唯一的区别是它会自动重置完成状态，直到收集到所需的帧数。

我们需要向其提供我们的探索策略。此外，为了像运行一个策略一样运行所有组的策略，我们将它们放在一个序列中。它们不会相互干扰，因为每个组在不同的位置写入和读取键。

# Put exploration policies from each group in a sequence
agents_exploration_policy = TensorDictSequential(*exploration_policies.values())

collector = SyncDataCollector(
    env,
    agents_exploration_policy,
    device=device,
    frames_per_batch=frames_per_batch,
    total_frames=total_frames,
)

回放缓冲区

回放缓冲区是非策略强化学习算法的一个常见构建块。存在许多类型的缓冲区，在本教程中，我们使用一个基本缓冲区来随机存储和采样 tensordict 数据。

此缓冲区使用 LazyMemmapStorage，它将数据存储在磁盘上。这允许使用磁盘内存，但由于需要将数据转换到训练设备上，因此可能会导致采样速度较慢。要将缓冲区存储在 GPU 上，您可以使用 LazyTensorStorage 并传入所需设备。这将加快采样速度，但受所选设备的内存限制。

replay_buffers = {}
scratch_dirs = []
for group, _agents in env.group_map.items():
    scratch_dir = tempfile.TemporaryDirectory().name
    scratch_dirs.append(scratch_dir)
    replay_buffer = ReplayBuffer(
        storage=LazyMemmapStorage(
            memory_size,
            scratch_dir=scratch_dir,
        ),  # We will store up to memory_size multi-agent transitions
        sampler=RandomSampler(),
        batch_size=train_batch_size,  # We will sample batches of this size
    )
    if device.type != "cpu":
        replay_buffer.append_transform(lambda x: x.to(device))
    replay_buffers[group] = replay_buffer

损失函数

为了方便起见，可以使用 DDPGLoss 类直接从 TorchRL 导入 DDPG 损失。这是使用 DDPG 的最简单方法：它隐藏了 DDPG 的数学运算及其控制流程。

对于每个组，也可以有不同的策略。

losses = {}
for group, _agents in env.group_map.items():
    loss_module = DDPGLoss(
        actor_network=policies[group],  # Use the non-explorative policies
        value_network=critics[group],
        delay_value=True,  # Whether to use a target network for the value
        loss_function="l2",
    )
    loss_module.set_keys(
        state_action_value=(group, "state_action_value"),
        reward=(group, "reward"),
        done=(group, "done"),
        terminated=(group, "terminated"),
    )
    loss_module.make_value_estimator(ValueEstimators.TD0, gamma=gamma)

    losses[group] = loss_module

target_updaters = {
    group: SoftUpdate(loss, tau=polyak_tau) for group, loss in losses.items()
}

optimisers = {
    group: {
        "loss_actor": torch.optim.Adam(
            loss.actor_network_params.flatten_keys().values(), lr=lr
        ),
        "loss_value": torch.optim.Adam(
            loss.value_network_params.flatten_keys().values(), lr=lr
        ),
    }
    for group, loss in losses.items()
}

训练工具

我们确实需要定义两个将在训练循环中使用的辅助函数。它们非常简单，不包含任何重要的逻辑。

def process_batch(batch: TensorDictBase) -> TensorDictBase:
    """
    If the `(group, "terminated")` and `(group, "done")` keys are not present, create them by expanding
    `"terminated"` and `"done"`.
    This is needed to present them with the same shape as the reward to the loss.
    """
    for group in env.group_map.keys():
        keys = list(batch.keys(True, True))
        group_shape = batch.get_item_shape(group)
        nested_done_key = ("next", group, "done")
        nested_terminated_key = ("next", group, "terminated")
        if nested_done_key not in keys:
            batch.set(
                nested_done_key,
                batch.get(("next", "done")).unsqueeze(-1).expand((*group_shape, 1)),
            )
        if nested_terminated_key not in keys:
            batch.set(
                nested_terminated_key,
                batch.get(("next", "terminated"))
                .unsqueeze(-1)
                .expand((*group_shape, 1)),
            )
    return batch

训练循环

现在我们已经具备了编写训练循环所需的所有部分。步骤包括：

• 收集所有组的数据

  • 遍历各组

    • 将组数据存储在组缓冲区中

    • 遍历 epoch

      • 从组缓冲区采样

      • 计算采样数据的损失

      • 反向传播损失

      • 优化

    • 重复

  • 重复

• 重复

pbar = tqdm(
    total=n_iters,
    desc=", ".join(
        [f"episode_reward_mean_{group} = 0" for group in env.group_map.keys()]
    ),
)
episode_reward_mean_map = {group: [] for group in env.group_map.keys()}
train_group_map = copy.deepcopy(env.group_map)

# Training/collection iterations
for iteration, batch in enumerate(collector):
    current_frames = batch.numel()
    batch = process_batch(batch)  # Util to expand done keys if needed
    # Loop over groups
    for group in train_group_map.keys():
        group_batch = batch.exclude(
            *[
                key
                for _group in env.group_map.keys()
                if _group != group
                for key in [_group, ("next", _group)]
            ]
        )  # Exclude data from other groups
        group_batch = group_batch.reshape(
            -1
        )  # This just affects the leading dimensions in batch_size of the tensordict
        replay_buffers[group].extend(group_batch)

        for _ in range(n_optimiser_steps):
            subdata = replay_buffers[group].sample()
            loss_vals = losses[group](subdata)

            for loss_name in ["loss_actor", "loss_value"]:
                loss = loss_vals[loss_name]
                optimiser = optimisers[group][loss_name]

                loss.backward()

                # Optional
                params = optimiser.param_groups[0]["params"]
                torch.nn.utils.clip_grad_norm_(params, max_grad_norm)

                optimiser.step()
                optimiser.zero_grad()

            # Soft-update the target network
            target_updaters[group].step()

        # Exploration sigma anneal update
        exploration_policies[group][-1].step(current_frames)

    # Stop training a certain group when a condition is met (e.g., number of training iterations)
    if iteration == iteration_when_stop_training_evaders:
        del train_group_map["agent"]

    # Logging
    for group in env.group_map.keys():
        episode_reward_mean = (
            batch.get(("next", group, "episode_reward"))[
                batch.get(("next", group, "done"))
            ]
            .mean()
            .item()
        )
        episode_reward_mean_map[group].append(episode_reward_mean)

    pbar.set_description(
        ", ".join(
            [
                f"episode_reward_mean_{group} = {episode_reward_mean_map[group][-1]}"
                for group in env.group_map.keys()
            ]
        ),
        refresh=False,
    )
    pbar.update()

结果

我们可以绘制每个 episode 获得的平均奖励。

要使训练持续更长时间，请增加 n_iters 超参数。

在本地运行此脚本时，您可能需要关闭打开的窗口才能继续屏幕的其余部分。

fig, axs = plt.subplots(2, 1)
for i, group in enumerate(env.group_map.keys()):
    axs[i].plot(episode_reward_mean_map[group], label=f"Episode reward mean {group}")
    axs[i].set_ylabel("Reward")
    axs[i].axvline(
        x=iteration_when_stop_training_evaders,
        label="Agent (evader) stop training",
        color="orange",
    )
    axs[i].legend()
axs[-1].set_xlabel("Training iterations")
plt.show()

渲染

渲染说明适用于 VMAS，即使用 use_vmas=True 运行时。

TorchRL 提供了一些用于录制和保存渲染视频的工具。您可以在此处详细了解这些工具。

在以下代码块中，我们附加了一个转换，它将调用 VMAS 包装环境中 render() 方法，并将帧堆栈保存到由自定义记录器 video_logger 确定的位置的 mp4 文件。请注意，此代码可能需要一些外部依赖项，例如 torchvision。

if use_vmas and not is_sphinx:
    # Replace tmpdir with any desired path where the video should be saved
    with tempfile.TemporaryDirectory() as tmpdir:
        video_logger = CSVLogger("vmas_logs", tmpdir, video_format="mp4")
        print("Creating rendering env")
        env_with_render = TransformedEnv(env.base_env, env.transform.clone())
        env_with_render = env_with_render.append_transform(
            PixelRenderTransform(
                out_keys=["pixels"],
                # the np.ndarray has a negative stride and needs to be copied before being cast to a tensor
                preproc=lambda x: x.copy(),
                as_non_tensor=True,
                # asking for array rather than on-screen rendering
                mode="rgb_array",
            )
        )
        env_with_render = env_with_render.append_transform(
            VideoRecorder(logger=video_logger, tag="vmas_rendered")
        )
        with set_exploration_type(ExplorationType.DETERMINISTIC):
            print("Rendering rollout...")
            env_with_render.rollout(100, policy=agents_exploration_policy)
        print("Saving the video...")
        env_with_render.transform.dump()
        print("Saved! Saved directory tree:")
        video_logger.print_log_dir()

结论与下一步

在本教程中，我们已经了解了：

• 如何在 TorchRL 中创建一个具有竞争力的多组多代理环境，其 specs 如何工作，以及它如何与库集成；

• 如何在 TorchRL 中为多个组创建多代理网络架构；

• 如何使用 tensordict.TensorDict 来承载多代理多组数据；

• 如何在多代理多组 MADDPG/IDDPG 训练循环中将所有库组件（收集器、模块、回放缓冲区和损失函数）联系起来。

现在您已精通多代理 DDPG，可以查看 GitHub 存储库中所有 TorchRL 多代理实现。这些是许多 MARL 算法的纯代码脚本，例如本教程中介绍的算法、QMIX、MADDPG、IQL 等等！

另外，请务必查看我们的教程：使用 TorchRL 进行多代理强化学习 (PPO) 教程。

最后，您可以修改本教程的参数，尝试许多其他配置和场景，成为一名 MARL 大师。

PettingZoo 和 VMAS 包含更多场景。以下是一些您可以在 VMAS 中尝试的场景的视频。

VMAS 中可用的场景