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

• 如何使用 TorchRL 中的平行 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 网络产生的梯度来优化确定性策略。有关更多信息，请参阅 Deep Deterministic Policy Gradients 论文。此类算法通常是离线训练的。有关离线学习的更多信息，请参阅 *Sutton, Richard S., and Andrew G. Barto. Reinforcement learning: An introduction. MIT press, 2018*。

离线学习

这种方法已扩展到多智能体学习，如 Multi-Agent Actor-Critic for Mixed Cooperative-Competitive Environments，它引入了多智能体 DDPG (MADDPG) 算法。在多智能体设置中，情况略有不同。我们现在有多个策略 \(\mathbf{\pi}\)，每个智能体一个。策略通常是局部和分散的。这意味着单个智能体的策略将仅根据其观察来输出该智能体的动作。在 MARL 文献中，这被称为 **分散执行**。另一方面，对于 critic，存在不同的表述，主要是

• 在 MADDPG 中，critic 是中心化的，并将全局状态和系统全局动作作为输入。全局状态可以是全局观察，也可以仅仅是智能体观察的连接。全局动作是智能体动作的连接。MADDPG 可用于执行 **中心化训练** 的场景，因为它需要访问全局信息。

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

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

本教程的结构如下：

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 在环境规范方面要求非常严格，但您可以轻松检查您的环境规范是否足够。在我们的示例中，模拟器包装器负责为您的 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)（1 用于 agents，2 用于 adversaries），这意味着该组的张量数据将始终具有该前导形状（组内的智能体数据已堆叠）。

查看 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)

变换 (Transforms)

我们可以附加任何需要的 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() 函数运行一个小的回滚，并将其输出与环境规范进行比较。如果没有引发错误，我们可以确信规范已正确定义。

check_env_specs(env)

回滚 (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)

我们可以看到我们的滚动具有 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 中（因为它们在 reset 时和 step 之后都可用）。Action 仅在 root 中可用（因为 step 没有产生 action），reward 仅在 next 中可用（因为 reset 时没有 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。我们还提供了此分布的最小值和最大值，这些值来自环境规范。

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 参数** 的决定。总的来说，参数共享将实现更快的训练收敛，但有一些重要的考虑因素：

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

• 在分散式训练场景中，如果不进行额外的基础设施来同步参数，则无法进行共享。

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

这也是我们必须在 **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() 方法，并将帧堆栈保存到 mp4 文件中，该文件的位置由自定义日志记录器 video_logger 确定。请注意，此代码可能需要一些外部依赖项，例如 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 中创建竞争性多组多智能体环境、其规范如何工作以及它如何与库集成；

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

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

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

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

另外，请务必查看我们的教程：TorchRL 教程：多智能体强化学习 (PPO)。

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

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

VMAS 中可用的场景 VMAS