# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
from __future__ import annotations
from dataclasses import dataclass, MISSING
from typing import Optional, Sequence, Type
import torch
from tensordict import TensorDictBase
from torch import nn
from torchrl.modules import MLP, MultiAgentMLP
from benchmarl.models.common import Model, ModelConfig
[docs]
class Mlp(Model):
"""Multi layer perceptron model.
Args:
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.
layer_class (Type[nn.Module]): class to be used for the linear layers;
activation_class (Type[nn.Module]): activation class to be used.
activation_kwargs (dict, optional): kwargs to be used with the activation class;
norm_class (Type, optional): normalization class, if any.
norm_kwargs (dict, optional): kwargs to be used with the normalization layers;
num_feature_dims: number of dimensions to be considered as features.
"""
def __init__(
self,
**kwargs,
):
self.num_feature_dims = kwargs.pop("num_feature_dims", 1)
super().__init__(
input_spec=kwargs.pop("input_spec"),
output_spec=kwargs.pop("output_spec"),
agent_group=kwargs.pop("agent_group"),
input_has_agent_dim=kwargs.pop("input_has_agent_dim"),
n_agents=kwargs.pop("n_agents"),
centralised=kwargs.pop("centralised"),
share_params=kwargs.pop("share_params"),
device=kwargs.pop("device"),
action_spec=kwargs.pop("action_spec"),
model_index=kwargs.pop("model_index"),
is_critic=kwargs.pop("is_critic"),
)
self.input_features = sum(
[
torch.prod(torch.tensor(spec.shape[-self.num_feature_dims :])).item()
for spec in self.input_spec.values(True, True)
]
)
self.output_features = self.output_leaf_spec.shape[-1]
if self.input_has_agent_dim:
self.mlp = MultiAgentMLP(
n_agent_inputs=self.input_features,
n_agent_outputs=self.output_features,
n_agents=self.n_agents,
centralised=self.centralised,
share_params=self.share_params,
device=self.device,
**kwargs,
)
else:
self.mlp = nn.ModuleList(
[
MLP(
in_features=self.input_features,
out_features=self.output_features,
device=self.device,
**kwargs,
)
for _ in range(self.n_agents if not self.share_params else 1)
]
)
def _perform_checks(self):
super()._perform_checks()
input_shape = None
for input_key, input_spec in self.input_spec.items(True, True):
if (
self.input_has_agent_dim
and len(input_spec.shape) == self.num_feature_dims + 1
) or (
not self.input_has_agent_dim
and len(input_spec.shape) == self.num_feature_dims
):
if input_shape is None:
input_shape = input_spec.shape[: -self.num_feature_dims]
else:
if input_spec.shape[: -self.num_feature_dims] != input_shape:
raise ValueError(
f"MLP inputs should all have the same shape up to the last {self.num_feature_dims} dimensions, got {self.input_spec}"
)
else:
raise ValueError(
f"MLP input value {input_key} from {self.input_spec} has an invalid shape, maybe you need a CNN or more feature dimensions?"
)
if self.input_has_agent_dim:
if input_shape[-1] != self.n_agents:
raise ValueError(
"If the MLP input has the agent dimension,"
f" the second to last spec dimension should be the number of agents, got {self.input_spec}"
)
if (
self.output_has_agent_dim
and self.output_leaf_spec.shape[-2] != self.n_agents
):
raise ValueError(
"If the MLP output has the agent dimension,"
" the second to last spec dimension should be the number of agents"
)
[docs]
def _forward(self, tensordict: TensorDictBase) -> TensorDictBase:
# Gather in_key and flatten the last self.num_feature_dims dimensions
input = torch.cat(
[
torch.flatten(tensordict.get(in_key), start_dim=-self.num_feature_dims)
for in_key in self.in_keys
],
dim=-1,
)
# Has multi-agent input dimension
if self.input_has_agent_dim:
res = self.mlp.forward(input)
if not self.output_has_agent_dim:
# If we are here the module is centralised and parameter shared.
# Thus the multi-agent dimension has been expanded,
# We remove it without loss of data
res = res[..., 0, :]
# Does not have multi-agent input dimension
else:
if not self.share_params:
res = torch.stack(
[net(input) for net in self.mlp],
dim=-2,
)
else:
res = self.mlp[0](input)
tensordict.set(self.out_key, res)
return tensordict
[docs]
@dataclass
class MlpConfig(ModelConfig):
"""Dataclass config for a :class:`~benchmarl.models.Mlp`."""
num_cells: Sequence[int] = MISSING
layer_class: Type[nn.Module] = MISSING
activation_class: Type[nn.Module] = MISSING
activation_kwargs: Optional[dict] = None
norm_class: Type[nn.Module] = None
norm_kwargs: Optional[dict] = None
num_feature_dims: int = 1
[docs]
@staticmethod
def associated_class():
return Mlp