注意

转到末尾下载完整的示例代码。

使用 TorchRL 教程进行多智能体强化学习 (PPO)¶

作者: Matteo Bettini

另请参阅

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

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

为了便于使用，本教程将遵循以下已提供的通用结构：使用 TorchRL 教程进行强化学习 (PPO)。建议但不强制要求在开始本教程之前熟悉该教程。

在本教程中，我们将使用 VMAS 中的 Navigation 环境，这是一个多机器人模拟器，也基于 PyTorch，可在设备上运行并行批量模拟。

在 Navigation 环境中，我们需要训练多个机器人（在随机位置生成）导航到它们的目标（也在随机位置），同时使用激光雷达传感器以避免彼此碰撞。

主要学习内容

如何在 TorchRL 中创建多智能体环境，其规范如何工作，以及它如何与库集成；
如何在 TorchRL 中使用 GPU 向量化环境；
如何在 TorchRL 中创建不同的多智能体网络架构（例如，使用参数共享、中心化评论家）
如何使用 tensordict.TensorDict 来携带多智能体数据；
如何将所有库组件（收集器、模块、回放缓冲区和损失）绑定到多智能体 MAPPO/IPPO 训练循环中。

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

!pip3 install torchrl
!pip3 install vmas
!pip3 install tqdm

近端策略优化 (PPO) 是一种策略梯度算法，其中收集一批数据并直接消耗以训练策略，从而在给定某些邻近约束的情况下最大化预期回报。您可以将其视为 REINFORCE 的复杂版本，REINFORCE 是基础策略优化算法。有关更多信息，请参阅近端策略优化算法论文。

这种类型的算法通常以在线策略方式训练。这意味着，在每次学习迭代中，我们都有一个采样和一个训练阶段。在迭代 \(t\) 的采样阶段，使用当前策略 \(\mathbf{\pi}_t\) 从智能体在环境中的交互中收集 rollout。在训练阶段，所有收集的 rollout 都会立即馈送到训练过程以执行反向传播。这会导致策略更新，然后再次用于采样。在循环中执行此过程构成在线策略学习。

在 PPO 算法的训练阶段，使用评论家来估计策略所采取行动的优劣。评论家学习近似特定状态的值（平均折扣回报）。然后，PPO 损失将策略获得的实际回报与评论家估计的回报进行比较，以确定所采取行动的优势并指导策略优化。

在多智能体设置中，情况略有不同。我们现在有多个策略 \(\mathbf{\pi}\)，每个智能体一个。策略通常是本地的和去中心化的。这意味着单个智能体的策略将仅基于其观察结果输出该智能体的动作。在 MARL 文献中，这被称为去中心化执行。另一方面，对于评论家存在不同的公式，主要是

在 MAPPO 中，评论家是中心化的，并将系统的全局状态作为输入。这可以是全局观察或智能体观察的简单连接。MAPPO 可用于执行中心化训练的上下文中，因为它需要访问全局信息。
在 IPPO 中，评论家仅将相应智能体的观察作为输入，就像策略一样。这允许去中心化训练，因为评论家和策略都只需要本地信息来计算其输出。

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

本教程的结构如下

首先，我们将定义一组我们将使用的超参数。
接下来，我们将使用 TorchRL 的 VMAS 模拟器包装器创建一个向量化多智能体环境。
接下来，我们将设计策略和评论家网络，讨论各种选择对参数共享和评论家中心化的影响。
接下来，我们将创建采样收集器和回放缓冲区。
最后，我们将运行我们的训练循环并分析结果。

如果您在 Colab 或带有 GUI 的机器上运行此程序，您还可以选择渲染和可视化您自己训练的策略在训练之前和之后的效果。

让我们导入我们的依赖项

# Torch
import torch

# Tensordict modules
from tensordict.nn import TensorDictModule
from tensordict.nn.distributions import NormalParamExtractor
from torch import multiprocessing

# Data collection
from torchrl.collectors import SyncDataCollector
from torchrl.data.replay_buffers import ReplayBuffer
from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
from torchrl.data.replay_buffers.storages import LazyTensorStorage

# Env
from torchrl.envs import RewardSum, TransformedEnv
from torchrl.envs.libs.vmas import VmasEnv
from torchrl.envs.utils import check_env_specs

# Multi-agent network
from torchrl.modules import MultiAgentMLP, ProbabilisticActor, TanhNormal

# Loss
from torchrl.objectives import ClipPPOLoss, ValueEstimators

# Utils
torch.manual_seed(0)
from matplotlib import pyplot as plt
from tqdm import tqdm

定义超参数¶

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

# 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")
)
vmas_device = device  # The device where the simulator is run (VMAS can run on GPU)

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

# Training
num_epochs = 30  # Number of optimization steps per training iteration
minibatch_size = 400  # Size of the mini-batches in each optimization step
lr = 3e-4  # Learning rate
max_grad_norm = 1.0  # Maximum norm for the gradients

# PPO
clip_epsilon = 0.2  # clip value for PPO loss
gamma = 0.99  # discount factor
lmbda = 0.9  # lambda for generalised advantage estimation
entropy_eps = 1e-4  # coefficient of the entropy term in the PPO loss

环境¶

多智能体环境模拟多个智能体与世界的交互。TorchRL API 允许集成各种类型的多智能体环境风格。一些示例包括具有共享或单独智能体奖励、完成标志和观察的环境。有关多智能体环境 API 在 TorchRL 中如何工作的更多信息，您可以查看专门的文档部分。

特别是，VMAS 模拟器使用单独的奖励、信息、观察和动作来建模智能体，但使用集体的完成标志。此外，它使用向量化来批量执行模拟。这意味着其所有状态和物理都是 PyTorch 张量，第一个维度表示批处理中并行环境的数量。这允许利用 GPU 的单指令多数据 (SIMD) 范例，并通过利用 GPU warp 中的并行化显着加速并行计算。这也意味着，在 TorchRL 中使用它时，模拟和训练都可以在设备上运行，而无需将数据传递到 CPU。

我们今天将解决的多智能体任务是 Navigation（参见上面的动画图）。在 Navigation 中，随机生成的智能体（带有周围点的圆圈）需要导航到随机生成的目标（较小的圆圈）。智能体需要使用激光雷达（它们周围的点）来避免彼此碰撞。智能体在具有阻力和弹性碰撞的 2D 连续世界中行动。它们的动作是 2D 连续力，决定了它们的加速度。奖励由三个项组成：碰撞惩罚、基于到目标距离的奖励以及当所有智能体都到达目标时给出的最终共享奖励。基于距离的项计算为智能体及其目标在两个连续时间步长之间的相对距离差异。每个智能体观察其位置、速度、激光雷达读数以及相对于其目标的相对位置。

我们现在将实例化环境。在本教程中，我们将 episodes 限制为 max_steps，在此之后设置完成标志。此功能已在 VMAS 模拟器中提供，但 TorchRL StepCount 变换也可以替代使用。我们还将使用 num_vmas_envs 向量化环境，以利用批量模拟。

max_steps = 100  # Episode steps before done
num_vmas_envs = (
    frames_per_batch // max_steps
)  # Number of vectorized envs. frames_per_batch should be divisible by this number
scenario_name = "navigation"
n_agents = 3

env = VmasEnv(
    scenario=scenario_name,
    num_envs=num_vmas_envs,
    continuous_actions=True,  # VMAS supports both continuous and discrete actions
    max_steps=max_steps,
    device=vmas_device,
    # Scenario kwargs
    n_agents=n_agents,  # These are custom kwargs that change for each VMAS scenario, see the VMAS repo to know more.
)

环境不仅由其模拟器和变换定义，还由一系列元数据定义，这些元数据描述了在执行期间可以预期什么。出于效率目的，TorchRL 在环境规范方面非常严格，但您可以轻松检查您的环境规范是否足够。在我们的示例中，VmasEnv 负责为您的环境设置正确的规范，因此您不必担心这一点。

有四个规范需要查看

action_spec 定义动作空间；
reward_spec 定义奖励域；
done_spec 定义完成域；
observation_spec 定义来自环境步骤的所有其他输出的域；

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

action_spec: Composite(
    agents: Composite(
        action: BoundedContinuous(
            shape=torch.Size([60, 3, 2]),
            space=ContinuousBox(
                low=Tensor(shape=torch.Size([60, 3, 2]), device=cpu, dtype=torch.float32, contiguous=True),
                high=Tensor(shape=torch.Size([60, 3, 2]), device=cpu, dtype=torch.float32, contiguous=True)),
            device=cpu,
            dtype=torch.float32,
            domain=continuous),
        device=cpu,
        shape=torch.Size([60, 3])),
    device=cpu,
    shape=torch.Size([60]))
reward_spec: Composite(
    agents: Composite(
        reward: UnboundedContinuous(
            shape=torch.Size([60, 3, 1]),
            space=ContinuousBox(
                low=Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, contiguous=True),
                high=Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, contiguous=True)),
            device=cpu,
            dtype=torch.float32,
            domain=continuous),
        device=cpu,
        shape=torch.Size([60, 3])),
    device=cpu,
    shape=torch.Size([60]))
done_spec: Composite(
    done: Categorical(
        shape=torch.Size([60, 1]),
        space=CategoricalBox(n=2),
        device=cpu,
        dtype=torch.bool,
        domain=discrete),
    terminated: Categorical(
        shape=torch.Size([60, 1]),
        space=CategoricalBox(n=2),
        device=cpu,
        dtype=torch.bool,
        domain=discrete),
    device=cpu,
    shape=torch.Size([60]))
observation_spec: Composite(
    agents: Composite(
        observation: UnboundedContinuous(
            shape=torch.Size([60, 3, 18]),
            space=ContinuousBox(
                low=Tensor(shape=torch.Size([60, 3, 18]), device=cpu, dtype=torch.float32, contiguous=True),
                high=Tensor(shape=torch.Size([60, 3, 18]), device=cpu, dtype=torch.float32, contiguous=True)),
            device=cpu,
            dtype=torch.float32,
            domain=continuous),
        info: Composite(
            pos_rew: UnboundedContinuous(
                shape=torch.Size([60, 3, 1]),
                space=ContinuousBox(
                    low=Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, contiguous=True),
                    high=Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, contiguous=True)),
                device=cpu,
                dtype=torch.float32,
                domain=continuous),
            final_rew: UnboundedContinuous(
                shape=torch.Size([60, 3, 1]),
                space=ContinuousBox(
                    low=Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, contiguous=True),
                    high=Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, contiguous=True)),
                device=cpu,
                dtype=torch.float32,
                domain=continuous),
            agent_collisions: UnboundedContinuous(
                shape=torch.Size([60, 3, 1]),
                space=ContinuousBox(
                    low=Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, contiguous=True),
                    high=Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, contiguous=True)),
                device=cpu,
                dtype=torch.float32,
                domain=continuous),
            device=cpu,
            shape=torch.Size([60, 3])),
        device=cpu,
        shape=torch.Size([60, 3])),
    device=cpu,
    shape=torch.Size([60]))

使用刚刚显示的命令，我们可以访问每个值的域。这样做我们可以看到，除了完成规范外，所有规范都有前导形状 (num_vmas_envs, n_agents)。这表示这些值将存在于每个个体环境中的每个智能体。另一方面，完成规范具有前导形状 num_vmas_envs，表示完成在智能体之间共享。

TorchRL 有一种方法可以跟踪哪些 MARL 规范是共享的，哪些不是。实际上，具有附加智能体维度（即，它们对于每个智能体都不同）的规范将包含在内部“agents”键中。

如您所见，奖励和动作规范呈现“agent”键，这意味着属于这些规范的 tensordict 中的条目将嵌套在“agents” tensordict 中，将所有每个智能体的值分组在一起。

为了快速访问 tensordict 中每个值的键，我们可以简单地向环境询问各自的键，我们将立即了解哪些是每个智能体的，哪些是共享的。此信息对于告诉所有其他 TorchRL 组件在哪里找到每个值非常有用

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

action_keys: [('agents', 'action')]
reward_keys: [('agents', 'reward')]
done_keys: ['done', 'terminated']

变换¶

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

例如，在这种情况下，我们将实例化一个 RewardSum 变换，它将对整个 episode 的奖励求和。我们将告诉此变换在哪里找到奖励键以及在哪里写入求和的 episode 奖励。变换后的环境将继承包装环境的设备和元数据，并根据其包含的变换序列来变换这些。

env = TransformedEnv(
    env,
    RewardSum(in_keys=[env.reward_key], out_keys=[("agents", "episode_reward")]),
)

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

check_env_specs(env)

Rollout¶

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

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

rollout of three steps: TensorDict(
    fields={
        agents: TensorDict(
            fields={
                action: Tensor(shape=torch.Size([60, 5, 3, 2]), device=cpu, dtype=torch.float32, is_shared=False),
                episode_reward: Tensor(shape=torch.Size([60, 5, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                info: TensorDict(
                    fields={
                        agent_collisions: Tensor(shape=torch.Size([60, 5, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                        final_rew: Tensor(shape=torch.Size([60, 5, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                        pos_rew: Tensor(shape=torch.Size([60, 5, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
                    batch_size=torch.Size([60, 5, 3]),
                    device=cpu,
                    is_shared=False),
                observation: Tensor(shape=torch.Size([60, 5, 3, 18]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([60, 5, 3]),
            device=cpu,
            is_shared=False),
        done: Tensor(shape=torch.Size([60, 5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                agents: TensorDict(
                    fields={
                        episode_reward: Tensor(shape=torch.Size([60, 5, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                        info: TensorDict(
                            fields={
                                agent_collisions: Tensor(shape=torch.Size([60, 5, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                                final_rew: Tensor(shape=torch.Size([60, 5, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                                pos_rew: Tensor(shape=torch.Size([60, 5, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
                            batch_size=torch.Size([60, 5, 3]),
                            device=cpu,
                            is_shared=False),
                        observation: Tensor(shape=torch.Size([60, 5, 3, 18]), device=cpu, dtype=torch.float32, is_shared=False),
                        reward: Tensor(shape=torch.Size([60, 5, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
                    batch_size=torch.Size([60, 5, 3]),
                    device=cpu,
                    is_shared=False),
                done: Tensor(shape=torch.Size([60, 5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                terminated: Tensor(shape=torch.Size([60, 5, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([60, 5]),
            device=cpu,
            is_shared=False),
        terminated: Tensor(shape=torch.Size([60, 5, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([60, 5]),
    device=cpu,
    is_shared=False)
Shape of the rollout TensorDict: torch.Size([60, 5])

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

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

在根目录中（通过运行 rollout.exclude("next") 访问）我们将找到在第一个时间步调用重置后可用的所有键。我们可以通过索引 n_rollout_steps 维度来查看它们在 rollout 步骤中的演变。在这些键中，我们将找到 rollout["agents"] tensordict 中每个智能体不同的键，其批处理大小为 (num_vmas_envs, n_rollout_steps, n_agents)，表示它存储了额外的智能体维度。此智能体 tensordict 之外的键将是共享的键（在本例中仅为 done）。
在 next 中（通过运行 rollout.get("next") 访问）。我们将找到与根目录相同的结构，但对于仅在步骤之后可用的键。

在 TorchRL 中，约定是 done 和 observations 将同时存在于根目录和 next 中（因为这些在重置时间和步骤后都可用）。Action 将仅在根目录中可用（因为没有来自步骤的动作），而 reward 将仅在 next 中可用（因为重置时没有奖励）。此结构遵循 Reinforcement Learning: An Introduction (Sutton and Barto) 中的结构，其中根目录表示时间 \(t\) 的数据，而 next 表示世界步骤的时间 \(t+1\) 的数据。

渲染随机 rollout¶

如果您在 Google Colab 上，或在具有 OpenGL 和 GUI 的机器上，您实际上可以渲染随机 rollout。这将使您了解随机策略在此任务中将实现什么，以便将其与您自己训练的策略进行比较！

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

策略¶

PPO 利用随机策略来处理探索。这意味着我们的神经网络必须输出分布的参数，而不是与所采取的动作相对应的单个值。

由于数据是连续的，我们使用 Tanh-Normal 分布来尊重动作空间边界。TorchRL 提供了这样的分布，我们唯一需要关心的是构建一个神经网络，该网络输出正确数量的参数。

在这种情况下，每个智能体的动作将由一个 2 维独立正态分布表示。为此，我们的神经网络必须为每个动作输出一个均值和一个标准差。因此，每个智能体将具有 2 * n_actions_per_agents 个输出。

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

我们分三个步骤设计策略。

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

为此，我们使用 MultiAgentMLP，这是一个专为多智能体设计的 TorchRL 模块，具有许多可用的自定义选项。

share_parameters_policy = True

policy_net = torch.nn.Sequential(
    MultiAgentMLP(
        n_agent_inputs=env.observation_spec["agents", "observation"].shape[
            -1
        ],  # n_obs_per_agent
        n_agent_outputs=2 * env.action_spec.shape[-1],  # 2 * n_actions_per_agents
        n_agents=env.n_agents,
        centralised=False,  # the policies are decentralised (ie each agent will act from its observation)
        share_params=share_parameters_policy,
        device=device,
        depth=2,
        num_cells=256,
        activation_class=torch.nn.Tanh,
    ),
    NormalParamExtractor(),  # this will just separate the last dimension into two outputs: a loc and a non-negative scale
)

第二步：将神经网络包装在 TensorDictModule 中

这只是一个模块，它将从 tensordict 读取 in_keys，将其馈送到神经网络，并将输出就地写入 out_keys。

请注意，我们使用 ("agents", ...) 键，因为这些键表示具有附加 n_agents 维度的数据。

policy_module = TensorDictModule(
    policy_net,
    in_keys=[("agents", "observation")],
    out_keys=[("agents", "loc"), ("agents", "scale")],
)

第三步：将 TensorDictModule 包装在 ProbabilisticActor 中

我们现在需要从正态分布的位置和尺度构建一个分布。为此，我们指示 ProbabilisticActor 类从位置和尺度参数构建 TanhNormal。我们还提供了此分布的最小值和最大值，我们从环境规范中收集这些值。

in_keys 的名称（因此 TensorDictModule 上面的 out_keys 的名称）必须以 TanhNormal 分布构造函数关键字参数（loc 和 scale）结尾。

policy = ProbabilisticActor(
    module=policy_module,
    spec=env.unbatched_action_spec,
    in_keys=[("agents", "loc"), ("agents", "scale")],
    out_keys=[env.action_key],
    distribution_class=TanhNormal,
    distribution_kwargs={
        "low": env.unbatched_action_spec[env.action_key].space.low,
        "high": env.unbatched_action_spec[env.action_key].space.high,
    },
    return_log_prob=True,
    log_prob_key=("agents", "sample_log_prob"),
)  # we'll need the log-prob for the PPO loss

评论家网络¶

评论家网络是 PPO 算法的关键组成部分，即使它在采样时未使用。此模块将读取观察结果并返回相应的值估计。

与之前一样，应该仔细考虑共享评论家参数的决定。通常，参数共享将加快训练收敛速度，但有一些重要的注意事项需要考虑

当智能体具有不同的奖励函数时，不建议共享，因为评论家将需要学习为同一状态分配不同的值（例如，在混合合作-竞争设置中）。
在去中心化训练设置中，如果不使用额外的基础设施来同步参数，则无法执行共享。

在奖励函数（要与奖励区分开）对于所有智能体都相同的所有其他情况下（如当前场景中），共享可以提供改进的性能。这可能会以智能体策略的同质性为代价。通常，了解哪种选择更可取的最佳方法是快速试验这两种选择。

这里也是我们必须在 MAPPO 和 IPPO 之间做出选择的地方

使用 MAPPO，我们将获得一个具有完全可观察性的中心评论家（即，它将所有连接的智能体观察作为输入）。我们可以这样做，因为我们处于模拟器中，并且训练是中心化的。
使用 IPPO，我们将拥有一个本地去中心化的评论家，就像策略一样。

在任何情况下，评论家输出都将具有形状 (..., n_agents, 1)。如果评论家是中心化的和共享的，则沿 n_agents 维度的所有值都将相同。

share_parameters_critic = True
mappo = True  # IPPO if False

critic_net = MultiAgentMLP(
    n_agent_inputs=env.observation_spec["agents", "observation"].shape[-1],
    n_agent_outputs=1,  # 1 value per agent
    n_agents=env.n_agents,
    centralised=mappo,
    share_params=share_parameters_critic,
    device=device,
    depth=2,
    num_cells=256,
    activation_class=torch.nn.Tanh,
)

critic = TensorDictModule(
    module=critic_net,
    in_keys=[("agents", "observation")],
    out_keys=[("agents", "state_value")],
)

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

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

print("Running policy:", policy(env.reset()))
print("Running value:", critic(env.reset()))

Running policy: TensorDict(
    fields={
        agents: TensorDict(
            fields={
                action: Tensor(shape=torch.Size([60, 3, 2]), device=cpu, dtype=torch.float32, is_shared=False),
                episode_reward: Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                info: TensorDict(
                    fields={
                        agent_collisions: Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                        final_rew: Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                        pos_rew: Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
                    batch_size=torch.Size([60, 3]),
                    device=cpu,
                    is_shared=False),
                loc: Tensor(shape=torch.Size([60, 3, 2]), device=cpu, dtype=torch.float32, is_shared=False),
                observation: Tensor(shape=torch.Size([60, 3, 18]), device=cpu, dtype=torch.float32, is_shared=False),
                sample_log_prob: Tensor(shape=torch.Size([60, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                scale: Tensor(shape=torch.Size([60, 3, 2]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([60, 3]),
            device=cpu,
            is_shared=False),
        done: Tensor(shape=torch.Size([60, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        terminated: Tensor(shape=torch.Size([60, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([60]),
    device=cpu,
    is_shared=False)
Running value: TensorDict(
    fields={
        agents: TensorDict(
            fields={
                episode_reward: Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                info: TensorDict(
                    fields={
                        agent_collisions: Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                        final_rew: Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                        pos_rew: Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
                    batch_size=torch.Size([60, 3]),
                    device=cpu,
                    is_shared=False),
                observation: Tensor(shape=torch.Size([60, 3, 18]), device=cpu, dtype=torch.float32, is_shared=False),
                state_value: Tensor(shape=torch.Size([60, 3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([60, 3]),
            device=cpu,
            is_shared=False),
        done: Tensor(shape=torch.Size([60, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        terminated: Tensor(shape=torch.Size([60, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([60]),
    device=cpu,
    is_shared=False)

数据收集器¶

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

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

collector = SyncDataCollector(
    env,
    policy,
    device=vmas_device,
    storing_device=device,
    frames_per_batch=frames_per_batch,
    total_frames=total_frames,
)

回放缓冲区¶

回放缓冲区是离线策略 RL 算法的常见构建块。在在线策略上下文中，每次收集一批数据时都会重新填充回放缓冲区，并且其数据会重复使用一定数量的 epoch。

对于 PPO，使用回放缓冲区不是强制性的，我们可以简单地在线使用收集的数据，但使用这些类使我们能够以可重现的方式轻松构建内部训练循环。

replay_buffer = ReplayBuffer(
    storage=LazyTensorStorage(
        frames_per_batch, device=device
    ),  # We store the frames_per_batch collected at each iteration
    sampler=SamplerWithoutReplacement(),
    batch_size=minibatch_size,  # We will sample minibatches of this size
)

损失函数¶

为了方便起见，可以使用 ClipPPOLoss 类直接从 TorchRL 导入 PPO 损失。这是利用 PPO 的最简单方法：它隐藏了 PPO 的数学运算和随之而来的控制流。

PPO 需要计算一些“优势估计”。简而言之，优势是一个值，它反映了在处理偏差/方差权衡时的回报值的期望值。要计算优势，只需 (1) 构建优势模块，该模块利用我们的值运算符，以及 (2) 在每个 epoch 之前将每批数据传递给它。GAE 模块将使用新的 "advantage" 和 "value_target" 条目更新输入 TensorDict。"value_target" 是一个无梯度的张量，表示值网络应使用输入观察表示的经验值。这两者都将由 ClipPPOLoss 使用以返回策略和值损失。

loss_module = ClipPPOLoss(
    actor_network=policy,
    critic_network=critic,
    clip_epsilon=clip_epsilon,
    entropy_coef=entropy_eps,
    normalize_advantage=False,  # Important to avoid normalizing across the agent dimension
)
loss_module.set_keys(  # We have to tell the loss where to find the keys
    reward=env.reward_key,
    action=env.action_key,
    sample_log_prob=("agents", "sample_log_prob"),
    value=("agents", "state_value"),
    # These last 2 keys will be expanded to match the reward shape
    done=("agents", "done"),
    terminated=("agents", "terminated"),
)


loss_module.make_value_estimator(
    ValueEstimators.GAE, gamma=gamma, lmbda=lmbda
)  # We build GAE
GAE = loss_module.value_estimator

optim = torch.optim.Adam(loss_module.parameters(), lr)

训练循环¶

我们现在拥有编码训练循环所需的所有部分。步骤包括

收集数据
- 计算优势
  
  循环遍历 epoch
  
  循环遍历 minibatch 以计算损失值
  
  反向传播
  
  优化
  
  重复
  
  重复
- 重复
重复

pbar = tqdm(total=n_iters, desc="episode_reward_mean = 0")

episode_reward_mean_list = []
for tensordict_data in collector:
    tensordict_data.set(
        ("next", "agents", "done"),
        tensordict_data.get(("next", "done"))
        .unsqueeze(-1)
        .expand(tensordict_data.get_item_shape(("next", env.reward_key))),
    )
    tensordict_data.set(
        ("next", "agents", "terminated"),
        tensordict_data.get(("next", "terminated"))
        .unsqueeze(-1)
        .expand(tensordict_data.get_item_shape(("next", env.reward_key))),
    )
    # We need to expand the done and terminated to match the reward shape (this is expected by the value estimator)

    with torch.no_grad():
        GAE(
            tensordict_data,
            params=loss_module.critic_network_params,
            target_params=loss_module.target_critic_network_params,
        )  # Compute GAE and add it to the data

    data_view = tensordict_data.reshape(-1)  # Flatten the batch size to shuffle data
    replay_buffer.extend(data_view)

    for _ in range(num_epochs):
        for _ in range(frames_per_batch // minibatch_size):
            subdata = replay_buffer.sample()
            loss_vals = loss_module(subdata)

            loss_value = (
                loss_vals["loss_objective"]
                + loss_vals["loss_critic"]
                + loss_vals["loss_entropy"]
            )

            loss_value.backward()

            torch.nn.utils.clip_grad_norm_(
                loss_module.parameters(), max_grad_norm
            )  # Optional

            optim.step()
            optim.zero_grad()

    collector.update_policy_weights_()

    # Logging
    done = tensordict_data.get(("next", "agents", "done"))
    episode_reward_mean = (
        tensordict_data.get(("next", "agents", "episode_reward"))[done].mean().item()
    )
    episode_reward_mean_list.append(episode_reward_mean)
    pbar.set_description(f"episode_reward_mean = {episode_reward_mean}", refresh=False)
    pbar.update()

episode_reward_mean = 0:   0%|          | 0/10 [00:00<?, ?it/s]
episode_reward_mean = -0.4579917788505554:  10%|█         | 1/10 [00:05<00:51,  5.70s/it]
episode_reward_mean = 0.23260341584682465:  20%|██        | 2/10 [00:11<00:45,  5.64s/it]
episode_reward_mean = 1.1713813543319702:  30%|███       | 3/10 [00:16<00:39,  5.62s/it]
episode_reward_mean = 1.386345624923706:  40%|████      | 4/10 [00:22<00:33,  5.61s/it]
episode_reward_mean = 1.8939578533172607:  50%|█████     | 5/10 [00:28<00:27,  5.60s/it]
episode_reward_mean = 2.2214083671569824:  60%|██████    | 6/10 [00:33<00:22,  5.59s/it]
episode_reward_mean = 2.1770293712615967:  70%|███████   | 7/10 [00:39<00:16,  5.60s/it]
episode_reward_mean = 2.6274709701538086:  80%|████████  | 8/10 [00:44<00:11,  5.62s/it]
episode_reward_mean = 2.73148250579834:  90%|█████████ | 9/10 [00:50<00:05,  5.65s/it]
episode_reward_mean = 2.737316608428955: 100%|██████████| 10/10 [00:56<00:00,  5.68s/it]

结果¶

让我们绘制每个 episode 获得的平均奖励

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

plt.plot(episode_reward_mean_list)
plt.xlabel("Training iterations")
plt.ylabel("Reward")
plt.title("Episode reward mean")
plt.show()

渲染¶

如果您在具有 GUI 的机器上运行此程序，则可以通过运行以下命令来渲染训练后的策略

with torch.no_grad():
   env.rollout(
       max_steps=max_steps,
       policy=policy,
       callback=lambda env, _: env.render(),
       auto_cast_to_device=True,
       break_when_any_done=False,
   )

如果您在 Google Colab 中运行此程序，则可以通过运行以下命令来渲染训练后的策略

!apt-get update
!apt-get install -y x11-utils
!apt-get install -y xvfb
!pip install pyvirtualdisplay

import pyvirtualdisplay
display = pyvirtualdisplay.Display(visible=False, size=(1400, 900))
display.start()
from PIL import Image

def rendering_callback(env, td):
    env.frames.append(Image.fromarray(env.render(mode="rgb_array")))
env.frames = []
with torch.no_grad():
   env.rollout(
       max_steps=max_steps,
       policy=policy,
       callback=rendering_callback,
       auto_cast_to_device=True,
       break_when_any_done=False,
   )
env.frames[0].save(
    f"{scenario_name}.gif",
    save_all=True,
    append_images=env.frames[1:],
   duration=3,
   loop=0,
)

from IPython.display import Image
Image(open(f"{scenario_name}.gif", "rb").read())

结论和下一步¶

在本教程中，我们看到了

如何在 TorchRL 中创建多智能体环境，其规范如何工作，以及它如何与库集成；
如何在 TorchRL 中使用 GPU 向量化环境；
如何在 TorchRL 中创建不同的多智能体网络架构（例如，使用参数共享、中心化评论家）
如何使用 tensordict.TensorDict 来携带多智能体数据；
如何将所有库组件（收集器、模块、回放缓冲区和损失）绑定到多智能体 MAPPO/IPPO 训练循环中。

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

您还可以查看我们的另一个关于如何在 PettingZoo/VMAS 中使用多个智能体组训练竞争性 MADDPG/IDDPG 的多智能体教程：使用 TorchRL 教程进行竞争性多智能体强化学习 (DDPG)。

如果您有兴趣在 TorchRL 中创建或包装自己的多智能体环境，您可以查看专门的文档部分。

最后，您可以修改本教程的参数以尝试许多其他配置和场景，从而成为 MARL 大师。以下是一些您可以在 VMAS 中尝试的可能场景的视频。

脚本的总运行时间： (1 分钟 52.899 秒)

估计内存使用量： 357 MB

由 Sphinx-Gallery 生成的图库