1 Introduction↩︎

Reinforcement learning (RL) describes the field of machine learning concerned with learning a policy to solve a task through stochastic trial-and-error experience in an environment. The default setting typically assumed in RL is of online access to the environment; i.e., the learned policy may collect new trial-and-error experience directly from the environment. However, in many practical scenarios, where deploying a new policy to interact with the live environment is expensive or associated with risks or safety concerns [1], it is more common to have only offline access to the environment. That is, the trial-and-error experience available for learning a task-solving policy is a static, offline dataset of experience collected by some other behavior policy. This setting is known as offline RL and has attracted a significant amount of interest in recent years [2], [3].

Many of the recent approaches to offline RL utilize some form of behavior regularization, in which the learned policy is compelled to stay close to the data-generating behavior policy [4]. For example, in model-free actor critic algorithms, behavior regularization is typically done by augmenting the actor loss with a penalty measuring the divergence of the learned policy from the behavior policy [4]–[6], reminiscent of KL-control methods in related literature [7], [8]. While straightforward, a disadvantage of this approach is that it does little to regularize the critic itself, and thus, it is common for the critic to take on wildly extrapolated values on actions unseen in the training data, which may dominate any behavior regularization applied to the actor, as we demonstrate in this work.

In this work, we propose an alternative approach to encouraging the learned policy to stay close to the offline data. Focusing on the critic, we propose parameterizing the critic values as the logits of the behavior policy plus an additional offset term. When the actor (the learned policy) is then trained to choose actions which maximize the critic value, it will be compelled to stay close to the behavior policy as long as the offset term is suitably ‘small’. Thus, behavior regularization in this setting corresponds to augmenting the standard Bellman error critic loss with an appropriate regularization on the offset term.

What should this regularization term be? After noting that, in continuous control, the offset term’s effect on the learned policy is via the gradients of the offset with respect to actions, we propose regularizing the offset term with a gradient penalty. While this may appear heuristic at first, we present mathematical derivations establishing a connection between this gradient penalty and the Fisher divergence, which appears in the score matching and energy-based generative model literature [9], [10], interpreting the critic values as the energy function of a Boltzmann distribution.

We thus term our newly proposed actor critic algorithm Fisher-BRC (behavior regularized critic). To aid conceptual understanding of Fisher-BRC, we analyze its training dynamics in a simple toy setting, highlighting the advantage of its implicit Fisher divergence regularization as opposed to the more explicit divergence penalties imposed by alternative offline RL methods. We then present an extensive evaluation of Fisher-BRC on standard offline RL benchmarks. We find that Fisher-BRC yields state-of-the-art performance compared to a variety of existing model-free and model-based RL methods. We further show that while Fisher-BRC learns better policies, it is also much more computationally efficient than more sophisticated offline RL algorithms. Overall, Fisher-BRC presents a new approach to behavior regularization in offline RL with compelling practical benefits.

2 Related Work↩︎

Our work adds to the rich literature on behavior regularization methods in offline RL [4], which propose a number of regularizations in RL training that compel the learned policy to stay close to the offline data. These regularizers have appeared as divergence penalties [4]–[6], implicitly through appropriate network initializations [11], or more explicitly through careful parameterization of the policy [12]. Another way to apply behavior regularizers is via modification of the critic learning objective as in [13], [14]. Our work is unique in applying behavior regularization through a parameterization of the critic, and empirically, we find this to yield much better performance.

In this work, we demonstrate that the specific parameterization we employ has connections to Fisher divergence regularization, establishing connections with energy-based models. Prior methods have proposed using energy-based models for policies to express more multi-modal distributions [15], [16]. Meanwhile, while our work focuses on reinforcement learning – i.e., learning a return-maximizing policy – other works have established connections between energy-based models and inverse RL or imitation learning [17].

In practice, our regularization reduces to a gradient penalty applied to the offset term in the critic. Gradient penalties have appeared in previous work, arguably first popularized in machine learning by Wasserstein generative-adversarial networks [18], but also used in imitation learning [19], cross-domain disentanglement [20], and uncertainty estimation [21].

3 Background↩︎

We introduce the notation and assumptions used in this paper, as well as provide an in-depth review of the methods most closely related to ours.

3.1 Reinforcement Learning↩︎

In this work, we consider environments that can be represented as a Markov Decision Process (MDP) defined by a tuple \((\mathcal{S}, \mathcal{A}, p_0, p, r, \gamma)\), where \(\mathcal{S}\) is a state space, \(\mathcal{A}\) is an action space, \(p_0(s)\) is a distribution of initial states, \(p(s'|s,a)\) is a stochastic dynamics model, \(r: \mathcal{S}\times \mathcal{A}\rightarrow \mathbb{R}\) is a reward function and \(\gamma \in [0, 1)\) is a discount. We restrict our work to continuous action spaces; i.e., \(\mathcal{A}\subset \mathbb{R}^{d}\) for some \(d\). The goal of reinforcement learning is to find a policy \(\pi(a|s)\) that maximizes the cumulative discounted returns \(\mathop{\mathbb{E}}_\pi[\sum_{t=0}^\infty\gamma^t r(s_t, a_t) | s_0 \sim p_0(\cdot), a_t \sim \pi(\cdot|s_t), s_{t+1} \sim p(\cdot|s_t, a_t) ]\). In online reinforcement learning, it is usually assumed that an agent learns based on experience generated by the agent interacting with the learning environment.

3.2 Offline Reinforcement Learning↩︎

In this work, we focus on the offline setting, in which the agent cannot generate new experience data and learns based on a provided dataset \(\mathcal{D}\) of \((s, a, r, s')\) tuples generated by some other policy interacting with the environment. We call the policy that generated this dataset the behavior policy and denote it as \(\mu\).

Offline datasets usually do not provide complete state-action coverage. That is, the set \(\{(s,a) ~|~ (s,a,r,s')\in\mathcal{D}\}\) is typically a small subset of the full space \(\mathcal{S}\times\mathcal{A}\). Standard reinforcement learning methods such as SAC [22] or DDPG [23] suffer when applied to these datasets due to extrapolation errors [6], [12].

Behavior regularization is a prominent approach to offline reinforcement learning that aims to address this problem by using appropriate regularizers to compel the learned policy to stay close to the data. There are two common ways to incorporate behavior regularization into the actor-critic framework – via policy regularization or via a critic penalty. Since our own approach is related to both techniques, in the following sections we describe the two approaches as well as problems associated with them.

3.3 Policy Regularization↩︎

Policy regularization can be imposed either during critic or policy learning. First, we describe a family of approaches based on applying behavior constraints to training policies. These constraints can be applied in a hard fashion, by restricting the policy action space to the actions seen in the offline dataset as in BCQ [12]: \[\begin{align} \begin{aligned} \pi(s) \vcentcolon=\mathop{\mathrm{arg\,max}}_{a_i + \xi_\phi(s,a_i,\Phi)} Q_\theta(s, a_i + \xi_\phi(s,a_i,\Phi)),\\ \{a_i\sim \mu(\cdot|s)\}_{i=1}^n \end{aligned} \end{align}\] where \(\pi(s)\) is a deterministic policy, \(\xi_\phi(s,a,\Phi)\) is a perturbation model on actions constrained to the interval \([-\Phi, \Phi]\), \(\mu(a|s)\) is a behavioral policy contracted by fitting a density model to the offline dataset. The hyperparameter \(n\) controls the number of sampled actions. In other words, one samples \(n\) perturbed actions from the behavior policy and chooses from the action with the largest critic-approximated value. One of the limitations of this approach is that a large number of sampled actions might be required for competitive performance  [24]. Using a divergence penalty is an alternative approach to policy regularization. Instead of having hard constraints on the training policy, one can regularize the policy with an appropriately chosen probability divergence such as the KL-divergence [4], [5]: \[\label{eq:kl95control} \max_{\pi} \mathop{\mathbb{E}}_{s \sim \mathcal{D}}\left[\mathop{\mathbb{E}}_{a \sim \pi(\cdot|s)} [Q_\theta(s,a)] - \alpha D_{KL}(\pi(\cdot|s)\|\mu(\cdot|s)) \right].\tag{1}\]

Although these approaches demonstrate impressive performance on some tasks, they share a common problem. The Q-function, \(Q_\theta\), learned via standard TD-error minimization on \(\mathcal{D}\) receives no learning signal for actions not observed in the replay buffer, while this same Q-function is nevertheless queried on out-of-distribution actions during the policy update – i.e., when \(Q_\theta(s, a)\) is evaluated on \(a\sim\pi(\cdot|s)\) in 1 – and for bootstrapping critic targets in the squared TD-loss: \[\label{eq:tdloss} J(Q_\theta) \vcentcolon=\mathop{\mathbb{E}}_{\substack{(s,a,s')\sim\mathcal{D}\\ a' \sim \pi(\cdot|s')}}[(r(s,a) + \gamma Q_{\hat{\theta}}(s',a')-Q_\theta(s,a))^2].\tag{2}\] Thus, issues with critic extrapolation can still dominate divergence regularizers applied to the policy.

3.4 Critic Penalty↩︎

Other works, such as AlgaeDICE  [13] and CQL [14] attempt to incorporate some divergence regularization into the critic. In particular, AlgaeDICE introduces a term that pushes \(Q\)-values down for actions sampled from the training policy while minimizing TD-error via residual learning: \[\begin{align} \min_{\theta} \alpha(1-\gamma)&\mathop{\mathbb{E}}_{\substack{s_0\sim\pi_0(\cdot)\\a_0\sim\pi(s_0)}}[Q_\theta(s_0,a_0)] +\\ \mathop{\mathbb{E}}_{(s,a)\sim\mathcal{D}}[(r(s,a) + \gamma &\mathop{\mathbb{E}}_{\substack{s'\sim p(\cdot|s,a)\\a'\sim\pi(\cdot|s')}}[Q_\theta(s',a')]-Q_\theta(s,a))^2]. \end{align}\] This formulation can be generalized by replacing the squared function by some convex function \(f\). The choice of \(f\) may be shown to correspond to an implicit \(f\)-divergence regularization on the learned policy with respect to state-action distributions. For example, a choice of \(f(x) = \exp(x - 1)\) or \(f(x) = \log \mathop{\mathbb{E}}_{\mathcal{D}} \exp(x)\) corresponds to an implicit KL-divergence.

Based on a similar idea, CQL [14] extends the standard critic loss \(J(Q_\theta)\) in 2 with additional terms that minimize \(Q\)-values sampled from a policy and maximize values of the dataset actions: \[\begin{align} \label{eqn:cql} \min_{\theta} J(Q_\theta) + \lambda\mathop{\mathbb{E}}_{(s, a) \sim \mathcal{D}}[\log \Sigma_a \exp(Q_\theta(s,a)) - Q_\theta(s,a)]. \end{align}\tag{3}\] Although both of these methods provide learning signal to the critic \(Q\)-values on the entire action space, AlgaeDICE is based on residual learning with slower convergence than fitted TD-learning [25]. On the other hand, although CQL can achieve better empirical performance, the log-sum-exp term that appears in its formulation is not tractable for continuous actions and must be computed via numerical integration. The authors of CQL propose doing this via Monte-Carlo sampling with importance weights, where the current training policy is used to draw samples for the procedure. This process can add a significant computational burden to actor critic training. In contrast, our proposed Fisher-BRC will be shown to achieve favorable empirical performance while maintaining computational efficiency closer to standard actor critic.

4 Fisher-BRC↩︎

We now continue to describe our own approach to behavior regularization in offline RL settings, which aims to circumvent the issues associated with competing methods described in the previous section, namely (1) lack of well-defined critic values on out-of-distribution actions and (2) computational inefficiency of critic penalty approaches. We begin with a conceptual derivation of our method, before presenting a more formal connection to Fisher divergence regularization. See Algorithm 1 for a sketch of the full algorithm.

4.1 Conceptual Derivation↩︎

We start from the observation that we can represent the entropy smoothed \(Q\)-values of the behavior policy \(\mu(\cdot|s)\) as \[Q(s,a) = V(s) + \log \mu(a|s).\] This decomposition of \(Q\)-values into state-value and log-policy is popular in the entropy-regularized online RL literature [26], [27], where the \(\mu\) in our notation is treated as the policy \(\pi\) to be learned. Our own setting is markedly different from these previous approaches in that \(\mu\) is the behavior policy and is fixed. Nevertheless, what would happen if we were to parameterize \(Q\)-values in this way, i.e, as \[\label{eq:q-soft} Q_\theta(s,a) := V_\theta(s) + \log \mu(a|s),\tag{4}\] and then learn via standard TD error minimization 2 ? Well, there is an advantage, but also a disadvantage.

First, the main advantage of this formulation is that, in contrast to a \(Q\)-function parameterized by its own neural network function approximator, the density \(\mu(a|s)\) is well-defined for all actions, and thus \(Q\) is more likely to have a well-behaved landscape even though its training may only cover a small subset of the full action space. Accordingly, the learned policy \(\pi\), when trained in the standard way to choose actions which maximize the \(Q\)-values, is thus encouraged to stay close to \(\mu\), without the need for an explicit divergence penalty as in 1 . In practice, knowledge of \(\mu\) is not explicitly provided, and one usually resorts to behavioral cloning on the offline dataset to approximate \(\mu\) [4]. Nevertheless, the advantage of the formulation in 4 still holds, since even if \(\mu\) is trained on a sparse subset of all possible actions, the fact that it is a normalized probability distribution means that \(\mu(a|s)\) assigns low probabilities to actions outside of \(\mathcal{D}\). Still, in practice some density models might fail to generalize to out-of-distribution data [28]. For this reason, we parameterize the approximate density \(\mu\) as a mixture density model [29].

As for the disadvantage, it is clear that the formulation of \(Q\) in 4 is too restrictive. If this representation is used for training a new policy \(\pi\), the new policy will be limited to copying the behavior policy \(\mu\), regardless of \(V_\theta\), and this will lead to suboptimal performance. In order to address this issue and enable the learned \(\pi\) to generalize beyond just mimicking \(\mu\), we propose replacing the value function \(V_\theta(s)\) with a state-action value offset function \(O_\theta(s,a)\): \[\label{eq:qrepr} Q_\theta(s,a) \vcentcolon=O_\theta(s,a) + \log \mu(a|s).\tag{5}\] With this representation one can learn a richer representation of \(Q\)-values. However, this parameterization can potentially put us back in the fully-parameterized \(Q_\theta\) regime of vanilla actor critic. It is clear that the offset term must be suitably constrained or regularized to find an appropriate middle-ground between the overly-restrictive \(V_\theta(s)\) and the fully-parameterized alternative that is liable to extrapolation errors and policy divergence.

To motivate what an appropriate regularization on \(O_\theta\) should be, we begin by dissecting exactly how the offset term impacts the learned policy. Let’s take a closer look at the policy updates in standard continuous-control actor critic [22], [23]. These updates are based on the chain-rule gradient computation below: \[\begin{gather} \label{eq:inc95update} \nabla_\phi Q_\theta(s, \pi_\phi(s)) = \\ \big[\nabla_a O_\theta(s,a) + \nabla_a \log \mu(a|s)\big]_{a = \pi_\phi(s)} \nabla_\phi \pi_\phi(s). \end{gather}\tag{6}\] From this expression, it is clear that potential extrapolation issues with the actor arise via the gradient \(\nabla_a O_\theta(s,a)\). That is, without appropriate constraints or regularizers on \(O_\theta(s,a)\), its gradients might dominate over the gradients of the behavior policy \(\nabla_a \log \mu(a|s)\). Therefore, as a way to control the trade-off between over-constraining \(\pi\) to be close to the behavior policy and learning more rich representations of \(Q\)-values, we propose using a gradient penalty regularizer of the form \(\|\nabla_a O_\theta(s,a)\|^2\). Accordingly, the full critic optimization objective is as follows: \[\begin{align} \label{eq:bfrc} \min_{\theta} J(O_\theta + \log\mu) + \lambda \mathop{\mathbb{E}}_{\substack{s \sim \mathcal{D}\\a\sim\pi_\phi(\cdot|s)}}[\|\nabla_a O_\theta(s,a) \|^2], \end{align}\tag{7}\] where \(J(\cdot)\) is defined as in 2 and \(\mu\) is not updated during critic learning. In this objective, \(\lambda\) is a hyperparameter that controls the contribution of the gradient penalty term. Unless otherwise noted, we set \(\lambda=0.1\) as the regularization coefficient.

A keen reader may note that the gradients in 6 look similar to gradients of 1 , with the difference that we take gradients of the offset function instead of the critic function. However, in our setting the critic loss is substantially different, since the offset term plus the log behavior density learns to predict the unmodified \(Q\)-values of the training policy instead of \(Q\)-values augmented with a KL-divergence term. For example, if the rewards of the MDP already naturally compel the learned policy to match the behavior policy, the offset term in Fisher-BRC can vanish, whereas the explicit divergence penalty in BRAC will bias the learned \(Q\)-values unnecessarily.

4.2 Fisher Divergence Derivation↩︎

We now show how the same objective in 7 may be derived from the perspective of Fisher divergence regularization. We begin by introducing the idea of Boltzmann policies – essentially policies expressed as a Boltzmann distribution with energy function given by a set of \(Q\)-values. We then present the Fisher divergence, and show how a Fisher divergence regularizer between a Boltzmann policy and the behavior policy reduces to the gradient penalty proposed in 7 . Finally, we elaborate on connections to CQL, which we show can be interpreted as a KL divergence regularization between the Boltzmann policy and the behavior policy, and this insight may be of independent interest, since the original derivation of CQL is via a very different route.

4.2.0.1 Boltzmann Policies

For a given \(Q\)-value function, the associated Boltzmann policy is given by the following expression: \[\label{eqn:bolt} \pi_{ebm}(a|s) \vcentcolon=\cfrac{\exp(Q(s,a))}{\sum \exp(Q(s,a))}.\tag{8}\] In an actor critic setting, the actor will recover this same policy if an entropy regularizer is added to the actor loss, as is commonly done [22]. While there are some works which use this representation of a policy more explicitly [15], [27], [30], a main disadvantage is that the normalization term may not be tractable for continuous actions, and so computing the policy distribution \(\pi_{ebm}\) explicitly requires performing computationally expensive numerical integration. Thus, it is more common to only recover the policy via entropy regularization on the actor loss.

4.2.0.2 Fisher Divergence

In order to avoid issues with computing the normalization term, we consider the Fisher divergence, or Fisher information distance [31]: \[\label{eqn:fisher} F(p(\cdot), q(\cdot)) = \mathop{\mathbb{E}}_{x \sim p(\cdot)} \left[ \|\nabla_x \log p(x) - \nabla_x \log q(x) \|^2 \right].\tag{9}\] As one can see from the formulation, in order to compute the Fisher divergence between two distributions, we need only have sampling access to \(p(x)\) and the ability to compute \(\nabla_x \log p(x),\nabla_x \log q(x)\). Crucially, the computation of either \(\nabla_x \log p(x)\) or \(\nabla_x \log q(x)\) does not require normalized distributions, since the normalization term (which is a constant with respect to \(x\)) disappears from \(\log p(x),\log q(x)\) due to the differentiation. Thus, we can avoid computing the normalization term in 8 .

Since the Fisher divergence is amenable to Boltzmann representations of policies, let’s consider an optimization objective that consists of a squared TD-loss and a Fisher divergence term between the Boltzmann distribution and the behavior policy \(\mu\): \[\begin{align} \begin{aligned} \label{eqn:fisher95td} &J(Q_\theta) + \lambda \mathop{\mathbb{E}}_{s\sim \mathcal{D}} \left[F\left( \cfrac{\exp(Q(s,\cdot))}{\sum_a \exp(Q(s,a))},\mu(\cdot|s)\right)\right] =\\ &J(Q_\theta) + \lambda\mathop{\mathbb{E}}_{\substack{s \sim \mathcal{D}\\ a \sim \pi_{emb}(\cdot|s)}} \left[\left\|\nabla_a \log \mu(a|s)- \nabla_a Q(s,a) \right\|^2 \right]. \end{aligned} \end{align}\tag{10}\] The coefficient \(\lambda\) controls the strength of the Fisher divergence term. We can further simplify this objective by using the representation of \(Q\) proposed in 5 : \[\begin{align} &J(O_\theta + \log\mu) + \lambda\mathop{\mathbb{E}}_{\substack{s \sim \mathcal{D}\\ a \sim \pi_{ebm}(\cdot|s)}} \left[\left\|\nabla_a O_\theta(s,a) \right\|^2 \right]. \label{eq:fisher-ebm} \end{align}\tag{11}\] The only difference with 7 is that actions are sampled from the training policy \(\pi_\phi\), while in this case the actions are sampled from the Bolzman policy \(\pi_{ebm}\). In practice, sampling from \(\pi_{ebm}\) can be just as computationally expensive as computing the normalization term in 3 . However, as mentioned earlier, the Boltzmann policy \(\pi_{ebm}\) may be recovered by the actor via incorporation of an entropy regularizer in the actor loss. Thus, we propose to train our actor \(\pi_\phi\) exactly in this manner (which is already popular in model-free actor critic algorithms [22]), and then use it in 11 as a plug-in approximation of \(\pi_{ebm}\). In this way, we have arrived again to the same objective first defined in [sec:fbrc].

3pt

The connection of our proposed objective to the Fisher divergence recalls similar quantities in the score matching and energy-based generative model literature. In fact, the Fisher divergence is a popular metric in these literatures exactly because it avoids an expensive computation of a log-normalizer, which is necessary for other common divergences [9], [10]. Moreover, many works in the generative model literature employ gradient penalties, even if they do not explicitly make a connection to the Fisher divergence. This provides a further empirical advantage to our method, as these gradient penalties – due to their popularity – may be efficiently implemented using many modern machine learning libraries. We note that although we use a soft penalty, one can enforce hard constraints as in Spectral Normalized GANs [34], but we leave this for future work.

Due the connection of our method to Fisher Divergence, we dub our method Fisher-BRC (Fisher Behavior Regularized Critic).

4.2.0.3 Connections to CQL

The CQL objective, although originally derived in [14] from a very different perspective, can also be motivated as a regularizer on a Boltzmann policy. In the case of CQL, the regularizer is the common KL-divergence: \[\label{eqn:cql95kl}J(Q_\theta) + \lambda\mathop{\mathbb{E}}_{s\sim\mathcal{D}} \left[D_{KL}\left(\mu(\cdot|s)\middle|\cfrac{\exp(Q(s,\cdot))}{\sum_a \exp(Q(s,a))}\right)\right].\tag{12}\] Expanding the KL-Divergence term yields, \[\begin{align} D_{KL}(\mu(\cdot|s)|\pi_{emb}(\cdot|s)) = \mathop{\mathbb{E}}_{a\sim\mu(\cdot|s)}\left[\log \cfrac{\mu(a|s)}{\pi_{emb}(a|s)}\right]= \\ \mathop{\mathbb{E}}_{a\sim\mu(\cdot|s)}\left[\log \sum_a \exp(Q(s,a)) - Q(s,a) + \log \mu(a|s)\right], \end{align}\] and this is equivalent to the familiar form of CQL from 3 , since \(\log \mu(a|s)\) is a constant with respect to \(Q\).

As mentioned earlier, for continuous distributions the normalization term in CQL is not tractable and necessitates expensive numerical integration. In CQL this integration is calculated by sampling from the training policy and importance weighting. Thus, our own method enjoys a significant computational advantage. Our use of a novel critic representation and the Fisher divergence allows us to circumvent this practical issue.

We note that in this derivation of CQL, the direction of divergence in 12 is switched compared to 10 in Fisher-BRC. One may derive a variant of Fisher-BRC using this same direction of the divergence, and this will result in the same expression in 7 , only that the expectation of the gradient penalty is switched to be over \((s,a) \sim \mathcal{D}\). Anecdotally, in initial experiments we did not observe large empirical differences when training with this alternative objective. However, due to the closer connection to the actor loss when using \(a\sim \pi_\phi(\cdot|s)\) (see 4.1), we stick to the formulation originally presented in 7 .

5 Experiments↩︎

We present empirical demonstrations of Fisher-BRC in a variety of settings. We start with a simple continuous bandit experiment that illustrates the difference between our method and more common behavior regularization techniques based on explicit divergence penalties. Then we evaluate our method against state-of-the-art offline RL model-based and model-free algorithms on the D4RL benchmark datasets. Finally, we analyze the effect of gradient penalty regularization and provide statistics on the computational advantage of Fisher-BRC compared to CQL.

5.1 Toy Continuous Bandit Problem↩︎

We begin with a simple conceptual demonstration comparing Fisher-BRC to the similar and common alternative of explicit divergence penalties applied to the learned policy. Specifically, we consider the loss in 1 , which corresponds to the policy update used in BRAC [4].

For the purpose of this experiment, we consider a continuous bandit with one-dimensional action space given by \([-1,1]\). The rewards are given by \[r(a)= \begin{cases} |a| - 0.125,& \text{if } a \in [-0.25, 0.25] \\ -\infty, & \text{otherwise.} \end{cases}\] The offline training dataset is collected by sampling 1000 actions from a uniform distribution, \(a \sim \mathcal{U}(-0.25, 0.25)\), and recording the corresponding rewards \(r(a)\). Thus, the distribution of actions exhibits poor coverage of the full action space, and the rewards for \(a \in [-1, -0.25] \cup [0.25, 1]\) are not observed in the dataset.

Both BRAC and Fisher-BRC require a fitted behavior policy. To do so, we fit a behavior policy \(\mu(a)\) parameterized as a Laplace distribution.² Subsequently, we fit the critic for BRAC and F-BRC. In BRAC, we fit a critic using mean squared error to match the rewards: \(\mathop{\mathbb{E}}_{(a, r)\sim\mathcal{D}}[(R_\theta(a) - r)^2]\). For Fisher-BRC, we use the representation and regularization from 7 , assuming an initialization of \(\pi\) to \(\mathcal{U}\): \[\begin{align} \mathop{\mathbb{E}}_{(a, r)\sim\mathcal{D}}[(O_\theta(a) + \log\mu(a) - r)^2] + \\ \lambda \mathop{\mathbb{E}}_{a_{reg}\sim \mathcal{U}(-1, 1)} \|\nabla_a O_\theta(a_{reg})\|^2. \end{align}\] These critics then determine the objective landscape for the learned policy. We plots these landscapes in 2. Specifically, we plot \(R_\theta(a) + \alpha \log\mu(a)\) for BRAC and \(O_\theta(a) + \log \mu(a)\) for Fisher-BRC for a variety of choices of \(\alpha\) and \(\lambda\). Recall that the policy will be learned to choose actions which maximize these values. Thus, ideally these objective landscapes should possess optima around the globally optimal actions \(\{-0.25, 0.25\}\).

We observe that it is hard to pick the KL-coefficient \(\alpha\) for the policy loss landscape induced by BRAC to avoid either over generalizing (with optima outside of \([-0.25,0.25]\)) or over regularizing (with optima far from the optimal actions \(\{-0.25, 0.25\}\)); see 2, left.

On the other hand, Fisher-BRC correctly constrains the learned value function – in fact, the value function is unmodified for in-distribution samples – and is robust to the choice of the regularization hyperparameter (2, right).

5.2 Deep Offline RL Benchmarks↩︎

We continue to present Fisher-BRC on more complex environments. We compare our method to prior work on the OpenAI Gym MuJoCo tasks using D4RL datasets [32]. We consider the following baselines: BRAC-vp and BRAC-pr [4], due to a similar policy learning objective, MBOP [33], the state-of-state in offline model-based reinforcement learning, and CQL [14], due to a similar connection with energy based learning. Our implementation for Fisher-BRC follows the standard SAC implementation, only that we use a 3-layer network as in CQL. Additional implementation details are in the appendix.

5.2.0.1 Overall performance

The results of our method and all considered baselines are presented in Table 1. Our method performs comparably or surpasses prior work on most of the tasks. Notably, many of the baseline algorithms exhibit inconsistent performance across the tasks, achieving good performance on some while poor performance on others. In contrast, our Fisher-BRC exhibits consistent and good performance across almost all tasks.

5.2.0.2 Effect of gradient penalty

As a way of investigating the effect of gradient penalty vs. the offset parameterization in Fisher-BRC, we evaluate different values of gradient penalty. We present results of \(\lambda \in \{0.0, 0.1, 1.0\}\) in 3. We note that the gradient penalty is an important component of Fisher-BRC, since when \(\lambda=0.0\) performance degrades dramatically. On the other hand, when \(\lambda\) is set too high (\(\lambda=1.0\)), we see that the learned policy is over-constrained; i.e., we see that performance on the expert datasets is improved while performance is limited on the medium datasets, since the behavior policy in these datasets is highly sub-optimal. This is expected, since a high \(\lambda\) corresponds to a large gradient penalty on the regularization on the offset, compelling the offset to be near-constant with respect to actions.

5.2.0.3 Convergence speed

We also evaluate the wall-clock training time of our method compared with CQL, which demonstrates comparable policy performance on the D4RL tasks but significantly different computational efficiency. The total training time for 1 million steps for Fisher-BRC is 1.4 hours of behavioral cloning pretraining followed by 6.2 hours of policy training. CQL total training time is 16.3 hours (which does not require pre-training of a behavior density policy). Our method converges faster not only in terms of gradient updates (see 4), but it is also computationally faster due to omitting the expensive numerical integration to compute the log-sum-exp term of CQL. These experiments were carried out on a Google cloud instance containing an AMD EPYC 7B12 CPU at 2.25GHz (using 8 of 64 available cores) and 32GB of RAM.

6 Conclusions↩︎

We have introduced Fisher-BRC, a simple critic representation and regularization technique for offline reinforcement learning. Our derivations highlight connections between our training objective and Fisher divergence regularization from score matching and energy-based model literature. Our method is easy to implement and highly performant. Compared to existing offline RL algorithms, Fisher-BRC exhibits better and more consistent performance across a variety of domains.

Acknowledgements↩︎

We thank George Tucker for reviewing a draft of this paper and providing valuable feedback. We also thank Aviral Kumar for discussions and help with CQL results.

7 Implementation Details↩︎

7.0.0.1 Behavior policy

We fit the behavior model using a conditional Mixture of Gaussians [29] with tanh squashing [15]. We use 5 mixture components. We train the density model with Adam optimizer [35] for \(10^6\) steps and starting from learning rate \(10^{-3}\) and decreasing it by 10 at \(8\cdot10^5\) and \(9\cdot10^5\) gradient update steps. Similar to BRAC we train the behavior actor with SAC-style entropy regularization with the same target entropy. We parameterize the model as a 3 layer MLP with relu activations and 256 hidden units.

7.0.0.2 Actor and critic learning

Our implementation is based on Soft Actor Critic [22]. As in CQL we do not add entropy to the rewards and we modify the critic loss to accommodate the additional regularization term. We use default SAC hyper parameters without additional tuning, in contrast to CQL and BRAC which tune policy learning rate. Following CQL we increased network size for the actor and the critic to 3 layer MLP with 256 hidden units.

7.0.0.3 Survival bonus

The linear term used in CQL can be seen as adding a survival bonus for the environments with early early termination. The derivation is included in the Appendix. Adding a positive constant to the rewards does not have an effect on the optimal policy in infinite horizon MDPs, but in practice Q-targets for terminal states are replaced with \(0\) that leads to having either a survival bonus or step penalty. For this reason, we add a reward bonus to our implementation as well for fair comparison. We choose the same value \(\lambda_{cql}=5\) as in CQL.

In particular, one can verify that \[\begin{align} \nabla_\theta [-\lambda_{cql} &Q_\theta(s,a) + (\gamma Q_{\hat{\theta}}(s',a') + r(s,a) - Q_\theta(s,a))^2]= \\ -\lambda_{cql}\nabla_\theta &Q_\theta(s,a) - (\gamma Q_{\hat{\theta}}(s',a') + r(s,a) - Q_\theta(s,a)) \nabla_\theta Q_\theta(s,a)=\\ - (\gamma Q_{\hat{\theta}}(s',&a') + [r(s,a) + \lambda_{cql}] - Q_\theta(s,a)) \nabla_\theta Q_\theta(s,a)=\\ \nabla_\theta (\gamma &Q_{\hat{\theta}}(s',a') + [r(s,a) + \lambda_{cql}] - Q_\theta(s,a))^2. \end{align}\]

8 Gradient Penalty Ablation↩︎

We also present a full set of results for the ablations in 3.

9 Critic Regularization Ablation↩︎

We also evaluate the effect of gradient penalty of standard Soft Actor Critic without the critic representation introduced in this paper.

References↩︎