Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems — part 8/12

Cleaned-text code points [102985, 117996). Bibliography omitted here; source.md includes it.

The basic idea behind policy constraint methods for offline reinforcement learning via dynamic programming is to ensure, explicitly or implicitly, that regardless of the target values in Equation ( 3 ), the distribution over actions under which we compute the target value, π ⁡ ( 𝐚 ′ | 𝐬 ′ ) \pi(\mathbf{a}^{\prime}|\mathbf{s}^{\prime}) , is “close” to the behavior distribution π β ​ ( 𝐚 ′ | 𝐬 ′ ) {\pi_{\beta}}(\mathbf{a}^{\prime}|\mathbf{s}^{\prime}) . This ensures that the Q-function is never queried on out-of-distribution actions, which may introduce errors into the computation of the Bellman operator. If all states and actions fed into the Q-function for target value calculations are always in-distribution with respect to the Q-function training set, errors in the Q-function should not accumulate, and standard generalization results from supervised learning should apply. Since the Q-function is evaluated on the same states as the ones on which it is trained, only the action inputs can be out of distribution, when we compute 𝔼 𝐚 ′ ∼ π ⁡ ( 𝐚 ′ | 𝐬 ′ ) ​ [ Q ⁡ ( 𝐬 ′ , 𝐚 ′ ) ] \mathbb{E}_{\mathbf{a}^{\prime}\sim\pi(\mathbf{a}^{\prime}|\mathbf{s}^{\prime})}[Q(\mathbf{s}^{\prime},\mathbf{a}^{\prime})] , and therefore it is sufficient to keep π ⁡ ( 𝐚 ′ | 𝐬 ′ ) \pi(\mathbf{a}^{\prime}|\mathbf{s}^{\prime}) close to π β ​ ( 𝐚 ′ | 𝐬 ′ ) {\pi_{\beta}}(\mathbf{a}^{\prime}|\mathbf{s}^{\prime}) . Of course, in practice, the distributions do need to deviate in order for the learned policy to improve over the behavior policy, but by keeping this deviation small, errors due to out-of-distribution action inputs can be kept under control. The different methods in this family differ in terms of the probability metric used to define “closeness” and how this constraint is actually introduced and enforced. We can broadly categorize these methods along these two axes. We will discuss explicit f f -divergence constraints , which directly add a constraint to the actor update to keep the policy π \pi close to π β {\pi_{\beta}} in terms of an f f -divergence (typically the KL-divergence), implicit f f -divergence constraints , which utilize an actor update that, by construction, keeps π \pi close to π β {\pi_{\beta}} , and integral probability metric (IPM) constraints , which can express constraints with more favorable theoretical and empirical properties for offline RL. Furthermore, the constraints can be enforced either as direct policy constraints on the actor update, or via a policy penalty added to the reward function or target Q-value.

Formally, we can express the family of policy iteration methods with policy constraints as a fixed point iteration involving iterative (partial) optimization of the following objectives: Q ^ k + 1 π \displaystyle\hat{Q}^{\pi}_{k+1} ← arg ⁡ min Q ​ 𝔼 ( 𝐬 , 𝐚 , 𝐬 ′ ) ∼ 𝒟 ​ [ ( Q ⁡ ( 𝐬 , 𝐚 ) − ( r ⁡ ( 𝐬 , 𝐚 ) + γ ​ 𝔼 𝐚 ′ ∼ π k ​ ( 𝐚 ′ | 𝐬 ′ ) ​ [ Q ^ k π ​ ( 𝐬 ′ , 𝐚 ′ ) ] ) ) 2 ] \displaystyle\leftarrow\arg\min_{Q}\mathbb{E}_{(\mathbf{s},\mathbf{a},\mathbf{s}^{\prime})\sim\mathcal{D}}\left[\left(Q(\mathbf{s},\mathbf{a})-\left(r(\mathbf{s},\mathbf{a})+\gamma\mathbb{E}_{\mathbf{a}^{\prime}\sim\pi_{k}(\mathbf{a}^{\prime}|\mathbf{s}^{\prime})}[\hat{Q}^{\pi}_{k}(\mathbf{s}^{\prime},\mathbf{a}^{\prime})]\right)\right)^{2}\right] π k + 1 \displaystyle\pi_{k+1} ← arg ⁡ max π ​ 𝔼 𝐬 ∼ 𝒟 ​ [ 𝔼 𝐚 ∼ π ⁡ ( 𝐚 | 𝐬 ) ​ [ Q ^ k + 1 π ​ ( 𝐬 , 𝐚 ) ] ] ​ s.t. ​ D ​ ( π , π β ) ≤ ϵ . \displaystyle\leftarrow\arg\max_{\pi}\mathbb{E}_{\mathbf{s}\sim\mathcal{D}}\left[\mathbb{E}_{\mathbf{a}\sim\pi(\mathbf{a}|\mathbf{s})}[\hat{Q}^{\pi}_{k+1}(\mathbf{s},\mathbf{a})]\right]\text{ s.t. }D(\pi,{\pi_{\beta}})\leq\epsilon. When the min \min and max \max optimizations are not performed to convergence, but instead for a limited number of gradient steps, we recover the general actor-critic method in Algorithm 3 , with the exception of the constraint D ⁡ ( π , π β ) ≤ ϵ D(\pi,{\pi_{\beta}})\leq\epsilon on the policy update. A number of prior methods instantiate this approach with different choices of D D ( Kumar et al.,, 2019 ; Fujimoto et al.,, 2018 ; Siegel et al.,, 2020 ) . We will refer to this class of algorithms as policy constraint methods.

In policy penalty methods, the actor-critic algorithm is modified to incorporate the constraint into the Q-values, which forces the policy to not only avoid deviating from π β ​ ( 𝐚 | 𝐬 ) {\pi_{\beta}}(\mathbf{a}|\mathbf{s}) at each state, but to also avoid actions that may lead to higher deviation from π β ​ ( 𝐚 | 𝐬 ) {\pi_{\beta}}(\mathbf{a}|\mathbf{s}) at future time steps. This can be accomplished by adding a penalty term α D ( π ( ⋅ | 𝐬 ) , π β ( ⋅ | 𝐬 ) ) \alpha D(\pi(\cdot|\mathbf{s}),{\pi_{\beta}}(\cdot|\mathbf{s})) to the reward function r ⁡ ( 𝐬 , 𝐚 ) r(\mathbf{s},\mathbf{a}) leading to the modified reward function r ¯ ( 𝐬 , 𝐚 ) = r ( 𝐬 , 𝐚 ) − α D ( π ( ⋅ | 𝐬 ) , π β ( ⋅ | 𝐬 ) ) \bar{r}(\mathbf{s},\mathbf{a})=r(\mathbf{s},\mathbf{a})-\alpha D(\pi(\cdot|\mathbf{s}),{\pi_{\beta}}(\cdot|\mathbf{s})) . The most well-known formulation of policy penalty methods stems from the linearly solvable MDP framework ( Todorov,, 2006 ) , or equivalently the control as inference framework ( Levine,, 2018 ) , where D D is chosen to be the KL-divergence. An equivalent formulation adds the penalty term α D ( π ( ⋅ | 𝐬 ) , π β ( ⋅ | 𝐬 ) ) \alpha D(\pi(\cdot|\mathbf{s}),{\pi_{\beta}}(\cdot|\mathbf{s})) directly to the target Q-values and the actor objective ( Wu et al., 2019a, ; Jaques et al.,, 2019 ) , resulting in the following algorithm: Q ^ k + 1 π \displaystyle\hat{Q}^{\pi}_{k+1} ← arg ⁡ min Q \displaystyle\leftarrow\arg\min_{Q} 𝔼 ( 𝐬 , 𝐚 , 𝐬 ′ ) ∼ 𝒟 [ ( Q ( 𝐬 , 𝐚 ) − ( r ( 𝐬 , 𝐚 ) + γ 𝔼 𝐚 ′ ∼ π k ​ ( 𝐚 ′ | 𝐬 ′ ) [ Q ^ k π ( 𝐬 ′ , 𝐚 ′ ) ] − α γ D ( π k ( ⋅ | 𝐬 ′ ) , π β ( ⋅ | 𝐬 ′ ) ) ) ) 2 ] \displaystyle\mathbb{E}_{(\mathbf{s},\mathbf{a},\mathbf{s}^{\prime})\sim\mathcal{D}}\!\left[\left(Q(\mathbf{s},\mathbf{a})\!-\!\left(r(\mathbf{s},\mathbf{a})\!+\!\gamma\mathbb{E}_{\mathbf{a}^{\prime}\sim\pi_{k}(\mathbf{a}^{\prime}|\mathbf{s}^{\prime})}[\hat{Q}^{\pi}_{k}(\mathbf{s}^{\prime},\mathbf{a}^{\prime})]\!-\!\alpha\gamma\!D(\pi_{k}(\cdot|\mathbf{s}^{\prime}),{\pi_{\beta}}(\cdot|\mathbf{s}^{\prime}))\right)\right)^{2}\right] π k + 1 \displaystyle\pi_{k+1} ← arg max π 𝔼 𝐬 ∼ 𝒟 [ 𝔼 𝐚 ∼ π ⁡ ( 𝐚 | 𝐬 ) [ Q ^ k + 1 π ( 𝐬 , 𝐚 ) ] − α D ( π ( ⋅ | 𝐬 ) , π β ( ⋅ | 𝐬 ) ) ] . \displaystyle\leftarrow\arg\max_{\pi}\mathbb{E}_{\mathbf{s}\sim\mathcal{D}}\left[\mathbb{E}_{\mathbf{a}\sim\pi(\mathbf{a}|\mathbf{s})}[\hat{Q}^{\pi}_{k+1}(\mathbf{s},\mathbf{a})]-\alpha D(\pi(\cdot|\mathbf{s}),{\pi_{\beta}}(\cdot|\mathbf{s}))\right]. While the basic recipe for policy constraint and policy penalty methods is similar, the particular choice of how the constraints are defined and how they are enforced can make a significant difference in practice. We will discuss these choices next, as well as their tradeoffs.

#### Explicit f f -divergence constraints.

For any convex function f f , the corresponding f f -divergence is defined as: D f ( π ( ⋅ | 𝐬 ) , π β ( ⋅ | 𝐬 ) ) = 𝔼 𝐚 ∼ π ( ⋅ | 𝐬 ) [ f ( π ⁡ ( 𝐚 | 𝐬 ) π β ​ ( 𝐚 | 𝐬 ) ) ] . D_{f}(\pi(\cdot|\mathbf{s}),{\pi_{\beta}}(\cdot|\mathbf{s}))=\mathbb{E}_{\mathbf{a}\sim\pi(\cdot|\mathbf{s})}\left[f\left(\frac{\pi(\mathbf{a}|\mathbf{s})}{{\pi_{\beta}}(\mathbf{a}|\mathbf{s})}\right)\right]. (19) KL-divergence, χ 2 \chi^{2} -divergence, and total-variation distance are some commonly used members of the f f -divergence family, corresponding to different choices of function f f ( Nowozin et al.,, 2016 ) . A variational form for the f f -divergence can also be written as D f ( π ( ⋅ | 𝐬 ) , π β ( ⋅ | 𝐬 ) ) = max x : S × A → ℝ 𝔼 𝐚 ∼ π ( ⋅ | 𝐬 ) [ x ( 𝐬 , 𝐚 ) ] − 𝔼 𝐚 ′ ∼ π β ( ⋅ | 𝐬 ) [ f ∗ ( x ( 𝐬 , 𝐚 ′ ) ) ] , D_{f}(\pi(\cdot|\mathbf{s}),{\pi_{\beta}}(\cdot|\mathbf{s}))=\max_{x:S\times A\rightarrow\mathbb{R}}\mathbb{E}_{\mathbf{a}\sim\pi(\cdot|\mathbf{s})}\left[x(\mathbf{s},\mathbf{a})\right]-\mathbb{E}_{\mathbf{a}^{\prime}\sim{\pi_{\beta}}(\cdot|\mathbf{s})}\left[f^{*}(x(\mathbf{s},\mathbf{a}^{\prime}))\right], (20) where f ∗ f^{*} is the convex conjugate for the convex function f f . Both the primal form (Equation 19 ) and the dual variational form (Equation 20 ) of the f f -divergence has been used to specify policy constraints. In the dual form, an additional neural network is used to represent the function x x . The standard form of “passive dynamics” in linearly solvable MDPs ( Todorov,, 2006 ) or the action prior in control as inference ( Levine,, 2018 ) corresponds to the KL-divergence (primal form), which has also been used in recent work that adapts such a KL-divergence penalty for offline reinforcement learning ( Jaques et al.,, 2019 ; Wu et al., 2019a, ) . The KL-divergence, given by D KL ( π , π β ) = 𝔼 𝐚 ∼ π ( ⋅ | 𝐬 ) [ log π ( 𝐚 | 𝐬 ) − log π β ( 𝐚 | 𝐬 ) ] D_{\mathrm{KL}}(\pi,{\pi_{\beta}})=\mathbb{E}_{\mathbf{a}\sim\pi(\cdot|\mathbf{s})}[\log\pi(\mathbf{a}|\mathbf{s})-\log{\pi_{\beta}}(\mathbf{a}|\mathbf{s})] , can be computed by sampling action samples 𝐚 ∼ π ( ⋅ | 𝐬 ) \mathbf{a}\sim\pi(\cdot|\mathbf{s}) , and then computing sample-wise estimates of the likelihoods inside the expectation. It is commonly used together with “policy penalty” algorithms, by simply adding an estimate of − α ​ log ⁡ π β ​ ( 𝐚 | 𝐬 ) -\alpha\log{\pi_{\beta}}(\mathbf{a}|\mathbf{s}) to the reward function, and employing an entropy regularized reinforcement learning algorithm, which analytically adds ℋ ( π ( ⋅ | 𝐬 ) ) \mathcal{H}(\pi(\cdot|\mathbf{s})) to the actor objective ( Levine,, 2018 ) . One significant disadvantage of this approach is that it requires explicit estimation of the behavior policy (e.g., via behavioral cloning) to fit log ⁡ π β ​ ( 𝐚 | 𝐬 ) \log{\pi_{\beta}}(\mathbf{a}|\mathbf{s}) .

Additionally, the sub-family of asymmetrically-relaxed f f -divergences can be used for the policy constraint. For any chosen convex function f f , these divergences modify the expression for D f D_{f} to integrate over only those 𝐚 \mathbf{a} such that the density ration π ( ⋅ | 𝐬 ) / π β ( 𝐚 | 𝐬 ) \pi(\cdot|\mathbf{s})/{\pi_{\beta}}(\mathbf{a}|\mathbf{s}) is larger than a pre-defined threshold, thereby not penalizing small density ratios. Wu et al., 2019a () briefly discuss this divergence sub-family, and we refer readers to Wu et al., 2019b () for a detailed description.

#### Implicit f f -divergence constraints.

The KL-divergence constraint can also be enforced implicitly, as in the case of AWR ( Peng et al.,, 2019 ) , AWAC ( Nair et al.,, 2020 ) , and ABM ( Siegel et al.,, 2020 ) . These methods first solve for the optimal next policy iterate under the KL-divergence constraint, indicated as π ¯ k + 1 \bar{\pi}_{k+1} , non-parameterically, and then project it onto the policy function class via supervised regression, implementing the following procedure: π ¯ k + 1 ​ ( 𝐚 | 𝐬 ) \displaystyle\bar{\pi}_{k+1}(\mathbf{a}|\mathbf{s}) ← 1 Z ​ π β ​ ( 𝐚 | 𝐬 ) ​ exp ⁡ ( 1 α ​ Q k π ​ ( 𝐬 , 𝐚 ) ) \displaystyle\leftarrow\frac{1}{Z}{\pi_{\beta}}(\mathbf{a}|\mathbf{s})\exp\left(\frac{1}{\alpha}Q^{\pi}_{k}(\mathbf{s},\mathbf{a})\right) π k + 1 \displaystyle\pi_{k+1} ← arg ⁡ min π ​ D KL ​ ( π ¯ k + 1 , π ) \displaystyle\leftarrow\arg\min_{\pi}D_{\mathrm{KL}}(\bar{\pi}_{k+1},\pi) In practice, the first step can be implemented by weighting samples from π β ​ ( 𝐚 | 𝐬 ) {\pi_{\beta}}(\mathbf{a}|\mathbf{s}) (i.e., the data in the buffer 𝒟 \mathcal{D} ) by importance weights proportional to exp ⁡ ( 1 α ​ Q k π ​ ( 𝐬 , 𝐚 ) ) \exp\left(\frac{1}{\alpha}Q^{\pi}_{k}(\mathbf{s},\mathbf{a})\right) , and the second step can be implemented via a weighted supervised learning procedure employing these weights. In this way, importance sampling effectively implements a KL-divergence constraint on the policy, with α \alpha corresponding to the Lagrange multiplier. The Q-function Q k π Q^{\pi}_{k} corresponding to the previous policy π k \pi_{k} can be estimated with any Q-value or advantage estimator. We refer the reader to related work for a full derivation of this procedure ( Peng et al.,, 2019 ; Nair et al.,, 2020 ) .

#### Integral probability metrics (IPMs).

Another choice of the policy constraint D D is an integral probability metric, which is a divergence measure with the following functional form dependent on a function class Φ \Phi : D Φ ( π ( ⋅ | 𝐬 ) , π β ( ⋅ | 𝐬 ) ) = sup ϕ ∈ Φ , ϕ : S × A → ℝ | 𝔼 𝐚 ∼ π ( ⋅ | 𝐬 ) [ ϕ ( 𝐬 , 𝐚 ) ] − 𝔼 𝐚 ′ ∼ π β ( ⋅ | 𝐬 ) [ ϕ ( 𝐬 , 𝐚 ′ ) ] | . D_{\Phi}(\pi(\cdot|\mathbf{s}),{\pi_{\beta}}(\cdot|\mathbf{s}))=\sup_{\phi\in\Phi,\phi:S\times A\rightarrow\mathbb{R}}\left|\mathbb{E}_{\mathbf{a}\sim\pi(\cdot|\mathbf{s})}[\phi(\mathbf{s},\mathbf{a})]-\mathbb{E}_{\mathbf{a}^{\prime}\sim{\pi_{\beta}}(\cdot|\mathbf{s})}[\phi(\mathbf{s},\mathbf{a}^{\prime})]\right|. (21) Different choices for the function class Φ \Phi give rise to different divergences. For example, when Φ \Phi consists of all functions with a unit Hilbert norm in a reproducing kernel Hilbert space (RKHS), D Φ D_{\Phi} represents the maximum mean discrepancy (MMD) distance, which can alternatively be computed directly from samples as following: MMD 2 ( π ( ⋅ | 𝐬 ) , π β ( ⋅ | 𝐬 ) ) = 𝔼 𝐚 ∼ π ( ⋅ | 𝐬 ) , 𝐚 ′ ∼ π ( ⋅ | 𝐬 ) [ k ( 𝐚 , 𝐚 ′ ) ] − 2 𝔼 𝐚 ∼ π ( ⋅ | 𝐬 ) , 𝐚 ′ ∼ π β ( ⋅ | 𝐬 ) [ k ( 𝐚 , 𝐚 ′ ) ] + 𝔼 𝐚 ∼ π β ( ⋅ | 𝐬 ) , 𝐚 ′ ∼ π β ( ⋅ | 𝐬 ) [ k ( 𝐚 , 𝐚 ′ ) ] , \text{MMD}^{2}(\pi(\cdot|\mathbf{s}),{\pi_{\beta}}(\cdot|\mathbf{s}))=\mathbb{E}_{\mathbf{a}\sim\pi(\cdot|\mathbf{s}),\mathbf{a}^{\prime}\sim\pi(\cdot|\mathbf{s})}\left[k(\mathbf{a},\mathbf{a}^{\prime})\right]-\\ 2\mathbb{E}_{\mathbf{a}\sim\pi(\cdot|\mathbf{s}),\mathbf{a}^{\prime}\sim{\pi_{\beta}}(\cdot|\mathbf{s})}\left[k(\mathbf{a},\mathbf{a}^{\prime})\right]+\mathbb{E}_{\mathbf{a}\sim{\pi_{\beta}}(\cdot|\mathbf{s}),\mathbf{a}^{\prime}\sim{\pi_{\beta}}(\cdot|\mathbf{s})}\left[k(\mathbf{a},\mathbf{a}^{\prime})\right], where k ⁡ ( ⋅ , ⋅ ) k(\cdot,\cdot) represents any radial basis kernel, such as the Gaussian or Laplacian kernel. As another example, when Φ \Phi consists of all functions ϕ \phi with a unit Lipschitz constant, then D Φ D_{\Phi} is equivalent to the first order Wasserstein ( W 1 W_{1} ) or Earth-mover’s distance: W 1 ( π ( ⋅ | 𝐬 ) , π β ( ⋅ | 𝐬 ) ) = sup f , ‖ f ‖ L ≤ 1 | 𝔼 𝐚 ∼ π ( ⋅ | 𝐬 ) [ f ( 𝐚 ) ] − 𝔼 𝐚 ∼ π β ( ⋅ | 𝐬 [ f ( 𝐚 ) ] | . W_{1}(\pi(\cdot|\mathbf{s}),{\pi_{\beta}}(\cdot|\mathbf{s}))=\sup_{f,||f||_{L}\leq 1}\left|\mathbb{E}_{\mathbf{a}\sim\pi(\cdot|\mathbf{s})}[f(\mathbf{a})]-\mathbb{E}_{\mathbf{a}\sim{\pi_{\beta}}(\cdot|\mathbf{s}}[f(\mathbf{a})]\right|. (22) These metrics are appealing because they can often be estimated with non-parametric estimators. We refer the readers to Sriperumbudur et al., (2009) for a detailed discussion on IPMs. BEAR ( Kumar et al.,, 2019 ) uses the MMD distance to represent the policy constraint, while Wu et al., 2019a () evaluates an instantiation of the first order Wasserstein distance.

#### Tradeoffs between different constraint types.