MOSDT: Self-Distillation-Based Decision Transformer for Multi-Agent Offline Safe Reinforcement Learning¶

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

TL;DR¶

MOSDT is the first algorithm for multi-agent offline safe reinforcement learning (MOSRL), combining policy self-distillation, full parameter sharing, and a cost binary embedding module within a Decision Transformer framework. It achieves SOTA returns on 14/18 tasks across MuJoCo, Safety Gym, and Isaac Gym while guaranteeing safety on all tasks, using only 65% of the execution parameter count of its single-agent baseline CDT.

Problem¶

Multi-agent offline safe RL (MOSRL) is unstudied: existing offline MARL methods ignore safety constraints, and offline safe RL methods handle only single agents. Combining the two introduces compounded challenges—training instability from dynamic multi-agent policy updates, scalability of knowledge distillation, and the failure of predetermined cost targets in causal transformers (which cause unsafe convergence and mask return-cost correlations due to the masked attention ordering).

Method¶

MOSDT builds a CTDE (centralized training, decentralized execution) framework on top of a causal Decision Transformer with three integrated components:

Policy Self-Distillation (PSD): Instead of two-stage KD (train teacher, then distill), student networks are embedded inside the teacher. Global features are reconstructed by summing per-agent observation features across all N agents (Eq. 2). Three losses—decentralized student supervised learning (DSSL), centralized teacher supervised learning (CTSL), and KL+feature-norm distillation (PSD loss)—are optimized jointly and synchronously, reducing training parameters by 35% and training time by 24%.
Full Parameter Sharing: All N student networks share identical parameters (πS¹ = … = πSᴺ); only teacher action heads and feature projectors are agent-specific. This stabilizes training (reduces absolute variation coefficient from 0.35 to 0.27) and cuts execution parameters by 58%.
Cost Binary Embedding (CBE): Cumulative costs Cᵢₜ are binarized to a safety signal Bᵢₜ ∈ {0,1} via a fixed threshold c. The safety signal feature is concatenated with the return feature before the causal transformer, giving the model persistent access to reward-cost correlation regardless of token ordering. At execution, Bᵢₜ = 0 is fixed throughout, preventing convergence to a specific cost target.

The MOSDB dataset is also introduced, collecting ~2×10⁷ tuples from 18 Safety Gymnasium tasks using two online MARL algorithms (MACPO, MAPPO-Lagrangian) across three cost thresholds.

Key Contributions¶

MOSDT: First MOSRL algorithm; CTDE architecture via synchronous self-distillation, eliminating two-stage KD overhead.
Full parameter sharing: First application of complete parameter sharing across all agents in offline MARL; stabilizes training and dramatically reduces parameter count.
Cost Binary Embedding (CBE): Plug-and-play module replacing specific cost targets with a binarized safety signal fused into the return feature, improving safety control and training stability.
MOSDB: First dataset and benchmark for MOSRL, covering MuJoCo/Safety Gym/Isaac Gym tasks with 18.7 GB of diverse trajectory data collected under three cost thresholds.

Results¶

SOTA on 14/18 MOSDB tasks (all environments) while maintaining zero safety violations across all 18 tasks; the only SOTA-safe method.
vs. CDT (single-agent SOTA baseline, CTCE framework with global information): MOSDT achieves better returns on 14/18 tasks with only partial observations and 65% of CDT's execution parameter count.
Full parameter sharing gains over no-parameter-sharing on multi-agent MuJoCo tasks: 38.4× on 3x1Hopper, 2.3× on 4x2Ant, 3.5× on 6x1HalfCheetah.
CBE improves returns on 14/18 tasks vs. CDT-style cost processing; safety signal embedding alone improves 12/18 tasks.
PSD reduces training parameters by 35% vs. conventional two-stage KD, and cuts training time by 24%.
Inference: 0.80 ms/agent vs. 0.86 ms for CDT.
All other baselines (BC, BCQ-Lag, BEAR-Lag, COptiDICE, CPQ) either fail to achieve SOTA or produce unsafe policies.

Limitations¶

Baselines are all single-agent methods adapted to the CTCE (centralized execution) setting; no true multi-agent offline safe RL baselines exist for comparison.
Safety violations still occur in some ablation configurations (e.g., removing PSD makes the model unsafe), indicating fragility of the safety guarantee to design choices.
Performance on 2-agent tasks (e.g., 2x1Swimmer, 2x3HalfCheetah) shows smaller gains vs. no-parameter-sharing, suggesting benefits of full sharing scale with agent count.
MOSDB is collected from only two online algorithms (MACPO, MAPPO-Lagrangian), potentially limiting trajectory diversity.
The global information reconstruction scheme (sum of agent observation features) is a simple aggregation; richer cross-agent interaction structure is not modeled.

Relevance to Agentic AI / LLM Agents¶

MOSDT addresses a core challenge for deploying multi-agent systems in safety-critical real-world settings—learning coordinated, constraint-respecting policies from logged data without online interaction, directly relevant to autonomous vehicle fleets, robot teams, and power grid control. The self-distillation architecture (teacher-student integration with synchronous training) is a pattern increasingly relevant to LLM agent ensembles that need to share knowledge under partial observability. The CBE module offers a generalizable technique for injecting safety constraints into sequence-model-based decision making, applicable wherever transformer-based agents must satisfy hard budget or constraint targets. The introduction of MOSDB fills a benchmark gap that will facilitate future research on safe multi-agent foundations.