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