Diffusion Models for Offline Multi-agent Reinforcement Learning with Safety Constraints*

For a nonlinear affine control system:

where $s\in S\subseteq\mathbb{R}^{n}$ is the system state, and $a\in A\subseteq\mathbb{R}^{m}$ is the admissible control input.

Definition 1 A set $C\subseteq\mathbb{R}^{n}$ is forward invariant for system (1) if the solutions for some $a\in A$ beginning at any $s_{0}\in C$ meet $x_{t}\in C$, $\forall t\geq 0$.

Definition 2 A function $h$ is a control barrier function (CBF) if there exists an extended class $\mathcal{K}_{\infty}$ function $\alpha$, i.e., $\alpha$ is strictly increasing and satisfies $\alpha(0)=0$, such that for the control system (1):

for all $s\in C$.

Theorem 1.(Ames et al. (2017)) Given a CBF $h(s)$ from Definition 2, if $s_{0}\in\mathcal{C}$, then any $a$ generated by a Lipschitz continuous controller that satisfies the constraint in (2), $\forall t\geq 0$ renders $\mathcal{C}$ forward invariant for system (1).

In single-agent reinforcement learning, for a control policy $\pi:S\rightarrow A$, CBF $h$, state space $s\in S$, action space $a\in A$ and let ${S}_{d}\subseteq{S}$ be the dangerous set, ${S}_{s}={S}\setminus{S}_{d}$ be the safe set, which contains the set of initial conditions ${S}_{0}\subseteq{S}_{s}$. It is proved in (Ames et al., 2014) that if these three conditions:

are satisfied with $a=\pi(s)$, then $s(t)\in\{s|h(s)\geq 0\}$ for $\forall t\in[0,\infty)$, which means the state is forward invariant according to Theorem 1, and it would never enter the dangerous set $S_{d}$ under $\pi$.

Diffusion Models generate data from the dataset $D:=\{x_{i}\}_{0\leq i\leq M}$. The forward diffusion process is defined as: $q(x_{k+1}|x_{k}):=\mathcal{N}(x_{k+1};\sqrt{\alpha_{k}}x_{k},(1-\alpha_{k})I)$ and the reverse diffusion process is $p_{\theta}(x_{k-1}|x_{k}):=\mathcal{N}(x_{k-1}|\mu_{\theta}(x_{k},k),\Sigma_{\theta}(x_{k},k))$, where $\mathcal{N}(\mu,\Sigma)$ is a Gaussian distribution with mean $\mu$ and variance $\Sigma$. Here, $\alpha$ is known as the “diffusion rate” and is precalculated using a “variance scheduler.” The term $I$ is an identity matrix. By predicting the parameters for the reverse diffusion process at each time step with a neural network, new samples that closely match the underlying data distribution are generated. The reverse diffusion process can be estimated by the loss function as follows (Ho et al., 2020):

The predicted noise $\epsilon_{\theta}(x_{k},k)$ is to estimate the noise $\epsilon\sim\mathcal{N}(0,I)$.

In the forward process, we use classifier-free guidance, which requires an additional condition $y$ to generate target synthetic data. In this work, we incorporated CBF into the diffusion model for the finite-time forward invariance and reward for the optimal policy. Classifier-free guidance modifies the original training setup to learn both a conditional $\epsilon_{\theta}(x_{k},y,k)$ and an unconditional conditional noise $\epsilon_{\theta}(x_{k},\emptyset,k)$ where a dummy value $\emptyset$ takes the place of $y$. The perturbed noise $\epsilon_{\theta}(x_{k},\emptyset,k)+\omega(\epsilon_{\theta}(x_{k},y,k)-\epsilon_{\theta}(x_{k},\emptyset,k))$ is used to later generate samples.

About diffusion decision-making in single-agent settings, diffusing over state trajectories only (Ajay et al., 2023), is claimed to be easier to model and can obtain better performance due to the less smooth nature of action sequences:

where $H$ is the trajectory length that the diffusion model generates and $t$ is the time a state was visited in trajectory $\tau$. However, sampling states from the diffusion model cannot get the corresponding action. To infer the policy, we could use the inverse dynamics model to generate the action by two consecutive states in the trajectory:

The safe MARL problem is normally formulated as a Constrained Markov Decision Process (CMDP) $\{N,S,O,A,p,\rho^{0},\gamma,R,h\}$. Here, $N=\{1,\ldots,n\}$ is the set of agents, the joint state space is $S=\{S_{1},S_{2},\ldots,S_{n}\}$ where $s^{i}_{t}\in S_{i}$ denotes the state of agent $i$ at time step $t$, $O$ is the local observation, $A=\prod_{i=1}^{n}A_{i}$ is the joint action space, $p:S\times A\rightarrow S$ is the probabilistic transition function, $\rho^{0}$ is the initial state distribution, $\gamma\in[0,1]$ is the discount factor, $R:S\times A\times S\rightarrow\mathbb{R}$ is the joint reward function, $h:S\rightarrow\mathbb{R}$ is the constraint function; in this paper, we use CBF as the constraint function. At time step $t$, the joint state at time $t$ is denoted by $s_{t}=\{s^{1}_{t},\ldots,s^{n}_{t}\}$, and every agent $i$ takes action $a^{i}_{t}$ according to its policy $\pi^{i}(a^{i}_{t}|s_{t})$. Together with other agents’ actions, it gives a joint action $a_{t}=(a^{1}_{t},\ldots,a^{n}_{t})$ and the joint policy $\pi(a_{t}|s_{t})=\prod_{i=1}^{n}\pi^{i}(a^{i}_{t}|s_{t})$. In offline settings, instead of collecting online data in environments, we only have access to a static dataset $D$ to learn the policies. The dataset $D$ generally comprises trajectories $\tau$, i.e., observation-action sequences.

For each agent $i$, we define ${N}^{i}_{t}$ as the set of its neighborhood agents at time $t$. Let $o^{i}_{t}\in O_{i}$ be the local observation of agent $i$, which is the states of ${N}^{i}_{t}$ neighborhood agents. Notice that the dimension of $o_{i}$ is not fixed and depends on the quantity of neighboring agents.

We assume the safety of agent $i$ is jointly determined by $s_{i}$ and $o_{i}$. Let ${O}_{i}$ be the set of all observations and $X_{i}:={S}_{i}\times{O}_{i}$ be the state-observation space that contains the safe set $X_{i,s}$, dangerous set $X_{i,d}$, and initial conditions $X_{i,0}\subseteq X_{i,s}$. Let $d$ describe the minimum distance from agent $i$ to other agents, relative speed $V$, deceleration $b$, and minimal stopped gap $\kappa_{s}$, then $d(s_{i},o_{i})<\frac{V^{2}}{2\cdot b}+\kappa_{s}$ implies a collision. Then $X_{i,s}=\{(s_{i},o_{i})|d(s_{i},o_{i})\geq\frac{V^{2}}{2\cdot b}+\kappa_{s}\}$ and $X_{i,d}=\{(s_{i},o_{i})|d(s_{i},o_{i})<\frac{V^{2}}{2\cdot b}+\kappa_{s}\}$. Since there is a surjection from ${S}$ to $X_{i}$, we may define $\bar{d}_{i}$ as the lifting of $d$ from $X_{i}$ to ${S}$. ${S}_{s}:=\{s\in{S}|\forall i=1,\ldots,N,\bar{d}_{i}(s)\geq\frac{V^{2}}{2\cdot b}+\kappa_{s}\}$ is then defined. Formally speaking, a multi-agent system’s safety can be described as follows:

Definition 3 If the minimum distance satisfies $d(s_{i},o_{i})\geq\frac{V^{2}}{2\cdot b}+\kappa_{s}$ for agent $i$ and $t$, then agent $i$ is safe at time $t$. If for $\forall i$, agent $i$ is safe at time $t$, then the multi-agent system is safe at time $t$, and $s\in\mathcal{S}_{s}$.

A simple CBF for a multi-agent dynamic system is a centralized function that accounts for the joint states of all agents. However, it may cause an exponential explosion in the state space; it is also difficult to define a safety constraint for the entire system while ensuring that the security of individual agents will not be violated.

By Definition 3, we consider decentralized CBF to guarantee the multi-agent system’s safety. From equation (3), we propose the following CBF:

where $\dot{o}_{i}(t)$ denotes the time derivative of observation, which depends on other agents’ actions. It can be assessed and included in the training process without an explicit expression. Here, the state $s_{i}$ and $o_{i}$ are the local state and observation of the corresponding agent $i$. We refer to conditions (3) as the decentralized CBF for agent $i$.

Proposition 1 If the decentralized CBF conditions in (7) are satisfied, then $\forall t$ and $\forall i$, $(s^{i}_{t},o^{i}_{t})\in\{(s_{i},o_{i})\mid h_{i}(s_{i},o_{i})\geq 0\}$, which implies the state would never enter $X_{i,d}$ for any agent $i$. Thus, the multi-agent system is safe by Definition (3).

According to Proposition 1, CBF can be a decentralized paradigm for every agent in the entire multi-agent system. Since state-observation satisfying $h_{i}(s_{i},o_{i})\geq 0$ is forward invariant, agent $i$ never gets closer than $\frac{V^{2}}{2\cdot b}+\kappa_{s}$ to all of its neighboring agents. According to the definition of $h_{i}$, $h_{i}(s_{i},o_{i})>0\Rightarrow\bar{d}_{i}(s)\geq\frac{V^{2}}{2\cdot b}+\kappa_{s}$. The multi-agent system is safe according to Definition 3 since $\forall i,h_{i}(s_{i},o_{i})\geq 0$ implies that $\forall i,\bar{d}_{i}(s)\geq\frac{V^{2}}{2\cdot b}+\kappa_{s}$.

Next, we need to formulate the control barrier function $h_{i}(s_{i},o_{i})$ to get a safe set from dataset $D$. Let $\tau_{i}=\{s_{i},o_{i}\}$ be a trajectory of the state and observation of agent $i$. Let $\mathcal{T}_{i}$ be the set of all possible trajectories of agent $i$. Let $\mathcal{H}_{i}$ and $\mathcal{V}_{i}$ be the function classes of $h_{i}$ and policy $\pi_{i}$. Define the function $y_{i}$: $\mathcal{T}_{i}\times\mathcal{H}_{i}\times\mathcal{V}_{i}\rightarrow\mathbb{R}$ as:

Notice that the third item on the right side of Equation (8) depends on both the policy and CBF, since $\dot{h}_{i}=\nabla_{s_{i}}h_{i}\cdot f_{i}(s_{i},u_{i})+\nabla_{o_{i}}h_{i}\cdot\dot{o}_{i}(t),u_{i}=\pi_{i}(s_{i},o_{i})$. It is clear that if we can find $h_{i}$ and $\pi_{i}(s_{i},o_{i})$ such that $y_{i}(\tau_{i},h_{i},\pi_{i})>0$ for $\forall\tau_{i}\in\mathcal{T}_{i}$ and $\forall i$, then the conditions in (7) are satisfied. We solve the objective:

For all $i$, find $h_{i}\in\mathcal{H}_{i}$ and $\pi_{i}\in\mathcal{V}_{i}$, such that $y_{i}(\tau_{i},h_{i},\pi_{i})\geq\gamma$, where $\gamma>0$ is a margin for the satisfaction of the CBF condition in (7).

We formulate the diffusion model as follows:

Our goal is to estimate $\tau$ conditioned on ${y}(\cdot)$ with $p_{\theta}$. In this paper, $\boldsymbol{y}(\tau)$ includes the CBF and the reward under the trajectory.

Given an offline dataset $D$ that consists of all agents’ trajectories data, our diffusion model also takes a decentralized manner to make it consistent with the decentralized CBF. The model is parameterized through the unified noise model $\epsilon_{\theta}$ and the inverse dynamics model $I_{\phi}$ of each agent $i$ with the reverse diffusion loss and the inverse dynamics loss:

However, there are some other gaps between methodology and practical implementation. First, equation (8) does not provide an exact way of designing loss functions. Second, the CBF and $\pi$ are coupled, where minor approximation errors can bootstrap across them and lead to severe instability; furthermore, $\dot{h}_{i}$ has term $\dot{o}_{i}$. Third, we do not have the loss function considering reward maximization.

Based on equation (8), we formulate the loss function: $\mathcal{L}^{c}=\sum_{i}\mathcal{L}^{c}_{i}$, where $\mathcal{L}^{c}_{i}$ is the loss function for agent $i$:

where $\gamma$ is the margin of satisfaction of CBF. We need to evaluate $\dot{o}_{i}$, which is the time derivative of the observation. Instead, we approximate $\dot{h}_{i}(s_{i},o_{i})={\nabla_{s_{i}}h_{i}^{\theta_{i}}\cdot f_{i}(s_{i},a_{i})+\nabla_{o_{i}}h_{i}^{\theta_{i}}\cdot\dot{o}_{i}}$ with the forward difference method $\dot{h}_{i}(s_{i},o_{i})=h_{i}[(s_{i}(t+\Delta t),o_{i}(t+\Delta t))-h_{i}(s_{i}(t),o_{i}(t))]/\Delta t.$ So, we only need $s_{i}$ and $o_{i}$ from the dataset, then the loss function becomes:

where

For the class-$\mathcal{K}$ function $\alpha(\cdot)$, we simply choose a linear function. Note here that $\mathcal{L}^{c}_{i}$ only proposes safety constraints. We incorporate the safety reward into its reward and denote this safe version. The safety reward is $r_{p}$ when the agent enters the dangerous set.

We propose the objective function:

Inspired by IQL, we do not explicitly learn the policy by a separate value function that approximates an expectile only concerning the action distribution:

We propose a novel framework for safe multi-agent reinforcement learning (MARL) using the following components:

• •

  Centralized Training with Decentralized Execution (CTDE): This architecture allows for centralized policy learning while enabling each agent to execute actions based on local observations.

• •

  Diffusion Model: Utilized for trajectory prediction, the diffusion model generates potential future states of the agents in the environment.

• •

  Control Barrier Functions (CBFs): These functions enforce safety constraints, ensuring that agents operate within safe bounds at all times.

The training procedure involves the following steps:

1. 1.

   Pre-training Diffusion Models: Initially, diffusion models are pre-trained on offline datasets to learn the underlying data distribution. Let $\mathcal{D}=\{x_{i}\}_{i=1}^{M}$ represent the dataset. The forward diffusion process is defined as:

   $$q(x_{k+1}|x_{k})=\mathcal{N}(\sqrt{\alpha_{k}}x_{k},(1-\alpha_{k})\mathbf{I}),$$

   (17)

   where $\alpha_{k}$ is the diffusion rate and $\mathbf{I}$ is the identity matrix. The reverse diffusion process is modeled as:

   $$p_{\theta}(x_{k-1}|x_{k})=\mathcal{N}(\mu_{\theta}(x_{k},k),\Sigma_{\theta}(x_{k},k)).$$

   (18)

   The parameters $\theta$ are learned by minimizing the loss function:

   $$\mathcal{L}(\theta)=\mathbb{E}_{k,x_{0},\epsilon}\left[\|\epsilon-\epsilon_{\theta}(x_{k},k)\|^{2}\right],$$

   (19)

   where $\epsilon\sim\mathcal{N}(0,\mathbf{I})$.

2. 2.

   Integration of CBFs: Control Barrier Functions are integrated into the diffusion model to guide the learning process towards safe trajectories. For a nonlinear affine control system:

   $$\dot{s}(t)=f(s(t),a(t)),$$

   (20)

   where $s\in S\subseteq\mathbb{R}^{n}$ is the system state and $a\in A\subseteq\mathbb{R}^{m}$ is the control input, a CBF $h(s)$ satisfies:

   $$\sup_{a\in A}[\nabla_{s}h(s)\cdot f(s,a)]\geq-\alpha(h(s)).$$

   (21)

3. 3.

   Multi-objective Loss Function: The combined model is trained using a loss function that includes both safety constraints and reward optimization. Let $\tau=\{(s_{t},a_{t},r_{t})\}_{t=1}^{T}$ denote a trajectory, the objective is to maximize:

   $$\max_{\theta,\phi}\mathbb{E}_{\tau\sim\mathcal{D}}\left[\sum_{t=1}^{T}\gamma^{t}r_{t}\right]+\lambda\sum_{i=1}^{n}\mathcal{L}_{\text{CBF}}(h_{i}),$$

   (22)

   where $\gamma$ is the discount factor, $\lambda$ is a weighting factor, and $\mathcal{L}_{\text{CBF}}(h_{i})$ is the loss associated with the control barrier function for agent $i$:

   	$\displaystyle\mathcal{L}_{\text{CBF}}(h_{i})=$	$\displaystyle\sum_{s_{i}\in X_{i,0}}\max(0,\gamma-h_{i}(s_{i}))$		(23)
   		$\displaystyle+\sum_{s_{i}\in X_{i,d}}\max(0,\gamma+h_{i}(s_{i})).$