Latent-Variable Advantage-Weighted Policy Optimization for Offline Reinforcement Learning
Anonymous Author(s)
AffiliationAddress email
Abstract
Offline reinforcement learning methods hold the promise of learning policies frompre-collected datasets without the need to query the environment for new samples.
This setting is particularly well-suited for continuous control robotic applicationsfor which online data collection based on trial-and-error is costly and potentiallyunsafe.
In practice, offline datasets are often heterogeneous , i.e., collected ina variety of scenarios, such as data from several human demonstrators or frompolicies that act with different purposes.
Unfortunately, such datasets often containaction distributions with multiple modes and, in some cases, lack a sufficientnumber of high-reward trajectories, which render offline policy training inefficient.
To address this challenge, we propose to leverage latent-variable generative modelto represent high-advantage state-action pairs leading to better adherence to datadistributions that contributes to solving the task, while maximizing reward via apolicy over the latent variable.
As we empirically show on a range of simulatedlocomotion, navigation, and manipulation tasks, our method referred to as latentvariable advantage-weighted policy optimization (LAPO), improves the averageperformance of the next best-performing offline reinforcement learning methodsby 49% on heterogeneous datasets, and by 8% on datasets with narrow and biaseddistributions.
Introduction
Offline reinforcement learning (RL), also known as batch RL [26], addresses the problem of learningan effective policy from a static fixed-sized dataset without interacting with the environment tocollect new data.
This formulation is especially important for robotics, as it avoids costly and unsafetrial-and-error and provides an alternative way of leveraging a pre-collected dataset.
However, inpractical settings, such offline datasets are often heterogeneous and are collected using differentpolicies, leading to a data distribution with multiple modes.
These data-collection policies may aim toaccomplish tasks that are not necessarily aligned with the target task or may accomplish the same taskbut provide conflicting solutions.
In contrast to the prior works that have focused on the distributionalshift problem [24, 9, 25], in this paper, we address the problem of learning from heterogeneous data.
The main challenges in learning from heterogeneous settings are the existence of conflicting actionsin the dataset, and the lack of sufficient high-return trajectories as the dataset is constructed byre-labeling state-action pairs from some different tasks.
Learning from heterogeneous datasets withconflicting actions is particularly challenging for implicit policy constraint methods that formulatea supervised learning objective function based on the forward Kullback–Leibler (KL)-divergencebetween the parametric policy being learned and the closed-form optimal policy found throughadvantage-weighted behavior cloning [39, 41, 44].
Figure 1 illustrates an example of offline heterogeneous settings in which state-conditional distributions learned directly over the action space using the forward KL-divergence objective may assign high probability to out-of-distribution (OOD)actions which results in sub-optimal policies.
Replacing the forward-KL-divergence with a reverse
KL-divergence can avoid this issue in theory, as it forces the policy to only learn a single mode ofthe distribution.
But in practice, it requires querying the behavior policy which is unknown, andusing an erroneous approximation of the behavior policy can negatively affect the performance([39]).
A similar problem has been reported when explicitly constraining the policy on multi-modaldistributions using maximum mean discrepancy (MMD) distances [57].
This problem however mayseem to be solved by using more expressive policies such as variational auto-encoders (VAEs) or
Gaussian mixture models (GMMs).
However, in practice this does not help in heterogeneous settingswith a low density of high-return trajectories.
In such settings, adhering to the data distribution usingpolicy constraint methods can result in sub-optimal performances.
We illustrate this problem with a toy navigation task in which an agent navigates to a goal positionwhile avoiding an obstacle.
In this example, the demonstrated expert actions have two prominentmodes at the initial state corresponding to moving to either the right or left side of the obstacle.
Weassume moving to the right leads to higher returns, but we have a lower density of such actions in thedataset.
In Figure 2, we show the action distribution learned by different offline RL algorithms, andin Table 1, we present their averaged performance.
Figure 2.a illustrates the histogram of actionssampled from a trained policy with AWAC [39], a policy constraint method.
The method fails inthis setting since the learned action distribution includes OOD actions which result in a collisionwith the obstacle.
Learning a GMM policy with AWAC can avoid the OOD action problem but stillassigns a high probability distribution over low-return actions (Figure 2.b).
Other offline RL methods,including BCQ [9] and PLAS [57], which learn a VAE policy, also fail to represent the high-returnactions in this setting.
PLAS can only model the low-return but high-density samples in the dataset(Figure 2.c).
Figure 1 also illustrates a potential solution to the problem of learning from heterogeneous datasets.
The intuition is to construct a state-conditioned latent space that represents high-return actions.
Inthis case, we can learn simple state-conditional distributions such as Gaussian distributions overthe latent space which, as shown in Figure 2.d, can capture task-relevant and high-reward actionswithout including out-of-distribution samples.
Based on this intuition, we propose to learn a latentspace policy by alternating between training the policy and maximizing the advantage-weightedlog-likelihood of data.
This biases the RL policy to choose actions supported by the training datawhile effectively solving the target task.
The main contributions of this work is the introduction of a new method, which we refer to it as latentvariable advantage-weighted policy optimization (LAPO), that can efficiently solve heterogeneousoffline RL tasks.
LAPO learns an advantage function and a state-conditional latent space in whichhigh-advantage action samples are generated by sampling from a prior distribution over the latentspace.
Furthermore, following a prior work, [57], we also train a latent policy that obtains stateconditioned latent values which result in higher reward outcomes compared to latent samples directly
[Figure]
[Table] drawn from a prior distribution. We compare LAPO to vanilla behavior cloning, BCQ [9], PLAS[57], AWAC [39], IQL [23] and CQL [25] on a variety of simulated locomotion, navigation andmanipulation tasks, provided heterogeneous offline datasets and also biased datasets with narrowdata distribution in standard offline RL benchmarks.
LAPO is the only method that yields goodperformance across all of these tasks and on average improves by 49% over the next best method forheterogeneous datasets, and by 8% on other offline RL tasks with narrow data distributions.
Preliminaries
The goal of reinforcement learning is to obtain a policy that maximizes a notion of accumulatedreward for a task as a Markov decision process (MDP).
In offline RL settings, we assume that weare given a dataset of tuples D
= { ( s t , a t , r t , s t +1 ) } , where, s t ∈ S and
a t ∈
A denote the stateand action at the time step t , r t = r
( s t , a t ) is the reward given by a bounded reward function, ands t +1 is sampled from a fixed transition probability distribution p ( s t +1 | s t , a t ) .
Note that actions cancome from a mixture of policies which is referred to as the behavior policy.
The policies may try toaccomplish different tasks and hence generate trajectories unrelated to the target task.
Given an initial state distribution µ 0 ( s ) and a discount factor γ ∈ (0 , 1) , the RL agent optimizes apolicy π ( a | s ) to maximize the expected cumulative reward
[Equation]
The state action value function (Q function), of the policy is defined as Q π ( s, a ) =
91 E π [ [Equation] ∞ t =0 γ t r t | s 0 = s, [Equation]
0 = a ] , which can be approximated with the Bellman equation:
[Equation]
where, ϕ denotes the parameters of the Q-function, V denotes the value function V ( s ) =
93 E π ( a | s ) [ Q ϕ ′ ( s, a )] and can be approximated by sampling actions from the policy distribution andaveraging the corresponding Q value.
ϕ
′ denotes the parameters of the Q-function at the previousiteration.
The advantage of a pair of state and action is then defined as A ( s, a ) =
Q ϕ ( s, a ) − V ( s ) .
Implicit Policy Constraints with KL
divergence
A central challenge for offline RL methods is to limit the state-conditional action distribution of thelearned policy to the empirical conditional action distribution of the dataset.
It is very likely foroff-policy algorithms to have overly optimistic values, as they tend to query the Q-function for action inputs outside the training distribution.
Policy constraints methods [24, 51] deal with this problemby keeping the learned policy close to the behavior policy using the following constrained policyoptimization formulation:
[Equation]
where, π β ( a | s ) denotes the unknown empirical conditional action distribution of the dataset, and Ddenotes a distance measure such as KL-divergence [39, 44] or maximum mean discrepancy (MMD)[24], and ϵ is a threshold parameter.
In case, we use the KL-divergence as the divergence in Equation3, the optimal π ∗ can be expressed as
[Equation]
where, λ is a temperature parameter that depends on the ϵ .
Prior work [39, 44, 41] suggested to incrementally solve Equation 4 by representing the optimalpolicy π ∗ ( a | s ) as a non-parametric policy, and then project it onto the parametric policy π θ ( a | s ) viasupervised regression:
[Equation]
where, the expectation is estimated by sampling from the dataset.
In practice, the non-parametricpolicy can be implemented by weighting samples from the dataset D using importance weightsω =
exp
( A ( s, a ) /λ ) , and the projection via KL-divergence can be done via weighted supervisedlearning using these weights [29].
Latent-Variable Advantage-Weighted Policy Optimization
In this work, we aim to address the problem of learning from heterogeneous offline datasets collectedby policies with different purposes.
The challenge arises from the fact that the dataset containsconflicting actions distributed with multiple modes.
Besides, the dataset often does not includea sufficient number of high-return trajectories to learn the task properly, as the dataset may beconstructed by re-labeling rewards for state-action pairs collected for some other tasks that notnecessarily aligned with the target task.
We argue that in this setting we not only need a moreexpressive policy class to represent the multi-modal action distributions, but also need a mechanismthat can select which samples to use for the given target task.
We present a novel method, named latent-space advantage-weighted policy training (LAPO), whichlearns a policy by maximizing the Q function over a latent space that is trained to approach thedistribution of actions that leads to higher returns in solving the task.
The overall objective of the LAPO can be defined as follows:
[Equation]
The objective consists of two main components: (1) an advantage-weighted log-likelihood of theactions in the dataset and, (2) a state-action value function. This objective is optimized iteratively byalternating between maximizing each component separately.
The weighted log-likelihood is maximized by training a latent-variable state-conditioned generative model p θ ( a | s ) =
E
z ∼ p
( z ) [ p θ ( a | s, z )] . Here, z denotes a latent variable, p ( z ) is the prior distribution, p θ ( a | s, z ) is a generative model(decoder), and θ ′ denotes the parameters of the decoder in the previous iteration.
This componentextends advantage-weighted regression (AWR, Equation 5) [41] using more expressive policy classes. Also, as illustrated in Figure 2c and d, weighting by advantages allows the generative model toeffectively represent high-return state-action transitions.
The state-action value function Q ϕ ( s, a ) isupdated in every iteration, and is then improved by updating the latent policy π ϑ ( z | s ) . The latentpolicy outputs a latent variable z which is then mapped to an action a ∈ A by the decoder p ′
θ ( a | s, z )with the parameters updated in the previous iteration.
The latent policy is trained with a standardactor-critical RL approach to maximize the state-action value function.
Next, we will describe eachof these components in detail.
Maximizing the weighted log-likelihood: LAPO incrementally learns to maximize the weightedlog-likelihood by alternating between two steps: (1) estimating the importance weight of each action, Input: D = { ( s t , a t , s t +1 ) i } , i = 1 , . . . , N . Initialize: ϕ, θ, ψ, ϑ randomly, θ ′ =
θ , and ω = 1 . 0 .
repeat
Update the decoder p θ using Equation 7. Update the Q-function Q ϕ using Equations 2 and 8. Estimate the advantage weights for every state-action pair. Update the latent policy π ϑ using Equation
Update θ ′
according to θ until M iterations completed and (2) regressing the actions in the dataset with the importance weight of each action.
In practice, tocharacterize the importance of high-return trajectories, we can use the exponential advantages as theimportance weight, ω = exp( A ( s, a ) /λ ) , where λ is a temperature parameter.
Following [21], wederive a weighted variational lower bound that maximizes the weighted log-likelihood of data:
[Equation]
where q ψ ( z | s, a ) is the variational posterior distribution, p ( z ) is the prior distribution over the latentvariable and typically modelled by a standard normal distribution, and β is a hyper-parameter tobalance the two loss terms [17].
Policy evaluation: As explained earlier, the estimated advantage of each state-action pair is used toupdate the decoder.
These advantage estimations are incrementally found through a policy evaluationprocess in which the overall policy, formed by combining the latent policy and the decoder, is beingevaluated.
The state-conditional action distribution of the overall policy π ϑ,θ ′ ( a | s ) , is found by firstsampling a latent value from the latent policy z ∼ π ϑ ( z | s ) , and then feeding the latent value to thedecoder to sample the action a ∼ p θ ′
( a | s, z ) .
In the policy evaluation step, we update the valuefunction by minimizing the squared temporal difference error using Equation 2.
To calculate theexpected value of a state, we sample actions from the overall policy, i.e.,
[Equation]
Since the latent values z will be used as the input to the decoder, an unbounded z can lead togenerating out-of-distribution actions.
Therefore, to better ensure sampling in-distribution actions,π ϑ ( z | s ) can be modeled as a truncated Gaussian [5].
For deterministic versions of the latent policy,the policy output can be limited to [ − z max , z max ] using a tanh function.
Policy improvement: LAPO optimizes the latent policy π ϑ ( z | s ) in every iteration to directlymaximize the return.
It is trained based on standard RL approaches such as DDPG [30] or TD3 [8] tomaximize the RL objective in Equation 1, by learning latent actions which result in high returns afterbeing converted to the original action space using the decoder.
The latent policy is updated by maximizing the Q-function over states sampled from the dataset,latent variables from the latent policy, and actions from the decoder p θ ′ :
[Equation]
A summary of our method is presented in Algorithm 1.
We randomly initialize the parameters of thestate-action Q function, the decoder, the amortized variational distribution, and the latent policy.
Inevery training iteration, we first update the decoder and the variational distribution q using Equation 7given the latest values of the advantage weights.
Then, we update the Q function using Equations 2and 8 provided the updated action and latent policies.
Then, for every state-action pair in the dataset,we compute the advantage using the updated Q function, the latent policy and the decoder, and thenestimate the weights.
Finally, we update the latent policy using Equation 9.
Related Works
Offline Reinforcement Learning: Offline RL methods generally address the problem of distributionshift between the behavior policy and the policy being learned [9, 24], which can cause issues due to out-of-distribution actions sampled from the learned policy and passed into the learnedcritic.
To address this issue, prior methods constrain the learned policy to stay close to the behaviorpolicy via explicit policy regularization [32, 18, 51, 24, 22, 13], via implicit policy constraints[44, 49, 42, 41, 39, 22, 23, 53, 35], by regularizing based on importance sampling [37, 31, 50, 52],by learning of conservative value functions [25, 46], by leveraging auxiliary behavioral cloninglosses [7, 38], and through model-based training with conservative penalties [56, 20, 3, 47, 36, 28, 55].
Compared to these prior works, we opt to develop a method that uses an implicit policy constraint.
However, prior implicit policy constraint methods have been largely limited to Gaussian policies,leading to OOD action samplings and sub-optimal performance on heterogeneous datasets, Weovercome this challenge by introducing a new method that leverages latent variable models.
As wewill find in Section 6, our method also outperforms state-of-the-art prior methods in other categories,particularly when learning from heterogeneous datasets.
Prior works [1, 19, 54] also studied the problem of learning from heterogeneous dataset in theoffline RL with heterogeneous dynamics setting [1] and multi-task offline RL with data sharingscenario [19, 54] respectively.
Both settings require knowledge of the agent/task identifier.
Here, westudy a more general problem setting including datasets with conflicting actions without assumingaccess to the ground-truth identifier of the source of the heterogeneity.
RL with Generative Models: Generative models have been used by prior works for improvingtraining performance [45, 11, 12], enabling transfer learning [4, 10], implementing hierarchical RL[2, 33, 34, 40], avoiding distributional shift in offline RL settings [9, 57], and learning dynamicsmodels [27, 16, 43, 15].
Our proposed method resembles PLAS [57], which learns a generativemodel and latent policy; but, unlike this prior work, our method trains an advantage-weightedgenerative model by alternating between learning the generative model, the advantage function,and the latent policy.
This allows LAPO to capture different high-reward solutions using a simple
Gaussian distribution in the latent space.
As we will find in Section 6, this distinction is crucial, as
LAPO significantly outperforms PLAS on heterogeneous datasets.
Experiments
Our experiments aim to answer the following questions: (1) How does LAPO compare to otheroffline RL methods on a set of standard offline RL tasks, including learning from heterogeneous andhomogeneous offline datasets?
(2) How does LAPO compare to prior methods implemented with GMM policies when learning behaviors offline from heterogeneous datasets?
and (3) In which settingdoes LAPO benefit from the latent policy training, versus only using the generative model?
Can LAPO learn without constraining the latent values?
We study the first question by evaluating LAPO and comparing it to several prior methods onoffline RL benchmarks with heterogeneous datasets.
We consider a range of simulated robotictasks, including navigation, locomotion, and manipulation, each with a corresponding static datasetwith a heterogeneous data distribution.
We also consider datasets with narrow and biased datadistributions, containing near-optimal or random trajectories.
To answer the second question, weinclude a comparison of AWAC [39] with a GMM policy, since AWAC trains using a similaradvantage-weighted objective as the actor policy in LAPO.
This comparison is only performed onheterogeneous datasets, since this is where we expect to see the most improvement from a multi-modalpolicy.
Finally, the questions in (3), we conduct two ablation studies in which we train LAPO withoutlearning the latent policy, i.e. evaluating only LAPO’s action policy, and without limiting the latentvalues to the range [ − z max , z max ] .
We compare this version of the method to the original LAPO onfour offline RL tasks that differ in the amount of high-performing task-relevant data contained in theoffline dataset.
6.1 Comparisons
We compare LAPO to 7 offline RL approaches: (1) vanilla behavioral cloning (BC), (2) the BCQmethod [9], which approximates the dataset distribution with a generative model and manually selectsthe action with the maximum Q value among a set of generated actions, (3) the PLAS method [57],which learns a generative model via maximum likelihood and a latent-space policy to maximize the
RL objective in the latent space, (4) the AWAC method [39], which performs advantage-weightedimitation learning with a forward-KL objective for the policy projection step, (5) the IQL method [Figure] [23], which learns the value function using expectile regression without an explicit policy, and thenextracts the policy using advantage weighted regression (AWR), (6) the CQL method [25] whichlearns a conservative Q-function, and (7) AWAC (GMM) which trains GMM policies using AWAC.
In our experiments, we train each method three times using three different random seeds each for onemillion steps, and report the return averaged over 10 test episodes from the trained policies.
We usedthe implementation provided by [48] for the BCQ and AWAC, the original implementation for the
PLAS, IQL and CQL, and our implementation for the BC method.
We report the implementationdetails in the Appendix C.
6.2 Tasks and Datasets
We compare LAPO to the introduced prior methods on learning from both heterogeneous andhomogeneous offline datasets.
We evaluate the methods on three simulated task domains, locomotion,navigation, and manipulation domains (Figure 3), with heterogeneous data distributions, and alsoa locomotion task with random, narrow and biased data distribution.
We consider two sources ofheterogeneous datasets: (1) data collected by policies accomplishing different tasks, and (2) datacollected by different policies that accomplish one task but in different ways.
For (1), we use themulti-task heterogeneous datasets introduced by [54], and for (2), we leverage heterogeneous datasetsfrom the standard D4RL benchmark [6].
More details about our tasks and datasets are provided inthe Appendix D.
6.3 Experimental Results
Heterogeneous datasets: Table 6 reports the results of our experiments on heterogeneous datasets,and in Figure 6 (Appendix B), we present the learning performance of LAPO with 95% confidenceinterval on Heterogeneous tasks.
Among the 12 tasks introduced in the previous section, LAPO(our method) achieves the best performance on nine tasks and the second-best performance on onetask.
Besides, on average, LAPO improves by 49% over the next best method on the heterogeneousdatasets.
The heterogeneous locomotion tasks can be accomplished more easily by the prior methods since still33% of the data comes from the target task.
However, from the prior works, AWAC and CQL are theonly methods that can accomplish the tasks.
LAPO performs well on all tasks and on average yieldsthe best performance on one out of the three tasks.
This shows that LAPO is capable of learning fromheterogeneity introduced by multi-task datasets.
The navigation tasks are in general more challenging especially for those with medium and large mapsizes.
The main challenge is to learn policies for long-horizon planning from datasets that do notcontain optimal trajectories, and there are only very few states with rewards.
LAPO significantlyoutperforms all prior works for all of the navigation tasks.
Similarly, the manipulation tasks arealso challenging, as they require the assembly of sub-trajectories related to completing a given taskconsisting of multiple sub-tasks.
In addition, the agent has access to fewer training samples whilehaving to learn to interact with complicated dynamics.
LAPO outperforms previous approaches bya large margin on Kitchen-partial and Kitchen-mixed for which only a few optimal trajectories areprovided by the dataset.
It also performs competitively on the Kitchen-complete task and yields thesecond-best performance.we observe that GMM policy training does not perform well on heterogeneous offline RL settings.
The AWAC method with GMM policies yields similar or even worse performance compared to theoriginal AWAC.
These results suggest that the VAE is more robust to model the high-reward regionsin heterogeneous data. [Table] [Table] [Table]
Random/Narrow datasets:
Table 7 reports the results of the locomotion tasks with random, narrowand biased datasets.
Our method yields the best performance for four out of nine tasks.
It alsoachieves competitive performance on the rest of the tasks, showing that LAPO is applicable to generaloffline RL settings and is not limited to learning from heterogeneous datasets.
However, the mostsignificant gain is when it is applied to offline settings with heterogeneous datasets.
Learning curvesand results with 95%-confidence interval are reported in Appendix B.
6.4 Ablations
Latent Policy Training :
To study the importance of the latent policy training in achieving goodperformances, we conduct an ablation study on 4 different tasks by eliminating the latent policytraining from LAPO.
The main motivation is that, as described in Section 4, the latent variablegenerative model is trained to incrementally learn the action distribution that lead to higher return,hence training the latent policy on top of that may seem redundant.
In this case, another option forsampling latent variables is to sample from the prior distribution, which changes Equation 8 to:
[Equation]
where, the latent z is now sampled from a fixed prior.
As shown in Fig.
4, training a latent policy does not significantly contribute to higher performancefor tasks such as the Walker2D-medium and Walker2d-mix-backward which contain sufficient task- [Figure]
[Figure] relevant data. However, latent policy training results in large performance gains when learning from datasets which contain fewer optimal task-related data such as Maze2D and Antmaze.
LAPOovercomes this challenge by explicitly optimizing the RL objective through learning the latent policy.
Limiting the latent values : We also study the importance of limiting the latent variable z to bewithin the range [ − z max , z max ] .
This is important to understand the limits of using the latent variabledecoder as a generative model of in-distribution actions.
Figure 5 illustrates the result of LAPOpolicy training using unbounded latent actions on four tasks.
The policy training performance issignificantly worse when we do not limit the latent values.
This suggests that the the action policycan generate in-distribution actions only when in-distribution latent values (close to samples drawnfrom the prior distribution p ( z ) ) are given as the input.
Conclusion
In this paper, we study an offline RL setup for learning from heterogeneous datasets where trajectoriesare collected using policies with different purposes, leading to a multi-modal data distribution, andin some cases, do not contain sufficient high-reward trajectories.
Through empirical analysis, wefind that in such cases, policies constrained methods may contain out-of-distribution actions andlead to suboptimal performance for continuous control tasks, especially for heterogeneous datasets.
To address this challenge, we present the latent-variable advantage-weighted policy optimization(LAPO) algorithm, which learns a latent variable model that generates high-advantage actions whensampling from a prior distribution over the latent space.
In addition, we train a latent policy thatobtains state-conditioned latent actions which result in higher reward outcomes compared to samplingfrom the prior distribution.
We compare our method to 6 prior methods on a variety of simulatedlocomotion, navigation, and manipulation tasks provided heterogeneous offline datasets, and also onstandard offline RL benchmarks with narrow and biased datasets.
We find that our proposed methodconsistently outperforms prior methods by a large margin on tasks with heterogeneous datasets, whilebeing competitive on other offline RL tasks with narrow data distributions.
For our future work,we will extend LAPO to multi-task offline RL settings in which an agent learns multiple RL tasksprovided a heterogeneous dataset of diverse behaviors.