Constrained Reinforcement Learning via Dissipative Saddle Flow Dynamics

Reinforcement learning (RL) studies sequential decision-making problems where the agent aims to maximize its expected total reward by interacting with an unknown environment over time. However, in many applications such as electric grids and robotics, the agent often deals with conflicting requirements [1], or has safety constraints during the learning process [2]. The constrained reinforcement learning (C-RL) framework is a natural way to embed all conflicting requirements efficiently and incorporate safety [3, 4, 5, 2, 6, 7, 8].

There are two major approaches to finding the optimal policy of a C-RL problem, where the first approach solves it in the occupancy measure space. The constrained Markov Decision Process (CMDP) framework is a standard, and well-studied formulation for reinforcement learning with constraints [3]. The agent aims to maximize the total reward function while satisfying requirements in secondary cumulative reward constraints. The CMDP problem can be equivalently written as a linear programming problem in occupancy measure space, and the optimal policy could be recovered from the optimal occupancy measure [3]. However, this approach requires knowledge of the transition kernel of the underlying dynamical system explicitly, which is not always available in many realistic applications.

An alternative approach is to apply the Lagrangian duality and solve the C-RL problem in policy space [6, 7, 8, 9, 10]. These approaches solve the min-max optimization problem using a sampling-based primal-dual algorithm or stochastic gradient descent-ascent (SGDA) algorithm, where the Lagrangian function is augmented with a possible regularization term, e.g., a KL divergence regularization. The primal variables and dual variables are updated iteratively, either using gradient information or solving a sub-optimization problem. The outcome of primal-dual algorithms is often subject to two cases: in the first case, the output of the primal-dual algorithm is a mixing policy, which is a weighted average of history outputs [6, 7, 8]. In the second case, instead of showing the output policy converges to the optimal policy, they present a regret analysis for objective functions, and constraints [9, 10]. In summary, a key limitation is that the policy often oscillates and does not converge to the optimal policy, i.e., there is a mismatch between the behavioral policy and the optimal one. In this paper, we aim to tackle the above limitations by introducing a novel SGDA algorithm leveraging recent results on regularized saddle flow dynamics. Some of the proofs are omitted due to space constraints.

The key insight that the above sampling-based primal-dual algorithms do not converge is that the Lagrangian function for the C-RL problem does not possess sufficient convexity. The Lagrangian function is bilinear in occupancy measure space and is non-convex-concave in policy space. Our proposed method is rooted in the study of saddle flow dynamics [11, 12]. By adding a carefully crafted augmented regularization, the dissipative saddle flow proposed in [11] makes minimal requirements on convexity-concavity and yet still guarantees asymptotic convergence to a saddle point.

Leveraging tools from this dissipative saddle flow framework, we propose a novel algorithm to solve the C-RL problem in occupancy measure space, where the dynamics asymptotically converge to the optimal occupancy measure and optimal policy. We further extend the continuous-time algorithm in a model-free setting, where the discretized SGDA algorithm is shown to be the stochastic approximation of the continuous-time saddle flow dynamic. We prove that the SGDA algorithm almost surely converges to the optimal solution of the C-RL problem. To the best of our knowledge, this work is the first attempt to solve the C-RL problem to converge to the optimal occupancy measure and policy.

Notation: Let $\mathcal{K}\subset\mathbb{R}^{n}$ be a closed convex set. Given a point $y\in\mathbb{R}^{n}$, $\Psi_{\mathcal{K}}[y]=\operatorname*{argmin}_{z\in\mathcal{K}}\|z-y\|$ denote the point-wise projection (nearest point) in $\mathcal{K}$ to $y$. Given $x\in\mathcal{K}$ and $v\in\mathbb{R}^{n}$, define the vector field projection of $v$ at $x$ with respect to $\mathcal{K}$ as: $\Pi_{\mathcal{K}}[x,v]=\lim_{\delta\to 0^{+}}\frac{\Psi_{\mathcal{K}}[x+\delta v]-x}{\delta}$

In the constrained reinforcement learning problem (C-RL), $\mathcal{S}$ denotes the finite state space, $\mathcal{A}$ denotes the finite action space, and $P:\mathcal{S}\times\mathcal{A}\to\triangle^{|\mathcal{S}|}$ gives the transition dynamics of the CMDP, where $P(\cdot|s,a)$ denotes the probability distribution of next state conditioned on the current state $s$ and action $a$. $r:\mathcal{S}\times\mathcal{A}\to[0,1]$ is the reward function, $g^{i}:\mathcal{S}\times\mathcal{A}\to[-1,1]$ denotes the $i^{th}$ constraint cost function. The scalar $\gamma$ denotes the discount factor, and $q$ denotes the initial distribution of the states. A stationary policy is a map $\pi:\mathcal{S}\to\triangle^{|\mathcal{A}|}$ from states to a distribution in the action space. The value functions for both reward and constraints’ cost following policy $\pi$ are given by:

The standard C-RL problem aims to maximize the total reward function while satisfying requirements in secondary cumulative reward constraints:

There exist two classes of approaches to solving the optimal policy of a constrained reinforcement learning problem. The constrained Markov Decision Process (CMDP) framework equivalently expresses the C-RL problem as a linear programming problem in occupancy measure space [3]. Given a policy $\pi$, define $\lambda^{\pi}:\mathcal{S}\times\mathcal{A}\to[0,1]$ as occupancy measure:

where $s_{0}\sim q$. By definition, the occupancy measure belongs to the probability simplex $\lambda^{\pi}\in\Delta$. Problem (II) can be equivalently written as a linear programming problem:

where $\lambda_{a}=[\lambda(1,a),\dots,\lambda(s,a)]^{T}\in\mathbb{R}^{|\mathcal{S}|}$ is the $a^{th}$ column of $\lambda^{\pi}$, $r_{a}=[r(1,a),\dots,r(s,a)]^{T}\in\mathbb{R}^{|\mathcal{S}|}$ denotes reward function associated with action $a$, $P_{a}$ denotes the transition matrix associated with action $a$. The optimal policy could be recovered by finding the optimal occupancy measure However, a key limitation in this approach is that it requires knowledge of the transition kernel of the underlying dynamical system explicitly, i.e., $P_{a},r_{a},g_{a}^{i}$.

Another approach is to apply the primal-dual algorithm to find the saddle points of the associated Lagrangian function of problem (II) in policy space:

Algorithms often augment Lagrangian function with a regularization term $\hat{L}(\pi,\mu)=L(\pi,\mu)+R(\pi,\mu)$, e.g., a KL divergence regularization, and update the policy and dual variable using one of the following rules:

Among the sampling-based primal-dual algorithms, several algorithms output a mixing policy of the form $\pi_{T}=\sum_{k=0}^{T-1}\eta_{k}\pi_{k}$, which is a weighted average of the history updates [6, 7, 8]. The output policy oscillates and does not converge to the optimal policy. On the other hand, several papers provide a regret analysis instead of showing the algorithm’s convergence. To summarize, the CMDP approach could directly solve the optimal occupancy measure and the optimal policy while requiring knowledge of the transition kernel. The sampling-based primal-dual algorithms often output a mixing policy of history and do not converge to the optimal policy. The key limitation is that the Lagrangian function for the C-RL problem does not possess sufficient convexity. Specifically, the Lagrangian function is bilinear in occupancy measure space and is nonconvex in policy space. In this paper, we aim to provide a novel algorithm that tackles the above difficulties.

Before introducing our algorithm, we would like to illustrate our key insight from saddle flow dynamics, which explains why the primal-dual algorithm oscillates and does not converge. For a min-max optimization problem, primal-dual algorithms require the Lagrangian $L(x,y)$ function to be strictly convex or concave on $x$ or $y$, respectively, to converge. Consider the following motivating example with bilinear Lagrangian function:

Our goal is to apply different dynamic laws that seek to converge to some saddle point $(x^{*},y^{*})=(0,0)$ of $L(x,y)$, which satisfies $L(x^{*},y)\leq L(x^{*},y^{*})\leq L(x,y^{*})$. In particular, consider the following classical primal-dual algorithm:

In Figure 1, (a) plots the time series trajectory of states $x$ and $y$, and (b) plots the vector field and corresponding phase portrait. We observe that the dynamical system oscillates and does not converge to the saddle point (0,0).

In [11], the authors introduce a regularization framework for saddle flow dynamics that guarantees asymptotic convergence to a saddle point based on mild assumptions. In this paper, we further extend the above framework to solve the C-RL problem. Specifically, consider the following constrained min-max optimization problem,

where $\mathcal{K}\subset\mathbb{R}^{n},\mathcal{V}\subset\mathbb{R}^{m}$ are bounded closed convex sets. We propose a regularized surrogate for $L(x,y)$ via the following augmentation:

The following projected and regularized saddle flow dynamics aim to find the saddle points of the regularized Lagrangian, which contains the saddle point of the original Lagrangian. The regularized saddle flow dynamics still preserve the same distribution structure, which can be implemented in a fully distributed fashion, and requires the same gradient information as the classical primal-dual algorithm:

The above theorem shows the projected and regularized saddle flow dynamics will asymptotically converge to the saddle point of the Lagrangian function, which requires mild assumptions on convexity. Additionally, the following result summarizes conditions under which the solutions of the projected system exist and are unique.

We now apply the regularized saddle flow dynamics to the bilinear Lagrangian function $L(x,y)=xy$.

According to Figure 2, the trajectories of the above saddle flow dynamics asymptotically converge to the saddle point $(0,0,0,0)$, even when the original Lagrangian function is bilinear.

A direct application of the above projected and regularized saddle flow dynamic is to solve the C-RL problem in occupancy measure space (2), where the Lagrangian function is also bilinear. Specifically, the Lagrangian function for (2) in occupancy measure space is:

where $\mu_{i}\geq 0$ is the Lagrange multiplier associated with the $i^{th}$ inequality constraint and $v$ is the Lagrange multiplier associated with the equality constraint. Therefore, motivated by the projected and regularized saddle flow dynamics framework, we propose a regularized surrogate for (III) via the following augmentation:

Slater’s condition for C-RL and the following Lemma establishes the boundedness of dual decision variables, which naturally provides a closed convex set for projection.

Therefore, we propose the following projected saddle flow dynamics to find the saddle points of (III), where $\mathcal{U}:=\{\mu|\mu\in\mathbb{R}^{I}_{\geq 0},\|\mu\|_{1}\leq\frac{2}{\psi}\},\mathcal{V}:=\{v|v\in\mathbb{R}^{s},\|v^{*}\|_{\infty}\leq\frac{1}{1-\gamma}+\frac{2}{(1-\gamma)\psi}\}$ are both closed convex polyhedrons.

The following theorem is a direct application of Theorem 1 and Proposition 2, which guarantees (III) asymptotically converge to the unique (optimal) saddle point of the C-RL problem (2). Then we could recover the optimal policy from the optimal occupancy measure $\lambda^{*}$.

In the following section, we aim to extend the proposed continuous-time saddle flow algorithm (III) to a model-free setting. Specifically, we propose a novel stochastic gradient descent-ascent algorithm, which does not require the knowledge of transition kernel. We show that the SGDA algorithm is a stochastic approximation of the continuous time saddle flow dynamics (III), which almost surely (w.p.1) converges to the unique saddle point of the C-RL problem.

In many optimization problems, the goal is to find some recursive numerical procedure that sequentially approximates a value of the decision variable $x$, which minimizes the objective function, e.g., $\dot{x}=h(x)$ or $x^{n+1}=x^{n}+\alpha^{n}h(x^{n})$. Stochastic approximations attempt to solve the problem when one cannot actually observe $h(x)$, but rather $h(x)$ plus some error or noise. Consider the following projection algorithm:

where $\mathcal{G}:=\{x:q_{i}(x)\leq 0,i\in[s]\}$ denotes the constraints and $\{\xi^{n}\}$ denotes a sequence of random variables. The goal is to generate a sequence $\{x^{n}\}$ estimate of the optimal value of $x$ when the actual observation has random noise $h(x^{n})+\xi^{n}$. In general, the projection $\Psi_{\mathcal{G}}[x]$ is easy to compute when the constraints are linear; i.e., when $\mathcal{G}$ is a polyhedron. We introduce the following list of standard assumptions for stochastic approximation

Under those standard assumptions for stochastic approximations, the sequence $\{x^{n}\}$ generated by the projection algorithm (8) will converge almost surely to a stable solution to the projected system.

Consider the following randomized primal-dual approach proposed in [7, 14], where we assume the presence of a generative model. For a given state action pair $(s,a)$, the generative model provides the next state $s^{\prime}$ and the reward functions $r(s,a),g^{i}(s,a)$ to train the policy. Consider the following stochastic approximation for the Lagrangian function (III) for a distribution $\xi$:

where $s_{0}\sim q,(s,a)\sim\xi$, and the next state $s^{\prime}\sim P(\cdot|s,a)$. The stochastic approximation $L^{\xi}(\lambda,\mu,v)$ (10) is an unbiased estimator for the Lagrangian function (III), i.e., $\mathbf{E}_{\xi,P(\cdot|s,a),q}\Bigr{[}L^{\xi}(\lambda,\mu,v)\Bigr{]}=L(\lambda,\mu,v)$. Using the proposed stochastic approximation of the Lagrangian function, consider the following projection algorithm for solving the C-RL problem in a model-free setting:

The following Theorem is a direct application of Theorem 5 and 4, which shows the sequence from (IV) almost surely converges to the optimal solution to the C-RL problem.

In this section, we illustrate the effectiveness of our proposed approach using a classical CMDP problem: flow and service control problem in a single-server queue [3]. Specifically, we consider a discrete-time single-server queue with a buffer of finite size $L$. We assume that, at most, one customer may join the system in a time slot. The state $s$ corresponds to the number of customers in the queue at the beginning of a time slot $(|\mathcal{S}|=L+1)$. The service action $a$ is selected from a finite subset $A$, and the flow action $b$ is selected from a finite subset $B$. Specifically, for two real numbers satisfying $0<a_{\min}\leq a_{\max}<1$, if the queue is non-empty and if the action of the server is $a\in A$, where $A$ is a finite subset of $[a_{\min},a_{\max}]$, then the service of a customer is successfully completed with probability $a$. Likewise, for two real numbers satisfying $0\leq b_{\min}\leq b_{\max}<1$, if the queue is not full and if the action of the server is $b\in B(s)$, where $B(s)$ is a finite subset of $[b_{\min},b_{\max}]$, then the probability of having one arrival during this time slot is equal to $b$. We assume that $0\in B(x)$ for all $x$; moreover, when the buffer is full, no arrivals are possible $(B(L)={0})$. The transition law $P(\cdot|s,a)$ is therefore given by:

The reward function $r(s,a,b)$ is a real-valued decreasing function that depends only on $s$, which can be interpreted as a holding cost. The reward function $g^{1}(s,a,b)$ corresponding to the service rate is assumed to be a decreasing function that depends only on $a$. It can be interpreted as a higher service success rate having a higher cost. The reward function $g^{2}(s,a,b)$ corresponding to the flow rate $b$ is assumed to be an increasing function that depends only on $b$. It can be interpreted as a higher flow rate is more desired.

Suppose we want to solve the optimal policy for C-RL problem (II), while satisfying constraints for service and flow. In the following numerical example, we compare the result generated by (IV) and the ground truth result by directly solving the linear programming 2, where we use the transition law stated above. Specifically, we choose $L=4,A=[0.2,0.3,0.5,0.6,0.8],B=[0.1,0.3,0.5,0.9,0]$. The initial distribution $q$ is set as uniform distribution. The reward functions are $r(s)=-s+5,g^{1}(a)=-10a+3,g^{2}(b)=10b-3$.

We compare the cumulative reward function, constraint functions, and output decision variables $\lambda,\mu,v$ with the ground truth result by directly solving the linear programming problem (2). Results show that the decision variables converge to the optimal solution while satisfying the constraints for flow and service.

In this work, we propose a novel SGDA algorithm to solve the C-RL problem in occupancy measure space leveraging tools from regularized saddle flow dynamics. Even when the Lagrangian function is bilinear, the continuous dynamics asymptotically converge to the optimal occupancy measure and policy. The discretized SGDA is a stochastic approximation of the continuous-time saddle flow dynamic. We further proved the SGDA algorithm almost surely converges to the optimal solution to the C-RL problem.