diff --git a/test/objectives/test_coma.py b/test/objectives/test_coma.py new file mode 100644 index 00000000000..15ec9b8a91d --- /dev/null +++ b/test/objectives/test_coma.py @@ -0,0 +1,234 @@ +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, +): + 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, + ) + + +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)) + + 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..e5dc17c9847 --- /dev/null +++ b/torchrl/objectives/multiagent/coma.py @@ -0,0 +1,269 @@ +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") + 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, + ) -> None: + super().__init__() + if n_step < 1: + raise ValueError(f"n_step must be >= 1, got {n_step}.") + 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 + + 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) + + 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)) + + 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(), + }, + 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 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