COCOA: Counterfactual Contribution Analysis

• COCOA is a family of causal attribution methods that define and quantify action and state contributions using counterfactual reasoning.
• It leverages techniques like counterfactual contribution coefficients, Shapley values, and intrinsic causal contributions to decompose total effects into agent- and state-specific impacts.
• COCOA reduces variance in credit assignment and enhances policy gradient estimators, proving effective in both long-horizon and multi-agent reinforcement learning settings.

Counterfactual Contribution Analysis (COCOA) is a family of model-based causal attribution methods developed to precisely measure and decompose the contributions of actions and states to future outcomes in both single-agent and multi-agent reinforcement learning (RL). COCOA algorithms are designed to answer counterfactual queries of the form: “Would the agent still have achieved this rewarding outcome if a different action had been taken?” They address the long-standing challenge of high-variance and biased credit assignment in long-horizon and multi-agent RL by quantifying principled, reward-centric contribution coefficients and by decomposing the total effect in multi-agent settings into agent- and state-specific components. Two major strands anchor the field: reward-focused credit assignment in single-agent RL (Meulemans et al., 2023) and causal decomposition in multi-agent sequential decision making (&&&1&&&).

1. Motivation and Problem Scope

Traditional methods for credit assignment in RL—such as REINFORCE, temporal-difference (TD) approaches, and Hindsight Credit Assignment (HCA)—struggle with high variance and bias when addressing long-term dependencies or delayed rewards. Specifically, Monte Carlo methods have variance that grows rapidly with horizon, while discounting methods trade signal strength for bias. HCA, which attributes credit based on achieved states, is susceptible to spurious contributions in environments where action trajectories uniquely determine states, thereby regressing to the variance of REINFORCE.

COCOA circumvents these failures by assigning credit not to every state but directly to rewarding outcomes—either the raw rewards or learned object-centric representations—using principled counterfactual reasoning. In multi-agent settings, the challenge escalates: understanding how an individual’s decision affects an outcome is confounded by inter-agent dynamics and non-trivial environment evolution. Here, COCOA introduces a formal decomposition of effects, allowing attribution both to individual agents (via Shapley values) and to state variables (via intrinsic causal contributions) (Meulemans et al., 2023, Triantafyllou et al., 2024).

2. Formalization and Core Definitions

In the single-agent case, COCOA operates within the standard Markov Decision Process (MDP) formalism:

• States $S$, actions $A$, transition kernel $p(s'|s,a)$, and policy $\pi(a|s; \theta)$.
• Rewarding outcome random variable $U'$ is fully predictive of $R$ if $p(R_k=r|S_0=s, A_0=a, U_k=u) = p(R=r|U=u)$ for all $k$.
• Counterfactual contribution coefficient:

$$w(s, a, u') = \frac{p^\pi(A_t=a|S_t=s,U'=u')}{\pi(a|s)} - 1.$$

In the multi-agent extension (Triantafyllou et al., 2024), an MMDP with a Structural Causal Model (SCM) is used. A factual trajectory $\tau$ and intervention $\doop(A_{i,t} := a')$ define:

• Total Counterfactual Effect (TCFE): $\Delta = \mathbb{E}[Y_{\doop(A_{i,t}=a')}] - \tau(Y)$,
• Decomposition: $$\Delta = \Delta_{\mathrm{actions}} - \Delta_{\mathrm{states}}$$,
  • $\Delta_{\mathrm{actions}}$ is the total agent-specific effect (tot-ASE), measuring the effect propagating via all future agents’ actions.
  • $\Delta_{\mathrm{states}}$ is the reverse state-specific effect (r-SSE), measuring the effect propagating via state transitions alone.

3. COCOA Policy-Gradient Estimators and Causal Decompositions

Single-Agent Policy-Gradient Estimator:

The COCOA gradient estimator introduces a correction term for each reward, leveraging action-outcome contribution coefficients:

$$\nabla_\theta V^\pi(s_0) = \mathbb{E}_\pi\Bigl[ \sum_{t \ge 0}\nabla_\theta\log\pi(A_t|S_t)R_t + \sum_{a \in A} \nabla_\theta\pi(a|S_t)\sum_{k \ge 1}w(S_t, a, U_{t+k})R_{t+k} \Bigr].$$

This estimator produces unbiased gradients with variance strictly lower (in single-reward settings) than HCA or REINFORCE, as per the variance ordering established in [(Meulemans et al., 2023), Theorem 3].

Multi-Agent Counterfactual Decomposition:

COCOA supports granular attribution in MMDPs by decomposing the effect of an agent's counterfactual action (Definition 1 in (Triantafyllou et al., 2024)):

• Action Path: Agent-specific effects are attributed with the Shapley value, reflecting each agent's marginal impact on the outcome through all possible agent interaction subsets.
• State Path: State variables’ effects are decomposed using the Intrinsic Causal Contribution (ICC), which quantifies the reduction in r-SSE variance by conditioning on individual state noise sources.

Table: Key Quantities in COCOA Decomposition

Quantity	Description	Attribution Mechanism
TCFE ($\Delta$)	Total counterfactual effect	Sum of action + state paths
tot-ASE ($\Delta_\text{actions}$)	Effect through future agents	Shapley value (ASE-SV)
r-SSE ($\Delta_\text{states}$)	Effect through state transitions	Intrinsic Causal Contribution (ICC)

4. Algorithms and Implementation

Single-Agent (Reward-based COCOA) (Meulemans et al., 2023):

1. Collect trajectories under the current policy.
2. For each $(s_t, a_t, u_{t+k})$, train a hindsight model $h_\phi$ to estimate $p^\pi(A=a|s_t, u_{t+k})$ (cross-entropy).
3. Compute $w(s_t, a, u_{t+k})$ as the normalized contribution coefficient.
4. Estimate the policy gradient via the COCOA formula, and update policy parameters $\theta$.

Multi-Agent Causal Attribution (Triantafyllou et al., 2024):

1. Simulate the factual trajectory.
2. Evaluate TCFE by intervention on the action of interest and simulation of the SCM.
3. Compute tot-ASE via interventions that propagate the action effect through all future agents.
4. Decompose tot-ASE into agent-wise effects via Shapley values (requiring $O(n 2^n)$ simulations in the worst case, approximated via sampling).
5. For r-SSE, partition its variance into state-wise ICCs, attributing contributions to pivotal state steps via repeated posterior sampling of noise variables.

Computationally, naive Shapley calculation scales exponentially with the number of agents, but sampling-based approximations reduce this to $O(n(\log n)^2)$. ICC computation is linear in horizon with grouping or logarithmic search (Triantafyllou et al., 2024).

5. Variance Reduction and Theoretical Insights

COCOA's fundamental advantage over HCA and model-free approaches is its principled targeting of rewards or meaningful reward features as the object of contribution analysis. This reward-centric grouping means that actions are credited for their genuine causal influence on outcomes of interest, rather than on arbitrary or environment-specific state sequences. Theoretical analysis demonstrates:

• The variance of COCOA-reward estimators is strictly minimized among the family (reward-, object-, state-based estimators); order: $$\operatorname{Var}[\hat\nabla^R] \preccurlyeq \operatorname{Var}[\hat\nabla^U] \preccurlyeq \operatorname{Var}[\hat\nabla^S] \preccurlyeq \operatorname{Var}[\hat\nabla^{\text{REINFORCE}}]$$ (Meulemans et al., 2023).
• Reward aliasing is a failure mode: if distractor events yield indistinguishable rewards, action contributions may not be recoverable unless object or causal features are used as targets.
• In multi-agent settings, the decomposition theorem ensures that all agent and state contributions sum exactly to the observed counterfactual effect (Triantafyllou et al., 2024).

6. Empirical Validation and Illustrative Applications

Single-Agent Tasks (Meulemans et al., 2023):

• Key-to-door and task-interleaving environments, with delayed and interleaved reward structures.
• Metrics: Signal-to-Noise Ratio (SNR), policy gradient bias/variance norm.
• COCOA matches oracle Q-critic sample efficiency (<100 episodes with ground-truth models), remains robust as causal distance increases (SNR remains flat as key-door separation grows), and maintains low bias and variance in long-horizon tasks.

Multi-Agent Scenarios (Triantafyllou et al., 2024):

• Gridworld with LLM-assisted agents: COCOA precisely attributes a single agent’s deviation as the exclusive causal factor in outcome change.
• Sepsis treatment simulation: COCOA’s agent decomposition shifts attributions between AI and clinician as the clinician’s trust changes, and ICC identifies bottleneck state steps that concentrate the r-SSE.

7. Limitations and Future Research Directions

COCOA’s effectiveness is contingent on several structural prerequisites:

• The necessity for an on-policy hindsight model $h_\phi$ and sufficient reward or object-feature observations—performance in sparse-reward domains may be limited.
• In N-step bootstrapping, state-based spurious contributions can re-enter, reintroducing variance, especially in highly stochastic or continuous settings.
• The accuracy of multi-agent decompositions depends on SCM assumptions: noise-monotonicity and the (first-order) Markov property.

Future directions highlighted include:

• Off-policy and model-based counterfactual queries, enabling generalization akin to SVG/Dreamer approaches for discrete actions.
• Information-bottlenecked representations for low-variance, efficient learning.
• Full counterfactual inference by conditioning on all exogenous noise for sharper attribution.
• Application to partially observable domains via disentangled belief-state representations (Meulemans et al., 2023, Triantafyllou et al., 2024).

COCOA is particularly applicable to domains requiring reliable long-term or multi-agent credit assignment, such as robotics with delayed rewards, strategic multi-agent games, dialogue systems, and sequential decision-making in critical systems.