使用 TorchRL 教程进行竞争性多智能体强化学习（DDPG）

作者： 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 中使用竞争性多智能体环境、它们的 spec 如何工作以及它们如何与库集成；

• 如何使用 TorchRL 中的 Parallel PettingZoo 和 VMAS 环境以及多个智能体组；

• 如何在 TorchRL 中创建不同的多智能体网络架构（例如，使用参数共享、集中式 critic）；

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

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

如果您在 Google Colab 中运行此代码，请确保安装以下依赖项

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

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

离策略学习

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

• 在 MADDPG 中，critic 是集中的，并将全局状态和系统全局动作作为输入。全局状态可以是全局观察，或者简单地是智能体观察的串联。全局动作是智能体动作的串联。MADDPG 可用于需要 **集中式训练** 的场景，因为它需要访问全局信息。

• 在 IDDPG 中，critic 只接收一个智能体的观察和动作作为输入。这允许 **去中心化训练**，因为 critic 和策略都只需要局部信息来计算它们的输出。

集中式 critics 有助于克服多个智能体并发学习的非平稳性，但另一方面，它们可能会受到其庞大输入空间的影响。在本教程中，我们将能够训练这两种形式，并讨论参数共享（跨智能体共享网络参数的做法）对它们的影响。

本教程的结构如下：

1. 最初，我们将建立一组超参数以供使用。

2. 随后，我们将构建一个多智能体环境，使用 TorchRL 的包装器来处理 PettingZoo 或 VMAS。

3. 在此之后，我们将制定策略和 critic 网络，讨论各种选择对参数共享和 critic 集中化的影响。

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）。逃避者的速度和加速度都高于追逐者。环境中还有障碍物（黑色圆圈）。智能体和障碍物根据均匀随机分布生成。智能体在 2D 连续世界中活动，具有阻力和弹性碰撞。它们的动作是 2D 连续力，决定了它们的加速度。每个智能体都观察到其位置、速度、与其他智能体和障碍物的相对位置以及逃避者的速度。

PettingZoo 和 VMAS 版本在奖励函数上略有不同，因为 PettingZoo 会惩罚逃避者越出边界，而 VMAS 则在物理上阻止其越界。这就是为什么您会发现 VMAS 中两个团队的奖励是相同的，只是符号相反，而在 PettingZoo 中，逃避者的奖励会更低。

我们现在将实例化环境。在本教程中，我们将把剧集限制在 max_steps，之后会设置 terminated 标志。这是 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 在环境 spec 方面非常严格，但您可以轻松检查您的环境 spec 是否足够。在我们的示例中，模拟器包装器负责为您的 base_env 设置正确的 spec，因此您不必担心这一点。

有四个 spec 需要查看：

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

使用上面显示的命令，我们可以访问每个值的域。

我们可以看到所有 spec 都被组织成一个字典，根始终包含组名。此结构将遵循来自环境的所有 tensordict 数据。此外，每个组的 spec 都具有前导形状 (n_agents_in_that_group)（智能体为 1，对手为 2），这意味着该组的张量数据将始终具有该前导形状（组内的智能体数据已堆叠）。

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

要快速访问这些值在 tensordict 中的键，我们可以直接向环境请求相应的键，并且我们立即就能理解哪些是每个智能体独有的，哪些是共享的。这些信息将有助于告诉所有其他 TorchRL 组件在哪里可以找到每个值。

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

变换 (Transforms)

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

例如，在此案例中，我们将实例化一个 RewardSum 转换，该转换将对整个剧集的奖励进行求和。我们将告诉此转换在何处查找每个奖励键的重置键。本质上，我们只是说明每个组的剧集奖励将在 "_reset" tensordict 键设置为 True 时重置，这意味着调用了 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() 运行一个小的滚动，并将其输出与环境 spec 进行比较。如果没有引发错误，我们可以确信 spec 已正确定义。

check_env_specs(env)

回滚 (Rollout)

为了好玩，让我们看看一个简单的随机滚动是什么样的。您可以调用 env.rollout(n_steps) 并获得环境输入和输出外观的概述。动作将自动从动作 spec 域中随机抽取。

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)

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

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

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

• 在 next（通过运行 rollout.get("next") 访问）。我们将找到与根相同的结构，但有一些小的区别如下：

在 TorchRL 中，约定是 done 和 observations 将同时出现在 root 和 next 中（因为它们在重置时和步骤之后都可用）。Action 仅在 root 中可用（因为没有由步骤产生的 action），Reward 仅在 next 中可用（因为在重置时没有 reward）。此结构遵循 **Reinforcement Learning: An Introduction (Sutton and Barto)** 中的结构，其中 root 代表时间 \(t\) 的数据，next 代表世界步骤时间 \(t+1\) 的数据。

渲染随机滚动

如果您在 Google Colab 或带有 OpenGL 和 GUI 的机器上，您实际上可以渲染一个随机滚动。这将让您了解随机策略在此任务中能达到什么效果，以便与您自己训练的策略进行比较！

要渲染滚动，请按照本教程末尾的“渲染”部分中的说明操作，然后从 env.rollout() 中删除 policy=agents_exploration_policy 行。

策略 (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。我们还提供了此分布的最小值和最大值，我们从环境 spec 中获取这些。

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

Critic 网络

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

和以前一样，应该仔细考虑 **共享 Critic 参数** 的决定。总的来说，参数共享将加快训练收敛速度，但有一些重要的注意事项需要考虑：

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

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

在所有其他情况下，如果组内所有智能体的奖励函数（与奖励区分开）相同（如当前场景），共享可以提高性能。这可能会以智能体策略的同质性为代价。总的来说，知道哪种选择更可取的方法是快速尝试这两种选择。

在这里，我们还必须在 **MADDPG 和 IDDPG** 之间进行选择：

• 使用 MADDPG，我们将获得一个具有完全可观测性的集中式 critic（即，它将接收所有串联的全局智能体观察和动作作为输入）。由于我们处于模拟器中且训练是集中的，因此我们可以做到这一点。

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

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

与策略一样，我们为每个组创建一个 critic 网络，并将它们存储在字典中。

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

让我们尝试我们的策略和 critic 模块。如前所述，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)),
    )

数据收集器 (Data collector)

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

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

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

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

回放缓冲区 (Replay buffer)

回放缓冲区是离策略 RL 算法的常见构建块。有许多类型的缓冲区，在本教程中，我们使用一个基本缓冲区来随机存储和采样 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

损失函数 (Loss function)

DDPG 损失可以直接从 TorchRL 导入，方便起见，可以使用 DDPGLoss 类。这是使用 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()

结果 (Results)

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

为了延长训练时间，请增加 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 中创建竞争性多组多智能体环境、其 spec 如何工作以及它如何与库集成；

• 如何在 TorchRL 中为多个组创建多智能体网络架构；

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

• 如何将所有库组件（收集器、模块、回放缓冲区和损失）绑定到一个多智能体多组 MADDPG/IDDPG 训练循环中。

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

别忘了也查看我们的教程：使用 TorchRL 教程进行多智能体强化学习（PPO）。

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

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

VMAS 中的可用场景