[Feature] Add per-agent advantage normalization support for IPPO/MAPPO

benchmarl/algorithms/ippo.py

@@ -5,7 +5,7 @@
 #
 
 from dataclasses import dataclass, MISSING
-from typing import Dict, Iterable, Tuple, Type
+from typing import Dict, Iterable, List, Tuple, Type
 
 import torch
 from tensordict import TensorDictBase
@@ -39,6 +39,13 @@ class Ippo(Algorithm):
         minibatch_advantage (bool): if ``True``, advantage computation is perfomend on minibatches of size
             ``experiment.config.on_policy_minibatch_size`` instead of the full
             ``experiment.config.on_policy_collected_frames_per_batch``, this helps not exploding memory usage
+        normalize_advantage (bool): if ``True``, the advantage will be normalized
+            before being used. Defaults to ``False``.
+            See ``normalize_advantage_exclude_dims`` for which dimensions are kept independent.
+        normalize_advantage_exclude_dims (list[int]): dimensions to exclude from the advantage
+            standardization. Negative dimensions are valid. This is useful in multiagent (or multiobjective) settings
+            where the agent (or objective) dimension may be excluded from the reductions. Default: ``(-2,)``
+            (``[-2]`` in YAML). ``()`` normalizes across all dimensions (including agent 1).
 
     """
 
@@ -53,6 +60,8 @@ def __init__(
         scale_mapping: str,
         use_tanh_normal: bool,
         minibatch_advantage: bool,
+        normalize_advantage: bool,
+        normalize_advantage_exclude_dims: List[int],
         **kwargs
     ):
         super().__init__(**kwargs)
@@ -66,6 +75,8 @@ def __init__(
         self.scale_mapping = scale_mapping
         self.use_tanh_normal = use_tanh_normal
         self.minibatch_advantage = minibatch_advantage
+        self.normalize_advantage = normalize_advantage
+        self.normalize_advantage_exclude_dims = normalize_advantage_exclude_dims
 
     #############################
     # Overridden abstract methods
@@ -82,7 +93,8 @@ def _get_loss(
             entropy_coeff=self.entropy_coef,
             critic_coeff=self.critic_coef,
             loss_critic_type=self.loss_critic_type,
-            normalize_advantage=False,
+            normalize_advantage=self.normalize_advantage,
+            normalize_advantage_exclude_dims=self.normalize_advantage_exclude_dims,
         )
         loss_module.set_keys(
             reward=(group, "reward"),
@@ -309,6 +321,8 @@ class IppoConfig(AlgorithmConfig):
     scale_mapping: str = MISSING
     use_tanh_normal: bool = MISSING
     minibatch_advantage: bool = MISSING
+    normalize_advantage: bool = MISSING
+    normalize_advantage_exclude_dims: List[int] = MISSING
 
     @staticmethod
     def associated_class() -> Type[Algorithm]:

benchmarl/algorithms/mappo.py

@@ -5,7 +5,7 @@
 #
 
 from dataclasses import dataclass, MISSING
-from typing import Dict, Iterable, Tuple, Type
+from typing import Dict, Iterable, List, Tuple, Type
 
 import torch
 from tensordict import TensorDictBase
@@ -43,6 +43,13 @@ class Mappo(Algorithm):
         minibatch_advantage (bool): if ``True``, advantage computation is perfomend on minibatches of size
             ``experiment.config.on_policy_minibatch_size`` instead of the full
             ``experiment.config.on_policy_collected_frames_per_batch``, this helps not exploding memory usage
+        normalize_advantage (bool): if ``True``, the advantage will be normalized
+            before being used. Defaults to ``False``.
+            See ``normalize_advantage_exclude_dims`` for which dimensions are kept independent.
+        normalize_advantage_exclude_dims (list[int]): dimensions to exclude from the advantage
+            standardization. Negative dimensions are valid. This is useful in multiagent (or multiobjective) settings
+            where the agent (or objective) dimension may be excluded from the reductions. Default: ``(-2,)``
+            (``[-2]`` in YAML). ``()`` normalizes across all dimensions (including agent 1).
 
     """
 
@@ -57,6 +64,8 @@ def __init__(
         scale_mapping: str,
         use_tanh_normal: bool,
         minibatch_advantage: bool,
+        normalize_advantage: bool,
+        normalize_advantage_exclude_dims: List[int],
         **kwargs
     ):
         super().__init__(**kwargs)
@@ -70,6 +79,8 @@ def __init__(
         self.scale_mapping = scale_mapping
         self.use_tanh_normal = use_tanh_normal
         self.minibatch_advantage = minibatch_advantage
+        self.normalize_advantage = normalize_advantage
+        self.normalize_advantage_exclude_dims = normalize_advantage_exclude_dims
 
     #############################
     # Overridden abstract methods
@@ -86,7 +97,8 @@ def _get_loss(
             entropy_coeff=self.entropy_coef,
             critic_coeff=self.critic_coef,
             loss_critic_type=self.loss_critic_type,
-            normalize_advantage=False,
+            normalize_advantage=self.normalize_advantage,
+            normalize_advantage_exclude_dims=self.normalize_advantage_exclude_dims,
         )
         loss_module.set_keys(
             reward=(group, "reward"),
@@ -332,6 +344,8 @@ class MappoConfig(AlgorithmConfig):
     scale_mapping: str = MISSING
     use_tanh_normal: bool = MISSING
     minibatch_advantage: bool = MISSING
+    normalize_advantage: bool = MISSING
+    normalize_advantage_exclude_dims: List[int] = MISSING
 
     @staticmethod
     def associated_class() -> Type[Algorithm]:

benchmarl/conf/algorithm/ippo.yaml

@@ -12,3 +12,5 @@ lmbda: 0.9
 scale_mapping: "biased_softplus_1.0"
 use_tanh_normal: True
 minibatch_advantage: False
+normalize_advantage: False
+normalize_advantage_exclude_dims: [-2]

benchmarl/conf/algorithm/mappo.yaml

@@ -13,3 +13,5 @@ lmbda: 0.9
 scale_mapping: "biased_softplus_1.0"
 use_tanh_normal: True
 minibatch_advantage: False
+normalize_advantage: False
+normalize_advantage_exclude_dims: [-2]

test/test_algorithm.py

@@ -4,12 +4,28 @@
 #  LICENSE file in the root directory of this source tree.
 #
 
-import pytest
+import copy
+import warnings
+from dataclasses import asdict
 
-from benchmarl.algorithms import algorithm_config_registry
+import pytest
+import torch
+from benchmarl.algorithms import (
+    algorithm_config_registry,
+    Ippo,
+    IppoConfig,
+    Mappo,
+    MappoConfig,
+)
 from benchmarl.algorithms.common import AlgorithmConfig
 from benchmarl.hydra_config import load_algorithm_config_from_hydra
+from benchmarl.models import MlpConfig
 from hydra import compose, initialize
+from tensordict import TensorDict
+from tensordict.nn import TensorDictModule
+from torch import nn
+from torchrl.data.tensor_specs import Categorical, Composite, OneHot, Unbounded
+from torchrl.objectives import ClipPPOLoss
 
 
 @pytest.mark.parametrize("algo_name", algorithm_config_registry.keys())
@@ -24,3 +40,96 @@ def test_loading_algorithms(algo_name):
         )
         algo_config: AlgorithmConfig = load_algorithm_config_from_hydra(cfg.algorithm)
         assert algo_config == algorithm_config_registry[algo_name].get_from_yaml()
+
+
+def _make_experiment(stub_algo_config):
+    from types import SimpleNamespace
+
+    model_config = MlpConfig(num_cells=[4], activation_class=nn.Tanh, layer_class=nn.Linear)
+    critic_model_config = copy.deepcopy(model_config)
+    critic_model_config.is_critic = True
+
+    return SimpleNamespace(
+        config=SimpleNamespace(
+            train_device="cpu",
+            buffer_device="cpu",
+            gamma=0.99,
+            share_policy_params=True,
+        ),
+        model_config=model_config,
+        critic_model_config=critic_model_config,
+        on_policy=True,
+        group_map={"agents": [0, 1]},
+        observation_spec=Composite({"agents": Composite({"observation": Unbounded(shape=(2, 3))}, shape=(2,))}),
+        action_spec=Composite({"agents": Composite({"action": Categorical(shape=(2,), n=2)}, shape=(2,))}),
+        state_spec=None,
+        action_mask_spec=None,
+        algorithm_config=stub_algo_config,
+    )
+
+def _make_ppo_config(algo_config_cls, normalize_advantage, exclude_dims):
+    return algo_config_cls(
+        share_param_critic=False,
+        clip_epsilon=0.2,
+        entropy_coef=0.0,
+        critic_coef=1.0,
+        loss_critic_type="l2",
+        lmbda=0.9,
+        scale_mapping="biased_softplus_1.0",
+        use_tanh_normal=True,
+        minibatch_advantage=False,
+        normalize_advantage=normalize_advantage,
+        normalize_advantage_exclude_dims=exclude_dims,
+    )
+
+@pytest.fixture(params=[(Ippo, IppoConfig), (Mappo, MappoConfig)], ids=["ippo", "mappo"])
+def ppo_setup(request):
+    algo_cls, algo_config_cls = request.param
+
+    def _make(normalize_advantage, exclude_dims):
+        algo_config = _make_ppo_config(algo_config_cls, normalize_advantage, exclude_dims)
+        experiment = _make_experiment(algo_config)
+        algo = algo_cls(**asdict(algo_config), experiment=experiment)
+        continuous = not isinstance(experiment.action_spec["agents", "action"], (Categorical, OneHot))
+        return algo, continuous
+    return _make
+
+
+@pytest.mark.parametrize("normalize_advantage,exclude_dims",[(True, [-2]), (True, []), (False, [-2])])
+def test_ppo_advantage_normalization_init(ppo_setup, normalize_advantage, exclude_dims):
+    algo, continuous = ppo_setup(normalize_advantage, exclude_dims)
+    loss_module, use_target = algo._get_loss(
+        "agents",
+        policy_for_loss=TensorDictModule(lambda td, **kwargs: td, in_keys=[], out_keys=[]),
+        continuous=continuous,
+    )
+    assert use_target is False
+    assert isinstance(loss_module, ClipPPOLoss)
+    assert loss_module.normalize_advantage is normalize_advantage
+    assert loss_module.normalize_advantage_exclude_dims == exclude_dims
+
+
+def test_ppo_advantage_normalization(ppo_setup):
+    algo, continuous = ppo_setup(True, (-2,))
+    policy = algo.get_policy_for_loss("agents")
+    loss_module, _ = algo._get_loss("agents", policy, continuous=continuous)
+
+    torch.manual_seed(0)
+    batch, n_agents = 64, 2
+    tensordict = TensorDict(batch_size=[batch])
+    tensordict["agents", "observation"] = torch.rand(batch, n_agents, 3)
+    tensordict["agents", "action"] = torch.randint(0, 2, (batch, n_agents))
+    tensordict["agents", "log_prob"] = torch.zeros(batch, n_agents)
+    tensordict["agents", "value_target"] = torch.zeros(batch, n_agents, 1)
+    tensordict["agents", "advantage"] = torch.cat([torch.rand(batch, 1, 1), 5.0 * torch.rand(batch, 1, 1)], dim=1)
+
+    def loss_with(exclude_dims):
+        loss_module.normalize_advantage_exclude_dims = exclude_dims
+        with warnings.catch_warnings():
+            return loss_module(tensordict.clone())["loss_objective"].item()
+
+    with torch.no_grad():
+        joint = loss_with([])
+        per_agent = loss_with((-2,))
+
+    assert per_agent != pytest.approx(joint, rel=1e-3)