Source code for torchrl.modules.models.multiagent
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import annotations
import abc
from copy import deepcopy
from textwrap import indent
from typing import Optional, Sequence, Tuple, Type, Union
import numpy as np
import torch
from tensordict import TensorDict
from torch import nn
from torchrl.data.utils import DEVICE_TYPING
from torchrl.modules.models import ConvNet, MLP
from torchrl.modules.models.utils import _reset_parameters_recursive
[docs]class MultiAgentNetBase(nn.Module):
"""A base class for multi-agent networks.
.. note:: to initialize the MARL module parameters with the `torch.nn.init`
module, please refer to :meth:`~.get_stateful_net` and :meth:`~.from_stateful_net`
methods.
"""
_empty_net: nn.Module
def __init__(
self,
*,
n_agents: int,
centralized: bool | None = None,
share_params: bool | None = None,
agent_dim: int | None = None,
vmap_randomness: str = "different",
**kwargs,
):
super().__init__()
# For backward compatibility
centralized = kwargs.pop("centralised", centralized)
if centralized is None:
raise TypeError("centralized arg must be passed.")
if share_params is None:
raise TypeError("share_params arg must be passed.")
if agent_dim is None:
raise TypeError("agent_dim arg must be passed.")
self.n_agents = n_agents
self.share_params = share_params
self.centralized = centralized
self.agent_dim = agent_dim
self._vmap_randomness = vmap_randomness
agent_networks = [
self._build_single_net(**kwargs)
for _ in range(self.n_agents if not self.share_params else 1)
]
initialized = True
for p in agent_networks[0].parameters():
if isinstance(p, torch.nn.UninitializedParameter):
initialized = False
break
self.initialized = initialized
self._make_params(agent_networks)
# We make sure all params and buffers are on 'meta' device
# To do this, we set the device keyword arg to 'meta', we also temporarily change
# the default device. Finally, we convert all params to 'meta' tensors that are not params.
kwargs["device"] = "meta"
with torch.device("meta"):
try:
self._empty_net = self._build_single_net(**kwargs)
except NotImplementedError as err:
if "Cannot copy out of meta tensor" in str(err):
raise RuntimeError(
"The network was built using `factory().to(device), build the network directly "
"on device using `factory(device=device)` instead."
)
# Remove all parameters
TensorDict.from_module(self._empty_net).data.to("meta").to_module(
self._empty_net
)
@property
def vmap_randomness(self):
if self.initialized:
return self._vmap_randomness
# The class _BatchedUninitializedParameter and buffer are not batched
# by vmap so using "different" will raise an exception because vmap can't find
# the batch dimension. This is ok though since we won't have the same config
# for every element (as one might expect from "same").
return "same"
def _make_params(self, agent_networks):
if self.share_params:
self.params = TensorDict.from_module(agent_networks[0], as_module=True)
else:
self.params = TensorDict.from_modules(*agent_networks, as_module=True)
@abc.abstractmethod
def _build_single_net(self, *, device, **kwargs):
...
@abc.abstractmethod
def _pre_forward_check(self, inputs):
...
@staticmethod
def vmap_func_module(module, *args, **kwargs):
def exec_module(params, *input):
with params.to_module(module):
return module(*input)
return torch.vmap(exec_module, *args, **kwargs)
[docs] def forward(self, *inputs: Tuple[torch.Tensor]) -> torch.Tensor:
if len(inputs) > 1:
inputs = torch.cat([*inputs], -1)
else:
inputs = inputs[0]
inputs = self._pre_forward_check(inputs)
# If parameters are not shared, each agent has its own network
if not self.share_params:
if self.centralized:
output = self.vmap_func_module(
self._empty_net, (0, None), (-2,), randomness=self.vmap_randomness
)(self.params, inputs)
else:
output = self.vmap_func_module(
self._empty_net,
(0, self.agent_dim),
(-2,),
randomness=self.vmap_randomness,
)(self.params, inputs)
# If parameters are shared, agents use the same network
else:
with self.params.to_module(self._empty_net):
output = self._empty_net(inputs)
if self.centralized:
# If the parameters are shared, and it is centralized, all agents will have the same output
# We expand it to maintain the agent dimension, but values will be the same for all agents
n_agent_outputs = output.shape[-1]
output = output.view(*output.shape[:-1], n_agent_outputs)
output = output.unsqueeze(-2)
output = output.expand(
*output.shape[:-2], self.n_agents, n_agent_outputs
)
if output.shape[-2] != (self.n_agents):
raise ValueError(
f"Multi-agent network expected output with shape[-2]={self.n_agents}"
f" but got {output.shape}"
)
return output
[docs] def get_stateful_net(self, copy: bool = True):
"""Returns a stateful version of the network.
This can be used to initialize parameters.
Such networks will often not be callable out-of-the-box and will require a `vmap` call
to be executable.
Args:
copy (bool, optional): if ``True``, a deepcopy of the network is made.
Defaults to ``True``.
If the parameters are modified in-place (recommended) there is no need to copy the
parameters back into the MARL module.
See :meth:`~.from_stateful_net` for details on how to re-populate the MARL model with
parameters that have been re-initialized out-of-place.
Examples:
>>> 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)
"""
if copy:
try:
net = deepcopy(self._empty_net)
except RuntimeError as err:
raise RuntimeError(
"Failed to deepcopy the module, consider using copy=False."
) from err
else:
net = self._empty_net
self.params.to_module(net)
return net
[docs] def from_stateful_net(self, stateful_net: nn.Module):
"""Populates the parameters given a stateful version of the network.
See :meth:`~.get_stateful_net` for details on how to gather a stateful version of the network.
Args:
stateful_net (nn.Module): the stateful network from which the params should be
gathered.
"""
params = TensorDict.from_module(stateful_net, as_module=True)
keyset0 = set(params.keys(True, True))
keyset1 = set(self.params.keys(True, True))
if keyset0 != keyset1:
raise RuntimeError(
f"The keys of params and provided module differ: "
f"{keyset1-keyset0} are in self.params and not in the module, "
f"{keyset0-keyset1} are in the module but not in self.params."
)
self.params.data.update_(params.data)
def __repr__(self):
empty_net = self._empty_net
with self.params.to_module(empty_net):
module_repr = indent(str(empty_net), 4 * " ")
n_agents = indent(f"n_agents={self.n_agents}", 4 * " ")
share_params = indent(f"share_params={self.share_params}", 4 * " ")
centralized = indent(f"centralized={self.centralized}", 4 * " ")
agent_dim = indent(f"agent_dim={self.agent_dim}", 4 * " ")
return f"{self.__class__.__name__}(\n{module_repr},\n{n_agents},\n{share_params},\n{centralized},\n{agent_dim})"
[docs] def reset_parameters(self):
"""Resets the parameters of the model."""
def vmap_reset_module(module, *args, **kwargs):
def reset_module(params):
with params.to_module(module):
_reset_parameters_recursive(module)
return params
return torch.vmap(reset_module, *args, **kwargs)
if not self.share_params:
vmap_reset_module(self._empty_net, randomness="different")(self.params)
else:
with self.params.to_module(self._empty_net):
_reset_parameters_recursive(self._empty_net)
[docs]class MultiAgentMLP(MultiAgentNetBase):
"""Mult-agent MLP.
This is an MLP that can be used in multi-agent contexts.
For example, as a policy or as a value function.
See `examples/multiagent` for examples.
It expects inputs with shape (*B, n_agents, n_agent_inputs)
It returns outputs with shape (*B, n_agents, n_agent_outputs)
If `share_params` is True, the same MLP will be used to make the forward pass for all agents (homogeneous policies).
Otherwise, each agent will use a different MLP to process its input (heterogeneous policies).
If `centralized` is True, each agent will use the inputs of all agents to compute its output
(n_agent_inputs * n_agents will be the number of inputs for one agent).
Otherwise, each agent will only use its data as input.
Args:
n_agent_inputs (int or None): number of inputs for each agent. If ``None``,
the number of inputs is lazily instantiated during the first call.
n_agent_outputs (int): number of outputs for each agent.
n_agents (int): number of agents.
centralized (bool): If `centralized` is True, each agent will use the inputs of all agents to compute its output
(n_agent_inputs * n_agents will be the number of inputs for one agent).
Otherwise, each agent will only use its data as input.
share_params (bool): If `share_params` is True, the same MLP will be used to make the forward pass
for all agents (homogeneous policies). Otherwise, each agent will use a different MLP to process
its input (heterogeneous policies).
device (str or toech.device, optional): device to create the module on.
depth (int, optional): depth of the network. A depth of 0 will produce a single linear layer network with the
desired input and output size. A length of 1 will create 2 linear layers etc. If no depth is indicated,
the depth information should be contained in the num_cells argument (see below). If num_cells is an
iterable and depth is indicated, both should match: len(num_cells) must be equal to depth.
default: 3.
num_cells (int or Sequence[int], optional): number of cells of every layer in between the input and output. If
an integer is provided, every layer will have the same number of cells. If an iterable is provided,
the linear layers out_features will match the content of num_cells.
default: 32.
activation_class (Type[nn.Module]): activation class to be used.
default: nn.Tanh.
**kwargs: for :class:`torchrl.modules.models.MLP` can be passed to customize the MLPs.
.. note:: to initialize the MARL module parameters with the `torch.nn.init`
module, please refer to :meth:`~.get_stateful_net` and :meth:`~.from_stateful_net`
methods.
Examples:
>>> 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)
>>> # instantiate a local network shared by all agents (e.g. a parameter-shared policy)
>>> mlp = MultiAgentMLP(
... n_agent_inputs=n_agent_inputs,
... n_agent_outputs=n_agent_outputs,
... n_agents=n_agents,
... centralized=False,
... share_params=True,
... depth=2,
... )
>>> print(mlp)
MultiAgentMLP(
(agent_networks): ModuleList(
(0): MLP(
(0): Linear(in_features=3, out_features=32, bias=True)
(1): Tanh()
(2): Linear(in_features=32, out_features=32, bias=True)
(3): Tanh()
(4): Linear(in_features=32, out_features=2, bias=True)
)
)
)
>>> assert mlp(obs).shape == (batch, n_agents, n_agent_outputs)
Now let's instantiate a centralized network shared by all agents (e.g. a centalised value function)
>>> mlp = MultiAgentMLP(
... n_agent_inputs=n_agent_inputs,
... n_agent_outputs=n_agent_outputs,
... n_agents=n_agents,
... centralized=True,
... share_params=True,
... depth=2,
... )
>>> print(mlp)
MultiAgentMLP(
(agent_networks): ModuleList(
(0): MLP(
(0): Linear(in_features=18, out_features=32, bias=True)
(1): Tanh()
(2): Linear(in_features=32, out_features=32, bias=True)
(3): Tanh()
(4): Linear(in_features=32, out_features=2, bias=True)
)
)
)
We can see that the input to the first layer is n_agents * n_agent_inputs,
this is because in the case the net acts as a centralized mlp (like a single huge agent)
>>> assert mlp(obs).shape == (batch, n_agents, n_agent_outputs)
Outputs will be identical for all agents.
Now we can do both examples just shown but with an independent set of parameters for each agent
Let's show the centralized=False case.
>>> 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,
... )
>>> print(mlp)
MultiAgentMLP(
(agent_networks): ModuleList(
(0-5): 6 x MLP(
(0): Linear(in_features=3, out_features=32, bias=True)
(1): Tanh()
(2): Linear(in_features=32, out_features=32, bias=True)
(3): Tanh()
(4): Linear(in_features=32, out_features=2, bias=True)
)
)
)
We can see that this is the same as in the first example, but now we have 6 MLPs, one per agent!
>>> assert mlp(obs).shape == (batch, n_agents, n_agent_outputs)
"""
def __init__(
self,
n_agent_inputs: int | None,
n_agent_outputs: int,
n_agents: int,
centralized: bool | None = None,
share_params: bool | None = None,
device: Optional[DEVICE_TYPING] = None,
depth: Optional[int] = None,
num_cells: Optional[Union[Sequence, int]] = None,
activation_class: Optional[Type[nn.Module]] = nn.Tanh,
**kwargs,
):
self.n_agents = n_agents
self.n_agent_inputs = n_agent_inputs
self.n_agent_outputs = n_agent_outputs
self.share_params = share_params
self.centralized = centralized
self.num_cells = num_cells
self.activation_class = activation_class
self.depth = depth
super().__init__(
n_agents=n_agents,
centralized=centralized,
share_params=share_params,
device=device,
agent_dim=-2,
**kwargs,
)
def _pre_forward_check(self, inputs):
if inputs.shape[-2] != self.n_agents:
raise ValueError(
f"Multi-agent network expected input with shape[-2]={self.n_agents},"
f" but got {inputs.shape}"
)
# If the model is centralized, agents have full observability
if self.centralized:
inputs = inputs.flatten(-2, -1)
return inputs
def _build_single_net(self, *, device, **kwargs):
n_agent_inputs = self.n_agent_inputs
if self.centralized and n_agent_inputs is not None:
n_agent_inputs = self.n_agent_inputs * self.n_agents
return MLP(
in_features=n_agent_inputs,
out_features=self.n_agent_outputs,
depth=self.depth,
num_cells=self.num_cells,
activation_class=self.activation_class,
device=device,
**kwargs,
)
[docs]class MultiAgentConvNet(MultiAgentNetBase):
"""Multi-agent CNN.
In MARL settings, agents may or may not share the same policy for their actions: we say that the parameters can be shared or not. Similarly, a network may take the entire observation space (across agents) or on a per-agent basis to compute its output, which we refer to as "centralized" and "non-centralized", respectively.
It expects inputs with shape ``(*B, n_agents, channels, x, y)``.
.. note:: to initialize the MARL module parameters with the `torch.nn.init`
module, please refer to :meth:`~.get_stateful_net` and :meth:`~.from_stateful_net`
methods.
Args:
n_agents (int): number of agents.
centralized (bool): If ``True``, each agent will use the inputs of all agents to compute its output, resulting in input of shape ``(*B, n_agents * channels, x, y)``. Otherwise, each agent will only use its data as input.
share_params (bool): If ``True``, the same :class:`~torchrl.modules.ConvNet` will be used to make the forward pass
for all agents (homogeneous policies). Otherwise, each agent will use a different :class:`~torchrl.modules.ConvNet` to process
its input (heterogeneous policies).
Keyword Args:
in_features (int, optional): the input feature dimension. If left to ``None``,
a lazy module is used.
device (str or torch.device, optional): device to create the module on.
num_cells (int or Sequence[int], optional): number of cells of every layer in between the input and output. If
an integer is provided, every layer will have the same number of cells. If an iterable is provided,
the linear layers ``out_features`` will match the content of ``num_cells``.
kernel_sizes (int, Sequence[Union[int, Sequence[int]]]): Kernel size(s) of the convolutional network.
Defaults to ``5``.
strides (int or Sequence[int]): Stride(s) of the convolutional network. If iterable, the length must match the
depth, defined by the num_cells or depth arguments.
Defaults to ``2``.
activation_class (Type[nn.Module]): activation class to be used.
Default to :class:`torch.nn.ELU`.
**kwargs: for :class:`~torchrl.modules.models.ConvNet` can be passed to customize the ConvNet.
Examples:
>>> import torch
>>> from torchrl.modules import MultiAgentConvNet
>>> batch = (3,2)
>>> n_agents = 7
>>> channels, x, y = 3, 100, 100
>>> obs = torch.randn(*batch, n_agents, channels, x, y)
>>> # Let's consider a centralized network with shared parameters.
>>> cnn = MultiAgentConvNet(
... n_agents,
... centralized = True,
... share_params = True
... )
>>> print(cnn)
MultiAgentConvNet(
(agent_networks): ModuleList(
(0): ConvNet(
(0): LazyConv2d(0, 32, kernel_size=(5, 5), stride=(2, 2))
(1): ELU(alpha=1.0)
(2): Conv2d(32, 32, kernel_size=(5, 5), stride=(2, 2))
(3): ELU(alpha=1.0)
(4): Conv2d(32, 32, kernel_size=(5, 5), stride=(2, 2))
(5): ELU(alpha=1.0)
(6): SquashDims()
)
)
)
>>> result = cnn(obs)
>>> # The final dimension of the resulting tensor would be determined based on the layer definition arguments and the shape of input 'obs'.
>>> print(result.shape)
torch.Size([3, 2, 7, 2592])
>>> # Since both observations and parameters are shared, we expect all agents to have identical outputs (eg. for a value function)
>>> print(all(result[0,0,0] == result[0,0,1]))
True
>>> # Alternatively, a local network with parameter sharing (eg. decentralized weight sharing policy)
>>> cnn = MultiAgentConvNet(
... n_agents,
... centralized = False,
... share_params = True
... )
>>> print(cnn)
MultiAgentConvNet(
(agent_networks): ModuleList(
(0): ConvNet(
(0): Conv2d(4, 32, kernel_size=(5, 5), stride=(2, 2))
(1): ELU(alpha=1.0)
(2): Conv2d(32, 32, kernel_size=(5, 5), stride=(2, 2))
(3): ELU(alpha=1.0)
(4): Conv2d(32, 32, kernel_size=(5, 5), stride=(2, 2))
(5): ELU(alpha=1.0)
(6): SquashDims()
)
)
)
>>> print(result.shape)
torch.Size([3, 2, 7, 2592])
>>> # Parameters are shared but not observations, hence each agent has a different output.
>>> print(all(result[0,0,0] == result[0,0,1]))
False
>>> # Or multiple local networks identical in structure but with differing weights.
>>> cnn = MultiAgentConvNet(
... n_agents,
... centralized = False,
... share_params = False
... )
>>> print(cnn)
MultiAgentConvNet(
(agent_networks): ModuleList(
(0-6): 7 x ConvNet(
(0): Conv2d(4, 32, kernel_size=(5, 5), stride=(2, 2))
(1): ELU(alpha=1.0)
(2): Conv2d(32, 32, kernel_size=(5, 5), stride=(2, 2))
(3): ELU(alpha=1.0)
(4): Conv2d(32, 32, kernel_size=(5, 5), stride=(2, 2))
(5): ELU(alpha=1.0)
(6): SquashDims()
)
)
)
>>> print(result.shape)
torch.Size([3, 2, 7, 2592])
>>> print(all(result[0,0,0] == result[0,0,1]))
False
>>> # Or where inputs are shared but not parameters.
>>> cnn = MultiAgentConvNet(
... n_agents,
... centralized = True,
... share_params = False
... )
>>> print(cnn)
MultiAgentConvNet(
(agent_networks): ModuleList(
(0-6): 7 x ConvNet(
(0): Conv2d(28, 32, kernel_size=(5, 5), stride=(2, 2))
(1): ELU(alpha=1.0)
(2): Conv2d(32, 32, kernel_size=(5, 5), stride=(2, 2))
(3): ELU(alpha=1.0)
(4): Conv2d(32, 32, kernel_size=(5, 5), stride=(2, 2))
(5): ELU(alpha=1.0)
(6): SquashDims()
)
)
)
>>> print(result.shape)
torch.Size([3, 2, 7, 2592])
>>> print(all(result[0,0,0] == result[0,0,1]))
False
"""
def __init__(
self,
n_agents: int,
centralized: bool | None = None,
share_params: bool | None = None,
*,
in_features: int | None = None,
device: DEVICE_TYPING | None = None,
num_cells: Sequence[int] | None = None,
kernel_sizes: Union[Sequence[Union[int, Sequence[int]]], int] = 5,
strides: Union[Sequence, int] = 2,
paddings: Union[Sequence, int] = 0,
activation_class: Type[nn.Module] = nn.ELU,
**kwargs,
):
self.in_features = in_features
self.num_cells = num_cells
self.strides = strides
self.kernel_sizes = kernel_sizes
self.paddings = paddings
self.activation_class = activation_class
super().__init__(
n_agents=n_agents,
centralized=centralized,
share_params=share_params,
device=device,
agent_dim=-4,
**kwargs,
)
def _build_single_net(self, *, device, **kwargs):
in_features = self.in_features
if self.centralized and in_features is not None:
in_features = in_features * self.n_agents
return ConvNet(
in_features=in_features,
num_cells=self.num_cells,
kernel_sizes=self.kernel_sizes,
strides=self.strides,
paddings=self.paddings,
activation_class=self.activation_class,
device=device,
**kwargs,
)
def _pre_forward_check(self, inputs):
if len(inputs.shape) < 4:
raise ValueError(
"""Multi-agent network expects (*batch_size, agent_index, x, y, channels)"""
)
if inputs.shape[-4] != self.n_agents:
raise ValueError(
f"""Multi-agent network expects {self.n_agents} but got {inputs.shape[-4]}"""
)
if self.centralized:
# If the model is centralized, agents have full observability
inputs = torch.flatten(inputs, -4, -3)
return inputs
class Mixer(nn.Module):
"""A multi-agent value mixer.
It transforms the local value of each agent's chosen action of shape (*B, self.n_agents, 1),
into a global value with shape (*B, 1).
Used with the :class:`torchrl.objectives.QMixerLoss`.
See `examples/multiagent/qmix_vdn.py` for examples.
Args:
n_agents (int): number of agents.
needs_state (bool): whether the mixer takes a global state as input.
state_shape (tuple or torch.Size): the shape of the state (excluding eventual leading batch dimensions).
device (str or torch.Device): torch device for the network.
Examples:
Creating a VDN mixer
>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModule
>>> from torchrl.modules.models.multiagent import VDNMixer
>>> n_agents = 4
>>> vdn = TensorDictModule(
... module=VDNMixer(
... n_agents=n_agents,
... device="cpu",
... ),
... in_keys=[("agents","chosen_action_value")],
... out_keys=["chosen_action_value"],
... )
>>> td = TensorDict({"agents": TensorDict({"chosen_action_value": torch.zeros(32, n_agents, 1)}, [32, n_agents])}, [32])
>>> td
TensorDict(
fields={
agents: TensorDict(
fields={
chosen_action_value: Tensor(shape=torch.Size([32, 4, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([32, 4]),
device=None,
is_shared=False)},
batch_size=torch.Size([32]),
device=None,
is_shared=False)
>>> vdn(td)
TensorDict(
fields={
agents: TensorDict(
fields={
chosen_action_value: Tensor(shape=torch.Size([32, 4, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([32, 4]),
device=None,
is_shared=False),
chosen_action_value: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([32]),
device=None,
is_shared=False)
Creating a QMix mixer
>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModule
>>> from torchrl.modules.models.multiagent import QMixer
>>> n_agents = 4
>>> qmix = TensorDictModule(
... module=QMixer(
... state_shape=(64, 64, 3),
... mixing_embed_dim=32,
... n_agents=n_agents,
... device="cpu",
... ),
... in_keys=[("agents", "chosen_action_value"), "state"],
... out_keys=["chosen_action_value"],
... )
>>> td = TensorDict({"agents": TensorDict({"chosen_action_value": torch.zeros(32, n_agents, 1)}, [32, n_agents]), "state": torch.zeros(32, 64, 64, 3)}, [32])
>>> td
TensorDict(
fields={
agents: TensorDict(
fields={
chosen_action_value: Tensor(shape=torch.Size([32, 4, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([32, 4]),
device=None,
is_shared=False),
state: Tensor(shape=torch.Size([32, 64, 64, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([32]),
device=None,
is_shared=False)
>>> vdn(td)
TensorDict(
fields={
agents: TensorDict(
fields={
chosen_action_value: Tensor(shape=torch.Size([32, 4, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([32, 4]),
device=None,
is_shared=False),
chosen_action_value: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.float32, is_shared=False),
state: Tensor(shape=torch.Size([32, 64, 64, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([32]),
device=None,
is_shared=False)
"""
def __init__(
self,
n_agents: int,
needs_state: bool,
state_shape: Union[Tuple[int, ...], torch.Size],
device: DEVICE_TYPING,
):
super().__init__()
self.n_agents = n_agents
self.device = device
self.needs_state = needs_state
self.state_shape = state_shape
def forward(self, *inputs: Tuple[torch.Tensor]) -> torch.Tensor:
"""Forward pass of the mixer.
Args:
*inputs: The first input should be the value of the chosen action of shape (*B, self.n_agents, 1),
representing the local q value of each agent.
The second input (optional, used only in some mixers)
is the shared state of all agents of shape (*B, *self.state_shape).
Returns:
The global value of the chosen actions obtained after mixing, with shape (*B, 1)
"""
if not self.needs_state:
if len(inputs) > 1:
raise ValueError(
"Mixer that doesn't need state was passed more than 1 input"
)
chosen_action_value = inputs[0]
else:
if len(inputs) != 2:
raise ValueError("Mixer that needs state was passed more than 2 inputs")
chosen_action_value, state = inputs
if state.shape[-len(self.state_shape) :] != self.state_shape:
raise ValueError(
f"Mixer network expected state with ending shape {self.state_shape},"
f" but got state shape {state.shape}"
)
if chosen_action_value.shape[-2:] != (self.n_agents, 1):
raise ValueError(
f"Mixer network expected chosen_action_value with last 2 dimensions {(self.n_agents,1)},"
f" but got {chosen_action_value.shape}"
)
batch_dims = chosen_action_value.shape[:-2]
if not self.needs_state:
output = self.mix(chosen_action_value, None)
else:
output = self.mix(chosen_action_value, state)
if output.shape != (*batch_dims, 1):
raise ValueError(
f"Mixer network expected output with same shape as input minus the multi-agent dimension,"
f" but got {output.shape}"
)
return output
def mix(self, chosen_action_value: torch.Tensor, state: torch.Tensor):
"""Forward pass for the mixer.
Args:
chosen_action_value: Tensor of shape [*B, n_agents]
Returns:
chosen_action_value: Tensor of shape [*B]
"""
raise NotImplementedError
[docs]class VDNMixer(Mixer):
"""Value-Decomposition Network mixer.
Mixes the local Q values of the agents into a global Q value by summing them together.
From the paper https://arxiv.org/abs/1706.05296 .
It transforms the local value of each agent's chosen action of shape (*B, self.n_agents, 1),
into a global value with shape (*B, 1).
Used with the :class:`torchrl.objectives.QMixerLoss`.
See `examples/multiagent/qmix_vdn.py` for examples.
Args:
n_agents (int): number of agents.
device (str or torch.Device): torch device for the network.
Examples:
>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModule
>>> from torchrl.modules.models.multiagent import VDNMixer
>>> n_agents = 4
>>> vdn = TensorDictModule(
... module=VDNMixer(
... n_agents=n_agents,
... device="cpu",
... ),
... in_keys=[("agents","chosen_action_value")],
... out_keys=["chosen_action_value"],
... )
>>> td = TensorDict({"agents": TensorDict({"chosen_action_value": torch.zeros(32, n_agents, 1)}, [32, n_agents])}, [32])
>>> td
TensorDict(
fields={
agents: TensorDict(
fields={
chosen_action_value: Tensor(shape=torch.Size([32, 4, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([32, 4]),
device=None,
is_shared=False)},
batch_size=torch.Size([32]),
device=None,
is_shared=False)
>>> vdn(td)
TensorDict(
fields={
agents: TensorDict(
fields={
chosen_action_value: Tensor(shape=torch.Size([32, 4, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([32, 4]),
device=None,
is_shared=False),
chosen_action_value: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([32]),
device=None,
is_shared=False)
"""
def __init__(
self,
n_agents: int,
device: DEVICE_TYPING,
):
super().__init__(
needs_state=False,
state_shape=torch.Size([]),
n_agents=n_agents,
device=device,
)
[docs] def mix(self, chosen_action_value: torch.Tensor, state: torch.Tensor):
return chosen_action_value.sum(dim=-2)
[docs]class QMixer(Mixer):
"""QMix mixer.
Mixes the local Q values of the agents into a global Q value through a monotonic
hyper-network whose parameters are obtained from a global state.
From the paper https://arxiv.org/abs/1803.11485 .
It transforms the local value of each agent's chosen action of shape (*B, self.n_agents, 1),
into a global value with shape (*B, 1).
Used with the :class:`torchrl.objectives.QMixerLoss`.
See `examples/multiagent/qmix_vdn.py` for examples.
Args:
state_shape (tuple or torch.Size): the shape of the state (excluding eventual leading batch dimensions).
mixing_embed_dim (int): the size of the mixing embedded dimension.
n_agents (int): number of agents.
device (str or torch.Device): torch device for the network.
Examples:
>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModule
>>> from torchrl.modules.models.multiagent import QMixer
>>> n_agents = 4
>>> qmix = TensorDictModule(
... module=QMixer(
... state_shape=(64, 64, 3),
... mixing_embed_dim=32,
... n_agents=n_agents,
... device="cpu",
... ),
... in_keys=[("agents", "chosen_action_value"), "state"],
... out_keys=["chosen_action_value"],
... )
>>> td = TensorDict({"agents": TensorDict({"chosen_action_value": torch.zeros(32, n_agents, 1)}, [32, n_agents]), "state": torch.zeros(32, 64, 64, 3)}, [32])
>>> td
TensorDict(
fields={
agents: TensorDict(
fields={
chosen_action_value: Tensor(shape=torch.Size([32, 4, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([32, 4]),
device=None,
is_shared=False),
state: Tensor(shape=torch.Size([32, 64, 64, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([32]),
device=None,
is_shared=False)
>>> vdn(td)
TensorDict(
fields={
agents: TensorDict(
fields={
chosen_action_value: Tensor(shape=torch.Size([32, 4, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([32, 4]),
device=None,
is_shared=False),
chosen_action_value: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.float32, is_shared=False),
state: Tensor(shape=torch.Size([32, 64, 64, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([32]),
device=None,
is_shared=False)
"""
def __init__(
self,
state_shape: Union[Tuple[int, ...], torch.Size],
mixing_embed_dim: int,
n_agents: int,
device: DEVICE_TYPING,
):
super().__init__(
needs_state=True, state_shape=state_shape, n_agents=n_agents, device=device
)
self.embed_dim = mixing_embed_dim
self.state_dim = int(np.prod(state_shape))
self.hyper_w_1 = nn.Linear(
self.state_dim, self.embed_dim * self.n_agents, device=self.device
)
self.hyper_w_final = nn.Linear(
self.state_dim, self.embed_dim, device=self.device
)
# State dependent bias for hidden layer
self.hyper_b_1 = nn.Linear(self.state_dim, self.embed_dim, device=self.device)
# V(s) instead of a bias for the last layers
self.V = nn.Sequential(
nn.Linear(self.state_dim, self.embed_dim, device=self.device),
nn.ReLU(),
nn.Linear(self.embed_dim, 1, device=self.device),
)
[docs] def mix(self, chosen_action_value: torch.Tensor, state: torch.Tensor):
bs = chosen_action_value.shape[:-2]
state = state.view(-1, self.state_dim)
chosen_action_value = chosen_action_value.view(-1, 1, self.n_agents)
# First layer
w1 = torch.abs(self.hyper_w_1(state))
b1 = self.hyper_b_1(state)
w1 = w1.view(-1, self.n_agents, self.embed_dim)
b1 = b1.view(-1, 1, self.embed_dim)
hidden = nn.functional.elu(
torch.bmm(chosen_action_value, w1) + b1
) # [-1, 1, self.embed_dim]
# Second layer
w_final = torch.abs(self.hyper_w_final(state))
w_final = w_final.view(-1, self.embed_dim, 1)
# State-dependent bias
v = self.V(state).view(-1, 1, 1)
# Compute final output
y = torch.bmm(hidden, w_final) + v # [-1, 1, 1]
# Reshape and return
q_tot = y.view(*bs, 1)
return q_tot
