Offline Reinforcement Learning with Behavioral Supervisor Tuning

1 Introduction

Reinforcement learning (RL) is a method of learning where an agent interacts with an environment to collect experiences and seeks to maximize the reward provided by the environment. This typically follows a repeating cycle of experience collecting and improvement [Sutton and Barto, 2018]. This is termed online RL due to the need for policy rollouts in the environment. Both on-policy and off-policy RL require some schedule of online interaction which, in some domains, can be infeasible due to experimental or environmental limitations [Mirowski et al., 2018; Yu et al., 2021]. With such constraints, a dataset may instead be collected that consists of demonstrations by arbitrary (potentially multiple, unknown) behavior policies [Lange et al., 2012] that may be suboptimal. Offline reinforcement learning algorithms are designed to recover optimal policies from such static datasets.

The primary challenge in offline RL is the evaluation of out-of-distribution (OOD) actions; offline datasets rarely offer support over the entire state-action space and neural networks overestimate values when extrapolating to OOD actions [Fujimoto et al., 2018; Gulcehre et al., 2020; Kumar et al., 2020; Kumar et al., 2019]. If trained using standard offpolicy methods, a policy will select any actions that maximize reward, which includes OOD actions. The difference between the rewards implied by the value function and the environment results in a distribution shift that can result in failure in real-world policy rollouts. Thus, offline RL algorithms must both maximize the reward and follow the behavioral policy, while having to potentially "stitch" together several suboptimal trajectories. The former requirement is usually satisfied by introducing a constraint on the actor to either penalize deviation from the behavior policy or epistemic uncertainty of the value function, or by regularizing the value function to directly minimize OOD action-values.

Many recent approaches to offline RL [Tarasov et al., 2023; Zhu et al., 2022; Li et al., 2022; Nikulin et al., 2023; Xu et al., 2023] demonstrate success in D4RL benchmarks [Fu et al., 2020], but demand the onerous task of per-dataset hyperparameter tuning [Zhang and Jiang, 2021]. Algorithms that require substantial offline fine-tuning can be infeasible in real-world applications [Tang and Wiens, 2021], hampering their adoption in favor of simpler, older algorithms [Emerson et al., 2023; Zhu et al., 2023]. These older methods [Fujimoto and Gu, 2021; Kumar et al., 2020; Kostrikov et al., 2021b] provide excellent "bang-for-buck" as their hyperparameters work well across a range of D4RL datasets.

Contributions In this paper, we show how a trained uncertainty model can be incorporated into the regularized policy objective as a behavioral supervisor to yield TD3 with behavioral supervisor tuning (TD3-BST). The key advantage of our method is the dynamic regularization weighting performed by the uncertainty network, which allows the learned policy to maximize Q-values around dataset modes. Evaluation on D4RL datasets demonstrates that TD3-BST achieves SOTA performance, and ablation experiments analyze the performance of the uncertainty model and the sensitivity of the parameters of the BST objective.

2 RelatedWork

Reinforcement learning is a framework for sequential decision making often formulated as a Markov decision process (MDP), M = {S,A,R, p, p0, γ} with state space S, action space A, a scalar reward dependent on state and action R(s, a), transition dynamics p, initial state distribution p0 and discount factor γ ∈ [0, 1) [Sutton and Barto, 2018]. RL aims to learn a policy π ∈ Π that executes action a = π(s) that will maximize the expected discounted reward J(π) = Eτ∼Pπ(τ) [∑Tt=0 γtR(st, at)] where Pπ(τ) = p0(s0) ∏Tt=0 π(at | st)p(st+1 | st, at) is the trajectory under π. Rather than rolling out an entire trajectory, a state-action value function (Q function) is often used: Qπ(s, a) = Eτ∼Pπ(τ ) [∑Tt=0 γtr(st, at) | s0 = s, a0 = a].

2.1 Offline Reinforcement Learning

Offline RL algorithms are presented with a static dataset D that consists of tuples {s, a, r, s′ } where r ∼ R(s, a) and s′ ∼ p(· | s, a).

D has limited coverage over S×A; hence, offline RL algorithms must constrain the policy to select actions within the dataset support. To this end, algorithms employ one of three approaches: 1) policy constraints; 2) critic regularization; or 3) uncertainty penalization.

Policy constraint Policy constraints modify the actor's objective only to minimize divergence from the behavior policy. Most simply, this adds a constraint term [Fujimoto and Gu, 2021; Tarasov et al., 2023] to the policy objective:

where α is a scalar controlling the strength of regularization, D(·, ·) is a divergence function between the policy π and the behavior policy πβ. In offline RL, we do not have access to πβ; some prior methods attempt to estimate it empirically

[Kostrikov et al., 2021a; Li et al., 2023] which is challenging when the dataset is generated by a mixture of policies. Furthermore, selecting the constraint strength can be challenging and difficult to generalize across datasets with similar environments [Tarasov et al., 2023; Kostrikov et al., 2021a].

Other policy constraint approaches use weighted BC [Nair et al., 2020; Kostrikov et al., 2021b; Xu et al., 2023] or (surrogate) BC constraints [Li et al., 2022; Wu et al., 2019; Li et al., 2023]. The former methods may be too restrictive as they do not allow OOD action selection, which is crucial to improve performance [Fu et al., 2022]. The latter methods may still require substantial tuning and in addition to training if using model-based score methods. Other methods impose architectural constraints [Kumar et al., 2019; Fujimoto et al., 2019] that parameterize separate BC and reward-maximizing policy models.

Critic Regularization Critic regularization methods directly address the OOD action-value overestimation problem by penalizing large values for adversarially sampled actions [Kostrikov et al., 2021a].

Ensembles Employing an ensemble of neural network estimators is a commonly used technique for prediction with a measure of epistemic uncertainty [Kondratyuk et al., 2020]. A family of offline RL methods employ large ensembles of value functions [An et al., 2021] and make use of the diversity of randomly initialized ensembles to implicitly reduce the selection of OOD actions or directly penalize the variance of the reward in the ensemble [Ghasemipour et al., 2022; Sutton and Barto, 2018].

Model-Based Uncertainty Estimation Learning an uncertainty model of the dataset is often devised analogously to exploration-encouraging methods used in online RL, but, employing these for anti-exploration instead [Rezaeifar et al., 2022]. An example is SAC-RND which directly adopts such an approach [Nikulin et al., 2023]. Other algorithms include DOGE [Li et al., 2022] which trains a model to estimate uncertainty as a distance to dataset action and DARL [Zhang et al., 2023] which uses distance to random projections of stateaction pairs as an uncertainty measure. As a whole, these methods optimize a distance d(·, ·) ≥ 0 that represents the uncertainty of an action.

2.2 Uncertainty Estimation

Neural networks are known to predict confidently even when presented with OOD samples [Nguyen et al., 2015; Goodfellow et al., 2014; Lakshminarayanan et al., 2017]. A classical approach to OOD detection is to fit a generative model to the dataset that produces a high probability for in-dataset samples and a low probability for OOD ones. These methods work well for simple, unimodal data but can become computationally demanding for more complex data with multiple modes. Another approach trains classifiers that are leveraged to become finer-grained OOD detectors [Lee et al., 2018]. In this work, we focus on Morse neural networks [Dherin et al., 2023], an approach that trains a generative model to produce an unnormalized density that takes on value 1 at the dataset modes.

3 Preliminaries

A Morse neural network produces an unnormalized density M(x) ∈ [0, 1] on an embedding space Re [Dherin et al., 2023]. A Morse network can produce a density in Re that attains a value of 1 at mode submanifolds and decreases towards 0 when moving away from the mode. The rate at which the value decreases is controlled by a Morse Kernel.

Definition 1 (Morse Kernel). A Morse Kernel is a positive definite kernel K. When applied in a space Z = Rk , the kernel K(z1, z2) takes values in the interval [0, 1] where K(z1, z2) = 1 iff z1 = z2.

All kernels of the form K(z1, z2) = e-D(z1,z2) where D(·, ·) is a divergence [Amari, 2016] are Morse Kernels. Examples include common kernels such as the Radial Basis Function (RBF) Kernel,

The RBF kernel and its derivatives decay exponentially, leading learning signals to vanish rapidly. An alternative is the ubiquitous Rational Quadratic (RQ) kernel:

where λ is a scale parameter in each kernel. The RQ kernel is a scaled mixture of RBF kernels controlled by κ and, for small κ, decays much more slowly [Williams and Rasmussen, 2006].

Consider a neural network that maps from a feature space into a latent space fϕ : X → Z, with parameters ϕ, X ∈ Rd and Z ∈ Rk . A Morse Kernel can impose structure on the latent space.

Definition 2 (Morse Neural Network). A Morse neural network is a function fϕ : X → Z in combination with a Morse Kernel on K(z, t) where t ⊂ Z is a target, chosen as a hyperparameter of the model. The Morse neural network is defined as Mϕ(x) = K(fϕ(x), t).

Using Definition 1 we see that Mϕ(x) ∈ [0, 1], and when Mϕ(x) = 1, x corresponds to a mode that coincides with the level set of the submanifold of the Morse neural network. Furthermore, Mϕ(x) corresponds to the certainty of the sample x being from the training dataset, so 1-Mϕ(x) is a measure of the epistemic uncertainty of x.

The function -logMϕ(x) measures a squared distance, d(·, ·), between fϕ(x) and the closest mode in the latent space at m:

where M is the set of all modes. This encodes information about the topology of the submanifold and satisfies the Morse-Bott non-degeneracy condition [Basu and Prasad, 2020].

The Morse neural network offers the following properties:

• 1 Mϕ(x) ∈ [0, 1].
• 2 Mϕ(x) = 1 at its mode submanifolds.
• 3 -logMϕ(x) ≥ 0 is a squared distance that satisfies the Morse-Bott non-degeneracy condition on the mode submanifolds.
• 4 As Mϕ(x) is an exponentiated squared distance, the function is also distance aware in the sense that as fϕ(x) → t, Mϕ(x) → 1.

Proof of each property is provided in the appendix.

4 Policy Constraint with a Behavioral Supervisor

We now describe the constituent components of our algorithm, building on the Morse network and showing how it can be incorporated into a policy-regularized objective.

4.1 Morse Networks for Offline RL

The target t is a hyperparameter that must be chosen. Experiments in [Dherin et al., 2023] use simple, toy datasets with classification problems that perform well for categorical t. We find that using a static label for the Morse network yields poor performance; rather than a labeling model, we treat fϕ as a perturbation model that produces an action fϕ(s, a) = ˆa such that ˆa = a if and only if s, a ∼ D.

An offline RL dataset D consists of tuples {s, a, r, s′ } where we assume {s, a} pairs are i.i.d. sampled from an unknown distribution. The Morse network must be fitted on N state-action pairs [{s1, a1, }, ..., {sN, aN}] such that Mϕ(si, aj) = 1, ∀i, j ∈ 1, ..., N ] only when i = j.

We fit a Morse neural network to minimize the KL divergence between unnormalized measures [Amari, 2016] following [Dherin et al., 2023], DKL(D(s, a) || Mϕ(s, a)):

With respect to ϕ, this amounts to minimizing the empirical loss:

where au is an action sampled from a uniform distribution over the action space Duni.

A learned Morse density is well suited to modeling ensemble policies [Lei et al., 2023], more flexibly [Dherin et al., 2023; Kostrikov et al., 2021a; Li et al., 2023] and without down-weighting good, in-support actions that have low density under the behavior policy [Singh et al., 2022] as all modes have unnormalized density value 1.

A Morse neural network can be expressed as an energybased model (EBM) [Goodfellow et al., 2016]:

Proposition 1. A Morse neural network can be expressed as an energy-based model: Eϕ(x) = e-logMϕ(x) where Mϕ : Rd → R.

Note that the EBM Eϕ is itself unnormalized. Representing the Morse network as an EBM allows analysis analogous to [Florence et al., 2022].

Theorem 1. For a set-valued function F(x) : x ∈ Rm → Rn \{∅}, there exists a continuous function g : Rm+n → R that is approximated by a continuous function approximator gϕ with arbitrarily small bounded error ϵ. This ensures that any point on the graph Fϕ(x) = argminy gϕ(x, y) is within distance ϵ of F .

We refer the reader to [Florence et al., 2022] for a detailed proof.

The theorem assumes that F(x) is an implicit function and states that the error at the level-set (i.e. the modes) of F(x) is small.

4.2 TD3-BST

We can use the Morse network to design a regularized policy objective. Recall that policy regularization consists of Q-value maximization and minimization of a distance to the behavior policy (Equation 1). We reconsider the policy regularization term and train a policy that minimizes uncertainty while selecting actions close to the behavior policy. Let C π (s, a) denote a measure of uncertainty of the policy action. We solve the following optimization problem:

This optimization problem requires an explicit behavior model, which is difficult to estimate and using an estimated model has historically returned mixed results [Kumar et al., 2019; Fujimoto et al., 2019]. Furthermore, this requires direct optimization through C π which may be subject to exploitation. Instead, we enforce this implicitly by deriving the solution to the constrained optimization to obtain a closed-form solution for the actor [Peng et al., 2019; Nair et al., 2020]. Enforcing the KKT conditions we obtain the Lagrangian:

Computing ∂L ∂π and solving for π yields the uncertainty minimizing solution πC∗ (a | s) ∝ πβ(a | s)e 1µ Cπ (s,a) . When learning the parametric policy πψ, we project the nonparametric solution into the policy space as a (reverse) KL divergence minimization of πψ under the data distribution D:

which is a weighted maximum likelihood update where the supervised target is sampled from the dataset D and C π (s, a) = 1 -Mϕ(s, πψ(s)). This avoids explicitly modeling the behavior policy and uses the Morse network uncertainty as a behavior supervisor to dynamically adjust the strength of behavioral cloning. We provide a more detailed derivation in the appendix.

Interpretation Our regularization method shares similarities with other weighted regression algorithms [Nair et al., 2020; Peng et al., 2019; Kostrikov et al., 2021b] which weight the advantage of an action compared to the dataset/replay buffer action. Our weighting can be thought of as a measure of disadvantage of a policy action in the sense of howOOD it is.

We make modifications to the behavioral cloning objective. From Morse network property 1 we know Mϕ ∈ [0, 1], hence 1 ≤ e 1µ Cπ ≤ e 1µ , i.e. the lowest possible disadvantage coefficient is 1. To minimize the coefficient in the mode, we require it to approach 0 when near a mode. We adjust the weighted behavioral cloning term and add Q-value maximization to yield the regularized policy update:

where µ is the Lagrangian multiplier that controls the magnitude of the disadvantage weight and ZQ = 1N ∑Nn=1 |Q(s, aπ)| is a scaling term detached from the gradient update process [Fujimoto and Gu, 2021], necessary as Q(s, a) can be arbitrarily large and the BC-coefficient is upper-bounded at e 1µ .

The value function update is given by:

with y = r(s, a) + γEs′∼¯π(s′ ) ¯Q(s′ , a′ ) where ¯Q and ¯π are target value and policy functions, respectively.

4.3 Controlling the Tradeoff Constraint

Tuning TD3-BST is straightforward; the primary hyperparameters of the Morse network consist of the choice and scale of the kernel, and the temperature µ. Increasing λ for higher dimensional actions ensures that the high certainty region around modes remains tight. Prior empirical work has demonstrated the importance of allowing some degree of OOD actions [An et al., 2021]; in the TD3-BST framework, this is dependent on λ. In Figure 2 we provide a didactic example of the effect of λ. We construct a dataset consisting of 2-dimensional actions in [-1, 1] with means at the four locations {[0.0, 0.8], [0.0, -0.8], [0.8, 0.0], [-0.8, 0.0]} and each with standard deviation 0.05. We sample M = 128 points, train a Morse network and plot the density produced by the Morse network for λ = { 1 10 , 12 , 1.0, 2.0}. A behavioral cloning policy learned using vanilla MLE where all targets are weighted equally results in an OOD action being selected. Training using Morse-weighted BC downweights the behavioral cloning loss for far away modes, enabling the policy to select and minimize error to a single mode.

Algorithm 1 TD3-BST Training Procedure Outline. The policy is updated once for every m = 2 critic updates, as is the default in TD3.

Input: Dataset D = {s, a, r, s′ }

Initialize: Initialize Morse networkMϕ.

Output: Trained Morse network Mϕ.

Let t = 0. for t = 1 to TM do Sample minibatch (s, a) ∼ D Sample random actions a ¯D ∼ Duni for each state s Update ϕ by minimizing Equation 6

end for

Initialize: Initialize policy network πψ, critic Qθ, target policy ¯ ψ ← ψ and target critic ¯ θ ← θ. Output: Trained policy π.

Let t = 0. for t = 1 to TAC do Sample minibatch (s, a, r, s′ ) ∼ D Update θ using Equation 14 if t mod m = 0 then Obtain aπ = π(s) Update ψ using Equation 13 Update target networks ¯ θ ← ρθ + (1 - ρ)¯ θ, ¯ ψ ← ρψ + (1 - ρ) ¯ ψ end if end for return π

4.4 Algorithm Summary

Fitting the Morse Network The TD3-BST training procedure is described in Algorithm 1. The first phase fits the Morse network for TM gradient steps.

Actor-Critic Training In the second phase of training, a modified TD3-BC procedure is used for TAC iterations with alterations highlighted in red.

We provide full hyperparameter details in the appendix.

5 Experiments

In this section, we conduct experiments that aim to answer the following questions:

• How does TD3-BST compare to other baselines, with a focus on comparing to newer baselines that use perdataset tuning?
• Can the BST objective improve performance when used with one-step methods (IQL) that perform in-sample policy evaluation?
• How well does the Morse network learn to discriminate between in-dataset and OOD actions?
• How does changing the kernel scale parameter λ affect performance?
• Does using independent ensembles, a second method of uncertainty estimation, improve performance?

We evaluate our algorithm on the D4RL benchmark [Fu et al., 2020], including the Gym Locomotion and challenging Antmaze navigation tasks.

5.1 Comparison with SOTAMethods

We evaluate TD3-BST against the older, well known baselines of TD3-BC [Fujimoto and Gu, 2021], CQL [Kumar et al., 2020], and IQL [Kostrikov et al., 2021b]. There are more recent methods that consistently outperform these baselines; of these, we include SQL [Xu et al., 2023], SAC-RND [Nikulin et al., 2023], DOGE [Li et al., 2022], VMG [Zhu et al., 2022], ReBRAC [Tarasov et al., 2023], CFPI [Li et al., 2023] and MSG [Ghasemipour et al., 2022] (to our knowledge, the best-performing ensemble-based method). It is interesting to note that most of these baselines implement policy constraints, except for VMG (graph-based planning) and MSG (policy constraint using a large, independent ensemble). We note that all the aforementioned SOTA methods (except SQL) report scores with per-dataset tuned parameters in stark contrast with the older TD3-BC, CQL, and IQL algorithms, which use the same set of hyperparameters in each D4RL domain. All scores are reported with 10 evaluations in Locomotion and 100 in Antmaze across five seeds.

We present scores for D4RL Gym Locomotion in Table 1. TD3-BST achieves best or near-best results compared to all previous methods and recovers expert performance on five of nine datasets. The best performing prior methods include SAC-RND and ReBRAC, both of which require per-dataset tuning of BRAC-variant algorithms [Wu et al., 2019].

We evaluate TD3-BST on the more challenging Antmaze tasks which contain a high degree of suboptimal trajectories and follow a sparse reward scheme that requires algorithms to stitch together several trajectories to performwell. TD3-BST achieves the best scores overall in Table 2, especially as the maze becomes more complex. VMG and MSG are the bestperforming prior baselines and TD3-BST is far simpler and more efficient in its design as a variant of TD3-BC. The authors of VMG report the best scores from checkpoints rather than from the final policy. MSG report scores from ensembles with both 4 and 64 critics of which the best scores included here are from the 64-critic variant.

We pay close attention to SAC-RND, which, among all baselines, is most similar in its inception to TD3-BST. SACRND uses a random and trained network pair to produce a dataset-constraining penalty. SAC-RND achieves consistent SOTA scores on locomotion datasets, but fails to deliver commensurate performance on Antmaze tasks. TD3-BST performs similarly to SAC-RND in locomotion and achieves SOTA scores in Antmaze.

5.2 Improving One-StepMethods

One-step algorithms learn a policy from an offline dataset, thus remaining on-policy [Rummery and Niranjan, 1994; Sutton and Barto, 2018], and using weighted behavioral cloning [Brandfonbrener et al., 2021; Kostrikov et al., 2021b]. Empirical evaluation by [Fu et al., 2022] suggests that advantageweighted BC is too restrictive and relaxing the policy objective to Equation 1 can lead to performance improvement. We use the BST objective as a drop-in replacement for the policy improvement step in IQL [Kostrikov et al., 2021b] to learn an optimal policy while retaining in-sample policy evaluation.

We reproduce IQL results and report scores for IQL-BST, both times using a deterministic policy [Tarasov et al., 2022] and identical hyperparameters to the original work in Table 3. Reproduced IQL closely matches the original results, with slight performance reductions on the -large datasets. Relaxing weighted-BC with a BST objective leads to improvements in performance, especially on the more difficult -medium and -large datasets. To isolate the effect of the BST objective, we do not perform any additional tuning.

5.3 Ablation Experiments

Morse Network Analysis We analyze how well the Morse network can distinguish between dataset tuples and samples from Dperm, permutations of dataset actions, and Duni. We plot both certainty (Mϕ) density and t-SNEs [Van der Maaten and Hinton, 2008] in Figure 3 which show that the unsupervised Morse network is effective in distinguishing between Dperm and Duni and assigning high certainty to dataset tuples.

Ablating kernel scale We examine sensitivity to the kernel scale λ. Recall that k = dim(A). We see in Figure 4 that the scale λ = k2 is a performance sweet-spot on the challenging Antmaze tasks. We further illustrate this by plotting policy deviations from dataset actions in Figure 5. The scale λ = 1.0 is potentially too lax a behavioral constraint, while λ = k is too strong, resulting in performance reduction. However, performance on all scales remains strong and compares well with most prior algorithms. Performance may be further improved by tuning λ, possibly with separate scales for each input dimension.

Independent or Shared Targets? Standard TD3 employs Clipped Double Q-learning (CDQ) [Hasselt, 2010; Fujimoto et al., 2018] to prevent value overestimation. On tasks with sparse rewards, this may be too conservative [Moskovitz et al., 2021]. MSG [Ghasemipour et al., 2022] uses large ensembles of fully independent Q functions to learn offline. We examine how independent double Q functions perform compared to the standard CDQ setup in Antmaze with 2 and 10 critics. The results in Figure 6 show that disabling CDQ with 2 critics is consistently detrimental to performance. Using a larger 10-critic ensemble leads to moderate improvements. This suggests that combining policy regularization with an efficient, independent ensemble could bring further performance benefits with minimal changes to the algorithm.

6 Discussion

Morse Network In [Dherin et al., 2023], deeper architectures are required even when training on simple datasets. This rings true for our application of Morse networks in this work, with low-capacity networks performing poorly. Training the Morse network for each locomotion and Antmaze dataset typically takes 10 minutes for 100000 gradient steps using a batch size of 1 024. When training the policy, using the Morse network increases training time by approximately 15%.

Optimal Datasets On Gym Locomotion tasks TD3-BST performance is comparable to newer methods, all of which rarely outperform older baselines. This can be attributed to a significant proportion of high-return-yielding trajectories that are easier to improve.

7 Conclusion

In this paper, we introduce TD3-BST, an algorithm that uses an uncertainty model to dynamically adjust the strength of regularization. Dynamic weighting allows the policy to maximize reward around individual dataset modes. Our algorithm compares well against prior methods on Gym Locomotion tasks and achieves the best scores on the more challenging Antmaze tasks, demonstrating strong performance when learning from suboptimal data.

In addition, our experiments show that combining our policy regularization with an ensemble-based source of uncertainty can improve performance. Future work can explore how to apply alternative uncertainty measures and how best to combine multiple sources of uncertainty.