Reinforcement Learning via Fenchel-Rockafellar Duality

Introduction

Reinforcement learning (RL) aims to learn behavior policies to optimize a long-term decision-making process in an environment. RL is thus relevant to a variety of real-world applications, such as robotics, patient care, and recommendation systems. When tackling problems associated with the RL setting, two main difficulties arise: first, the decision-making problem is inherently sequential, with early decisions made by the policy affecting outcomes both near and far in the future; second, the learner’s knowledge of the environment is only through sampled experience, i.e., previously sampled trajectories of interactions, while the underlying mechanism governing the dynamics of these trajectories is unknown.

The environment and its underlying dynamics are typically abstracted as a Markov decision process (MDP). This abstraction gives rise to the Bellman recurrence, which characterizes the optimal value function and behavior policy through a dynamic programming (DP) view of the RL problem (Bellman, 1966). Most of the effective existing RL algorithms are rooted in this dynamic programming paradigm, attempting to find approximate fixed-point solutions to the Bellman recurrence, leading to the family of temporal-difference (TD) algorithms including SARSA (Sutton, 1996), $Q$-learning (Watkins, 1989), and their deep learning variants (Mnih et al., 2015; Hasselt et al., 2016; Wang et al., 2015). While the TD-based algorithms are powerful, their training can oscillate or even diverge in settings where function approximation is used or the ability to sample additional interactions with the environment is limited (Sutton and Barto, 1998).

An alternative paradigm for RL is based on linear programming (LP). A number of RL problems, such as policy optimization and policy evaluation, can be expressed as an LP – i.e., an optimization problem involving a linear objective and linear constraints. The LP may then be transformed to a form more amenable to stochastic and large-scale optimization via the tools of LP duality. Although the LP perspective has existed for decades (e.g., Manne (1960); Denardo (1970)), it has recently received renewed interest for its potential ability to circumvent the optimization challenges of DP-based approaches in exchange for more mature and well-studied techniques associated with convex optimization (De Farias and Van Roy, 2003; Wang et al., 2007; Chen and Wang, 2016; Wang, 2017; Bas-Serrano and Neu, 2019).

In this article, we generalize the LP approach and describe a number of convex problems relevant to RL – i.e., formulations of RL problems as a convex objective and linear constraints. With a convex objective, one must appeal to the more general Fenchel-Rockafellar duality. Perhaps the most useful property of this generalization is that, when the original primal problem involves a strictly convex objective (unlike the LP setting), application of Fenchel-Rockafellar duality leads to a dual problem which is unconstrained. We show that the Fenchel-Rockafellar duality and its variants are at the heart of a number of recent RL algorithms, although many of these were originally presented through less generalizable derivations (e.g., the DICE-family of offline RL algorithms: Nachum et al. (2019a, b); Kostrikov et al. (2019); Zhang et al. (2020)). By providing a unified perspective on these results and the tools and tricks which lead to them, we hope to enable future researchers to better use the techniques of convex duality to make further progress in RL.

Aiming to provide a useful reference for any interested researcher, we begin by reviewing basic knowledge of convex duality (Section 2) and RL (Section 3). We then focus on the discounted policy evaluation problem in RL (Section 4), originally expressed in an LP form known as the $Q$-LP. We show how the tools of LP and convex duality may be used to derive a variety of useful re-formulations of policy evaluation. We then continue to show how these same techniques can be applied to the policy optimization problem, starting from either the $Q$-LP (Section 5) or the potentially more streamlined $V$-LP (Section 6). We then generalize these algorithms to undiscounted settings in Section 7. We conclude in Section 8 with a brief summary and promising future directions.

Convex Duality

The concept of duality is a basic and powerful tool in optimization and machine learning, especially in the field of convex analysis, allowing a researcher to easily reformulate optimization problems in alternative ways that are potentially more tractable. In this section, we provide a brief overview of a few key convex duality results which will play an important role in the RL algorithms derived in later sections.

A full and detailed introduction to convex analysis is beyond the scope of this report, and so most of our statements will be presented informally (for example, we use $\min$ and $\max$ as opposed to $\inf$ and $\sup$ and we use $\sum_{x}f(x)$ and $\int f(x)dx$ without ambiguity), although we will strive to properly qualify theoretical claims when appropriate. The curious reader may refer to a number of resources for a more complete and mathematically precise treatment of this subject; e.g., Rockafellar (1970); Boyd and Vandenberghe (2004); Borwein and Lewis (2010); Bauschke and Lucet (2012).

where $\left\langle\cdot,\cdot\right\rangle$ denotes the inner product defined on $\Omega$. This function is also referred to as the convex conjugate or Legendre–Fenchel transformation of $f$.

We say a function $f$ is proper when $\left\{x\in\Omega:f(x)<\infty\right\}$ is non-empty and $f(x)>-\infty$ for all $x\in\Omega$.

For a proper, convex, lower semi-continuous $f$, its conjugate function $f_{*}$ is also proper, convex, and lower semi-continuous. Moreover, one has the duality $f_{**}=f$; i.e.,

where $\Omega^{*}$ denotes the domain of $f_{*}$. From now on we will assume any use of $f$ is for a convex function. Furthermore, we will assume any declared convex function is also proper and lower semi-continuous. Table 1 provides a few common functions and their corresponding Fenchel conjugates.

One especially useful function is the indicator function $\delta_{C}(x)$, which is defined as,

If $C$ is a closed, convex set, it is easy to check that $\delta_{C}$ is proper, convex, and lower semi-continuous. The indicator function can be used as a way of expressing constraints. For example, the optimization problem $\min_{Ax=0}f(x)$ may be alternatively expressed as $\min_{x}f(x)+\delta_{\{0\}}(Ax)$. It may be readily shown that the conjugate of $\delta_{\{a\}}(x)$ is the linear function $\langle a,y\rangle$ and vice-versa.

1.2 f𝑓f-Divergences

The family of $f$-divergences, also known as Csiszár-Morimoto or Ali-Silvey divergences (Ali and Silvey, 1966), has been widely applied in many machine learning applications, including variational inference (Wainwright and Jordan, 2003), generative model estimation (Nowozin et al., 2016; Dai et al., 2019), imitation learning (Ke et al., 2019; Ghasemipour et al., 2019; Kostrikov et al., 2019), and reinforcement learning (Nachum et al., 2019a; Zhang et al., 2020; Nachum et al., 2019b).

For a convex function $f$ and a distribution $p$ over some domain $\mathcal{Z}$, the $f$-divergence is defined as,

On the other hand, if one considers the domain of $D_{f}$ to be $\Delta(\mathcal{Z})$, then one must solve a constrained version of (5), which can be difficult depending on the form of $f$.

2 Fenchel-Rockafellar Duality

Fenchel conjugates are indispensable when tackling a variety of optimization problems. In this section, we present one of the most general and useful tools associated with Fenchel conjugates, known as the Fenchel-Rockafellar duality (Rockafellar, 1970; Borwein and Lewis, 2010).

where we use $A_{*}$ to denote the adjoint linear operator of $A$; i.e., $A_{*}$ is the linear operator for which $\langle y,Ax\rangle=\langle A_{*}y,x\rangle$, for all $x,y$. In the common case of $A$ simply being a real-valued matrix, $A_{*}$ is the transpose of $A$.

Of course, in the presence of Fenchel-Rockafellar duality, the label of primal and dual is arbitrary. One can consider (11) the primal problem and (10) its dual, and in our derivations we will use these labels interchangeably.

The Fenchel-Rockafellar duality is general enough that it can be used to derive the Lagrangian duality. Consider the constrained optimization problem

By considering $f_{*}$ in terms of its Fenchel conjugate (equation (1)), we may write the problem as

Using the fact that $\langle y,Ax\rangle=\langle x,A_{*}y\rangle$ for any $A$ we may express this as

The expression $L(x,y)$ is known as the Lagrangian of the original problem in (14). One may further derive the well-known Lagrange duality:See Veinott (2005) (https://web.stanford.edu/class/msande361/handouts/nlpdual.pdf) for a brief derivation of this fact and Ekeland and Temam (1999)[Proposition 2.1] for more general cases.

Moreover, the optimal value of the Lagrangian is the optimal value of the original problem (14), and the optimal solutions (equilibrium points) $x^{*},y^{*}$ are the solutions to the original primal (14) and its dual (15).

2.2 LP Duality

respectively. By making the switch $y\to-y$, the dual (20) may be equivalently expressed in the more familiar form,

Fenchel-Rockafellar duality thus provides us with the strong LP duality theorem. Namely, if the primal problem (19) is feasible, then its result is the same as that of the dual (21).

Reinforcement Learning

In this work, we will show how the Fenchel-Rockafellar duality (and the LP and Lagrangian dualities) can be applied to solve a number of reinforcement learning (RL) problems. Before we present these algorithms, we use this section as a brief introduction to RL.

In RL, one wishes to learn a behavior policy $\pi$ to interact with an environment in an optimal way, where the typical meaning of ‘optimal’ is with respect to future discounted rewards (feedback) provided by the environment. The RL environment is commonly abstracted as a Markov decision process (MDP) (Puterman, 1994; Sutton and Barto, 1998), which is specificied by a tuple $\mathcal{M}=\langle S,A,R,T,\mu_{0},\gamma\rangle$, consisting of a state space, an action space, a reward function, a transition probability function, an initial state distribution, and a discount factor $\gamma\in(0,1]$, respectively. The policy $\pi$ is a function $S\to\Delta(A)$. The policy interacts with the environment iteratively, starting with an initial state $s_{0}\sim\mu_{0}$. At step $t=0,1,\dots$, the policy produces a distribution $\pi(\cdot|s_{t})$ over the actions $A$, from which an action $a_{t}$ is sampled and applied to the environment. The environment produces a scalar reward $r_{t}=R(s_{t},a_{t})$For simplicity we consider a deterministic reward function. Stochastic rewards are more typical, although the same derivations are usually applicable in either case. and subsequently transitions to a new state $s_{t+1}\sim T(s_{t},a_{t})$.

In summary, the RL setting is concerned with a policy which sequentially makes decisions, and the effects of those decisions are observed through a per-step reward feedback and a stochastic, Markovian state transition process. For simplicity, we will consider infinite-horizon (non-terminating) environments, which may be extended to finite-horizon environments by considering an extra terminal state which continually loops onto itself with zero reward.

2 Policy Evaluation and Optimization

The first question one may ask when presented with an MDP $\mathcal{M}$ and a policy $\pi$ is, what is the long-term value of $\pi$ when interacting with $\mathcal{M}$? The next question might be, what is the optimal policy $\pi^{*}$ maximizing this long-term value? These two questions constitute the policy evaluation and optimization problems, respectively.

To formalize these questions, we consider a discount factor $\gamma\in(0,1)$.See Section 7 for consideration of $\gamma=1$. The value of $\pi$ is defined as the expected per-step reward obtained by following the policy, averaging over time via $\gamma$-discounting:

3 Online vs. Offline RL

One of the main limitations when approaching either the policy evaluation or policy optimization problems is that one does not have explicit knowledge of the environment; i.e., one does not have explicit knowledge of the functions $R,T,\mu_{0}$. Rather, access to the environment is given in the form of experience $s_{t},a_{t},r_{t},s_{t+1},a_{t+1},r_{t+1},\dots$ gathered via interactions with the environment. The specific nature of this experience depends on the context of one’s considered problem. The most common forms of experience may be generally categorized into online and offline.

In the online setting, experience from the environment may be collected at any point via Monte-Carlo rollouts. With this type of access to the environment, the policy evaluation and optimization problems may be easily solved. For example, the value of the policy may be estimated by simply averaging the discounted reward of a large number of Monte-Carlo rollouts. For this reason, online RL research typically strives to find sample-efficient algorithms, which find approximate solutions to policy evaluation or optimization with as few interactions with the environment as possible.

In practice (e.g., consumer web recommendation systems or health applications), interaction with the environment during training is not available at all. More commonly, access to the environment is offline. That is, interactions with the environment are limited to a static dataset of (logged) experience $\mathcal{D}=\{(s^{(i)},a^{(i)},r^{(i)},s^{(i)\prime})\}_{i=1}^{N}$, where $(s^{(i)},a^{(i)})\sim d^{\mathcal{D}}$ for some unknown distribution $d^{\mathcal{D}}$, $r^{(i)}=R(s^{(i)},a^{(i)})$, and $s^{(i)\prime}\sim T(s^{(i)},a^{(i)})$. One also typically assumes access to samples $\mathcal{U}=\{s_{0}^{(i)}\}_{i=1}^{M}$ from $\mu_{0}$. In this report we will mostly focus on the offline setting, although we will relax it to assume that our offline experience is effectively unlimited ($N,M\to\infty$), and so will write our expectations in terms of $d^{\mathcal{D}}$, $T$, and $\mu_{0}$. Performing the appropriate finite-sample analysis for finite $N,M$ is outside the scope of this report. Even with effectively unlimited experience, the offline setting presents a challenge to RL algorithms, due to the mismatch between the experience distribution given by the offline dataset and the online distribution typically needed for policy evaluation or optimization.

We emphasize a subtle difference between the offline setting and what is commonly referred to in the literature as off-policy learning. Off-policy algorithms are designed to enable an RL agent to learn from historical samples collected by other policies. However, these algorithms are typically allowed to interact with the environment during training to collect new samples. On the other hand, in the offline setting, one’s access to the environment is exclusively via a fixed dataset of experience. In other words, while an offline RL algorithm is necessarily off-policy, an off-policy algorithm is not necessarily offline.

4 Q𝑄Q-values and State-Action Visitations

When evaluating or optimizing policies, both online and offline, the notions of $Q$-values and state-action visitations are useful. For a policy $\pi$, the Q-values $Q^{\pi}(s,a)$ denote the future discounted sum of rewards of following $\pi$ starting at $s,a$:

The $Q$-values satisfy the single-step Bellman recurrence

where $\mathcal{P}^{\pi}$ is the policy transition operator,

The state-action visitations $d^{\pi}$ of $\pi$ (also known as occupancies or density) may be defined similarly as,

That is, the visitation $d^{\pi}(s,a)$ measures how likely $\pi$ is to encounter $s,a$ when interacting with $\mathcal{M}$, averaging these encounters over time via $\gamma$-discounting. The visitations $d^{\pi}$ constitute a normalized distribution, and this distribution is referred to as the on-policy distribution.

Like the $Q$-values, the visitations satisfy the single-step transpose Bellman recurrence:

where $\mathcal{P}^{\pi}_{*}$ is the transpose (or adjoint) policy transition operator,

These recursions simply reflect the conservation of flow (probability mass) of a stationary distribution on a Markov process. Note that both $\mathcal{P}^{\pi}$ and $\mathcal{P}^{\pi}_{*}$ are linear operators and that the transpose policy transition operator $\mathcal{P}^{\pi}_{*}$ is indeed the mathematical transpose (or adjoint) of $\mathcal{P}^{\pi}$ in the sense that $\langle y,\mathcal{P}^{\pi}x\rangle=\langle\mathcal{P}^{\pi}_{*}y,x\rangle$ for any $x,y$.

Both the $Q$-values and the visitations are useful in RL. For example, the value of a policy may be expressed in two ways:

Also, when performing policy optimization, the policy gradient theorem (Sutton et al., 2000) utilizes the $Q$-values and visitations to express the gradient of $\operatorname{\rho}(\pi)$ as

It is thus standard in most RL algorithms to either have access to $Q^{\pi}$ and $d^{\pi}$ or have some mechanism to estimate these quantities. Typically, the $Q$-values are estimated by finding the fixed point of the Bellman recurrence (Sutton et al., 2008, 2009; Scherrer, 2010; Liu et al., 2015; Dai et al., 2016; Du et al., 2017); i.e., minimizing the (surrogate) squared difference between the LHS and RHS of (23), potentially with target networks. For the visitations, it is more typical to assume access to the distribution $d^{\pi}$ (for example, by simply performing Monte-Carlo rollouts enabled by online access), although instances exist in which $d^{\pi}$ is approximated by importance-weighting a different (i.e., offline) distribution (Precup et al., 2001; Sutton et al., 2014).

Policy Evaluation

We now move on to demonstrating applications of Fenchel-Rockafellar duality to RL. We begin by approaching the policy evaluation problem. Although the policy evaluation problem may appear to be simpler or less interesting (it is not!) than the policy optimization problem, in our case the same techniques will be used in either setting. Thus, we will use this section to provide more detailed derivations of a variety of techniques which will be referenced repeatedly in the following sections.

The equivalent formulations of $\operatorname{\rho}(\pi)$ in (27) in terms of either $Q^{\pi}$ or $d^{\pi}$ hint at a duality which is formally given by the following LP characterization of $\operatorname{\rho}(\pi)$, known as the $Q$-LP:

The optimal $Q^{*}$ of this LP satisfies $Q^{*}(s,a)=Q^{\pi}(s,a)$ for all $s,a$ reachable by $\pi$.

The dual of this LP provides us with the visitation perspective on policy evaluation:

The optimal $d^{*}$ of this LP is the state-action visitation $d^{\pi}$ of $\pi$. It is important to note that this dual LP is over-constrained. The $|S|\times|A|$ equality constraints in (33) uniquely determine $d$ regardless of the objective in (32). This fact will prove useful in a number of later derivations.

For detailed and complete derivations of these LP representations of $Q^{\pi}$ and $d^{\pi}$, please refer to Nachum et al. (2019b).

2 Policy Evaluation via the Lagrangian

The potentially large number of constraints in either the primal or dual forms of the $Q$-LP introduce a challenge to estimating $\operatorname{\rho}(\pi)$. We may instead derive a more tractable unconstrained optimization perspective on the policy evaluation problem using the Lagrangian of the $Q$-LP:

In practical settings where $\left|S\right|\times\left|A\right|$ is possibly infinite, it is not feasible to optimize the sum in (35) over $S\times A$. In an offline setting, where we only have access to a distribution $d^{\mathcal{D}}$, we may make a change-of-variables via importance sampling, i.e., $\zeta(s,a)=\frac{d(s,a)}{d^{\mathcal{D}}(s,a)}$. If $d^{\mathcal{D}}$ has sufficient support or coverage (Sutton et al., 2016), we may re-write (35) as

The optimal $\zeta^{*}$ of this problem satisfies $\zeta^{*}(s,a)=\frac{d^{\pi}(s,a)}{d^{\mathcal{D}}(s,a)}$. Thus, to estimate $\operatorname{\rho}(\pi)$, one may optimize this objective with respect to $Q,\zeta$ (requiring only access to samples from $\mu_{0}$, $d^{\mathcal{D}}$, and $\pi$) and return $L(\hat{Q}^{*},\hat{\zeta}^{*})$ as the final estimate.

This more practical, offline estimator also has a desirably property, known as the doubly robust property (Funk et al., 2011; Jiang and Li, 2015; Kallus and Uehara, 2019a). Specifically,

Thus, this estimator is robust to errors in at most one of $Q$ and $\zeta$.

Despite the desirable properties of this estimator, the optimization problem associated with it involves rewards $R(s,a)$ and learning $Q^{\pi}$-values with respect to these rewards. Learning $Q^{\pi}$-values using single-step transitions turns out to be difficult in practice without the use of a number of tricks developed over the years (e.g., target networks, ensembling). Moreover, the bilinear nature of the Lagrangian can lead to instability or poor convergence in optimization (Dai et al., 2017; Bas-Serrano and Neu, 2019). Rather than tackling these various issues head-on, a number of recent works propose an alternative approach, which we describe in the following subsection.

3 Changing the Problem Before Applying Duality

As mentioned in Section 4.1, the dual of the $Q$-LP is over-constrained, in the sense that the $|S|\times|A|$ constraints in (33) uniquely determine the state-action visitation $d^{\pi}$. Thus, one may replace the objective in (32) with $\max_{d}\,\,-h(d)$ for some $h$ without affecting the optimal solution $d^{*}=d^{\pi}$. Therefore, the main idea of a number of recent works is to choose an appropriate $h$ so that either the Lagrangian or the Fenchel-Rockafellar dual of this problem is more approachable and potentially avoids the instabilities associated with the original LP.

Although the problem is changed, the solution is unaffected, and once a solution is found it may be used to provide an estimate of $\operatorname{\rho}(\pi)$. Specifically, if the problem is re-written in terms of $\zeta(s,a)=\frac{d(s,a)}{d^{\mathcal{D}}(s,a)}$, then once the problem is optimized, we can derive an estimate for the value of $\pi$ via the approximate solution $\hat{\zeta}^{*}$:

If $h$ is taken to be the trivial function $h(d):=0$, the offline form of the Lagrangian optimization becomes,

The optimal solution $\zeta^{*}$ of this problem is $d^{\pi}/d^{\mathcal{D}}$, and once an approximate solution is found, it may be used to estimate $\operatorname{\rho}(\pi)$ according to (38). Unlike the previous form of the Lagrangian in (36), this optimization does not involve learning $Q$-values with respect to environment rewards, and in practice this distinction leads to much better optimization behavior (Uehara and Jiang, 2019). Still, the Lagrangian is linear in both $Q$ and $\zeta$. This can be remedied by choosing a strictly convex form of $h$, for example, by using an $f$-divergence.

The use of an $f$-divergence objective leads to the set of general off-policy estimation techniques outlined in the recent DualDICE paper (Nachum et al., 2019a). Specifically, the various estimators proposed by DualDICE correspond to applying either the Lagrange or Fenchel-Rockafellar dualities to the optimization problem,

Application of Lagrange duality to the above problem yields

We transform the transpose policy transition operator $\mathcal{P}^{\pi}_{*}$ to its transpose $\mathcal{P}^{\pi}$ by using the fact $\langle y,Ax\rangle=\langle x,A_{*}y\rangle$:

Now we make the change-of-variables $\zeta(s,a)=\frac{d(s,a)}{d^{\mathcal{D}}(s,a)}$ to yield,

Thus we have recovered the general saddle-point form of DualDICE, which proposes to optimize (45) and then use an approximate solution $\hat{\zeta}^{*}$ to estimate $\operatorname{\rho}(\pi)$ via (38).

Rather than applying Lagrange duality, application of Fenchel-Rockafellar duality to (40) more clearly reveals the wisdom of choosing $h(d)=D_{f}(d\|d^{\mathcal{D}})$. We write the problem in (40) as

where $g(-Ad)$ corresponds to the linear constraints (41) with respect to the adjoint Bellman operator; i.e.,

We can now see that the use of an $f$-divergence with respect to $d^{\mathcal{D}}$ naturally leads to an offline problem with expectations over $d^{\mathcal{D}}$, without an explicit change-of-variables. Furthermore, unlike previous dual problems, there are no constraints in this optimization, and so standard gradient-based techniques may be applied to find a solution $Q^{*}$ without appealing to the Lagrange duality, which would necessarily involve nested $\max$-$\min$ optimizations. The Fenchel-Rockafellar duality also provides us with a way to recover $d^{*}=d^{\pi}$ from a solution $Q^{*}$:

If we set $f(x)=\frac{1}{2}x^{2}$, we may recover what is perhaps the most intriguing result in Nachum et al. (2019a):

That is, if one optimizes $Q$-value functions to minimize squared Bellman residuals (with respect to zero reward) while minimizing initial $Q$-values, then the optimal Bellman residuals are exactly the density ratios between the on-policy and offline state-action distributions.

Interestingly, the derivations in Nachum et al. (2019a) do not explicitly use Lagrangian or Fenchel-Rockafellar duality, but rather focus on a cleverly chosen change-of-variables (the so-called DualDICE trick). It is clear from our own derivations that this trick essentially comes from the relationship between $\mathcal{P}^{\pi}$ and $\mathcal{P}^{\pi}_{*}$ and is simply another way of applying Fenchel-Rockafellar duality.

It is important to note that there is a trade-off introduced by the use of the Fenchel-Rockafellar duality as opposed to the Lagrangian. Namely, the objective (47) involves optimizing a convex function of an expectation under $\mathcal{P}^{\pi}$, and thus also the environment transition function $T$, whereas in practice one typically only has access to a single empirical sample $s^{\prime}\sim T(s,a)$. Many works ignore this problem, and simply consider the single empirical sample $s^{\prime}$ as the full distribution $T(s,a)$, and this introduces a bias into the optimization. See Antos et al. (2008); Dai et al. (2016) for potential remedies to this issue.

4 Summary

We briefly summarize the main takeaways from this section.

The policy evaluation problem may be expressed as an LP, known as the $Q$-LP, whose solution is $Q^{\pi}$.

The dual of this LP has solution $d^{\pi}$.

Taking the Lagrangian of the $Q$-LP can lead to a doubly robust estimator for $\operatorname{\rho}(\pi)$.

Changing the objective in the dual of the $Q$-LP does not change its solution $d^{\pi}$.

Changing the objective to an appropriate alternative is a powerful tool. This can lead to a (regularized) Fenchel-Rockafellar dual that is unconstrained, and thus more amenable to stochastic and offline settings.

These techniques may be used to derive the results provided in a number of recent works, such as Nachum et al. (2019a) and Uehara and Jiang (2019).

Policy Optimization

The previous section detailed a number of ways to estimate $\operatorname{\rho}(\pi)$. In this section, we show how similar techniques may be applied for the policy optimization problem, which is concerned with finding the optimal solution $\arg\max_{\pi}\operatorname{\rho}(\pi)$.

Considering the Lagrangian formulation of $\operatorname{\rho}(\pi)$ given in (35) can provide a simple derivation of the policy gradient theorem in (28). Let $L(Q,d;\pi)$ be the inner expression in (35). Danskin’s theorem (Bertsekas, 1999) tells us that

where $Q^{*},d^{*}$ are the solutions to $\min_{Q}\max_{d\geq 0}L(Q,d;\pi)=\max_{d\geq 0}\min_{Q}L(Q,d;\pi)$. We may compute the gradient of $L(Q^{*},d^{*};\pi)$ w.r.t. $\pi$ term-by-term. For the first term in (35) we have

Recall that $Q^{*}(s,a)=Q^{\pi}(s,a)$ for all $s,a$ for which $d^{\pi}(s,a)>0$ and that $d^{*}=d^{\pi}$. Furthermore,

Using this information, we may combine (52) and (53) to yield

which matches the policy gradient theorem as presented in (28).

2 Offline Policy Gradient via the Lagrangian

The original policy gradient theorem in Sutton et al. (2000) relies on on-policy settings, while in practice we would like to optimize the policy with only offline samples. As in Section 4.2, one may write $\operatorname{\rho}(\pi)$ in an offline manner via the variables $\zeta(s,a)=\frac{d^{\pi}(s,a)}{d^{\mathcal{D}}(s,a)}$. Using (36), we may write the policy optimization problem as

3 Fenchel-Rockafellar Duality for Regularized Optimization

As noted in previous sections, both the linear nature and $\min$-$\max$ form of the Lagrangian may lead to numerical instability in practice. In Section 4.3.2, we showed how regularizing the objective of the LP for $d^{\pi}$ in an appropriate manner can lead to a Fenchel-Rockafellar dual that is unconstrained (and thus avoids the need for the Lagrangian). In this previous section, we only cared about estimating $d^{*}=d^{\pi}$, and conveniently, the regularization did not affect the optimal solution. However, in our current setting of policy optimization, changing the objective will change the optimal policy. Still, the modified objective may be motivated as a regularization of the max-reward policy objective, and finding the optimal regularized policy is still desirable in many applications.

Thus, we consider regularizing the max-reward policy objective with the $f$-divergence $D_{f}(d\|d^{\mathcal{D}})$. Our modified problem becomes

Fenchel-Rockafellar duality yields the following dual formulation:

Thus, one may optimize $\pi$ to maximize the regularized objective by solving the $\max$-$\min$ optimization

and the optimal solution $\pi^{*}$ of this objective is the optimal solution to the regularized problem $\max_{\pi}\operatorname{\rho}(\pi)-D_{f}(d^{\pi}\|d^{\mathcal{D}})$. This recovers the offline formulation of regularized policy gradient known as Primal AlgaeDICE derived in Nachum et al. (2019b). As noted in Nachum et al. (2019b), when using $f(x)=\frac{1}{2}x^{2}$, the objective bears some resemblance to actor-critic learning but with offline samples, in which a $Q$-value is learned to minimize squared Bellman errors and a policy $\pi$ is learned to maximize $Q$-values.

Of course, one may also write the regularized problem (57) in its Lagrangian form, and this can recover Fenchel AlgaeDICE as derived in Nachum et al. (2019b).

Subsequent application of Fenchel-Rockafellar duality leads to the following offline policy optimization objective:

The beauty of this objective is revealed when one considers optimizing $\pi$ with respect to a gradient-based method. For a specific $Q$, the gradient of this objective with respect to $\pi$ is

Thus, we see that the use of KL-divergence leads to a dual formulation which bears similarities to max-likelihood policy learning, a common goal of a number of recent works (Abdolmaleki et al., 2018; Song et al., 2019; Peng et al., 2019). The use of max-likelihood policy learning in conjunction with a log-average-exp objective for the $Q$-value function also bears strong resemblance to the REPS algorithm (Peters et al., 2010). Still, the policy objective here only resembles max-likelihood learning and is not exactly equivalent. The connections hinted at here will be made more explicit in Section 6.

4 Imitation Learning

5 Summary

We briefly summarize the main takeaways from this section.

One can apply many of the same techniques used for policy evaluation to the policy optimization problem by simply putting a $\max_{\pi}$ around a derived estimator for $\operatorname{\rho}(\pi)$.

Since the optimal solution of the inner optimization will typically be either $d^{\pi}$ or $d^{\pi}/d^{\mathcal{D}}$, one can appeal to Danskin’s theorem to argue that the gradients of $\pi$ will be on-policy policy gradients, even if one only has access to offline data.

We again see the power of modifying an objective before appealing to duality. Appropriate regularization of the max-reward objective can lead to an unconstrained Fenchel-Rockafellar dual problem.

The same technique can also be exploited for offline imitation learning, which can be simply derived from the policy optimization objectives by ignoring the reward.

Depending on the exact form of regularization, application of Fenchel-Rockafellar duality leads to objectives which hint at connections to actor-critic via squared Bellman error minimization as well as max-likelihood policy learning. More interesting or useful objectives may be possible for other (yet undiscovered?) specially chosen regularizers.

RL with the Linear Programming Form of V𝑉V

The policy optimization approaches of the previous section may seem superficial. We essentially took formulations of $\operatorname{\rho}(\pi)$ from Section 4 and put a $\max_{\pi}$ around them. While this is valid, it leads to problems-within-problems, i.e. $\max$-$\min$ problems, for which stochastic optimizations can be difficult to theoretically motivate. Is there a better way?

There is a more direct way to frame the policy optimization problem as a convex problem, and it is given by the $V$-LP (Puterman, 1994; Bertsekas et al., 1995; Bertsekas and Tsitsiklis, 1996; Wang et al., 2007). We begin by introducing what is usually referred to as the dual of the $V$-LP, expressed in terms of state-action visitations:

where we use $\mathcal{T}_{*}$ to denote the transpose (or adjoint) transition operator,

As in (33), the constraints in (68) describe the conservation of flow (probability mass) of a stationary distribution on a Markovian process, although now the conservation is measured with respect to states as opposed to state-action pairs. Crucially, unlike the problem in (32), this problem is not over-constrained. The result of this problem is $\operatorname{\rho}(\pi^{*})$ for an optimal max-reward policy $\pi^{*}$ and the solution is $d^{*}=d^{\pi^{*}}$.

The dual of problem (67) is the more commonly seen $V$-LP (which is typically referred to as the primal):

where now $\mathcal{T}$ is the transition operator,

The optimal $V^{*}$ of this problem is the value function of an optimal policy $V^{\pi^{*}}$, where

This objective can be seen as an improvement over our previous derivation using the $Q$-LP, since, unlike before, this objective only involves a single optimization over $V,K$ as opposed to a $\max$-$\min$ optimization over $\pi$ and $Q$. However, the solution to this problem will give us $V^{*}$, the (regularized) value function of the optimal (regularized) policy, while what we really want is the policy itself!

To derive the optimal policy, we first note that Fenchel-Rockafellar duality provides us the optimal $d^{*}$ from $V^{*},K^{*}$ via

The solution $d^{*}$ is the visitation of an optimal regularized policy $d^{\pi^{*}}$. Using Bayes’ rule we may find the optimal policy to be

This way, one may recover the optimal policy $\pi^{*}$ from the solutions $V^{*},K^{*}$ of (76). This form of value learning (and $K$ learning) before recovering a policy from the optimal values is the same as derived by Belousov and Peters (2017).

We can see that the use of the Fenchel-Rockafellar dual of $V$-LP introduces a trade-off compared to starting from the $Q$-LP. We have avoided nested $\max$-$\min$ optimizations that arise from the $Q$-LP, but now our problems do not directly provide us with a policy. One must perform extra derivations to derive the policy. Depending on the specific regularization employed, deriving the optimal policy from $V^{*},K^{*}$ may be difficult in practical (stochastic) settings.

and this recovers the REPS objective (Peters et al., 2010).See also Neu et al. (2017).

The visitations of the optimal policy are now given by the $\operatorname{softmax}$ function:

The optimal policy thus has a similar form:

We may now see that recovering $\pi^{*}$ from $V^{*}$ may be done via max-likelihood learning:

where $Z(s)$ is an arbitrary normalization. This either recovers or can be used to motivate many past and recent works advocating for max-likelihood policy learning (Peters et al., 2010; Abdolmaleki et al., 2018; Song et al., 2019; Peng et al., 2019)

2 Policy Evaluation with the V𝑉V-LP

Although we originally introduced the $V$-LP as a way to perform more streamlined policy optimization, one can also use it for policy evaluation. To do so, we decompose $d(s,a)=\mu(s)\pi(a|s)$ for a fixed policy $\pi(a|s)$. Plugging this form of $d(s,a)$ into (67) leads to

When one fixes the policy $\pi(a|s)$ in this manner, the resulting LP is over-constrained, in the sense that the $\left|S\right|$ constraints in (83) uniquely determine the on-policy state visitation $\mu^{\pi}(s)$. Applying the Lagrangian to (82) will result in objectives similar to those that have appeared in Kallus and Uehara (2019b); Tang et al. (2019). Moreover, as we have done in Section 4.3, one can also replace the objective (82) before applying either Lagrangian or Fenchel-Rockafellar duality. One should note that the use of a $V$-LP for offline policy evaluation results in objectives that require knowledge of the data distribution policy $d^{\mathcal{D}}(a|s)$ in order to correct the offline action samples to on-policy samples, as needed by the use of $\mu\times\pi$ in (83). This is in contrast to the behavior-agnostic objectives yielded by the $Q$-LP.

Undiscounted Settings

So far, our treatment has only considered discounted settings, i.e., $\gamma\in(0,1)$. The undiscounted setting $\gamma=1$ presents an interesting challenge for many RL algorithms, as notions of $Q$-values and convergence of Bellman backup operators are harder to get a handle on. On the other hand, the techniques of Fenchel-Rockafeller duality and its variants may be readily applied to the undiscounted setting, with only minor modifications. In this section, we provide an extension of our previous derivations to the $\gamma=1$ case, leading to several practical algorithms.

When $\gamma=1$, the policy evaluation problem concerns estimating the average per-step reward of the policy:

Under certain conditions,In finite state and action spaces, the MDP must be ergodic. For other (continuous) spaces, the conditions are more involved. Please refer to Zhang et al. (2020) for the details. this quantity may be alternatively written as an expectation over the stationary distribution of $\pi$ in $\mathcal{M}$:

where the undiscounted on-policy distribution $d^{\pi}$ is defined as the normalized distribution satisfying

Thus, we can formulate $\rho(\pi)$ analogous to (32) as,

Note that the only difference to (32) is that this LP requires the constraint (89), which ensures that $d$ constitutes a normalized distribution. Still, as in (32), this problem is over-constrained. The objective in (88) may be modified without changing the solution $d^{*}=d^{\pi}$, and so the same techniques used in Section 4 may be applied.

For example, if we modify the objective to $\max_{d}-D_{f}(d\|d^{\mathcal{D}})$ and then take the Fenchel-Rockafellar dual we arrive at the problem

Given a solution $Q^{*},\lambda^{*}$ to this problem, the optimal $d^{*}=d^{\pi}$ may be recovered as

If one were to instead use the Lagrange duality with the objective $\max_{d}-D_{f}(d\|d^{\mathcal{D}})$, after making the change-of-variables $\zeta(s,a)=\frac{d(s,a)}{d^{\mathcal{D}}(s,a)}$ one would arrive at the nested optimization

The optimal solution of this optimization is

We refer the reader to Zhang et al. (2020) for an alternative approach to undiscounted policy evaluation. This approach essentially considers regularizing $Q$ and $\lambda$ in the Lagrangian of (88) with a square regularization, eventually yielding the objective,

2 Policy Optimization

As in Section 5, we may approach the policy optimization problem by simply putting a $\max_{\pi}$ around a reward-aware form of (91) or (94); e.g.,

Similar to Section 5, one may use Danskin’s theorem to argue that the gradient with respect to $\pi$ of this objective is the on-policy policy gradient (although with respect to regularized $Q$-values).

Alternatively, one may use the techniques in Section 6, writing the policy optimization problem analogous to (67) as

The optimal policy $\pi^{*}$ to the original problem may then be recovered from $V^{*}$ via the max-likelihood optimization

Conclusion

We have presented a variety of ways to apply the Fenchel-Rockafellar duality to problems appearing in RL. Although our settings and corresponding results are numerous, the techniques we used can be summarized succinctly as,

when presented with a problem that appears difficult to solve, consider writing the problem as a constrained convex optimization and solving its Fenchel-Rockafellar dual, or its Lagrangian form;

if the dual is still difficult to solve (e.g., when the primal objective is linear, yielding a dual with constraints), consider modifying the original objective, e.g., by applying an appropriate convex regularizer.

This simple protocol has been a recurring theme in our derivations, leading to several algorithms to tackle the policy evaluation, policy optimization, and imitation learning problems regardless of online or offline access to the environment and discounted or undiscounted rewards.

We hope the connections exposed here between RL and optimization can ignite progress and collaborations of researchers from both communities. For example, applying the same protocol outlined here to problems in other RL settings (e.g., multi-agent RL, safe RL, exploration for RL, etc.) and other ways to appropriately regularize these problems are promising RL research directions. At the same time, questions of how well these new formulations interact with algorithms for convex optimization (especially in stochastic and function approximation settings), and whether these duality-based formulations are more efficient than DP-based approaches, bring new challenges and problems for optimization research.

We thank Lihong Li, Dale Schuurmans, Ilya Kostrikov, Yinlam Chow, Sherry Yang, George Tucker, Matthieu Geist, and other members of the Google Brain team for insightful thoughts and discussions.

References