Technical Summary: Automata-Conditioned Cooperative Multi-Agent Reinforcement Learning (ACC-MARL)

1. Problem Statement

The paper addresses the challenge of learning multi-task, multi-agent policies for cooperative teams operating under centralized training and decentralized execution (CTDE). The specific setting involves agents assigned a priori unknown temporal tasks at runtime, which they must complete without explicit communication.

The core difficulty lies in generalizing to these dynamic tasks while satisfying complex, long-horizon objectives. Existing approaches using formal specifications (like Deterministic Finite Automata, DFAs) to represent tasks are often sample-inefficient and limited to single-task scenarios, requiring retraining for new objectives. Furthermore, in a multi-agent context, three primary challenges hinder feasibility:

1. History Dependency: Policies must track task progress (history) to make decisions, breaking the Markov property and leading to suboptimal policies.
2. Credit Assignment: The team-level reward is sparse (only awarded upon total task completion), making it difficult for agents to discern how individual actions contribute to the global goal.
3. Representation Bottleneck: Simultaneously learning latent representations of DFAs and control policies creates a scalability issue, particularly for teams larger than two agents.

2. Methodology: ACC-MARL Framework

The authors propose Automata-Conditioned Cooperative Multi-Agent Reinforcement Learning (ACC-MARL), a framework that conditions decentralized policies on task representations derived from DFAs. The approach addresses the three challenges through the following mechanisms:

A. Mitigating History Dependency via Markovian Reformulation

Instead of conditioning policies on the raw history of observations, the framework augments the state space with the current minimal DFA state for each agent.

• Mechanism: At each step, the environment state is labeled, and the assigned DFAs are progressed (updated) based on these labels. The agents' observations are augmented with the latest minimal DFA representations.
• Result: This transforms the non-Markovian game into a Markov game over the product space of environment states and DFA states ($S \times D^n$). The authors prove (Theorem 3.1) that maximizing the return in this reformulated Markov game is equivalent to maximizing the probability of satisfying the original temporal objectives as the discount factor $\gamma \to 1$.

B. Addressing Credit Assignment via Potential-Based Reward Shaping (PBRS)

To overcome sparse rewards, the authors apply Potential-Based Reward Shaping.

• Mechanism: A potential function $\Phi_i$ is defined for each agent $i$, returning 1 if their assigned DFA is satisfied and 0 otherwise. The reward is shaped by the difference in potential between the current and next state: $R_{PBRS} = R_{base} + \gamma \Phi_i(next) - \Phi_i(current)$.
• Result: Agents receive dense positive feedback for completing their specific sub-tasks while maintaining the theoretical guarantee that the optimal policy for the shaped reward remains optimal for the original objective.

C. Solving the Representation Bottleneck with RAD Embeddings

To avoid learning latent DFA representations from scratch during training (which causes performance degradation in larger teams), the framework utilizes pretrained RAD (ReachAvoidDerived) Embeddings.

• Mechanism: A frozen Graph Attention Network (GATv2) encoder maps DFAs to latent vectors (RAD Embeddings). These embeddings are trained to uniquely represent distinct minimized DFAs and encode semantic similarities across a large class of temporal tasks.
• Result: Policies condition on these fixed embeddings rather than raw DFAs or learned representations, decoupling representation learning from control learning. This enables efficient scaling to larger teams (e.g., 4 agents).

D. Optimal Task Assignment

The framework leverages the learned value functions to assign tasks optimally at test time. By evaluating the expected team performance (sum of value functions) for different permutations of task assignments, the system can select a Pareto-optimal assignment that maximizes overall team success, particularly in asymmetric environments.

3. Key Contributions

1. Framework Introduction: The proposal of ACC-MARL, a framework for learning multi-task, multi-agent, ego-centric policies conditioned on DFAs.
2. Theoretical Guarantees: Proof that the proposed Markovian reformulation and reward shaping preserve optimality with respect to the original non-Markovian problem.
3. Scalability Solution: Demonstration that using pretrained RAD Embeddings overcomes the representation bottleneck, enabling effective learning in teams with more than two agents.
4. Task Assignment: A method for using learned value functions to compute optimal task assignments at test time.
5. Toolchain: Implementation of a JAX-based toolchain (DFAx) for parallelized DFA operations and a Python package for integrating automata tasks into RL applications.

4. Experimental Results

The framework was evaluated in a discrete, fully observable environment called TokenEnv, featuring two- and four-agent variants with layouts involving buttons, doors, and tokens.

• Feasibility and Scaling (Q1 & Q2):
  • Ablation studies show that without the proposed solutions (specifically PBRS and RAD Embeddings), policies fail to converge or achieve low success rates (<0.5).
  • In 2-agent environments, both pretrained and untrained encoders perform similarly, though pretrained embeddings offer lower variance.
  • In 4-agent environments, policies using pretrained RAD Embeddings significantly outperform those learning representations from scratch, confirming that pretraining is crucial for scalability.
• Generalization (Q3):
  • Policies trained on RAD DFAs generalize to unseen task classes (Reach and ReachAvoid) and out-of-distribution (OOD) DFAs with more states than seen during training.
  • Pretrained embeddings specifically improve performance on tasks involving "Avoid" constraints (irrecoverable conditions).
• Task Assignment (Q4):
  • In symmetric environments (Buttons), task assignment order had negligible impact.
  • In asymmetric environments (Rooms), optimal task assignment derived from learned value functions significantly improved team success probabilities compared to random assignment.
• Emergent Behaviors (Q5):
  • Qualitative analysis reveals emergent cooperative skills, such as agents holding doors for teammates to save time and "short-circuiting" tasks (e.g., one agent unlocking a door to allow another to take a shorter path).

5. Significance and Claims

The paper claims to be the first work to utilize formal specifications (DFAs) in a multi-task cooperative MARL setting where tasks are assigned at runtime. Its significance lies in:

• Bridging Formal Methods and MARL: It extends automata-conditioned RL from single-agent to multi-agent settings, enabling the decomposition of complex team objectives into sub-tasks.
• Solving the Scalability Challenge: By decoupling representation learning via pretrained embeddings, it addresses the performance bottlenecks that previously limited automata-conditioned MARL to small teams.
• Optimality: It provides a theoretically grounded approach that guarantees optimal team behavior while enabling efficient credit assignment and generalization.

The authors note limitations, including the assumption of fixed labeling functions, full observability, and the computational cost of enumerating task assignments for very large teams. They suggest future work in partial observability, dynamic labeling, and continuous domains.