MultiAgentNetBase

class torchrl.modules.MultiAgentNetBase(*, n_agents: int, centralized: bool | None = None, share_params: bool | None = None, agent_dim: int | None = None, vmap_randomness: str = 'different', use_td_params: bool = True, **kwargs)

多智慧體網路的基類。

注意

要使用 torch.nn.init 模組初始化 MARL 模組引數，請參考 get_stateful_net() 和 from_stateful_net() 方法。

forward(*inputs: Tuple[Tensor]) → Tensor

定義每次呼叫時執行的計算。

所有子類都應覆蓋此方法。

注意

雖然前向傳播的邏輯需要在此函式中定義，但之後應該呼叫 Module 例項而不是直接呼叫此函式，因為前者會處理註冊的鉤子，而後者會默默忽略它們。

from_stateful_net(stateful_net: Module)

根據網路的有狀態版本填充引數。

有關如何獲取網路的有狀態版本的詳細資訊，請參閱 get_stateful_net()。

引數:

stateful_net (nn.Module) – 應從中獲取引數的有狀態網路。

get_stateful_net(copy: bool = True)

返回網路的有狀態版本。

這可用於初始化引數。

此類網路通常無法直接呼叫，需要透過 vmap 呼叫才能執行。

引數:

copy (bool, 可選) – 如果為 True，則對網路進行深複製。預設為 True。

如果引數是原地修改的（推薦），則無需將引數複製回 MARL 模組。有關如何使用非原地（out-of-place）重新初始化的引數重新填充 MARL 模型，請參閱 from_stateful_net() 的詳細資訊。

示例

>>> from torchrl.modules import MultiAgentMLP
>>> import torch
>>> n_agents = 6
>>> n_agent_inputs=3
>>> n_agent_outputs=2
>>> batch = 64
>>> obs = torch.zeros(batch, n_agents, n_agent_inputs)
>>> mlp = MultiAgentMLP(
...     n_agent_inputs=n_agent_inputs,
...     n_agent_outputs=n_agent_outputs,
...     n_agents=n_agents,
...     centralized=False,
...     share_params=False,
...     depth=2,
... )
>>> snet = mlp.get_stateful_net()
>>> def init(module):
...     if hasattr(module, "weight"):
...         torch.nn.init.kaiming_normal_(module.weight)
>>> snet.apply(init)
>>> # If the module has been updated out-of-place (not the case here) we can reset the params
>>> mlp.from_stateful_net(snet)

reset_parameters()

重置模型的引數。