競爭性多智慧體強化學習 (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 中可用的場景