From e95b81128eaf072186bb9f02c563ebbdb9e2a01d Mon Sep 17 00:00:00 2001
From: Iliamsou <iliam.souami@gmail.com>
Date: Mon, 24 Aug 2026 10:59:02 +0200
Subject: [PATCH 1/2] Add COMA multi-agent objective

---
 test/objectives/test_coma.py              | 266 ++++++++++++++++++
 torchrl/objectives/__init__.py            |   5 +-
 torchrl/objectives/multiagent/__init__.py |   3 +-
 torchrl/objectives/multiagent/coma.py     | 311 ++++++++++++++++++++++
 4 files changed, 582 insertions(+), 3 deletions(-)
 create mode 100644 test/objectives/test_coma.py
 create mode 100644 torchrl/objectives/multiagent/coma.py

diff --git a/test/objectives/test_coma.py b/test/objectives/test_coma.py
new file mode 100644
index 00000000000..34a8d557c17
--- /dev/null
+++ b/test/objectives/test_coma.py
@@ -0,0 +1,266 @@
+import torch
+from tensordict import TensorDict
+from tensordict.nn import TensorDictModule
+from torch import nn
+from torchrl.modules import OneHotCategorical, ProbabilisticActor
+
+from torchrl.objectives.multiagent import COMALoss
+
+from torchrl.objectives.multiagent.coma import (
+    add_action_without_self,
+    add_joint_observation,
+    add_masked_joint_action,
+)
+
+
+class FixedActionValue(nn.Module):
+    def __init__(self, values: torch.Tensor):
+        super().__init__()
+        self.values = nn.Parameter(values.clone())
+
+    def forward(self, observation, action_without_self):
+        del action_without_self
+        return self.values.expand(*observation.shape[:-1], -1)
+
+
+def _one_hot(index, n_actions=3):
+    return torch.nn.functional.one_hot(torch.as_tensor(index), n_actions).to(torch.float)
+
+
+def _make_loss(
+    gamma=0.5, qvalue_loss_coef=0.5, entropy_coef=0.0, n_step=1, normalize_advantage=False, clip_epsilon=None
+):
+    obs_dim = 4
+    n_actions = 3
+    actor_net = nn.Linear(obs_dim, n_actions)
+    nn.init.zeros_(actor_net.weight)
+    nn.init.zeros_(actor_net.bias)
+    actor_module = TensorDictModule(
+        actor_net,
+        in_keys=[("agents", "observation")],
+        out_keys=[("agents", "logits")],
+    )
+    actor = ProbabilisticActor(
+        module=actor_module,
+        in_keys=[("agents", "logits")],
+        out_keys=[("agents", "action")],
+        distribution_class=OneHotCategorical,
+        return_log_prob=True,
+    )
+    qvalue_module = TensorDictModule(
+        FixedActionValue(torch.tensor([1.0, 2.0, 4.0])),
+        in_keys=[("agents", "observation"), ("agents", "action_without_self")],
+        out_keys=[("agents", "action_value")],
+    )
+    return COMALoss(
+        actor_network=actor,
+        qvalue_network=qvalue_module,
+        gamma=gamma,
+        qvalue_loss_coef=qvalue_loss_coef,
+        entropy_coef=entropy_coef,
+        n_step=n_step,
+        normalize_advantage=normalize_advantage,
+        clip_epsilon=clip_epsilon,
+    )
+
+
+def test_add_action_without_self_flattens_other_agents_actions():
+    action = _one_hot([[0, 1, 2]], n_actions=3)
+    tensordict = TensorDict({("agents", "action"): action}, batch_size=[1])
+
+    add_action_without_self(tensordict)
+
+    expected = torch.tensor(
+        [
+            [
+                [0.0, 1.0, 0.0, 0.0, 0.0, 1.0],
+                [1.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+                [1.0, 0.0, 0.0, 0.0, 1.0, 0.0],
+            ]
+        ]
+    )
+    torch.testing.assert_close(tensordict.get(("agents", "action_without_self")), expected)
+
+
+def test_coma_loss_uses_counterfactual_baseline_and_qvalue_target():
+    loss = _make_loss(qvalue_loss_coef=0.25, entropy_coef=0.0)
+    action = _one_hot([[0, 2]], n_actions=3)
+    tensordict = TensorDict(
+        {
+            ("agents", "observation"): torch.zeros(1, 2, 4),
+            ("agents", "action"): action,
+            "value_target": torch.zeros(1, 2, 1),
+        },
+        batch_size=[1],
+    )
+    add_action_without_self(tensordict)
+
+    loss_values = loss(tensordict)
+
+    expected_advantage = torch.tensor([1.0, 4.0]).mean() - torch.tensor([1.0, 2.0, 4.0]).mean()
+    expected_qvalue_loss = ((torch.tensor([1.0, 4.0]) ** 2).mean()) * 0.25
+    torch.testing.assert_close(loss_values["advantage"], expected_advantage)
+    torch.testing.assert_close(loss_values["loss_qvalue"], expected_qvalue_loss)
+    assert set(loss.out_keys).issubset(set(loss_values.keys()))
+
+
+def test_compute_value_target_bootstraps_along_time_dimension():
+    loss = _make_loss(gamma=0.5)
+    actions = _one_hot([[[0, 0], [1, 1], [2, 2]]], n_actions=3)
+    tensordict = TensorDict(
+        {
+            ("agents", "observation"): torch.zeros(1, 3, 2, 4),
+            ("agents", "action"): actions,
+            "next": {
+                "agents": {
+                    "reward": torch.zeros(1, 3, 2, 1),
+                    "done": torch.tensor([[[[False], [False]], [[False], [False]], [[True], [True]]]]),
+                }
+            },
+        },
+        batch_size=[1, 3],
+    )
+    add_action_without_self(tensordict)
+
+    loss.compute_value_target(tensordict)
+
+    expected = torch.tensor([[[[1.0], [1.0]], [[2.0], [2.0]], [[0.0], [0.0]]]])
+    torch.testing.assert_close(tensordict["value_target"], expected)
+
+
+def test_compute_value_target_supports_nstep_returns():
+    """n_step=2 accumulates two rewards then bootstraps, stopping at done."""
+    loss = _make_loss(gamma=0.5, n_step=2)
+    actions = _one_hot([[[0, 0], [1, 1], [2, 2]]], n_actions=3)
+    tensordict = TensorDict(
+        {
+            ("agents", "observation"): torch.zeros(1, 3, 2, 4),
+            ("agents", "action"): actions,
+            "next": {
+                "agents": {
+                    "reward": torch.ones(1, 3, 2, 1),
+                    "done": torch.tensor([[[[False], [False]], [[False], [False]], [[True], [True]]]]),
+                }
+            },
+        },
+        batch_size=[1, 3],
+    )
+    add_action_without_self(tensordict)
+
+    loss.compute_value_target(tensordict)
+
+    # G2(t0) = r0 + g*r1 + g^2*Q(t2) = 1 + 0.5 + 0.25*4 = 2.5
+    # G2(t1) = r1 + g*r2 = 1.5 (done at t2 stops the bootstrap)
+    # G2(t2) = r2 = 1.0
+    expected = torch.tensor([[[[2.5], [2.5]], [[1.5], [1.5]], [[1.0], [1.0]]]])
+    torch.testing.assert_close(tensordict["value_target"], expected)
+
+
+def test_add_joint_observation_repeats_team_observation_per_agent():
+    observation = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]])
+    tensordict = TensorDict({("agents", "observation"): observation}, batch_size=[1])
+
+    add_joint_observation(tensordict)
+
+    expected = torch.tensor([[[1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 3.0, 4.0]]])
+    torch.testing.assert_close(tensordict.get(("agents", "joint_observation")), expected)
+
+
+def test_add_masked_joint_action_zeroes_own_action_block():
+    action = _one_hot([[0, 2]], n_actions=3)
+    tensordict = TensorDict({("agents", "action"): action}, batch_size=[1])
+
+    add_masked_joint_action(tensordict)
+
+    expected = torch.tensor(
+        [
+            [
+                [0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+                [1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+            ]
+        ]
+    )
+    torch.testing.assert_close(tensordict.get(("agents", "masked_joint_action")), expected)
+
+
+def test_normalize_advantage_zero_centres_the_actor_signal():
+    """With a uniform policy, log-probs are constant, so the actor loss equals
+    -log(1/3) * mean(A); standardised advantages have zero mean, so the
+    normalized actor loss must vanish while the raw one does not."""
+    action = _one_hot([[0, 2]], n_actions=3)
+    data = {
+        ("agents", "observation"): torch.zeros(1, 2, 4),
+        ("agents", "action"): action,
+        "value_target": torch.zeros(1, 2, 1),
+    }
+    raw = TensorDict(dict(data), batch_size=[1])
+    add_action_without_self(raw)
+    norm = TensorDict(dict(data), batch_size=[1])
+    add_action_without_self(norm)
+
+    raw_loss = _make_loss(normalize_advantage=False)(raw)
+    norm_loss = _make_loss(normalize_advantage=True)(norm)
+
+    assert abs(raw_loss["loss_actor"].item()) > 1e-3
+    torch.testing.assert_close(norm_loss["loss_actor"], torch.tensor(0.0), atol=1e-5, rtol=0)
+
+
+def test_diagnostics_report_q_contrast_measures():
+    """Flat-critic instrumentation: own-action spread, others-sensitivity,
+    and empirical target contrast by chosen action."""
+    loss = _make_loss()
+    action = _one_hot([[0, 2]], n_actions=3)
+    tensordict = TensorDict(
+        {
+            ("agents", "observation"): torch.zeros(1, 2, 4),
+            ("agents", "action"): action,
+            ("agents", "logits"): torch.zeros(1, 2, 3),
+            "value_target": torch.tensor([[[0.5], [2.5]]]),
+        },
+        batch_size=[1],
+    )
+    add_action_without_self(tensordict)
+
+    diag = loss.diagnostics(tensordict)
+
+    # FixedActionValue outputs [1, 2, 4] for every input:
+    # population std of [1, 2, 4] = sqrt(14/9)
+    expected_spread = torch.tensor([1.0, 2.0, 4.0]).std(correction=0)
+    torch.testing.assert_close(diag["action_value_spread"], expected_spread.expand(1, 2))
+    # the fixed critic ignores other agents' actions entirely
+    torch.testing.assert_close(diag["others_sensitivity"], torch.zeros(1, 2))
+    # targets grouped by chosen action: {action0: 0.5, action2: 2.5} -> pop std 1.0
+    torch.testing.assert_close(diag["target_contrast"], torch.tensor(1.0))
+
+
+def test_clipped_loss_uses_frozen_old_advantage_and_clips_ratio():
+    """PPO-COMA: A frozen on stored pi_old logits; ratio vs stored log-probs;
+    min(unclipped, clipped) freezes over-drifted samples, keeps corrections."""
+    loss = _make_loss(clip_epsilon=0.2)
+    action = _one_hot([[0, 2]], n_actions=3)
+    tensordict = TensorDict(
+        {
+            ("agents", "observation"): torch.zeros(1, 2, 4),
+            ("agents", "action"): action,
+            ("agents", "logits"): torch.zeros(1, 2, 3),  # stored pi_old: uniform
+            ("agents", "sample_log_prob"): torch.full((1, 2), 0.5).log(),
+            "value_target": torch.zeros(1, 2, 1),
+        },
+        batch_size=[1],
+    )
+    add_action_without_self(tensordict)
+
+    loss.compute_counterfactual_advantage(tensordict)
+    # baseline under uniform pi_old = mean([1,2,4]) = 7/3 -> A = [1-7/3, 4-7/3]
+    expected_advantage = torch.tensor([[[-4.0 / 3.0], [5.0 / 3.0]]])
+    torch.testing.assert_close(tensordict.get(("agents", "advantage_old")), expected_advantage)
+
+    out = loss(tensordict)
+
+    # current policy uniform (1/3) vs stored 0.5 -> ratio 2/3, outside [0.8, 1.2]
+    torch.testing.assert_close(out["ratio_mean"], torch.tensor(2.0 / 3.0))
+    torch.testing.assert_close(out["clip_fraction"], torch.tensor(1.0))
+    # agent1 (A<0, already drifted far in A's direction): clipped branch -> 0.8*(-4/3)
+    # agent2 (A>0, drifted the wrong way): unclipped branch stays -> (2/3)*(5/3)
+    expected_loss = -((0.8 * (-4.0 / 3.0) + (2.0 / 3.0) * (5.0 / 3.0)) / 2.0)
+    torch.testing.assert_close(out["loss_actor"], torch.tensor(expected_loss))
\ No newline at end of file
diff --git a/torchrl/objectives/__init__.py b/torchrl/objectives/__init__.py
index c2c5c1824d8..c0ee5b9ecf1 100644
--- a/torchrl/objectives/__init__.py
+++ b/torchrl/objectives/__init__.py
@@ -32,7 +32,7 @@
 )
 from torchrl.objectives.gail import GAILLoss
 from torchrl.objectives.iql import DiscreteIQLLoss, IQLLoss
-from torchrl.objectives.multiagent import IPPOLoss, MAPPOLoss, QMixerLoss
+from torchrl.objectives.multiagent import COMALoss, IPPOLoss, MAPPOLoss, QMixerLoss
 from torchrl.objectives.pilco import ExponentialQuadraticCost
 from torchrl.objectives.ppo import ClipPPOLoss, KLPENPPOLoss, PPOLoss
 from torchrl.objectives.redq import REDQLoss
@@ -63,6 +63,7 @@
     "CQLLoss",
     "DiffusionBCLoss",
     "ClipPPOLoss",
+    "COMALoss",
     "CrossQLoss",
     "DDPGLoss",
     "DQNLoss",
@@ -114,4 +115,4 @@
     "two_hot_cross_entropy",
     "two_hot_decode",
     "two_hot_encode",
-]
+]
\ No newline at end of file
diff --git a/torchrl/objectives/multiagent/__init__.py b/torchrl/objectives/multiagent/__init__.py
index 6e25271b326..a2b782b047c 100644
--- a/torchrl/objectives/multiagent/__init__.py
+++ b/torchrl/objectives/multiagent/__init__.py
@@ -5,5 +5,6 @@
 
 from .mappo import IPPOLoss, MAPPOLoss
 from .qmixer import QMixerLoss
+from .coma import COMALoss
 
-__all__ = ["IPPOLoss", "MAPPOLoss", "QMixerLoss"]
+__all__ = ["IPPOLoss", "MAPPOLoss", "QMixerLoss", "COMALoss"]
diff --git a/torchrl/objectives/multiagent/coma.py b/torchrl/objectives/multiagent/coma.py
new file mode 100644
index 00000000000..ba4f3676ea8
--- /dev/null
+++ b/torchrl/objectives/multiagent/coma.py
@@ -0,0 +1,311 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from tensordict import TensorDict, TensorDictBase, TensorDictParams
+from tensordict.nn import TensorDictModule
+from tensordict.utils import NestedKey
+from torchrl.objectives.common import LossModule
+
+
+class COMALoss(LossModule):
+    """Counterfactual multi-agent policy-gradient loss.
+
+    The actor is decentralised. The Q-value network is centralised through its
+    inputs and must emit one action value per agent action under
+    ``("agents", "action_value")``.
+    """
+
+    @dataclass
+    class _AcceptedKeys:
+        action: NestedKey = ("agents", "action")
+        action_value: NestedKey = ("agents", "action_value")
+        chosen_action_value: NestedKey = ("agents", "chosen_action_value")
+        logits: NestedKey = ("agents", "logits")
+        sample_log_prob: NestedKey = ("agents", "sample_log_prob")
+        advantage_old: NestedKey = ("agents", "advantage_old")
+        reward: NestedKey = ("agents", "reward")
+        done: NestedKey = ("agents", "done")
+        terminated: NestedKey = ("agents", "terminated")
+        value_target: NestedKey = "value_target"
+
+    tensor_keys: _AcceptedKeys
+    default_keys = _AcceptedKeys
+    out_keys = [
+        "loss_actor",
+        "loss_qvalue",
+        "loss_entropy",
+        "entropy",
+        "pred_value",
+        "target_value",
+        "advantage",
+    ]
+
+    actor_network: TensorDictModule
+    actor_network_params: TensorDictParams
+    target_actor_network_params: TensorDictParams
+    qvalue_network: TensorDictModule
+    qvalue_network_params: TensorDictParams
+    target_qvalue_network_params: TensorDictParams
+
+    def __init__(
+        self,
+        actor_network: TensorDictModule,
+        qvalue_network: TensorDictModule,
+        *,
+        gamma: float = 0.99,
+        qvalue_loss_coef: float = 0.5,
+        entropy_coef: float = 0.0,
+        n_step: int = 1,
+        normalize_advantage: bool = False,
+        clip_epsilon: float | None = None,
+    ) -> None:
+        super().__init__()
+        if n_step < 1:
+            raise ValueError(f"n_step must be >= 1, got {n_step}.")
+        if clip_epsilon is not None and not 0.0 < clip_epsilon < 1.0:
+            raise ValueError(f"clip_epsilon must be in (0, 1), got {clip_epsilon}.")
+        self.convert_to_functional(actor_network, "actor_network")
+        self.convert_to_functional(qvalue_network, "qvalue_network", create_target_params=True)
+        self.gamma = gamma
+        self.qvalue_loss_coef = qvalue_loss_coef
+        self.entropy_coef = entropy_coef
+        self.n_step = n_step
+        self.normalize_advantage = normalize_advantage
+        self.clip_epsilon = clip_epsilon
+
+    def forward(self, tensordict: TensorDictBase) -> TensorDict:
+        td_copy = tensordict.clone(False)
+
+        dist = self.actor_network.get_dist(td_copy)
+        with self.qvalue_network_params.to_module(self.qvalue_network):
+            self.qvalue_network(td_copy)
+
+        chosen_action_value = self._chosen_action_value(td_copy)
+        if self.clip_epsilon is not None:
+            # PPO-COMA: the advantage was computed once per batch against the
+            # collection policy (pi_old) and frozen; see
+            # compute_counterfactual_advantage. Anchoring A and the ratio on
+            # the same pi_old is what makes the clipped surrogate coherent.
+            if self.tensor_keys.advantage_old not in tensordict.keys(True):
+                raise KeyError(
+                    "clip_epsilon is set: call compute_counterfactual_advantage on the batch before the update passes."
+                )
+            advantage = tensordict.get(self.tensor_keys.advantage_old).detach()
+        else:
+            advantage = self._counterfactual_advantage(td_copy, chosen_action_value).detach()
+        if self.normalize_advantage:
+            # MAPPO-style per-batch standardisation: same counterfactual
+            # advantage, rescaled so the actor step size is batch-invariant.
+            advantage = (advantage - advantage.mean()) / advantage.std().clamp_min(1e-6)
+
+        log_prob = dist.log_prob(td_copy.get(self.tensor_keys.action))
+        extra_outputs: dict[str, torch.Tensor] = {}
+        if self.clip_epsilon is not None:
+            old_log_prob = tensordict.get(self.tensor_keys.sample_log_prob).detach()
+            ratio = (log_prob - old_log_prob.reshape(log_prob.shape)).exp()
+            flat_advantage = advantage.squeeze(-1)
+            unclipped = ratio * flat_advantage
+            clipped = ratio.clamp(1.0 - self.clip_epsilon, 1.0 + self.clip_epsilon) * flat_advantage
+            loss_actor = -torch.min(unclipped, clipped).mean()
+            outside = (ratio < 1.0 - self.clip_epsilon) | (ratio > 1.0 + self.clip_epsilon)
+            extra_outputs["ratio_mean"] = ratio.detach().mean()
+            extra_outputs["clip_fraction"] = outside.float().mean()
+        else:
+            loss_actor = -(log_prob * advantage.squeeze(-1)).mean()
+
+        target_value = td_copy.get(self.tensor_keys.value_target)
+        loss_qvalue = F.mse_loss(chosen_action_value, target_value) * self.qvalue_loss_coef
+
+        entropy = dist.entropy().mean()
+        loss_entropy = -self.entropy_coef * entropy
+
+        return TensorDict(
+            {
+                "loss_actor": loss_actor,
+                "loss_qvalue": loss_qvalue,
+                "loss_entropy": loss_entropy,
+                "entropy": entropy.detach(),
+                "pred_value": chosen_action_value.detach().mean(),
+                "target_value": target_value.detach().mean(),
+                "advantage": advantage.detach().mean(),
+                **extra_outputs,
+            },
+            batch_size=[],
+        )
+
+    def compute_value_target(
+        self,
+        tensordict: TensorDictBase,
+        params: TensorDictParams | None = None,
+    ) -> TensorDictBase:
+        """Write the n-step Q-value target before flattening rollout data.
+
+        The collector batch is expected to keep time in dimension 1. Targets
+        follow the recursion G_k(t) = r(t) + gamma * (1 - done(t)) * G_{k-1}(t+1)
+        with G_0 = target-network chosen-action Q-values, applied ``n_step``
+        times. ``n_step=1`` is the TD(0) target; ``n_step=10`` matches
+        EPyMARL's ``q_nstep: 10`` within episodes. Values past the end of the
+        batch are treated as zero, so the last ``n_step`` transitions of a
+        batch are biased low unless they end an episode.
+        """
+        if params is None:
+            params = self.target_qvalue_network_params
+
+        td_copy = tensordict.clone(False)
+        with params.to_module(self.qvalue_network):
+            self.qvalue_network(td_copy)
+
+        chosen_action_value = self._chosen_action_value(td_copy)
+        if chosen_action_value.ndim < 2:
+            raise ValueError("COMALoss.compute_value_target expects an environment/time rollout batch.")
+
+        reward = tensordict.get(("next",) + tuple(self.tensor_keys.reward))
+        done = tensordict.get(("next",) + tuple(self.tensor_keys.done)).to(chosen_action_value.dtype)
+
+        value_target = chosen_action_value
+        for _ in range(self.n_step):
+            next_value = torch.zeros_like(value_target)
+            next_value[:, :-1] = value_target[:, 1:]
+            value_target = reward + self.gamma * (1.0 - done) * next_value
+        tensordict.set(self.tensor_keys.value_target, value_target.detach())
+        return tensordict
+
+    def compute_counterfactual_advantage(self, tensordict: TensorDictBase) -> TensorDictBase:
+        """Write the frozen counterfactual advantage of the collection policy.
+
+        PPO-style anchoring: the baseline is computed from the *stored*
+        collection-time logits (pi_old) — the actor is deliberately not run
+        here — and the result is written once per batch under
+        ``advantage_old`` so that every reuse epoch optimises the same fixed
+        coefficient, coherent with the importance ratio's anchor.
+        """
+        td_copy = tensordict.clone(False)
+        with self.qvalue_network_params.to_module(self.qvalue_network):
+            self.qvalue_network(td_copy)
+        chosen_action_value = self._chosen_action_value(td_copy)
+        advantage = self._counterfactual_advantage(td_copy, chosen_action_value)
+        tensordict.set(self.tensor_keys.advantage_old, advantage.detach())
+        return tensordict
+
+    def diagnostics(self, tensordict: TensorDictBase) -> dict[str, torch.Tensor]:
+        """Return unreduced COMA quantities for trainer-side observability.
+
+        This deliberately leaves aggregation to a reusable trainer hook: callers
+        can log global summaries, per-agent values, or custom histograms without
+        changing the loss used for optimization.
+        """
+        td_copy = tensordict.clone(False)
+        with self.qvalue_network_params.to_module(self.qvalue_network):
+            self.qvalue_network(td_copy)
+        chosen_action_value = self._chosen_action_value(td_copy)
+        advantage = self._counterfactual_advantage(td_copy, chosen_action_value)
+        target_value = td_copy.get(self.tensor_keys.value_target)
+
+        # Q-contrast measures (flat-critic hypothesis): how differentiated are
+        # the Q outputs across own actions, how sensitive are they to the other
+        # agents' actions, and how contrasted are the empirical targets by
+        # chosen action (the reference the critic should match).
+        action_value = td_copy.get(self.tensor_keys.action_value)
+        action_value_spread = action_value.std(dim=-1, correction=0)
+
+        permuted = tensordict.clone(False)
+        for key in (("agents", "masked_joint_action"), ("agents", "action_without_self")):
+            if key in permuted.keys(True):
+                values = permuted.get(key)
+                index = torch.randperm(values.shape[0], device=values.device)
+                permuted.set(key, values[index])
+        with self.qvalue_network_params.to_module(self.qvalue_network):
+            self.qvalue_network(permuted)
+        others_sensitivity = (permuted.get(self.tensor_keys.action_value) - action_value).abs().mean(dim=-1)
+
+        flat_action = tensordict.get(self.tensor_keys.action).to(torch.float).reshape(-1, action_value.shape[-1])
+        flat_target = target_value.reshape(-1, 1)
+        counts = flat_action.sum(dim=0)
+        means = (flat_action * flat_target).sum(dim=0) / counts.clamp_min(1.0)
+        taken = counts > 0
+        if int(taken.sum()) > 1:
+            target_contrast = means[taken].std(correction=0)
+        else:
+            target_contrast = torch.zeros((), device=action_value.device)
+
+        return {
+            "advantage": advantage.detach(),
+            "td_error": (target_value - chosen_action_value).detach(),
+            "chosen_action_value": chosen_action_value.detach(),
+            "target_value": target_value.detach(),
+            "action_value_spread": action_value_spread.detach(),
+            "others_sensitivity": others_sensitivity.detach(),
+            "target_contrast": target_contrast.detach(),
+        }
+
+    def _chosen_action_value(self, tensordict: TensorDictBase) -> torch.Tensor:
+        action = tensordict.get(self.tensor_keys.action).to(torch.float)
+        action_value = tensordict.get(self.tensor_keys.action_value)
+        chosen_action_value = (action * action_value).sum(dim=-1, keepdim=True)
+        tensordict.set(self.tensor_keys.chosen_action_value, chosen_action_value)
+        return chosen_action_value
+
+    def _counterfactual_advantage(
+        self,
+        tensordict: TensorDictBase,
+        chosen_action_value: torch.Tensor,
+    ) -> torch.Tensor:
+        action_prob = tensordict.get(self.tensor_keys.logits).softmax(dim=-1)
+        action_value = tensordict.get(self.tensor_keys.action_value)
+        counterfactual_baseline = (action_prob * action_value).sum(dim=-1, keepdim=True)
+        return chosen_action_value - counterfactual_baseline
+
+
+def add_action_without_self(tensordict: TensorDictBase) -> TensorDictBase:
+    """Write each agent's other-agent actions under ``action_without_self``."""
+    if ("agents", "action") not in tensordict.keys(True):
+        return tensordict
+
+    actions = tensordict.get(("agents", "action"))
+    n_agents = actions.shape[-2]
+    if n_agents == 1:
+        tensordict.set(("agents", "action_without_self"), actions.new_empty(*actions.shape[:-2], 1, 0))
+        return tensordict
+
+    other_actions = torch.stack(
+        [torch.cat([actions[..., :i, :], actions[..., i + 1 :, :]], dim=-2) for i in range(n_agents)],
+        dim=-3,
+    )
+    tensordict.set(("agents", "action_without_self"), other_actions.reshape(*actions.shape[:-2], n_agents, -1))
+    return tensordict
+
+
+def add_joint_observation(tensordict: TensorDictBase) -> TensorDictBase:
+    """Write the concatenated team observation under ``joint_observation``.
+
+    Every agent row receives the same concatenation of all agents'
+    observations. This mirrors EPyMARL's Gym wrapper, where the global state
+    fed to the COMA critic is the concatenation of individual observations.
+    """
+    observation = tensordict.get(("agents", "observation"))
+    n_agents = observation.shape[-2]
+    joint = observation.reshape(*observation.shape[:-2], 1, -1).expand(
+        *observation.shape[:-2], n_agents, n_agents * observation.shape[-1]
+    )
+    tensordict.set(("agents", "joint_observation"), joint.clone())
+    return tensordict
+
+
+def add_masked_joint_action(tensordict: TensorDictBase) -> TensorDictBase:
+    """Write the joint one-hot action with each agent's own slot zeroed.
+
+    This mirrors EPyMARL's ``COMACritic._build_inputs``: agent ``i`` sees the
+    full joint action vector of size ``n_agents * n_actions`` with the block
+    corresponding to its own action masked to zero, so its Q-values are not
+    conditioned on the action being marginalised by the counterfactual
+    baseline.
+    """
+    actions = tensordict.get(("agents", "action")).to(torch.float)
+    n_agents, n_actions = actions.shape[-2], actions.shape[-1]
+    joint = actions.reshape(*actions.shape[:-2], 1, -1).expand(*actions.shape[:-2], n_agents, n_agents * n_actions)
+    own_block = torch.eye(n_agents, device=actions.device, dtype=actions.dtype).repeat_interleave(n_actions, dim=1)
+    tensordict.set(("agents", "masked_joint_action"), joint * (1.0 - own_block))
+    return tensordict
\ No newline at end of file

From 35879da4208d828be18a9b0391fb44e30be7e187 Mon Sep 17 00:00:00 2001
From: Iliamsou <iliam.souami@gmail.com>
Date: Mon, 24 Aug 2026 11:29:03 +0200
Subject: [PATCH 2/2] Remove PPO-COMA extensions

---
 test/objectives/test_coma.py          | 34 +-----------------
 torchrl/objectives/multiagent/coma.py | 52 +++------------------------
 2 files changed, 6 insertions(+), 80 deletions(-)

diff --git a/test/objectives/test_coma.py b/test/objectives/test_coma.py
index 34a8d557c17..15ec9b8a91d 100644
--- a/test/objectives/test_coma.py
+++ b/test/objectives/test_coma.py
@@ -28,7 +28,7 @@ def _one_hot(index, n_actions=3):
 
 
 def _make_loss(
-    gamma=0.5, qvalue_loss_coef=0.5, entropy_coef=0.0, n_step=1, normalize_advantage=False, clip_epsilon=None
+    gamma=0.5, qvalue_loss_coef=0.5, entropy_coef=0.0, n_step=1, normalize_advantage=False,
 ):
     obs_dim = 4
     n_actions = 3
@@ -60,7 +60,6 @@ def _make_loss(
         entropy_coef=entropy_coef,
         n_step=n_step,
         normalize_advantage=normalize_advantage,
-        clip_epsilon=clip_epsilon,
     )
 
 
@@ -233,34 +232,3 @@ def test_diagnostics_report_q_contrast_measures():
     torch.testing.assert_close(diag["target_contrast"], torch.tensor(1.0))
 
 
-def test_clipped_loss_uses_frozen_old_advantage_and_clips_ratio():
-    """PPO-COMA: A frozen on stored pi_old logits; ratio vs stored log-probs;
-    min(unclipped, clipped) freezes over-drifted samples, keeps corrections."""
-    loss = _make_loss(clip_epsilon=0.2)
-    action = _one_hot([[0, 2]], n_actions=3)
-    tensordict = TensorDict(
-        {
-            ("agents", "observation"): torch.zeros(1, 2, 4),
-            ("agents", "action"): action,
-            ("agents", "logits"): torch.zeros(1, 2, 3),  # stored pi_old: uniform
-            ("agents", "sample_log_prob"): torch.full((1, 2), 0.5).log(),
-            "value_target": torch.zeros(1, 2, 1),
-        },
-        batch_size=[1],
-    )
-    add_action_without_self(tensordict)
-
-    loss.compute_counterfactual_advantage(tensordict)
-    # baseline under uniform pi_old = mean([1,2,4]) = 7/3 -> A = [1-7/3, 4-7/3]
-    expected_advantage = torch.tensor([[[-4.0 / 3.0], [5.0 / 3.0]]])
-    torch.testing.assert_close(tensordict.get(("agents", "advantage_old")), expected_advantage)
-
-    out = loss(tensordict)
-
-    # current policy uniform (1/3) vs stored 0.5 -> ratio 2/3, outside [0.8, 1.2]
-    torch.testing.assert_close(out["ratio_mean"], torch.tensor(2.0 / 3.0))
-    torch.testing.assert_close(out["clip_fraction"], torch.tensor(1.0))
-    # agent1 (A<0, already drifted far in A's direction): clipped branch -> 0.8*(-4/3)
-    # agent2 (A>0, drifted the wrong way): unclipped branch stays -> (2/3)*(5/3)
-    expected_loss = -((0.8 * (-4.0 / 3.0) + (2.0 / 3.0) * (5.0 / 3.0)) / 2.0)
-    torch.testing.assert_close(out["loss_actor"], torch.tensor(expected_loss))
\ No newline at end of file
diff --git a/torchrl/objectives/multiagent/coma.py b/torchrl/objectives/multiagent/coma.py
index ba4f3676ea8..e5dc17c9847 100644
--- a/torchrl/objectives/multiagent/coma.py
+++ b/torchrl/objectives/multiagent/coma.py
@@ -24,8 +24,6 @@ class _AcceptedKeys:
         action_value: NestedKey = ("agents", "action_value")
         chosen_action_value: NestedKey = ("agents", "chosen_action_value")
         logits: NestedKey = ("agents", "logits")
-        sample_log_prob: NestedKey = ("agents", "sample_log_prob")
-        advantage_old: NestedKey = ("agents", "advantage_old")
         reward: NestedKey = ("agents", "reward")
         done: NestedKey = ("agents", "done")
         terminated: NestedKey = ("agents", "terminated")
@@ -60,13 +58,10 @@ def __init__(
         entropy_coef: float = 0.0,
         n_step: int = 1,
         normalize_advantage: bool = False,
-        clip_epsilon: float | None = None,
     ) -> None:
         super().__init__()
         if n_step < 1:
             raise ValueError(f"n_step must be >= 1, got {n_step}.")
-        if clip_epsilon is not None and not 0.0 < clip_epsilon < 1.0:
-            raise ValueError(f"clip_epsilon must be in (0, 1), got {clip_epsilon}.")
         self.convert_to_functional(actor_network, "actor_network")
         self.convert_to_functional(qvalue_network, "qvalue_network", create_target_params=True)
         self.gamma = gamma
@@ -74,7 +69,6 @@ def __init__(
         self.entropy_coef = entropy_coef
         self.n_step = n_step
         self.normalize_advantage = normalize_advantage
-        self.clip_epsilon = clip_epsilon
 
     def forward(self, tensordict: TensorDictBase) -> TensorDict:
         td_copy = tensordict.clone(False)
@@ -84,37 +78,17 @@ def forward(self, tensordict: TensorDictBase) -> TensorDict:
             self.qvalue_network(td_copy)
 
         chosen_action_value = self._chosen_action_value(td_copy)
-        if self.clip_epsilon is not None:
-            # PPO-COMA: the advantage was computed once per batch against the
-            # collection policy (pi_old) and frozen; see
-            # compute_counterfactual_advantage. Anchoring A and the ratio on
-            # the same pi_old is what makes the clipped surrogate coherent.
-            if self.tensor_keys.advantage_old not in tensordict.keys(True):
-                raise KeyError(
-                    "clip_epsilon is set: call compute_counterfactual_advantage on the batch before the update passes."
-                )
-            advantage = tensordict.get(self.tensor_keys.advantage_old).detach()
-        else:
-            advantage = self._counterfactual_advantage(td_copy, chosen_action_value).detach()
+
+        advantage = self._counterfactual_advantage(td_copy, chosen_action_value).detach()
+
         if self.normalize_advantage:
             # MAPPO-style per-batch standardisation: same counterfactual
             # advantage, rescaled so the actor step size is batch-invariant.
             advantage = (advantage - advantage.mean()) / advantage.std().clamp_min(1e-6)
 
         log_prob = dist.log_prob(td_copy.get(self.tensor_keys.action))
-        extra_outputs: dict[str, torch.Tensor] = {}
-        if self.clip_epsilon is not None:
-            old_log_prob = tensordict.get(self.tensor_keys.sample_log_prob).detach()
-            ratio = (log_prob - old_log_prob.reshape(log_prob.shape)).exp()
-            flat_advantage = advantage.squeeze(-1)
-            unclipped = ratio * flat_advantage
-            clipped = ratio.clamp(1.0 - self.clip_epsilon, 1.0 + self.clip_epsilon) * flat_advantage
-            loss_actor = -torch.min(unclipped, clipped).mean()
-            outside = (ratio < 1.0 - self.clip_epsilon) | (ratio > 1.0 + self.clip_epsilon)
-            extra_outputs["ratio_mean"] = ratio.detach().mean()
-            extra_outputs["clip_fraction"] = outside.float().mean()
-        else:
-            loss_actor = -(log_prob * advantage.squeeze(-1)).mean()
+
+        loss_actor = -(log_prob * advantage.squeeze(-1)).mean()
 
         target_value = td_copy.get(self.tensor_keys.value_target)
         loss_qvalue = F.mse_loss(chosen_action_value, target_value) * self.qvalue_loss_coef
@@ -131,7 +105,6 @@ def forward(self, tensordict: TensorDictBase) -> TensorDict:
                 "pred_value": chosen_action_value.detach().mean(),
                 "target_value": target_value.detach().mean(),
                 "advantage": advantage.detach().mean(),
-                **extra_outputs,
             },
             batch_size=[],
         )
@@ -173,22 +146,7 @@ def compute_value_target(
         tensordict.set(self.tensor_keys.value_target, value_target.detach())
         return tensordict
 
-    def compute_counterfactual_advantage(self, tensordict: TensorDictBase) -> TensorDictBase:
-        """Write the frozen counterfactual advantage of the collection policy.
 
-        PPO-style anchoring: the baseline is computed from the *stored*
-        collection-time logits (pi_old) — the actor is deliberately not run
-        here — and the result is written once per batch under
-        ``advantage_old`` so that every reuse epoch optimises the same fixed
-        coefficient, coherent with the importance ratio's anchor.
-        """
-        td_copy = tensordict.clone(False)
-        with self.qvalue_network_params.to_module(self.qvalue_network):
-            self.qvalue_network(td_copy)
-        chosen_action_value = self._chosen_action_value(td_copy)
-        advantage = self._counterfactual_advantage(td_copy, chosen_action_value)
-        tensordict.set(self.tensor_keys.advantage_old, advantage.detach())
-        return tensordict
 
     def diagnostics(self, tensordict: TensorDictBase) -> dict[str, torch.Tensor]:
         """Return unreduced COMA quantities for trainer-side observability.