ACPO: A Policy Optimization Algorithm for Average MDPs with Constraints

For a given discount factor $\gamma\in(0,1)$, the discounted reward objective is defined as

where $\tau$ refers to a sample trajectory of $(s_{0},a_{0},s_{1},\cdots)$ generated when following a policy $\pi$, that is, $a_{t}\sim\pi(\cdot|s_{t})$ and $s_{t+1}\sim P(\cdot|s_{t},a_{t})\,$; $d_{\pi,\gamma}$ is the discounted occupation measure that is defined by $d_{\pi,\gamma}(s)=(1-\gamma)\sum_{t=0}^{\infty}\gamma^{t}\underset{\tau\sim\pi}{P}(s_{t}=s)$, which essentially refers to the discounted fraction of time spent in state $s$ while following policy $\pi$.

The average-reward objective is given by:

where $d_{\pi}(s):=\lim_{N\to\infty}\frac{1}{N}\sum_{t=0}^{N-1}P_{\tau\sim\pi}(s_{t}=s)$ is the stationary state distribution under policy $\pi$. The limits in $J(\pi)$ and $d_{\pi}(s)$ are guaranteed to exist under our ergodic assumption. Since the MDP is aperiodic, it can also be shown that $d_{\pi}(s)=\lim_{t\to\infty}P_{\tau\sim\pi}(s_{t}=s)$. Since we have $\lim_{\gamma\to 1}d_{\pi,\gamma}(s)\to d_{\pi}(s),\forall s$, it can be shown that $\lim_{\gamma\to 1}(1-\gamma)J_{\gamma}(\pi)=J(\pi)$.

In the average setting, we seek to keep the estimate of the state value function unbiased and hence, introduce the average-reward bias function as

and the average-reward action-bias function as

Finally, define the average-reward advantage function as $\widebar{A}^{\pi}(s,a):=\widebar{Q}^{\pi}(s,a)-\widebar{V}^{\pi}(s)$.

A constrained Markov decision process (CMDP) is an MDP augmented with constraints that restrict the set of allowable policies for that MDP. Specifically, we augment the MDP with a set $C$ of auxiliary cost functions, $C_{1},\cdots,C_{m}$ (with each function $C_{i}:S\times A\times S\to{\mathbb{R}}$ mapping transition tuples to costs, just like the reward function), and bounds $l_{1},\cdots,l_{m}$. Similar to the value functions being defined for the average reward criterion, we define the average cost objective with respect to the cost function $C_{i}$ as

where $J_{C_{i}}$ will be referred to as the average cost for constraint $C_{i}$. The set of feasible stationary policies for a CMDP then is given by $\Pi_{C}:=\left\{\pi\in\Pi\;:\;J_{C_{i}}(\pi)\leq l_{i},{\;\;\forall\;}i\in\{1,\cdots,M\}\right\}$. The goal is to find a policy $\pi^{\star}$ such that $\pi^{\star}\in\arg\max_{\pi\in\Pi_{C}}J(\pi).$

However, finding an exact $\pi^{\star}$ is infeasible for large-scale problems. Instead, we aim to derive an iterative policy improvement algorithm that given a current policy, improves upon it by approximately maximizing the increase in the reward, while not violating the constraints by too much and not being too different from the current policy.

Lastly, analogous to $\widebar{V}^{\pi}$, $\widebar{Q}^{\pi}$, and $\widebar{A}^{\pi}$, we define similar quantities for the cost functions $C_{i}(\cdot)$, and denote them by $\widebar{V}_{C_{i}}^{\pi}$, $\widebar{Q}_{C_{i}}^{\pi}$, and $\widebar{A}_{C_{i}}^{\pi}$.

In many on-policy constrained RL problems, we improve policies iteratively by maximizing a predefined function within a local region of the current best policy as in (Tessler et al., 2019; Achiam et al., 2017; Yang et al., 2020; Song et al., 2020). (Achiam et al., 2017) derived a policy improvement bound for the discounted CMDP setting as:

where $A_{\gamma}^{\pi_{k}}$ is the discounted version of the advantage function, $\epsilon^{\pi_{k+1}}:=\max_{s}|{\mathbb{E}}_{a\sim\pi_{k+1}}[A_{\gamma}^{\pi_{k}}(s,a)]|$, and $D_{TV}(\pi_{k+1}||\pi_{k})[s]=(1/2)\sum_{a}\left|\pi_{k+1}(a|s)-\pi_{k}(a|s)\right|$ is the total variational divergence between $\pi_{k+1}$ and $\pi_{k}$ at $s$. These results laid the foundations for on-policy constrained RL algorithms as in(Wu et al., 2017; Vuong et al., 2019).

However, Equation (3) does not generalize to the average setting ($\gamma\to 1$) (see Appendix A.1). In the next section, we will derive a policy improvement bound for the average case and present an algorithm based on trust region methods, which will generate almost-monotonically improving iterative policies. Proofs of theorems and lemmas, if not already given, are available in Appendix A.

Let $\pi^{\prime}$ be the policy obtained via some update rule from the current policy $\pi$. Analogous to the discounted setting of a CMDP, we would like to characterize the performance difference $J(\pi^{\prime})-J(\pi)$ by an expression which depends on $\pi$ and some divergence metric between the two policies.

Note that this difference depends on the stationary state distribution obtained from the new policy, $d_{\pi^{\prime}}$. This is computationally impractical as we do not have access to this $d_{\pi^{\prime}}$. Fortunately, by use of the following lemma we can show that if $d_{\pi}$ and $d_{\pi^{\prime}}$ are “close” with respect to some metric, we can approximate Eq. (4) using samples from $d_{\pi}$.

, where $\epsilon=\max_{s}\big{|}\underset{\begin{subarray}{c}a\sim\pi^{\prime}\end{subarray}}{{\mathbb{E}}}[\widebar{A}^{\pi}(s,a)]\big{|}$. See Appendix A.1 for proof. Lemma 3.2 implies $J(\pi^{\prime})\approx J(\pi)+\underset{\begin{subarray}{c}\end{subarray}}{{\mathbb{E}}}[\widebar{A}^{\pi}(s,a)]$ when $d_{\pi}$ and $d_{\pi^{\prime}}$ are “close”. Now that we have established this approximation, we need to study the relation of how the actual change in policies affects their corresponding stationary state distributions. For this, we turn to standard analysis of the underlying Markov chain of the CMDP.

Under the ergodic assumption, we have that $P_{\pi}$ is irreducible and hence its eigenvalues $\{\lambda_{\pi,i}\}_{i=1}^{|S|}$ are such that $\lambda_{\pi,1}=1$ and $\lambda_{\pi,i\neq 1}<1$. For our analysis, we define $\sigma^{\pi}=\max_{i\neq 1}\,(1-\lambda_{\pi,i})^{-1/2}$, and from (Levene & Loizou, 2002) and (Doyle, 2009), we connect $\{\lambda_{\pi,i}\}_{i=1}^{|S|}$ to the sensitivity of the stationary distributions to changes in the policy using the result below.

, where $\sigma^{\star}=\max_{\pi}\sigma^{\pi}$. See Appendix A.1 for proof. This bound is tighter and easier to compute than the one given by (Zhang et al., 2021), which replaces $\sigma^{\star}$ by $\kappa^{\star}=\max_{\pi}\kappa^{\pi}$, where $\kappa^{\pi}$ is known as Kemeny’s constant from (Kemeny & Snell, 1960). It is interpreted as the expected number of steps to get to any goal state, where the expectation is taken with respect to the stationary-distribution of those states.

Combining the bounds in Lemma 3.2 and Lemma 3.3 gives us the following result:

It is interesting to compare the inequalities of Equation (7) to Equation (4). The term $\underset{\begin{subarray}{c}\end{subarray}}{{\mathbb{E}}}[\widebar{A}^{\pi}(s,a)]$ in Prop. 3.4 is somewhat of a surrogate approximation to $J(\pi^{\prime})-J(\pi)$, in the sense that it uses $d_{\pi}$ instead of $d_{\pi^{\prime}}$. As discussed before, we do not have access to $d_{\pi^{\prime}}$ since the trajectories of the new policy are not available unless the policy itself is updated. This surrogate is a first order approximation to $J(\pi^{\prime})-J(\pi)$ in the parameters of $\pi^{\prime}$ in a neighborhood around $\pi$ as in (Kakade & Langford, 2002). Hence, Eq. (7) can be viewed as bounding the worst-case approximation error.

Extending this discussion to the cost function of our CMDP, similar expressions follow immediately.

Until now, we have been dealing with bounds given with regards to the TV divergence of the policies. However, in practice, bounds with respect to the KL divergence of policies is more commonly used as in (Schulman et al., 2015, 2016; Ma et al., 2021). From Pinsker’s and Jensen’s inequalities, we have that

We can thus use Eq. (9) in the bounds of Proposition 3.4 and Corollary 3.5 to make policy improvement guarantees, i.e., if we find updates such that $\pi_{k+1}\in\operatorname*{arg\,max}_{\pi}L_{\pi_{k}}^{-}(\pi)$, then we will have monotonically increasing policies as, at iteration $k$, $\underset{\begin{subarray}{c}s\sim d_{\pi_{k}},a\sim\pi\end{subarray}}{{\mathbb{E}}}[\widebar{A}^{\pi_{k}}(s,a)]=0$, $\underset{\begin{subarray}{c}s\sim d_{\pi_{k}}\end{subarray}}{{\mathbb{E}}}[D_{\text{KL}}\left(\pi\middle\|\pi_{k}\right)[s]]=0$ for $\pi=\pi_{k}$, implying that $J(\pi_{k+1})-J(\pi_{k})\geq 0$. However, this sequence does not guarantee constraint satisfaction at each iteration, so we now turn to trust region methods to incorporate constraints, do policy improvement and provide safety guarantees.

For large or continuous state and action CMDPs, solving for the exact optimal policy is impractical. However, trust region-based policy optimization algorithms have proven to be effective for solving such problems as in (Schulman et al., 2015, 2016, 2017; Achiam, 2017). For these approaches, we usually consider some parameterized policy class $\Pi_{\Theta}=\{\pi_{\theta}:\theta\in\Theta\}$ for tractibility. In addition, for CMDPs, we also require the policy iterates to be feasible, so instead of optimizing just over $\Pi_{\Theta}$, we optimize over $\Pi_{\Theta}\cap\Pi_{C}$. However, it is much easier to solve the above problem if we introduce hard constraints, rather than limiting the set to $\Pi_{\Theta}\cap\Pi_{C}$. Therefore, we now introduce the ACPO algorithm, which is inspired by the trust region formulations above as the following optimization problem:

where $\bar{D}_{\text{KL}}(\pi\parallel\pi_{\theta_{k}}):=\underset{\begin{subarray}{c}s\sim d_{\pi_{\theta_{k}}}\end{subarray}}{{\mathbb{E}}}[D_{\text{KL}}\left(\pi\middle\|\pi_{\theta_{k}}\right)[s]]$, $\widebar{A}^{\pi_{\theta_{k}}}(s,a)$ is the average advantage function defined earlier, and $\delta>0$ is a step size. We use this form of updates as it is an approximation to the lower bound given in Proposition 3.4 and the upper bound given in Corollary 3.5.

In most cases, the trust region threshold for formulations like Eq. (10) are heuristically motivated. We now show that it is quantitatively motivated and comes with a worst case performance degradation and constraint violation. Proof is in Appendix A.2.

Remark 3.7. Note that if the constraints are ignored (by setting $V_{max}=0$), then this bound is tighter than given in (Zhang et al., 2021) for the unconstrained average-reward setting.

However, the update rule of Eq. (10) is difficult to implement in practice as it takes steps that are too small, which degrades convergence. In addition, it requires the exact knowledge of $\widebar{A}^{\pi_{\theta_{k}}}(s,a)$ which is computationally infeasible for large-scale problems. In the next section, we will introduce a specific sampling-based practical algorithm to alleviate these concerns.

Since we are working with a parameterized class, we shall now overload notation to use $\theta_{k}$ as the policy at iteration $k$, i.e., $\theta_{k}\equiv\pi_{\theta_{k}}$. In addition, we use $g$ to denote the gradient of the advantage function objective, $a_{i}$ to denote the gradient of the advantage function of the cost $C_{i}$, $H$ as the Hessian of the KL-divergence. Formally,

In addition, let $c_{i}:=J_{C_{i}}(\theta_{k})-l_{i}$. The approximation to the problem in Eq. (10) is:

This is a convex optimization problem in which strong duality holds, and hence it can be solved using a Lagrangian method. The update rule for the dual problem then takes the form

where $\lambda^{\star}$ and $\mu^{\star}$ are the Lagrange multipliers satisfying the dual

with $r:=g^{T}H^{-1}A$, $S:=A^{T}H^{-1}A$, $A:=[a_{1},\cdots,a_{m}]$, and $c:=[c_{1},\cdots,c_{m}]^{T}$.

The approximation regime described in Eq. (13) requires $H$ to be invertible. For large parametric policies, $H$ is computed using the conjugate gradient method as in (Schulman et al., 2015). However, in practice, using this approximation along with the associated statistical sampling errors, there might be potential violations of the approximate constraints leading to infeasible policies.

To rectify this, for the case where we only have one constraint, one can recover a feasible policy by applying a recovery step inspired by the TRPO update on the cost surrogate as:

where $t\in[0,1]$. Contrasting with the policy recovery update of (Achiam et al., 2017) which only uses the cost advantage function gradient $a$, we introduce the reward advantage function gradient $g$ as well. This choice is to ensure recovery while simultaneously balancing the “regret” of not choosing the best (in terms of the objective value) policy $\pi_{k}$. In other words, we wish to find a policy $\pi_{k+1/2}$ as close to $\pi_{k}$ in terms of their objective function values. We follow up this step with a simple linesearch to find feasible $\pi_{k+1}$. Based on this, Algorithm 1 provides a basic outline of ACPO. For more details of the algorithm, see Appendix A.3.

For the Gather and Circle tasks we test two distinct agents: a point-mass ($S\subseteq{\mathbb{R}}^{9},A\subseteq{\mathbb{R}}^{2}$), and an ant robot ($S\subseteq{\mathbb{R}}^{32},A\subseteq{\mathbb{R}}^{8}$). The agent in the Bottleneck task in $S\subseteq{\mathbb{R}}^{71},A\subseteq{\mathbb{R}}^{16}$, and for the Grid task is $S\subseteq{\mathbb{R}}^{96},A\subseteq{\mathbb{R}}^{4}$. We use two hidden layer neural networks to represent Gaussian policies for the tasks. For Gather and Circle, size is (64,32) for both layers, and for Grid and Bottleneck the layer sizes are (16,16) and (50,25). We set the step size $\delta$ to $10^{-4}$, and for each task, we conduct 5 runs to get the mean and standard deviation for reward objective and cost constraint values during training. We train CPO, PCPO, and PPO with the discounted objective, however, evaluation and comparison with BVF-PPO, ATRPO and ACPO¹¹1Code of the ACPO implementation will be made available on GitHub. is done using the average reward objective (this is a standard evaluation scheme as in (Schulman et al., 2015; Wu et al., 2017; Vuong et al., 2019)).

For each environment, we train an agent for $10^{5}$ steps, and for every $10^{3}$ steps, we instantiate 10 evaluation trajectories with the current (deterministic) policy. For each of these trajectories, we calculate the trajectory average reward for the next $10^{3}$ steps and finally report the total average-reward as the mean of these 10 trajectories. Learning curves for the algorithms are compiled in Figure 1 (for Point-Circle, Point-Gather, and Ant-Circle see Appendix A.6).

Since there are two objectives (rewards in the objective and costs in the constraints), we show the plots which maximize the reward objective while satisfying the cost constraint. See Appendix A.4 and A.5 for more details.

From Figure 1, we can see that ACPO is able to improve the reward objective while having approximate constraint satisfaction on all tasks. In particular, ACPO is the only algorithm that best learns almost-constraint-satisfying maximum average-reward policies across all tasks: in a simple Gather environment, ACPO is able to almost exactly track the cost constraint values to within the given threshold $l$; however, for the high dimensional Grid and Bottleneck environments we have more constraint violations due to complexity of the policy behavior. Regardless, in these environments, ACPO still outperforms all other baselines.

In Equation (16) we introduced a hyperparameter $t$, which provides for an intuitive trade-off as follows: either we purely decrease the constraint violations ($t=1$), or we decrease the average-reward ($t=0$), which consequently decreases the constraint violation. The latter formulation is principled in that if we decrease rewards, we are bound to decrease constraints violation due to the nature of the environments. Figure 2 shows the experiments we conducted with varying $t$. With $t=1$, we obtain the same recovery scheme as that of (Achiam et al., 2017). Our results show that this scheme does not lead to the best performance, and that $t=0.75$ and $t=1$ perform the best across all tasks. See Appendix A.6 for a detailed study.