Reward is Enough for Convex MDPs
Year: 2022

Abstract:
Maximising a cumulative reward function that is Markov and stationary, i.e., defined over state-action pairs and independent of time, is sufficient to capture many kinds of goals in a Markov decision process (MDP). However, not all goals can be captured in this manner. In this paper we study convex MDPs in which goals are expressed as convex functions of the stationary distribution and show that they cannot be formulated using stationary reward functions. Convex MDPs generalize the standard reinforcement learning (RL) problem formulation to a larger framework that includes many supervised and unsupervised RL problems, such as apprenticeship learning, constrained MDPs, and so-called ‘pure exploration’. Our approach is to reformulate the convex MDP problem as a min-max game involving policy and cost (negative reward) ‘players’, using Fenchel duality. We propose a meta-algorithm for solving this problem and show that it unifies many existing algorithms in the literature.

heading: 1 Introduction
text: In reinforcement learning (RL), an agent learns how to map situations to actions so as to maximize a cumulative scalar reward signal. The learner is not told which actions to take, but instead must discover which actions lead to the most reward [64]. Mathematically, the RL problem can be written as finding a policy whose state occupancy has the largest inner product with a reward vector [55], i.e., the goal of the agent is to solve
RL: max dπ∈K ∑ s,a r(s, a)dπ(s, a), (1)
where dπ is the state-action stationary distribution induced by policy π and K is the set of admissible stationary distributions (see Definition 1). A significant body of work is dedicated to solving the RL problem efficiently in challenging domains [45, 62]. However, not all decision making problems of interest take this form. In particular we consider the more general convex MDP problem,
Convex MDP: min dπ∈K f(dπ), (2)
where f : K → R is a convex function. Sequential decision making problems that take this form include Apprenticeship Learning (AL), pure exploration, and constrained MDPs, among others; see Table 1. In this paper we prove the following claim:
We can solve Eq. (2) by using any algorithm that solves Eq. (1) as a subroutine.
In other words, any algorithm that solves the standard RL problem can be used to solve the more general convex MDP problem. More specifically, we make the following contributions.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Firstly, we adapt the meta-algorithm of Abernethy and Wang [3] for solving Eq. (2). The key idea is to use Fenchel duality to convert the convex MDP problem into a two-player zero-sum game between the agent (henceforth, policy player) and an adversary that produces rewards (henceforth, cost player) that the agent must maximize [3, 6]. From the agent’s point of view, the game is bilinear, and so for fixed rewards produced by the adversary the problem reduces to the standard RL problem with non-stationary reward (Fig. 1).
Secondly, we propose a sample efficient policy player that uses a standard RL algorithm (eg, [35, 60]), and computes an optimistic policy with respect to the non-stationary reward at each iteration. In other words, we use algorithms that were developed to achieve low regret in the standard RL setup, to achieve low regret as policy players in the min-max game we formulate to solve the convex MDP. Our main result is that the average of the policies produced by the policy player converges to a solution to the convex MDP problem (Eq. (2)). Inspired by this principle, we also propose a recipe for using deep-RL
(DRL) agents to solve convex MDPs heuristically: provide the agent non-stationary rewards from the cost player. We explore this principle in our experiments.
Finally, we show that choosing specific algorithms for the policy and cost players unifies several disparate branches of RL problems, such as apprenticeship learning, constrained MDPs, and pure exploration into a single framework, as we summarize in Table 1.

heading: 2 Reinforcement Learning Preliminaries
text: In RL an agent interacts with an environment over a number of time steps and seeks to maximize its cumulative reward. We consider two cases, the average reward case and the discounted case. The Markov decision process (MDP) is defined by the tuple (S,A, P,R) for the average reward case and by the tuple (S,A, P,R, γ, d0) for the discounted case. We assume an infinite horizon, finite state-action problem where initially, the state of the agent is sampled according to s0 ∼ d0, then at each time t the agent is in state st ∈ S, selects action at ∈ A according to some policy π(st, ·), receives reward rt ∼ R(st, at) and transitions to new state st+1 ∈ S according to the probability distribution P (·, st, at). The two performance metrics we consider are given by
Javgπ = lim T→∞
1 T E T∑ t=1 rt, J γ π = (1− γ)E ∞∑ t=1 γtrt, (3)
for the average reward case and discounted case respectively. The goal of the agent is to find a policy that maximizes Javgπ or J γ π . Any stationary policy π induces a state-action occupancy measure dπ,
which measures how often the agent visits each state-action when following π. Let Pπ(st = ·) be the probability measure over states at time t under policy π, then
davgπ (s, a) = lim T→∞
1 T E T∑ t=1 Pπ(st = s)π(s, a), dγπ(s, a) = (1− γ)E ∞∑ t=1 γtPπ(st = s)π(s, a),
for the average reward case and the discounted case respectively. With these, we can rewrite the RL objective in Eq. (3) in terms of the occupancy measure using the following well-known result, which for completeness we prove in Appendix B. Proposition 1. For both the average and the discounted case, the agent objective function Eq. (3) can be written in terms of the occupancy measure as Jπ = ∑ s,a r(s, a)dπ(s, a).
Given an occupancy measure it is possible to recover the policy by setting π(s, a) = dπ(s, a)/ ∑ a dπ(s, a) if ∑ a dπ(s, a) > 0, and π(s, a) = 1/|A| otherwise. Accordingly, in this paper we shall formulate the RL problem using the state-action occupancy measure, and both the standard RL problem (Eq. (1)) and the convex MDP problem (Eq. (2)) are convex optimization problems in variable dπ. For the purposes of this manuscript we do not make a distinction between the average and discounted settings, other than through the convex polytopes of feasible occupancy measures, which we define next. Definition 1 (State-action occupancy’s polytope [55]). For the average reward case the set of admissible state-action occupancies is
Kavg = {dπ | dπ ≥ 0, ∑ s,a dπ(s, a) = 1, ∑ a dπ(s, a) = ∑ s′,a′ P (s, s′, a′)dπ(s ′, a′) ∀s ∈ S},
and for the discounted case it is given by Kγ = {dπ | dπ ≥ 0, ∑ a dπ(s, a) = (1− γ)d0(s) + γ ∑ s′,a′ P (s, s′, a′)dπ(s ′, a′) ∀s ∈ S}.
We note that being a polytope implies that K is a convex and compact set. The convex MDP problem is defined for the tuple (S,A, P, f) in the average cost case and (S,A, P, f, γ, d0) in the discounted case. This tuple is defining a state-action occupancy’s polytope K (Definition 1), and the problem is to find a policy π whose state occupancy dπ is in this polytope and minimizes the function f (Eq. (2)).

heading: 3 A Meta-Algorithm for Solving Convex MDPs via RL
text: To solve the convex MDP problem (Eq. (2)) we need to find an occupancy measure dπ (and associated policy) that minimizes the function f . Since both f : K → R and the set K are convex this is a convex optimization problem. However, it is a challenging one due to the nature of learning about the environment through stochastic interactions. In this section we show how to reformulate the convex MDP problem (Eq. (2)) so that standard RL algorithms can be used to solve it, allowing us to harness decades of work on solving vanilla RL problems. To do that we will need the following definition. Definition 2 (Fenchel conjugate). For a function f : Rn → R ∪ {−∞,∞}, its Fenchel conjugate is denoted f∗ : Rn → R ∪ {−∞,∞} and defined as f∗(x) := supy x · y − f(y). Remark 1. The Fenchel conjugate function f∗ is always convex (when it exists) even if f is not. Furthermore, the biconjugate f∗∗ := (f∗)∗ equals f if and only if f is convex and lower semicontinuous.
Using this we can rewrite the convex MDP problem (Eq. (2)) as
fOPT = min dπ∈K f(dπ) = min dπ∈K max λ∈Λ (λ · dπ − f∗(λ)) = max λ∈Λ min dπ∈K (λ · dπ − f∗(λ)) (4)
where Λ is the closure of (sub-)gradient space {∂f(dπ)|dπ ∈ K}, which is a convex set [3, Theorem 4]. As both sets are convex, this is a convex-concave saddle-point problem and a zero-sum two-player game [54, 49], and we were able to swap the order of minimization and maximization using the minimax theorem [71].
With this we define the Lagrangian as L(dπ, λ) := λ · dπ − f∗(λ). For a fixed λ ∈ Λ, minimizing the Lagrangian is a standard RL problem of the form of Eq. (1), i.e., equivalent to maximizing a reward r = −λ. Thus, one might hope that by producing an optimal dual variable λ? we could simply solve d?π = argmindπ∈K L(·, λ
?) for the optimal occupancy measure. However, the next lemma states that this is not possible in general. Lemma 1. There exists an MDP M and convex function f for which there is no stationary reward r ∈ RS×A such that arg maxdπ∈K dπ · r = arg mindπ∈K f(dπ).
To see this note that for any reward r there is a deterministic policy that optimizes the reward [55], but for some choices of f no deterministic policy is optimal, eg, when f is the negative entropy function. This result tells us that even if we have access to an optimal dual-variable we cannot simply use it to recover the stationary distribution that solves the convex MDP problem in general.
To overcome this issue we develop an algorithm that generates a sequence of policies {πk}k∈N such that the average converges to an optimal policy for Eq. (2), i.e., (1/K) ∑K k=1 d k π → d?π ∈ arg mindπ∈K f(dπ). The algorithm we develop is described in Algorithm 1 and is adapted from the meta-algorithm described in Abernethy and Wang [3]. It is referred to as a meta-algorithm since it relies on supplied sub-routine algorithms Algπ and Algλ. The reinforcement learning algorithm Algπ takes as input a reward vector and returns a state-action occupancy measure dπ. The cost algorithm Algλ can be a more general function of the entire history. We discuss concrete examples of Algπ and Algλ in Section 4.
Algorithm 1: meta-algorithm for convex MDPs 1: Input: convex-concave payoff L : K × Λ→ R, algorithms Algλ,Algπ , K ∈ N 2: for k = 1, . . . ,K do 3: λk = Algλ(d 1 π, . . . , d k−1 π ;L)
4: dkπ = Algπ(−λk) 5: end for 6: Return d̄Kπ = 1 K ∑K k=1 d k π, λ̄ K = 1K ∑K k=1 λ k
In order to analyze this algorithm we will need a small detour into online convex optimization (OCO). In OCO, a learner is presented with a sequence of K convex loss functions `1, `2, . . . , `K : K → R and at each round k must select a point xk ∈ K after which it suffers a loss of `k(xk). At time period k the learner is assumed to have perfect knowledge of the loss functions `1, . . . , `k−1. The learner wants to minimize its average regret, defined as
R̄K := 1
K ( K∑ k=1 `k(xk)−min x∈K K∑ k=1 `k(x) ) .
In the context of convex reinforcement learning and meta-algorithm 1, the loss functions for the cost player are `kλ = −L(·, λk), and for the policy player are `kπ = L(dkπ, ·), with associated average regrets R̄πK and R̄ λ K . This brings us to the following theorem. Theorem 1 (Theorem 2, [3]). Assume that Algπ and Algλ have guaranteed average regret bounded as R̄πK ≤ K and R̄λK ≤ δK , respectively. Then Algorithm 1 outputs d̄Kπ and λ̄K satisfying mindπ∈K L(dπ, λ̄K) ≥ fOPT − K − δK and maxλ∈Λ L(d̄Kπ , λ) ≤ fOPT + K + δK .
This theorem tells us that so long as the RL algorithm we employ has guaranteed low-regret, and assuming we choose a reasonable low-regret algorithm for deciding the costs, then the meta-algorithm will produce a solution to the convex MDP problem (Eq. (2)) to any desired tolerance, this is because fOPT ≤ f(d̄Kπ ) = maxλ L(d̄Kπ , λ) ≤ fOPT + K +δK . For example, we shall later present algorithms that have regret bounded as K = δK ≤ O(1/ √ K), in which case we have
f(d̄Kπ )− fOPT ≤ O(1/ √ K). (5)
Non-Convex f . Remark 1 implies that the game maxλ∈Λ mindπ∈K (λ · dπ − f∗(λ)) is concaveconvex for any function f , so we can solve it with Algorithm 1, even for a non-convex f . From weak duality the value of the Lagrangian on the output of Algorithm 1, L(d̄π, λ̄), is a lower bound on the optimal solution fOPT. In addition, since f(dπ) is always an upper bound on fOPT we have both an upper bound and a lower bound on the optimal value: L(d̄π, λ̄) ≤ fOPT ≤ f(d̄π).

heading: 4 Policy and Cost Players for Convex MDPs
text: In this section we present several algorithms for the policy and cost players that can be used in Algorithm 1. Any combination of these algorithms is valid and will come with different practical and theoretical performance. In Section 6 we show that several well known methods in the literature correspond to particular choices of cost and policy players and so fall under our framework.
In addition, in this section we assume that
λmax = max λ∈Λ max s,a |λ(s, a)| <∞,
which holds when the set Λ is compact. One way to guarantee that Λ is compact is to consider functions f with Lipschitz continuous gradients (which implies bounded gradients since the set K is compact). For simplicity, we further assume that λmax ≤ 1. By making this assumption we assure that the non stationary rewards produced by the cost player are bounded by 1 as is usually done in RL.

heading: 4.1 Cost Player
text: Follow the Leader (FTL) is a classic OCO algorithm that selects λk to be the best point in hindsight. In the special case of convex MDPs, as defined in Eq. (4), FTL has a simpler form:
λk = arg max λ∈Λ ∑k−1 j=1 L(djπ, λ) = arg max λ∈Λ ( λ · ∑k−1 j=1 djπ −Kf∗(λ) ) = ∇f(d̄k−1π ), (6)
where d̄k−1π = ∑k−1 j=1 d j π and the last equality follows from the fact that (∇f∗)−1 = ∇f [56]. The average regret of FTL is guaranteed to be R̄K ≤ c/ √ K under some assumptions [29]. In some cases, and specifically when the set K is a polytope and the function f is strongly convex, FTL can enjoy logarithmic or even constant regret; see [32, 29] for more details.
Online Mirror Descent (OMD) uses the following update [47, 9]:
λk = arg max λ∈Λ
( (λ− λk−1) · ∇λL(dk−1π , λk−1) + αkBr(λ, λk−1) ) ,
where αk is a learning rate and Br is a Bregman divergence [14]. For Br(x) = 0.5||x||22, we get online gradient descent [79] and for Br(x) = x · log(x) we get multiplicative weights [23] as special cases. We also note that OMD is equivalent to a linearized version of Follow the Regularized Leader (FTRL) [43, 28]. The average regret of OMD is R̄K ≤ c/ √ K under some assumptions, see, for example [28].

heading: 4.2 Policy Players

heading: 4.2.1 Best Response
text: In OCO, the best response is to simply ignore the history and play the best option on the current round, which has guaranteed average regret bound of R̄K ≤ 0 (this requires knowledge of the current loss function, which is usually not applicable but is in this case). When applied to Eq. (4), it is possible to find the best response dkπ using standard RL techniques since
dkπ = arg min dπ∈K Lk(dπ, λk) = arg min dπ∈K dπ · λk − f∗(λk) = arg max dπ∈K dπ · (−λk),
which is an RL problem for maximizing the reward (−λk). In principle, any RL algorithm that eventually solves the RL problem can be used to find the best response, which substantiates our claim in the introduction. For example, tabular Q-learning executed for sufficiently long and with a suitable exploration strategy will converge to the optimal policy [72]. In the non-tabular case we could parameterize a deep neural network to represent the Q-values [45] and if the network has sufficient capacity then similar guarantees might hold. We make no claims on efficiency or tractability of this approach, just that in principle such an approach would provide the best-response at each iteration and therefore satisfy the required conditions to solve the convex MDP problem.

heading: 4.2.2 Approximate Best Response
text: The caveat in using the best response as a policy player is that in practice, it can only be found approximately by executing an RL algorithm in the environment. This leads to defining an approximate best response via the Probably Approximately Correct (PAC) framework. We say that a policy player is PAC( , δ), if it finds an -optimal policy to an RL problem with probability of at least 1 − δ. In addition, we say that a policy π′ is -optimal if its state occupancy d′π is such that
max dπ∈K
dπ · (−λk)− d′π · (−λk) ≤ .
For example, the algorithm in [40] can find an -optimal policy to the discounted RL problem after seeing O
( SA (1−γ)3 2 log( 1 δ ) ) samples; and the algorithm in [36] can find an -optimal policy for the
average reward RL problem after seeing O ( t2mixSA 2 log( 1 δ ) )
samples, where tmix is the mixing time (see, eg, [42, 76] for a formal definition). The following Lemma analyzes the sample complexity of Algorithm 1 with an approximate best response policy player for the average reward RL problem [36]. The result can be easily extended to the discounted case using the algorithm in [40]. Other relaxations to the best response for specific algorithms can be found in [65, 44, 33, 30]. Lemma 2 (The sample complexity of approximate best response in convex MDPs with average occupancy measure). For a convex function f , running Algorithm 1 with an oracle cost player with regret R̄λK = O(1/K) and an approximate best response policy player that solves the average reward RL problem in iteration k to accuracy k = 1/k returns an occupancy measure d̄Kπ that satisfies f(d̄Kπ )− fOPT ≤ with probability 1− δ after seeing O(t2mixSA log(2K/ δ)/ 3δ3) samples. Similarly, for R̄λK = O(1/ √ K), setting k = 1/ √ k requires O(t2mixSA log(2K/ δ)/ 4δ4) samples.

heading: 4.2.3 Non-Stationary RL Algorithms
text: We now discuss a different type of policy players; instead of solving an MDP to accuracy , these algorithms perform a single RL update to the policy, with cost −λk. In our setup the reward is known and deterministic but non-stationary, while in the standard RL setup it is unknown, stochastic, and stationary. We conjecture that any RL algorithm can be adapted to the known non-stationary reward setup we consider here. In most cases both Bayesian [51, 48] and frequentist [8, 35] approaches to the stochastic RL problem solve a modified (eg, by adding optimism) Bellman equation at each time period and swapping in a known but non-stationary reward is unlikely to present a problem.
To support this conjecture we shall prove that this is exactly the case for UCRL2 [35]. UCRL2 is an RL algorithm that was designed and analyzed in the standard RL setup, and we shall show that it is easily adapted to the non-stationary but known reward setup that we require. To make this claim more general, we will also discuss a similar result for the MDPO algorithm [61] that was given in a slightly different setup.
UCRL2 is a model based algorithm that maintains an estimate of the reward and the transition function as well as confidence sets about those estimates. In our case the reward at time k is known, so we only need to consider uncertainty in the dynamics. UCRL2 guarantees that in any iteration k, the true transition function is in a confidence set with high probability, i.e., P ∈ Pk for confidence set Pk. If we denote by JP,Rπ the value of policy π in an MDP with dynamics P and reward R then the optimistic policy is π̃k = arg maxπ maxP ′∈Pk J P ′,−λk π . Acting according to this policy is guaranteed to attain low regret. In the following results for UCRL2 we will use the constant D, which denotes the diameter of the MDP, see [35, Definition 1] for more details. In the supplementary material (Appendix E), we provide a proof sketch that closely follows [35]. Lemma 3 (Non stationary regret of UCRL2). For an MDP with dynamics P, diameter D, an arbitrary sequence of known and bounded rewards { ri : maxs,a |ri(s, a)| ≤ 1 }K i=1
, such that the optimal average reward at episode k, with respect to P and rk is J?k , then with probability at least 1−δ, the average regret of UCRL2 is at most R̄K = 1K ∑K k=1 J ? k−J π̃k k ≤ O(DS √ A log(K/δ)/K).
Next, we give a PAC( , δ) sample complexity result for the mixed policy π̄K , that is produced by running Algorithm 1 with UCRL2 as a policy player.
Lemma 4 (The sample complexity of non-stationary RL algorithms in convex MDPs). For a convex function f, running Algorithm 1 with an oracle cost player with regret R̄λK ≤ c0/ √ K and UCRL2 as a policy player returns an occupancy measure d̄Kπ that satisfies f(d̄ K π )− fOPT ≤ with probability
1− δ after K = O ( D2S2A δ2 2 log( 2DSA δ ) ) steps.
MDPO. Another optimistic algorithm is Mirror Descent Policy Optimization [60, MDPO]. MDPO is a model free RL algorithm that is very similar to popular DRL algorithms like TRPO [58] and MPO [2]. In [24, 59, 5], the authors established the global convergence of MDPO and in [15, 60], the authors showed that MDPO with optimistic exploration enjoys low regret.
The analysis for MDPO is given in a finite horizon MDP with horizon H , which is not the focus of our paper. Nevertheless, to support our conjecture that any stochastic RL algorithm can be adapted to the known non-stationary reward setup, we quickly discuss the regret of MDPO in this setup. We also note that MDPO is closer to practical DRL algorithms [70]. In a finite horizon MDP with horizon H and known, non-stationary and bounded rewards, the regret of MDPO is bounded by R̄K ≤ O(H2S √ A/K) [61, Lemma 4] with high probability.
To compare this result with UCRL2, we refer to a result from [57], which analyzed UCRL2 in the adversarial setup, that includes our setup as a special case. In a finite horizon MDP with horizon H it was shown that setting δ = SA/K with probability 1 − δ its regret is bounded by R̄K ≤ O(HS √ A log(K)/K) [57, Corollary 5], which is better by a factor of H than MDPO.
Discussion. Comparing the results in Lemma 4 with Lemma 2 suggests that using an RL algorithm with non stationary reward as a policy player requires O(1/ 2) samples to find an −optimal policy, while using an approximate best response requires O(1/ 3). In first glance, this results also improves the previously best known result of Hazan et al. [30] for approximate Frank-Wolfe (FW) that requires O(1/ 3) samples. However, there are more details that have to be considered as we now discuss.
Firstly, Lemma 4 and Lemma 2 assume access to an oracle cost player with some regret and do not consider how to implement such a cost player. The main challenge is that the cost player does not have access to the true state occupancy and must estimate it from samples. If we do not reuse samples from previous policies to estimate the state occupancy of the current policy we will require O(1/ 3) trajectories overall [30]. A better approach would use the samples from previous episodes to learn the transition function. Then, given the estimated transition function and the policy, we can compute an approximation of the state occupancy. We conjecture that such an approach would lead to a O(1/ 2) sample complexity, closing the gap with standard RL.
Secondly, while our focus is on the dependence in , our bound Lemma 4 is not tight in δ, i.e., it scales with 1/δ2 where it should be possible to achieve a log(1/δ) scaling. Again we conjecture an improvement in the bound is possible; see, eg, [38, Appendix F.].

heading: 5 Convex Constraints
text: We have restricted the presentation so far to unconstrained convex problems, in this section we extend the above results to the constrained case. The problem we consider is
min dπ∈K
f(dπ) subject to gi(dπ) ≤ 0, i = 1, . . .m,
where f and the constraint functions gi are convex. Previous work focused on the case where both f and gi are linear [7, 67, 12, 68, 18, 16, 11]. We can use the same Fenchel dual machinery we developed before, but now taking into account the constraints. Consider the Lagrangian
L(dπ, µ) = f(dπ) + ∑m
i=1 µigi(dπ) = max ν (ν · dπ − f∗(ν)) + ∑m i=1 µi max vi (dπvi − g∗i (vi)) .
over dual variables µ ≥ 0, with new variables vi and ν. At first glance this does not look convexconcave, however we can introduce new variables ζi = µivi to obtain
L(dπ, µ, ν, ζ1, . . . , ζm) = ν · dπ − f∗(ν) + ∑m
i=1 (dπζi − µig∗i (ζi/µi)) . (7)
This is convex (indeed affine) in dπ and concave in (ν, µ, ζ1, . . . , ζm), since it includes the perspective transform of the functions gi [13]. The Lagrangian involves a cost vector, ν + ∑m i=1 ζi, linearly
interacting with dπ, and therefore we can use the same policy players as before to minimize this cost. For the cost player, it is possible to use OMD on Eq. (7) jointly for the variables ν, µ and ζ. It is more challenging to use best-response and FTL for the cost-player variables as the maximum value of the Lagrangian is unbounded for some values of dπ. Another option is to treat the problem as a three-player game. In this case the policy player controls dπ as before, one cost player chooses (ν, ζ1, . . . , ζm) and can use the algorithms we have previously discussed, and the other cost player chooses µ with some restrictions on their choice of algorithm. Analyzing the regret in that case is outside the scope of this paper.

heading: 6 Examples
text: In this section we explain how existing algorithms can be seen as instances of the meta-algorithm for various choices of the objective function f and the cost and policy player algorithms Algλ and Algπ . We summarized the relationships in Table 1.

heading: 6.1 Apprenticeship Learning
text: In apprenticeship learning (AL), we have an MDP without an explicit reward function. Instead, an expert provides demonstrations which are used to estimate the expert state occupancy measure dE . Abbeel and Ng [1] formalized the AL problem as finding a policy π whose state occupancy is close to that of the expert by minimizing the convex function f(dπ) = ||dπ − dE ||. The convex conjugate of f is given by f∗(y) = y · dE if ||y||∗ ≤ 1 and∞ otherwise, where || · ||∗ denotes the dual norm. Plugging f∗ into Eq. (4) results in the following game:
min dπ∈K ||dπ − dE || = min dπ∈K max ||λ||∗≤1 λ · dπ − λ · dE . (8)
Inspecting Eq. (8), we can see that the norm in the function f that is used to measure the distance from the expert induces a constraint set for the cost variable, which is a unit ball in the dual norm.
Algλ=OMD, Algπ=Best Response/RL. The Multiplicative Weights AL algorithms [65, MWAL] was proposed to solve the AL problem with f(dπ) = ||dπ − dE ||∞. It uses the best response as the policy player and multiplicative weights as the cost player (a special case of OMD). MWAL has also been used to solve AL in contextual MDPs [10] and to find feasible solutions to convex-constrained MDPs [44]. We note that in practice the best response can only be solved approximately, as we discussed in Section 4. Instead, in online AL [61] the authors proposed to use MDPO as the policy player, which guarantees a regret bound of R̄K ≤ c/ √ K. They showed that their algorithm is equivalent to Wasserstein GAIL [73, 78] and in practice tends to perform similarly to GAIL.
Algλ=FTL, Algπ=Best Response. When the policy player plays the best response and the cost player plays FTL, Algorithm 1 is equivalent to the Frank-Wolfe algorithm [22, 3] for minimizing f (Eq. (2)). Pseudo-code for this is included in the appendix (Algorithm 3). The algorithm finds a point dkπ ∈ K that has the largest inner-product (best response) with the negative gradient (i.e., FTL). Abbeel and Ng [1] proposed two algorithms for AL, the projection algorithm and the max margin algorithm. The projection algorithm is essentially a FW algorithm, as was suggested in the supplementary [1] and was later shown formally in [75]. Thus, it is a projection free algorithm in the sense that it avoids projecting dπ into K, despite the name. In their case the gradient is given by ∇f (dπ) = dπ − dE . Thus, finding the best response is equivalent to solving an MDP whose reward is dE − dπ . In a similar fashion, FW can be used to solve convex MDPs more generally [30]. Specifically, in [30], the authors considered the problem of pure exploration, which they defined as finding a policy that maximizes entropy.
Fully Corrective FW. The FW algorithm has many variants (see [33] for a survey) some of which enjoy faster rates of convergence in special cases. Concretely, when the constraint set is a polytope, which is the case for convex MDPs (Definition 1), some variants achieve a linear rate of convergence [34, 75]. One such variant is the Fully corrective FW, which replaces the learning rate update (see line 4 of Algorithm 3 in the supplementary), with a minimization problem over the convex hull of occupancy measures at the previous time-step. This is guaranteed to be at least as good as the learning rate update. Interestingly, the second algorithm of Abbeel and Ng [1], the max margin algorithm, is
exactly equivalent to this fully corrective FW variant. This implies that the max-margin algorithm enjoys a better theoretical convergence rate than the ‘projection’ variant, as was observed empirically in [1].

heading: 6.2 GAIL and DIAYN: Algλ=FTL, Algπ=RL
text: We now discuss the objectives of two popular algorithms, GAIL [31] and DIAYN [20], which perform AL and diverse skill discovery respectively. Our analysis suggests that GAIL and DIAYN share the same objective function. In GAIL, this objective function is minimized, which is a convex MDP, however, in DIAYN it is maximized, which is therefore not a convex MDP. We start the discussion with DIAYN and follow with a simple construction showing the equivalence to GAIL.
DIAYN. Discriminative approaches [26, 20] rely on the intuition that skills are diverse when they are entropic and easily discriminated by observing the states that they visit. Given a probability space (Ω,F ,P), state random variables S : Ω→ S and latent skills Z : Ω→ Z with prior p, the key term of interest being maximized in DIAYN [20] is the mutual information:
I(S;Z) = Ez∼p;s∼dzπ [log p(z|s)− log p(z)], (9)
where dzπ is the stationary distribution induced by the policy π(a | s, z). For each skill z, this corresponds to a standard RL problem with (conditional) policy π(a | s, z) and reward function r(s|z) = log p(z|s) − log p(z). The first term encourages the policy to visit states for which the underlying skill has high-probability under the posterior p(z | s), while the second term ensures a high entropy distribution over skills. In practice, the full DIAYN objective further regularizes the learnt policy by including entropy terms − log π(a | s, z). For large state spaces, p(z|s) is typically intractable and Eq. 9 is replaced with a variational lower-bound, where the true posterior is replaced with a learned discriminator qφ(z|s). Here, we focus on the simple setting where z is a categorical distribution over |Z| outcomes, yielding |Z| policies πz , and qφ is a classifier over these |Z| skills with parameters φ.
We now show that a similar intrinsic reward can be derived using the framework of convex MDPs. We start by writing the true posterior as a function of the per-skill state occupancy dzπ = p(s | z), and using Bayes rules, p(z|s) = d
z π(s)p(z)∑ k d k π(s)p(k) . Combing this with Eq. (9) yields:
Ez∼p(z),s∼dzπ [log p(z|s)− p(z)] = ∑ z p(z) ∑ s dzπ(s) [ log ( dzπ(s)p(z)∑ k d k π(s)p(k) ) − log p(z) ] = ∑ z p(z)KL(dzπ|| ∑ k p(k)dkπ) = EzKL(dzπ||Ekdkπ), (10)
where KL denotes the Kullback–Leibler divergence [39].
Intuitively, finding a set of policies π1, . . . , πz that minimize Eq. (10) will result in finding policies that visit similar states, measured using the KL distance between their respective state occupancies d1π, . . . , d z π. This is a convex MDP because the KL-divergence is jointly convex in both arguments [13, Example 3.19]. We will soon show that this is the objective of GAIL. On the other hand, a set of policies that maximize Eq. (10) is diverse, as the policies visit different states, measured using the KL distance between their respective state occupancies d1π, . . . , d z π .
We follow on with deriving the FTL player for the convex MDP in Eq. (10). We will then show that this FTL player is producing an intrinsic reward that is equivalent to the intrinsic reward used in GAIL and DIAYN (despite the fact that DIAYN is not a convex MDP). According to Eq. (6), the FTL cost player will produce a cost λk at iteration k given by
∇dzπKL(d z π|| ∑ k p(k)dkπ) = E z∼p(z) [ log dzπ∑ k d k πp(k) + 1− d z πp(z)∑ k d k πp(k) ] = E z∼p(z) [ log(p(z|s))− log(p(z))︸ ︷︷ ︸
Mutual Information
+1− p(z|s)︸ ︷︷ ︸ Gradient correction
] , (11)
where the equality follows from writing the posterior as a function of the per-skill state occupancy dzπ = p(s | z), and using Bayes rules, p(z|s) = dzπ(s)p(z)∑ k d k π(s)p(k) . Replacing the posterior p(z|s) with
a learnt discriminator qφ(z|s) recovers the mutual-information rewards of DIAYN, with additional terms 1− p(z | s) which we refer to as “gradient correction” terms. Inspecting the common scenario of a uniform prior over the latent variables, p(z) = 1/|Z|, we get that the expectation of the gradient correction term ∑ z p(z)(1− p(z|s)) = 1− 1/|Z| in each state. From the perspective of the policy player, adding a constant to the reward does not change the best response policy, nor the optimistic policy. Therefore, the gradient correction term does not have an effect on the optimization under a uniform prior, and we retrieved the reward of DIAYN. These algorithms differ however for more general priors p(z), which we explore empirically in Appendix F.
GAIL. We further show how Eq. (10) extends to GAIL [31] via a simple construction. Consider a binary latent space of size |Z| = 2, where z = 1 corresponds to the policy of the agent and z = 2 corresponds to the policy of the expert which is fixed. In addition, consider a uniform prior over the latent variables, i.e., p(z = 1) = 12 . By removing the constant terms in Eq. (11), one retrieves the GAIL [31] algorithm. The cost log(p(z|s)) is the probability of the discriminator to identify the agent, and the policy player is MDPO (which is similar to TRPO in GAIL).

heading: 7 Discussion
text: In this work we reformulated the convex MDP problem as a convex-concave game between the agent and another player that is producing costs (negative rewards) and proposed a meta-algorithm for solving it.
We observed that many algorithms in the literature can be interpreted as instances of the metaalgorithm by selecting different pairs of subroutines employed by the policy and cost players. The Frank-Wolfe algorithm, which combines best response with FTL, was originally proposed for AL [1, 75] but can be used for any convex MDP problem as was suggested in [30]. Zhang et al. [77], unified the problems of RL, AL, constrained MDPs with linear constraints and maximum entropy exploration under the framework of convex MDPs. We extended the framework to allow convex constraints (Section 5) and explained the objective of GAIL as a convex MDP (Section 6.2). We also discussed non convex objectives (Section 3) and analyzed unsupervised skill discovery via the maximization of mutual information (Section 6.2) as a special case. Finally, we would like to point out a recent work by Geist et al. [25], which was published concurrently to ours, and studies the convex MDP problem from the viewpoint of mean field games.
There are also algorithms for convex MDPs that cannot be explained as instances of Algorithm 1. In particular, Zhang et al. [77] proposed a policy gradient algorithm for convex MDPs in which each step of policy gradient involves solving a new saddle point problem (formulated using the Fenchel dual). This is different from our approach since we solve a single saddle point problem iteratively, and furthermore we have much more flexibility about which algorithms the policy player can use. Moreover, for the convergence guarantee [77, Theorem 4.5] to hold, the saddle point problem has to be solved exactly, while in practice it is only solved approximately [77, Algorithm 1], which hinders its sample efficiency. Fenchel duality has also been used in off policy evaluation (OPE) in [46, 74]. The difference between these works and ours is that we train a policy to minimize an objective, while in OPE a target policy is fixed and its value is estimated from data produced by a behaviour policy.
In order to solve a practical convex MDP problem in a given domain it would be prudent to use an RL algorithm that is known to be high performing for the vanilla RL problem as the policy player. From the theoretical point of view this could be MDPO or UCRL2, which we have shown come with strong guarantees. From the practical point of view using a high performing DRL algorithm, which may be specific to the domain, will usually yield the best results. For the cost player using FTL, i.e., using the gradient of the objective function, is typically the best choice.
Acknowledgments and Disclosure of Funding
We would like to thank Yasin Abbasi-Yadkorie, Vlad Mnih, Jacob Abernethy, Lior Shani and Doina Precup for their comments and discussion on this work. Work done at DeepMind, the authors received no specific funding for this work.

Review summaries:
This paper addresses an extended reinforcement learning setting with general utilities, where the agent's objective can be any convex function of the state-action distribution. Especially, it recasts the problem as a min-max game between a policy player and a cost player, and it describes a principled methodology to solve the game. Finally, it provides ways to instantiate the methodology with known algorithms, and a brief experimental evaluation in simple domains.

The authors propose to consider a convex optimization problem over the polytope formed by the state-visitation frequency of an MDP. They constructed algorithms based on existing tools of sub-gradient descents, such as Frank-Wolfe, FTL and OMD, in conjunction with an oracle that (approximate)-solves an MDP with scalar rewards. Altogether, they provide a meta-algorithm that translates a scalar reward MDP oracle to an algorithm for the convex MDP. The authors also highlight several applications, such as the Apprenticeship Learning problem and the constrained MDP problem, of their formulation. 

The paper studies the convex RL problem, which is defined as finding policy $\pi$ minimizing $f(d_\pi)$ where $f$ is a convex function and $d_\pi$ is the occupancy measure of policy $\pi$. Using Fenchel duality, this formulation can be written as a min-max and be solved using the well-known framework of repeated game playing first introduced by Freund and Schapire (1999). The authors show that their framework unifies some well-known approaches in RL by appropriately choosing the no-regret algorithm for the policy and cost player.

The paper discusses the problem of minimizing a convex function of the stationary distribution of a system with Markov dynamics, subject to the (convex) Bellman flow constraint and optionally subject to other convex constraints. The paper provides a family of algorithms designed to tackle this problem. It also proves results about sample complexity. The proposed framework is a significant innovation since it unifies several apprenticeship learning and imitation learning algorithms, which have been proved to be useful.