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 [25], 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 [23, 9, 24], 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 [37, 39, 42].
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([37]).
A similar problem has been reported when explicitly constraining the policy on multi-modaldistributions using maximum mean discrepancy (MMD) distances [53].
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 actionsin the dataset.
Figure 2.a illustrates the histogram of actions sampled from a trained policy with AWAC [37], a policy constraint method.
The method fails in this setting since the learned actiondistribution includes OOD actions which result in a collision with the obstacle.
Learning a GMMpolicy with AWAC can avoid the OOD action problem but still assigns a high probability distributionover low-return actions (Figure 2.b).
Other offline RL methods, including BCQ [9] and PLAS [53],which learn a VAE policy, also fail to represent the high-return actions in this setting.
PLAS can onlymodel 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, [53], we also train a latent policy that obtains stateconditioned latent values which result in higher reward outcomes compared to latent samples directlydrawn from a prior distribution.
We compare LAPO to vanilla behavior cloning, BCQ [9], PLAS[53], AWAC [37], IQL [22], and CQL [24] on a variety of simulated locomotion, navigation, and
[Figure]
manipulation 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 ) =
89 E π [ (cid:80) ∞ t =0 γ t r t | s 0 =
s, a 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 ) =
91 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 ) .
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.
At a high-level, the overall objective of the LAPO can be represented as follows:
[Equation]
LAPO has two key components.
On the right term, it learns a latent variable state-conditioned generative model p θ ( a | s ) =
E
z ∼ p
( z ) [ p θ ( a | s, z )] to regress the actions in the dataset with an importanceweight ω . Here p ( z ) is a specified prior distribution and p θ ( a | s, z ) is a decoder.
The importanceweight ω indicates how significant an action is to solving the task, which can be computed in a similarway as other weighted-imitation learning methods [37, 39, 42].
On the left term of the objectivefunction, LAPO learns a latent policy π ϑ ( z | s ) to generate latent values that can be decoded intoin-distribution actions that maximize the state-action Q values, via the decoder p θ ′ ( a | s, z ) .
Theparameter θ ′ denotes the parameter of the decoder in earlier versions, which do not provide gradient information and therefore do not interfere with the optimization of the first component. Allcomponents are trained iteratively, and the parameter θ ′ is updated after one or more iterations. Overall, the latent variable generative model learns to represent the multi-modal distribution includedin the dataset, while actively prioritize regions containing actions that lead to higher returns throughthe importance weights.
Using the decoder, the latent policy will naturally avoid sampling OODactions, and can learn to optimize Q values without being further constrained by the data distribution,leading to better performance (in Figure. 2.d).
Next, we will describe each of these components indetail.
Maximizing the weighted log-likelihood: To ensure conservatism to the dataset while learning toadhere to high return actions.
LAPO incrementally learns to maximize the weighted log-likelihood byalternating between two steps: (1) estimating the importance weight of each action, and (2) regressingthe actions in the dataset with the importance weight of each action.
In practice, to characterizethe importance of high-return trajectories, we can use the exponential advantages as the importanceweight, ω = exp( A ( s, a ) /λ ) .
Following [20], we derive a weighted variational lower bound thatmaximizes 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 [16].
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 [28] 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]
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 4.
Update the Q-function Q ϕ using Equations 2 and 5.
Estimate the advantage weights for every state-action pair using ω = exp( Q ϕ ( s,a ) − V ( s ) λ )
Update the latent policy π ϑ using Equation
Update θ ′
according to θ until M iterations completed
A summary of our method is presented in Algorithm 1.
We randomly initialize the parameters ofthe state-action Q function, the action policy, the amortized variational distribution, and the latentpolicy.
In every training iteration, we first update the decoder and the variational distribution q using Equation 4 given the latest values of the advantage weights.
Then, we update the Q function using Equations 2 and 5 provided the updated action and latent policies.
Then, for every state-action pairin the dataset, we compute the advantage using the updated Q function, the latent policy and thedecoder, and then estimate the weights as the exponential of the scaled advantages, with the scalefactor λ .
Finally, we update the latent policy using Equation 6.
Related Works
Offline Reinforcement Learning: Offline RL methods generally address the problem of distributionshift between the behavior policy and the policy being learned [9, 23], which can cause issuesdue 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 [30, 17, 48, 23, 21, 13], via implicit policy constraints[42, 47, 40, 39, 37, 21, 22, 49, 33], by regularizing based on importance sampling [35, 29], by learningof conservative value functions [24, 44], by leveraging auxiliary behavioral cloning losses [7, 36],and through model-based training with conservative penalties [52, 19, 3, 45, 34, 27, 51].
Comparedto 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, We overcome thischallenge by introducing a new method that leverages latent variable models.
As we will find in Section 5, our method also outperforms state-of-the-art prior methods in other categories, particularlywhen learning from heterogeneous datasets.
Prior works [1, 18, 50] 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 [18, 50] 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 [43, 11, 12], enabling transfer learning [4, 10], implementing hierarchical RL[2, 31, 32, 38], avoiding distributional shift in offline RL settings [9, 53], and learning dynamicsmodels [26, 15, 41, 14].
Our proposed method resembles PLAS [53], 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 5, 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 [Figure] and homogeneous offline datasets?
(2) How does LAPO compare to prior methods implementedwith GMM policies when learning behaviors offline from heterogeneous datasets?
and (3) In whichsetting does LAPO benefit from the latent policy training, versus only using the tenerative 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 [37] 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.
5.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 [53],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 [37], which performs advantage-weightedimitation learning with a forward-KL objective for the policy projection step, (5) the IQL method[22], 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 [24] 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 [46] 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 supplementary materials. [Figure]
5.2 Tasks and Datasets
[Figure]
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 4), 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 [50], and for (2), we leverage heterogeneous datasetsfrom the standard D4RL benchmark [6].
More details about our tasks and datasets are provided inthe supplementary materials.
5.3 Experimental Results
Heterogeneous datasets: Table 1 reports the results of our experiments on heterogeneous datasets,and in Figure 3, we present the learning performance of LAPO with 95% confidence interval 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 one task.
Besides,on average, LAPO improves by 49% over the next best method on the heterogeneous datasets.
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. Surprisingly, we observe that GMM policy training does not perform well on heterogeneous offline RLsettings.
The AWAC method with GMM policies yields similar or even worse performance comparedto the original AWAC.
These results suggest that latent variable models are better candidates toeffectively model the high-reward regions of data especially when given data with multiple modes.
Random/Narrow datasets:
Table 2 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.
5.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 policy [Table] [Table] training from LAPO.
The main motivation is that, as described in Section 3, 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 5 to:
[Equation]
where, the latent z is now sampled from a fixed prior.
As shown in Fig.
5, training a latent policy does not significantly contribute to higher performancefor tasks such as the Walker2D-medium and Walker2d-mix-backward which contain sufficient taskrelevant data.
However, latent policy training results in large performance gains when learningfrom 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 6 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. [Figure] [Figure]
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.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] Our method out perform other
methods by a large margin on heterogeneous dataset, the performance on random/biased
datasets is not the best compared to other methods.
(c) Did you discuss any potential negative societal impacts of your work?
[No]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e
g., data splits, hyperparameters, how theywere chosen)?
[Yes]
(c) Did you report error bars (e
g., with respect to the random seed after running experi-
ments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e
g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[No]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]