LAPO: Latent-Variable Advantage-Weighted Policy Optimization for Offline Reinforcement Learning
Authors: Xi Chen
Ali Ghadirzadeh
Tianhe Yu
Jianhao Wang
Yuan Gao
Wenzhe Li
Liang Bin
Chelsea Finn
Chongjie Zhang
Venue: NeurIPS 2022 Accept (2022)
Keywords: Offline reinforcement learning
imitation learning

Abstract: Offline reinforcement learning methods hold the promise of learning policies from pre-collected datasets without the need to query the environment for new samples. This setting is particularly well-suited for continuous control robotic applications for which online data collection based on trial-and-error is costly and potentially unsafe. In practice, offline datasets are often heterogeneous, i.e., collected in a variety of scenarios, such as data from several human demonstrators or from policies that act with different purposes. Unfortunately, such datasets often contain action distributions with multiple modes and, in some cases, lack a sufficient number of high-reward trajectories, which render offline policy training inefficient. To address this challenge, we propose to leverage latent-variable generative model to represent high-advantage state-action pairs leading to better adherence to data distributions that contributes to solving the task, while maximizing reward via a policy over the latent variable. As we empirically show on a range of simulated locomotion, navigation, and manipulation tasks, our method referred to as latent-variable advantage-weighted policy optimization (LAPO), improves the average performance of the next best-performing offline reinforcement learning methods by 49\% on heterogeneous datasets, and by 8\% on datasets with narrow and biased distributions.

LAPO: Latent-Variable Advantage- Policy Optimization for Offline Reinforcement Learning
Xi Chen', Ali Ghadirzadeh", Tianhe Yu², Jianhao Wang', Yuan Gao Wenzhe Li", Bin Liang', Chelsea Finn², Chongjie Zhang' Tsinghua University, China, Stanford University, USA 3 Shenzhen Institute Artificial Intelligence and Roboties for Society, China
Abstract
Offline reinforcement learning methods hold the promise of learning policies from pre-collected datasets without the need to query the environment for new samples. This setting is particularly well-suited for continuous control robotic applications for which online data collection based on trial-and-error is costly and potentially unsafe. In practice, offline datasets are often heterogeneous, i.c., collected variety of scenarios. such data from several human demonstrators or from policies that act with different purposes. Unfortunately, such datasets often contain action distributions with multiple modes and, in some cases, lack sufficient number of high-reward trajectories, which render offline policy training inefficient. To address this challenge, propose to leverage latent-variable generative model to represent high-advantage state-action pairs leading to better adherence to data distributions that contributes solving the task, while maximizing reward via policy over the latent variable. As we empirically show on a range of simulated locomotion, navigation, and manipulation lasks, our method referred as latent- variable advantage-we weighted policy optimization (LAPO), improves the average performance of the next persfoeeforming offline reinforcement learning methods by 49% on heterogeneous datasets, and by 8% on datasets with narrow and biased distributions.
Introduction
Offline reinforcement learning (RL), also known as batch RL [26]. addresses the problem of learning an effective policy from static fixed-sized dataset without interacting with the environment to collect new data. This formulation is especially important for robotics, it avoids costly and unsafe trial-and-eroor and provides an alternative way of leveraging a pre-collected dataset. However, in practical settings, such offline datasets often heterogeneous and are collected using different policies, leading to a data distribution with multiple modes. These data-collection policies may aim to accomplish tasks that but provide conflicting solutions. In contrast to the prior works that have focused on the distributional shift problem (24 9 (25), this paper, we address the problem of learning from heterogeneous data.
necessarily aligned with the target task may accomplish the same task
The main challenges in learning from heterogeneous settings are the existence of conflicting actions in the dataset, and the lack of sufficient high-return trajectories as the dataset is constructed by re-labeling state-action pairs from some different tasks. Learning from heterogeneous datasets with conflicting actions is particularly challenging for implicit policy constraint methods that formulate a supervised learning objective function based on the forward Bullback-Leiblle (KL)-divereence between the parametric policy being learned and the closed-form optimal policy found through advantage-weighter behavior cloning (39 41 44 Figure T illustrates an example of offline het- erogeneous settings in which state-conditional distributions learned directly over the action space
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Learning from heterogenous datasets
Action Space Policy Leorning
An cut-of-disttiutoon sctimn smppe
Pashing
Palaxe
Latent Variable Advaninge Weighted Policy Optimization (LAPO)
Grasplng
Openinng Box
Afruango
Succo
VAE kiest space
Advastgg-nigtiaal lakasi qUAt
Task-immi actions
Task-mlevami actions
Prior dierrbution
Posterior distribution
Target Tisk: Object relocation
Trak-imeleent Istent actions
Task-rllearr latent actions
An example scenario of learning from data with heterogeneous action distributions. (left) The dataset includes data from three different tasks, pushing, grasping, and opening boxes. (right) The learned task is to relocate an object to a goal position. (middle up) Learning policies directly in the action space can result in sampling from fail the learning task. (middle bottom) LAPO constructs a state-conditional latent space and learns a latent policy only select in distribution high-return actions.
Figure 1:
actions which subsequently can
using the forward KL-divergence objective may assign high probability to out-of-distribution (OOD) actions which results in sub-optimal policies. Replacing the Sewaar-KLLdivvegeee with reverse KL-divergence can avoid this issue theory, as it forces the policy only learn single mode of the distribution. But in practice, it requires querying the behavior policy which is unknown, and using an erroneous approximation of the behavior policy can negatively affect the performance (339). A similar problem has been reported when explicitly constraining the policy on multi-modal distributions using maximum mean discrepancy (MMD) distances [57]. This problem however may seem 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 heterogeneous settings with low density of high-return trajectories. In such settings, adhering to the data distribution using policy 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 position while avoiding an obstacle. In this example, the demonstrated expert actions have two prominent modes at the initial state corresponding to moving to either the right or left side of the obstacle. We assume moving to the right leads to higher returns, but we have a lower density of such actions the dataset. In Figure we show the action distribution learned by different offline RL algorithms, and in Table [ we present their averaged performance. Figure illustrates the histogram of actions sampled from trained policy with AWAC [39], a policy constraint method. The method fails in this setting since the learned action distribution includes OOD actions which result collision with the obstacle. Learning GMM policy with AWAC can avoid the OOD action problem but still assigns a high probability distribution over low-return actions (Figure 2
b). Other offline RL methods, including BCQ [9] and PLAS (57), which learn VAE policy, also fail to represent the high-return actions in this setting. PLAS can only model the low-return but high-density samples in the dataset (Figure c).
Figure 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. In this case, we can learn simple state-conditional distributions such as Gaussian distributions over the latent space which, as shown in Figure dd, can capture task-relevant and high-reward actions without including out-of-distribution samples. Based on this intuition, we propose to learn a latent- space policy by alternating between training the policy and maximizing the advantage-weighter log-like ligeliiel of data. This biases the RL policy to choose actions supported by the training data while effectively solving the target task.
The main contributions of this work is the introduction of a new method, which we refer to it as latent- variable advantage- weighted policy optimization (LAPO), that can efficiently solve heterogeneous offline RL tasks. LAPO learns an advantage function and a state-conditional latent space in which high-advantage action samples are generated by sampling from a prior distribution over the latent space. Furthermore, following a prior work, (57), we also train latent policy that obtains state- conditioned latent values which result in higher reward outcomes compared to latent samples directly

---

600
Datenet
Ortaget
400
Detaset policy
400
Dotaset
LARO, jahent polrc
Actions with
Actions with Shigh returss
300
300
low returnd
200
200
100
100
(3)
(b)
(c)
The histogram of actions for the toy navigation task at the initial state in the offline dataset (in blue), and the histogram of actions sampled from (a) the AWAC Gaussian policy (in red), (b) the AWAC GMM policy (in red), (c) the PLAS generative model (in yellow) and the PLAS latent policy (in pink), and (d) the LAPO generative model (in green), and the LAPO latent policy (in pink).
Figure 2:
Table 1: The performance of policies the toy navigation task learned by AWAC, PLAS and LAPO. The return is averaged by 1, 000 samples. The agent receives reward of +1 to move left, a E [-1.5, 0.5), + +10 to move right, a [0.5, 1.5), and 0 otherwise.
AWAC Gaussian
PLAS
LAPO
GMM
VAE
policy
VAE
latent policy
Average return
2.43
4.39
1.03
1.00
4.24
10.00
drawn from prior distribution. We compare LAPO to vanilla behavior cloning, BCQ 19, PLAS [57], AWAC [39], IQL [23], and CQL [25] on a variety of simulated locomotion, navigation, and manipulation tasks, provided heterogeneous offline datasets and also biased datasets with narrow data distribution in standard offline RL benchmarks. LAPO is the only method that yields good performance across all of these tasks and on average improves by 49% over the next best method for heterogeneous datasets, and by 8% on other offline RL tasks with narrow data distributions.
2 Preliminaries
The goal of reinforcement learning is obtain policy that maximizes a notion of accumulated reward for a task as a Markov decision process (MDP). In offline RL settings, we assume that we {(st, dr Ft, St+1)]. where, St E S and a, E denote the state
are given a dataset of tuples and action at the time step f. = r(st, a.) is the reward given by bounded reward function, and 8141 is sampled from fixed transition probability distribution p(5t=1 at). Note that actions can come from a mixture of policies which is referred to as the behavior policy. The policies may try to accomplish different tasks and hence generate trajectories unrelated the target task.
Given an initial state distribution (s) and a discount factor y E (0, 1), the RL agent optimizes a policy क (als) to maximize the expected cumulative reward
J (T) E 2o (p(lssstt))))
Σ (r((n
The state action value function (Q function), of the policy is defined Q* E, 1> 7'rel = a0 = a). which can be approximated with the Bellman equation:
6² argmin B( [ra + (V(()))) - (s)()))))))
where, o denotes the parameters of the Q-function, V denotes the value function V(s) En(a]a (s, a]] and can be approximated by sampling actions from the policy distribution and averaging the corresponding Q value. o' denotes the parameters of the Q-function at the previous iteration. The advantage of pair of state and action then defined as A(s, a) Qs(s,a) vs)...
3 Implicit Policy Constraints with KL-divergence
A central challenge for offline RL methods is to limit the state-conditional action distribution of the learned policy the empirical conditional action distribution of the dataset. is very likely for off- off-policy algorithms to have overly optimistic values, as they tend to query the Q-function for action
3

---

inputs outside the training distribution. Policy constraints methods 22,51 deal with this problem by keeping the learned policy close to the behavior policy using the following constrained policy optimization formulation:
argmax EDD E [ s... EDn mp(als)) <<
(3)
where, (ala)) denotes the unknown empirical conditional action distribution of the dataset, and D
denotes a distance measure such as KL-ddvergence [39 44 or maximum mean discrepancy (MMD)
24 and is threshold parameter. In case, we use the KI.-divergence the divergence in Equation
3 the optimal x can be expressed as
n" (als) (a(as) xxp(A(s,))AA),
where, À a temperature parameter that depends on the e.
suggested to incrementally solve Equation 4 by representing the optimal policy n² (a s) as non- parammetree policy, and then project it onto the parametric policy (e(a)s) via supervised regression:
Prior work 39
argmin ED|DKL ( (alss To (a|s))) = arginaxpp (ala) [log Te [e(ass)]],
(5)
where, the expectation is estimated by sampling from the dataset. In practice, the non-parametric policy can be implemented by weighting samples from the dataset D using importance weights = (s,, A)) and the projection via KL-divergence can be done via weighted supervised learning using these weights (29].
4 Variable Advantage- Weighted Policy Optimization
In this work, we aim to address the problem of learning from heterogencous offline datasets collected by policies with different purposes. The challenge arises from the fact that the dataset contains conflicting actions distributed with multiple modes. Besides, the dataset often does not include a sufficient number of high-return trajectories to learn the task properly, as the dataset may be constructed by re-labeling rewards for state-action pairs collected for some other tasks that not necessarily aligned with the target task. We argue that this setting we not only need a more expressive policy class to represent the multi-modal action distributions, but also need mechanism that can select which samples to use for the given target task.
We present a novel method, named latent-space advantage-weighhee learns a policy by maximizing the Q function over a latent space that trained to approach the distribution of actions that leads to higher returns in solving the task. The overall objective of the LAPO can be defined as follows:
policy training (LAPO). which
argmax
E 2))0((1)) [Qo(s,a) (aGa))D ( * (als))]
0,0
am (18,5)
(6)
DnD
The objective consists of two main components: (1) advantage-weighted log-likelihood of the actions in the dataset and, (2) a state-action value function. This objective is optimized iteratively by alternating between maximizing each component separately. The weighted log-likelihood is maxi- mized by training a latent-variable state-conditioned generative model pa (a|s) = Ez-n(=) [Po(a) s, =)]. Here, z denotes latent variable, p(=) is the prior distribution, po(a|s, 88 =) is a generative model (decoder), and 0' denotes the parameters of the decoder in the previous iteration. This component extends advintage-weighted regression (AWR, Equation 4T using more expressive policy classes. Also, as illustrated in Figure 2 and d. weighting by advantages allows the generative model to effectively represent high-return state-action transitions. The state-action value function Qo a) is updated in every iteration, and is then improved by updating the latent policy (z(ss. The latent policy outputs a latent variable 2 which is then mapped to an action a E A by the decoder (a) s, with the parameters updated in the previous iteration. The latent policy is trained with standard actor-critical RL approach to maximize the state-action value function. Next, we will describe each of these components in detail.
Maximizing the weighted log-likelihood: LAPO incrementally learns to maximize the weighted ligeeiiodoood by alternating between two steps: (1) estimating the importance weight of each action,

---

Algorithm Latent- Variable advantage-Weighted Policy Optimization
Input:
D {(t,, at, St41),,, 1, I,...,N. N.
Initialize: Φ, 0, v,, 0 randomly, 6' = 8, and w = 1.0. repeat
Update the decoder pe using Equation 7
Update the Q-function Qs using Equations and
Estimate the advantage weights for every state-action pair.
Update the latent policy my using Equation Update 0' according to until M iterations completed
and (2) regressing the actions the dataset with the importance weight of each action. In practice, to use the exponential advantages the
characterize the importance of high-return trajectories, importance weight, w = exp(AA a)) where À is temperature parameter. Following (21], we derive a weighted variational lower bound that maximizes the weighted log-likelihood of data:
max / R. qo log(po (a|s, z)) (y (zls, a) [[)())
where qu (2|8, a) is the variational posterior distribution, p(2) is the prior distribution over the latent variable and typically modelled by standard normal distribution, and / is hyper-pante prper-prraeete to balance the two loss terms [17].
Policy evaluation: As explained earlier, the estimated advantage of each state-action pair is used to update the decoder. These advantage estimations are incrementally found through a policy evaluation process in which the overall policy, formed by combining the latent policy and the decoder, is being evaluated. The state-conditional action distribution of the overall policy To, (a|s). is found by first sampling a latent value from the latent policy z To (2(2|s), and then feeding the latent value the decoder to sample the action ~ pø (a|s, z). In the policy evaluation step, we update the value
function by minimizing the squared temporal difference error using Equation 2 To calculate the expected value of a state, we sample actions from the overall policy, i.e.,
V(s) B- (116,)).)~0 [Q (s.
Since the latent values will be used as the input to the decoder, an unbounded can lead to generating out-of-distribution actions. Therefore, to better ensure sampling in-distribution actions, (=(s))) can be modeled as truncated Gaussian [5]. For deterministic versions of the latent policy,
the policy output can be limited to | 2maax: Zmax using a tanh function.
Policy improvement: LAPO optimizes the latent policy Th (=(ss) in every iteration to directly maximize the return. trained based on standard RI. approaches such as DDPG (30 or TD3 X to maximize the RI. objective in Equation] [ by learning latent actions which result in high returns after being 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 pot
argmax Eapppr (-19,)), (aleea/alll)).ae. aD (ooii-i@@(ee QQ s)].
A summary of our method is presented in Algorithm [ We randomly initialize the parameters of the state-action Q function, the decoder, the amortized variational distribution, and the latent policy. In every training iteration, given the latest values of the advantage weights. Then, we update the Q function using Equations 2 and provided the updated action and latent policies. Then, for every state-action pair the dataset, we compute the advantage using the updated Q function, the latent policy and the decoder, and then estimate the weights. Finally, we update the latent policy using Equation
first update the decoder and the variational distribution g using Equation 7
5 Related Works
Offline Reinforcement I Learning: Offline RI. methods generally address the problem of distribution shift between the behavior policy and the policy being learned 9 24], which cause issues
5

---

actions sampled from the learned policy and passed into the learned critic. To address this issue. prior methods constrain the learned policy stay close the behavior
due to
policy via explicit policy regularization 32 18
51
24
22
13. via implicit policy constraints
39 22 23
5 35 by regularizing based on importance sampling [37 31 50 (2),
by learning of conservative value functions 225] 46 by leveraging auxiliary behavioral cloning
losses 7 38), and through model-based training with conservative penalties [56] 20 47 36 28 55
Compared to these prior works, we opt develop 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 this challenge by introducing a new method that leverages latent variable models. As we will find in Section [ our method also outperforms oftthe------ttttttt prior methods in other categories, particularly when learning from heterogeneous datasets.
Prior works 1 19 also studied the problem of learning from heterogeneous dataset in the offline RL with heterogeneous dynamics setting [ and multi-task offline RL. with data sharing scenario 19 54 respectively. Both settings require knowledge of the agent/task identifier. Here, we study a more general problem setting including datasets with conflicting actions without assuming access to the ground-truth identifier of the source of the heterogeneity.
RL with Generative Models: Generative models have been used by prior works for improving training performance [1] enabling transfer learning [4] [10]. implementing hierarchical RL
2 33
34 40). avoiding distributional shift in offline RL settings
57] and learning dynamics
57 which learns a generative model and latent policy; but, unlike this prior work, our method trains an advantage-weighter generative 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.
models 27 16 15). Our proposed method resembles PLAS
Experiments
Our experiments aim to answer the following questions: (1) How does LAPO compare other offline RL methods on a set of standard offline RL tasks, including learning from heterogeneous and homogeneous offline datasets? (2) How does LAPO compare prior methods implemented with GMM policies when learning behaviors offline from heterogeneous datasets? and (3) In which setting does 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 on offline RL benchmarks with heterogeneous datasets. We consider a range of simulated robotic tasks, including navigation, locomotion, and manipulation, each with corresponding static dataset with heterogeneous data distribution. We also consider datasets with narrow and biased data distributions, containing near-optimal or random trajectories. To answer the second question, we include comparison of AWAC (39 with a GMM policy, since AWAC trains using a similar advantage- weighted objective as the actor policy in LAPO. This comparison is only performed on heterogeneous datasets, since this is where we expect to see the most improvement from a multi-modal policy. Finally, the questions in (3), we conduct two ablation studies in which we train LAPO without learning the latent policy, i.c. evaluating only LAPO's s action policy, and without limiting the latent values to the range [- -maax 2. We compare this version of the method the original LAPO on four offline RL tasks that differ in the amount of high-performing tesk-relevant data contained in the offline dataset. Our implementation is available at httpps: http::://gttthn com wwwp///ccewiiiaaaa... LLAP -ww.AAPPo-ooffoomomma offlienRL
6.1 Comparisons
We compare LAPO to 7 offline RL approaches: (1) vanilla behavioral cloning (BC), (2) the BCQ method [9], which approximates the dataset distribution with generative model and manually selects the action with the maximum Q value among 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-weighh imitation learning with a forward-Ke KL objective for the policy projection step, (5) the IQL method

---

Figure 3: D4RL environments used in evaluations: (from left to right) Walker 2D, Hopper, Halfcheetah, Antmaze, Maze2d, and Kitchen environments.
23], which learns the value function using expectile regression without an explicit policy, and then extracts the policy using advantage weighted regression (AWR), (6) the CQL method [25] which learns a conservative Q-function, and (7) AWAC (смм; which trains GMM policies using AWAC.
In our experiments, we train each method three times using three different random seeds each for one million steps, and report the return averaged over 10 test episodes from the trained policies. We used the 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 implementation details in the Appendix E
6.2 Tasks and Datasets
We compare LAPO to the introduced prior methods on learning from both heterogeneous and homogeneous offline datasets. We evaluate the methods on three simulated task domains, locomotion, navigation, and manipulation domains (Figure B), with heterogeneous data distributions, and also locomotion task with random, narrow and biased data distribution. We consider two sources of heterogeneous datasets: (1) data collected by policies accomplishing different tasks, and (2) data collected by different policies that accomplish one task but in different ways. For (1), we use the multi-task heterogeneous datasets introduced by (54), and for (2), we leverage heterogeneous datasets from the standard D4RL benchmark [6]. More details about our tasks and datasets are provided in the Appendix
6.3 Experimental Results
Heterogeneous datasets: Table Sooleerepee reports the results of our experiments on heterogeneous datasets, we present the learning performance of I.APO 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, average, LAPO improves by 49% over the next best method on the heterogeneous datasets.
and in Figure (Appendix C). (Appenii
The heterogeneous locomotion tasks can be accomplished more easily by the prior methods since still 33% of the data comes from the target task. However, from the prior works, AWAC and CQL are the only methods that can accomplish the tasks. LAPO performs well on all tasks and on average yields the best performance on one out of the three tasks. This shows that LAPO capable of learning from heterogeneity introduced by multi-task datasets.
The navigation tasks are in general more challenging especially for those with medium and large map sizes. The main challenge is to learn policies for long-horizon planning from datasets that do not contain optimal trajectories, and there are only very few states with rewards. LAPO significantly outperforms all prior works for all of the navigation tasks. Similarly, the manipulation tasks are also challenging, they require the assembly of sub-trajectories related completing a given task consisting of multiple sub-tasks In addition, the agent has access to fewer training samples while having to learn to interact with complicated dynamics. LAPO outperforms previous approaches by a large margin on Kitchen- partial and Kitchen mixed for which only a few optimal trajectories are provided by the dataset. It also performs competitively on the Kitchen-complete task and yields the second-best performance.
we observe that GMM policy training does perform well heterogeneous offline RL settings. The AWAC method with GMM policies yields similar or even worse performance compared to the original AWAC. These results suggest that the VAE is more robust model the high-reward regions in heterogeneous data.
7

---

Table 2: The normalized performance of all methods tasks with heterogeneous dataset. 0 represents the performance of a random policy and 100 represents the performance of an expert policy. The scores averaged over the final 10 evaluations and 3 seeds. LAPO achieves the best performance on 9 tasks and achieves competitive performance on the rest 3 tasks. Results with 95%-ccnffienne
interval are reported in Appendix C.2
Task ID
BC
BCQ
PLAS
AWAC
AWACaam
IQL
CQL
LAPO
Walker2d-mis feier22dmmix-fowwacc forward-vl
-5.65
0.91
22.91
71.89
78.09
28.25
102.75
74.17
Walker2d-min mieaddmii-baaaae wwadd-iiii-bbb--------
-84.6
-1.37
-27.58
-7.95
71.33
-46.25
66,64
99.22
Walker2d-min jix-pp---l
-72.92
-41.33
0.56
51.08
28.01
-46.41
37.28
43.20
Maze2d-umaze-vl
0.99
18.91
80.12
94.53
19,45
51.00
22.86
118.86
Maze2d-mediun-vl
3.34
12.79
5.19
31.40
46.53
33.26
12.25
142.75
Maze2d-large-vl
-1.14
27.17
45.80
43.85
9.04
64.30
7.00
200.56
Autmarc-umarr--iiveree----
60.00
62.00
7.00
72.00
0.00
69.33
16.71
91.33
Autmareemmdduumdddvvres-w
0.00
11.33
8.67
0.33
0.00
73.00
1.00
85.67
Autmare-lare-diveere--
0.00
0.67
1,33
0.00
0.00
48.00
11.89
61.67
comple--o
4.50
9.08
38.08
3.83
1.08
66.67
4.67
53.17
-artial-pattill-O0 Kitchen-mixed-v()
31.67
17.58
27.00
0.25
0.42
32.33
0.55
53,67
30.00
11.50
29.92
0.00
3.92
49.92
1.86
62.42
Table 3: The normalized performance of all methods on locomotion tasks with random, narrow and biased dataset. 0 represents the performance of a random policy and 100 represents the performance of an expert policy. The scores are averaged over the final 10 evaluations and seeds. LAPO achieves the best performance on 4 tasks and achieves competitive performance on the rest of the tasks. Results with 95% 95% interval are reported in Appendix C.2
Task ID
BCQ
PLAS
AWAC
IQL
CQL
LAPO
Hopper-random-22 Walker24-random--2
2.23
7.80
6.68
8.01
7.89
8.33
23.46
1.11
4.87
9.17
0.42
5,41
-0,23
1.28
Halicheetab-randdmm-22
2.25
2,25
26.45
15.18
13.11
22.20
30.55
Hopper-medium medium-ve -22
49.23
56.44
50.96
69.55
65.75
71.59
51.63
Walker2d meeeddd-ddim--y
47.11
73.72
76.47
84,02
77.89
82.10
80.75
halfcheetah-medium--e
37.84
47.22
44.54
48.13
47.47
49.76
45.97
Hopper-expert-v2
76.16
68.86
107.05
109.32
109.36
102.27
106.76
Walker2d
79.22
110.51
109.56
110.46
109.93
108.76
112.27
Halicheetaah-experttv2
85.63
93.15
93.79
14.01
94.98
87.40
95.93
Randon/Narrow datasets: Table reports the results of the locomotion tasks with random, narrow and biased datasets. Our method yields the best performance for four out of nine tasks. It also achieves competitive performance
the rest of the tasks, showing that LAPO is applicable to general offline RL settings and is not limited to learning from heterogeneous datasets. However, the most significant gain is when it is applied to offline settings with heterogencous datasets. Learning curves and results with 95%-ccnffieecc interval are reported Appendix C
6.4 Ablations
Latent Policy Training: To study the importance of the latent policy training in achieving good performances, we conduct an ablation study on different tasks by eliminating the latent policy training from LAPO. The main motivation is that, as described in Section 4 the latent variable generative 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 for sampling latent variables to sample from the prior distribution, which changes Equation 8 to:
((s) En-ppr (ls.)))=p(=)112 1Qar (s,
where, the latent is now sampled from fixed prior.
As shown in Fig. training a latent policy does not significantly contribute to higher performance for tasks such as the Walker2D -eddium and moikerrdd-ii--uae which contain sufficient task-

---

mmzeedd mezezi-medium-vv1
diverse-v
walker2d medium
walkmr26.mix -itrr/2d-mii-hccwwrd--u
00
100
Nom Score
70
00
50
40
75-
20
20
40
0-
25
-20
-40
30
150 Without Istent policy
100
0
50
120
150
0
50
100
150
50
100
150
With istent policy
Gradient steps (~5000)
Figure 4: Comparison of LAPO with and without the latent policy on Maze2D-meddimm, Antmaze- medium-diverse, Walker2DD-medium, and Walker2D-mix barkwadd task. The shaded area represents +95%--cofiidene intervals, computed using 10 evaluations and random seeds. For tasks with less optimal task-relevant data like Maze2D-medium, autmaze-medium-diiee the learned latent policy results in gegh-performance gains.
maze2d-medium-/1
100
walleer24-medium-2
walkeedd-mix-baaeeeee
Normalzee Soore
150
80
60
-200
100
100
40
30
460
-200
1000
1.0
100
Without limlingg
With lmiting
Gradient (-5000)
Figure 5: Comparison of LAPO with and without limiting the latent values. The policy performance significantly worse when the latent values are not limited to the range [-Zmax zmax
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. LAPO overcomes 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 2 to be within the range [-2max: /max]. This is important to understand the limits of using the latent variable decoder as generative model of in-distribution actions. Figure 5 illustrates the result of LAPO policy training using unbounded latent actions four tasks. The policy training performance is significantly worse when we do not limit the latent values. This suggests that the the action policy can generate in-distribution actions only when in-distribution latent values (close to samples drawn from the prior distribution p(z)) are given as the input.
Conclusion
In this paper, study offline RL setup for learning from heterogeneous datasets where trajectories are collected using policies with different purposes, leading to a multi-modal data distribution, and in some cases, do not contain sufficient high-reward trajectories. Through empirical analysis, find that in such cases, policies constrained methods may contain -utto--distribution actions and lead to suboptimal performance for continuous control tasks, especially for heterogeneous datasets. advantage- advantagge-eiihted weighted policy optimization
To address this challenge, we present the (LAPO) algorithm, which learns latent variable model that generates adggtadd actions when sampling from prior distribution over the latent space. In addition, we train latent policy that obtains state-conditioned latent actions which result higher reward outcomes compared to sampling from the prior distribution. We compare our method to 6 prior methods on variety of simulated locomotion, navigation, and manipulation tasks provided heterogeneous offline datasets, and also on standard offline RL benchmarks with narrow and biased datasets. We find that our proposed method consistently outperforms prior methods by large margin on tasks with heterogeneous datasets, while being competitive on other offline RL tasks with narrow data distributions. For our future work, we will extend LAPO to multi-task offline RL settings which an agent learns multiple RL tasks provided heterogeneous dataset of diverse behaviors.

---

Acknowledgments
This work is supported in part by Science and Technology Innovation 2030 "New Generation Artifi- cial Intelligence" Major Project (No. 2018AAA0100904) and National Natural Science Foundation of China (62176135).
References
[1] Anish Agarwal, Abdullah Alomar, Varkey Alumootil, Devavrat Shah, Dennis Shen, Zhi Xu, and Cindy Yang. Porsim: Data-efficient offiine reinforcement learning with heterogeneous agents via personalized simulators. arXiv preprini arXi::2 2021.
[2] Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofrr Nachum. Opal: Offline primitive discovery for accelerating offline reinforcement learning. arXiv preprint arXiv:2010. 1361 2020.
[3] Arthur Argenson and Gabriel Dulac-Amold. mwXiv:2008.0555:0 2020.
Model-based offline planning.
arXiv preprint
[4] Karol Ardtt Murtaza Hazara, Ali Ghadirzadeh, and Ville Kyrki. Meta reinforcement learning for sim- to-raal domain adaptation. In 2020 IEFE International Conference on Robotics and Automation (ICRA), pages 2725-2222222 1233.. IEEE, 2020.
[5] Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv: 1809./ 2018.
[6] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint 07X1: 2020.
[7] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. arXiv preprint arXiv:2 /06. 068606 2021.
[8] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation etror in actor-critic methods. In International Conference on Machine Learning, pages 1587-1596. PMLR, 2018.
[9] Scott Fujimoto, David Meger, and Doina Precup. Off-policy dcep reinforcement learning without explo- ration. In International Conference on Machine Learning, pages 2052-2062. PMLR, 2019.
[10] Ali Ghadirzadeh, Xi Chen, Petra Poklukar, Chelsea Finn, Mårten Björkman, and Danica Kragic. Bayesian meta-leaming for few-shot policy adaptation across robotic platforms. arXiv preprint arXiv:2103.03697, 2021.
[11] Ali Ghadirzadeh, Atsuto Maki, Danica Kragic, and Mårten Björkman. Deep predictive policy training using reinforcement learning. In 2017 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 2351-2358. IEEE, 2017.
[12] Ali Ghadirzadeh, Petra Poklukar, Ville Kyrki, Danica Kragic, and Märten Björkman. Data-efficient visuomotor policy training using reinforcement lcarning and generative models. arXiv preprint 1rXiv:2007.13134. 2020. 2020.
[13] Seyed Kamyar Seyed Ghasemipour, Dale Schuurmans, and Shixiang Shane Gu. Emaq: Expected-mal g-learming operator for simple yet effective offiine and online rl. In International Conference on Machine Learning, pages 3682-3691. PMLR. 2021.
[14] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018.
[15] Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream control: Learning behaviors by latent imagination. arXiv preprint 1rX 2019.
[16] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pages 3365.. PMLR, 2019.
[17] Irina Higgins, Loie Mathey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. 2016.
10

---

[18] Natasha Jaques, Asma Ghandeharioun, Judy Hanwen Shen, Craig Ferguson, Agata Lapedriza, Noah Jones. Shixiang Gu, and Rosalind Picard. Way off-policy batch deep reinforcement learning of implicit human preferences in dialog. arXiv preprint arXiv: 00456,, 2019.
[19] Dmitry Kalashnikov, Jacob Varley, Yevgen Chebotar, Benjamin Swanson, Rico Jonschkowski. Chelsea Finn, Sergey Levine. and Karol Hausman. Mt-opt: Continuous multi-task robotic reinforcement learning at scale. arXiv preprint arXiv:2104.
2021.
[20] Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. arXiv preprint arXiv:2005.055 0,0551 2020.
[21] Diederik Kingma and Max Welling, Auto-encoding variational bayes. arXiv preprint urXiv:1312.6114. 2013.
[22] Ilya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. In International Conference on Machine Learning, pages 5774-5783 PMLR, 2021.
[23] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint 16000:2110.06169, 2021.
[24] Aviral Kumar, Justin Fu, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via boot- strapping error reduction. arXiv preprint wrXiv:1906.009/0 2019.
[25] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. arXiv preprint wrXiv:2006004779 2020.
[26] Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning, pages 45-73. Springer. 2012.
[27] Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with latent variable model. arXiv preprint arXiv: 190700090,,, 2019.
[28] Byung-Jun Lee, Jongmin Lee, and Kee-Eung Kim. Representation balancing offline model-based rein- forcement learning. In International Conference on Learning Representations, 2020.
[29] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu, Offine reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint urXiv:2005.0164 2020.
[30] Timothy P Lillicrap. Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint 1009v:1509.000 02971, 2015.
[31] Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Off-policy policy gradient with state distribution correction. arXiv preprint
2019.
[32] Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Provably good batch reinforcement learning without great exploration. arXiv preprint urXiv:2007.08202, 2020.
[33] Ying-Sheng Luo, Jonathan Hans Soeseno, Trista Pei-Chun Chen, and Wei-Chao Chen. Carl: Controllable agent with reinforcement learning for quadruped locomotion. ACM Transactions on Graphics (TOG), 39(4):38-1, 2020.
[34] Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. In Conference on Robot Learning, pages 1113-1132. PMLR, 2020.
[35] Xianteng Ma, Yiqin Yang, Hao Hu, Qihan Liu, Jun Yang, Chongjic Zhang, Qianchuan Zhao, and Bin Liang. Offline reinforcement learning with value-based episodic memory. arXiv preprint arXiv:2110.09796 2021.
[36] Tatsuya Matsushima, Hiroki Furuta, Yutaka Matsuo, Ofir Nachum, and Shixiang Gu. Deployment-efficient reinforcement learning via model-based offline optimization. arXiv preprint arXiv:2006. 03640, 2020.
[37] Ofir Nachum, Yinlam Chow, Bo Dai, and Lihong Li. Dualdice: Behavior-agnostic estimation of discounted stationary distribution corrections. arXiv preprini 146.:3190..44733, 2019.
[38] Ashvin Nair, Murtaza Dalal, Abhishek Gupta, and Sergey Levine. Accelerating online reinforcement learning with offline datascts. arXiv preprint
2020.

---

[39] Ashvin Nair, Murtaza Dalal, Abhishek Gupta, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. 2020.
[40] Xue Bin Peng, Michael Chang. Grace Zhang, Pieter Abbeel, and Sergey Levine. Mep: Learning composable hierarchical control with multiplicative compositional policies. Advances in Neural Information Processing Systems, 32:3686-3693 36976 2019.
[41] Xue Bin Peng. Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage- regression: Simple and scalable off pol-py reinforcement learning. arXiv preprint arXiv:1910.00177, 2019.
[42] Jan Peters, Katharina Mulling, and Yasemin Altun. Relative entropy policy search. In Twenty-Fourth AAAI Conference on Artificial Intelligence. 2010.
[43] Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Offline reinforcement learning from images with latent space models. In Learning for Dynamics and Control. pages 1154-1168, PMLR, 2021.
[44] Noah Y Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, Michael Neunert, Thomas Lampe, Roland Hafner, Nicolas Heess, Martin Riedmiller. Keep doing what worked: Behav- ioral modelling priors for offline reinforcement learning. arXiv preprint arXiv:2002.0839% 2020.
[45] Avi Singh, Huihan Liu, Gaoyue Zhou, Albert Yu, Nicholas Rhinehart, and Sergey Levine. Parrot: Data- driven behavioral priors for reinforcement learning, arXiv preprint arXiv:2011.1000 2020.
[46] Samarth Sinha, Ajay Mandlekar, and Animesh Garg. S4rl: Surprisingly simple self-supervision for offline reinforcement learning. arXiv preprint arXiv:2103.0m 100.0632, 2021.
[47] Phillip Swazinna, Steffen Udluft, and Thomas Runkler. Overcoming model bias for robust offline deep reinforcement learning. Engineering Applications of Artificial Intelligence, 104:104366 2021.
[48] Michita Imai Takuma Seno. d3rlpy: An offline deep reinforcement library. In NeurlPS 202/ Offline Reinforcement Learning Workshop, December 2021.
[49] Qing Wang, Jiechao Xiong. Lei Han, Peng Sun, Han Liu, and Tong Zhang. Exponentially weighted imitation learning for batched historical data. In NewrIPS, pages 6291-6300 2018.
[50] Yunke Wang, Chang Xu, Bo Du, and Honglak Lec. Learning to weight imperfect demonstrations. In International Conference on Machine Learning, pages 10961-10970. PMLR, 2021.
[51] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:191 I. 1161,, 2019.
[52] Yue Wu, Shuangfci Zhai, Nitish Srivastava, Joshua Susskind, Jian Zhang, Ruslan Salakhutdinov, and Hanlin Goh. Uncertainty weighted actor-critic for offline reinforcement learning. arXiv preprint 2005 2021.
[53] Yiqin Yang, Xiaoteng Ma, Li Chenghao, 7ewu Zheng, Qiyuan Zhang, Gao Huang, Jun Yang, and Qianchuan Zhao. Believe what you see: Implicit constraint approach for offline multi-agen reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021.
[54] Tianhe Yu, Aviral Kumar, Yevgen Chebotar, Karol Hausman, Sergey Levine, and Chelsea Finn. Conserva- tive data sharing for multi-task offline reinforcement leamning Advances in Neural Information Processing Systems, 34, 2021.
[55] Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Fian. Combo: Conservative offline masel-based policy optimization. arXiv preprint arXiv:2102.08363, 2021.
[56] Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-basedd offline policy optimization. Advances in Neural Information Processing Systems, 33:14 13:14129-14142. 2020.
[57] Wenxuan Zhou, Sujay Bajracharya, and David Held, Plas: Latent action space for offline reinforcement learning. arXiv preprint 2rX 2020.
12


Reviews:
[{"review_id": "ox7Z108hFkr", "reviewer": "Reviewer_eXjC", "summary": "The authors propose LAPO, an offline-RL algorithm that is designed to efficiently handle the scenario where the provided dataset has heterogeneity. In this setup, the dataset is composed of trajectories from policies solving tasks that could be different from the target task, and/or trajectories from policies solving the target task but in behaviorally diverse ways. There are two main models in LAPO that are jointly trained – 1.) a latent-variable generative model that maximizes the weighted data log-likelihood, where the weight is computed using the advantage function; 2.) a latent policy that is learned to maximize the long-term discounted cumulative rewards (Q-value). Through experiments on the standard D4RL benchmarks (locomotion, navigation, manipulation), the authors show that LAPO compared favorably to the baseline offline RL methods, especially when the dataset is heterogeneous.", "questions": "1. The crucial points of differentiation between LAPO and PLAS are: a.) advantage weighting that is applied in equation 4; b.) the iterative training of the generative model. As the authors mention in related-work, the proposal of a latent-variable generative model and a latent-policy was done in PLAS. When these concepts are introduced in Section 3, I find the corresponding references to PLAS missing. \nFor quantitative comparisons to PLAS, since LAPO is compared with the “original” PLAS code, could there be other differences in the model implementations (decoder, Q-function, policy) that are having an impact? Is this a reasonable way to ensure parity – run the LAPO code with $\\omega=1$ and move the decoder update out of the training loop, keeping everything else the same? \n\n2. Could the authors provide some intuition as to why the latent policy can avoid OOD actions by using the VAE decoder? The approach is understandable when the action distribution is uni-modal. However, with multi-modal action distributions, could using VAEs – that tend to be mode-covering – potentially lead to sampling OOD actions?\n", "limitations": "Yes, the authors have mentioned the limitations in the checklist. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The problem setting considered in this paper – offline-RL with heterogeneous datasets – seems to be under-explored in the community. The paper is well-written, well-motivated (e.g., Figure 2), and the experiments/ablations look thorough. In terms of originality, I am inclined to rate the paper as “incremental” as the proposed method looks to be a derivative of PLAS and prior advantage-weighted BC methods.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "fl83CkWgVqE", "reviewer": "Reviewer_cCVP", "summary": "In offline RL, the static datasets are often heterogeneous. \nThe standard policies constrained methods can be insufficient to learn from such datasets, since they may artificially assign high probability to out-of-distribution actions, leading to sub-optimal policies.\nTo address this challenge, this paper proposes to learn a latent variable model that generates high-advantage actions when sampling from a prior distribution on the latent space.\nFurthermore, this paper trains a latent policy to obtain state-conditional latent actions, which empirically results in better outcomes compared to sampling from a fixed prior distribution. \nThe proposed method, LAPO, performs well on several tested datasets.", "questions": "\n1. L30-31: why learning from heterogeneous settings have the problem of \"lack of sufficient high-return trajectories?\" Why the dataset \"is constructed by re-labeling state-action pairs from some different tasks?\" I don't think the tested Maze2d and Antmaze datasets has this relabelling process.\n2. It seems from Figure 2 that AWAC_GMM perform competitive to LAPO_VAE? Could you provide some intuition why it underperforms the LAPO in many of the D4RL datasets?\n3. L60: How to \"construct a state-conditioned latent space that represents high-return actions\" if at some state the dataset only contain low-return actions? This can be exacerbated when the dataset only contains few high-return actions, which is what this paper wants to address.\n4. For a given targeted task, if the goal is to \"select which (of the multiple conflicting action) samples to use,\" why bother using a \"more expressive policy class to represent the multi-modal action distributions?\" Since we only need one action choice at each state, why couldn't we just use a deterministic policy + advantage weighted learning? \\\nAs a further discussion, this paper aims at the following two goals: **(1)** represent the multi-modal action distributions, where the multiple modes come from the distribution of conflicting actions; and **(2)** learn a mechanism to select which samples (actions) to use. These two goals themselves seems conflicting and I am confused why we would like to simultaneously achieve both goals. In particular, since this paper use advantage weighted learning, the authors seem to implicitly assume that one action-mode is significantly better the others. Then why do we want to model other inferior action modes?\n5. L113-114: why the latent values will be decoded into in-distribution actions? Why couldn't the decoder generate out-of-distribution actions?\n6. I am confused by this sentence in L114-116 \"The parameter $\\theta'$ 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.\"\\\nIf the $p_{\\theta'}$ network does not provide gradient information, how can we train the $\\pi_v$ network? \\\nA follow up discussion, the use of $a$ in the first term of Eq. (3) seems unnecessary due the use of $\\theta'$. Can we define the Q function as $Q(s,z)$ so that the gradient information can be back-propagated directly to the $\\pi_v$ network? I think this one has better semantic meaning that the latent policy, which maximizes the Q-function, does not rely on the action decoder during training. The first term in Eq. (3) is confusing to me since the training of $\\pi_v$, and hence the generation of the latent values, entwists with the (lagged) decoder. Will the latent policy change as the decoder change? If so, could you explain the intuition behind this design? \\\nA even further discussion, since the decoder training already favors high advantage actions, I am not sure if the latent policy can actually *learn* \"latent actions which result in high returns after being converted to the original action space using the decoder\", as claimed in L148-149. \n7. L119-120, the sentence \"while actively prioritize regions ... through the importance weights\" is confusing. As shown in Eq. (3), the training of the latent policy $\\pi_v$ does not involve the \"importance weight\". Could you explain this? \n8. L120-121, the sentence \"Using the decoder, the latent policy will naturally avoid sampling OOD actions\" is confusing. Why the the OOD actions can be *naturally* avoided? How do you define the OOD actions in the latent space, the (function) range of the latent policy?\n9. I think Eq. (4) is wrong. Could you double check it with the Eq. (4) of PLAS [1]?\n10. I am still unsure how the estimation of the advantage function can be exempted from the OOD actions? Since the Q-function is trained by the TD-error Eq. (2), a wield $a'$ can propagate the error. If so, the training of the decoder $p_\\theta$ will be affected by the inaccurate advantage function.\n\n[1] Zhou, Wenxuan, Sujay Bajracharya, and David Held. \"PLAS: Latent action space for offline reinforcement learning.\" arXiv preprint arXiv:2011.07213 (2020).", "limitations": "The authors briefly discuss the limitations but not potential negative societal impacts.\n", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n1. It is interesting to see a method to learn a  policy more flexible than the standard deterministic / Gaussian policy.\n\n2. The empirical results are promising.\n\n### Weaknesses\n1. Novelty: while the whole algorithm is novel, it seems to be a combination of ideas in PLAS and AWAC. Thus, the novelty is incremental compared with prior work.\n\n2. Presentation: explanation to the proposed method is insufficient. In particular, it will be beneficially to explain the intuition/motivation behind the design of each component. \\\nPlease check the ***Questions*** section for questions I have for this paper.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "ESn4RRNgAE", "reviewer": "Reviewer_AQNa", "summary": "This paper investigates the heterogeneous datasets problem in offline reinforcement learning. The major idea of the proposed method LAPO can be summarized as two points. First, the authors construct the policy with a latent variable to better model multi-modal actions. Second, the advantage-weighted scheme is conducted to highlight the importance of high-reward trajectories during the training. The experimental result shows that LAPO can outperform other compared methods by a clear margin with heterogeneous datasets.", "questions": "See Weaknesses.", "limitations": "The authors have addressed the limitations of the paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n* The paper extends traditional offline reinforcement learning into a practical setting, where the offline datasets are heterogeneous and of diverse quality. This is good for making IL methods applicable in real-world tasks.\n\n* The motivation is clear and the proposed method is technically sound. Moreover, the experiments demonstrate the effectiveness of the proposed method.\n\nWeaknesses:\n\n* The idea of using latent policy to represent multi-modal actions and advantage weighting seems to derive from existing offline RL methods.\n\n* Some related work can be added.\n  \n  1. The method reminds me of another exponential advantage-weighted method [1] in imitation learning, which also addresses the diverse demonstrations problem. What is the difference in the advantage-weighting scheme between LAPO and [1]?\n\n  2. [2] also addresses the problem of OOD data by weighting in offline reinforcement learning. However, [2] assigns higher weights not to the data *with high advantage* but to the data that is *close to the trained policy distribution*. Is there any contradiction between these two weighting schemes and which one is better?\n\n  [1] Learning to Weight Imperfect Demonstrations, ICML 2021\n\n  [2] Uncertainty Weighted Actor-Critic for Offline Reinforcement Learning, ICML 2021", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "-G5TQAO8_PD", "reviewer": "Reviewer_A7nB", "summary": "The paper presents a method, latent-variable advantage-weighted policy optimization (LAPO), that learns a policy on top of an advantage-weighted latent space. In summary, the policy outputs a truncated Gaussian distribution over latent actions, which are decoded by a learned decoder. The key idea is that with this method, the agent can learn good policies from multi-modal data, or data collected by a number of demonstrators possibly completing different tasks. The paper also presents experiments suggesting LAPO does in fact perform well on heterogenous data compared to other offline algorithms.", "questions": "- Was I correct in assuming that multi-walker and walker-mix are the same environment?\n- Can you clarify the purpose of the experiments mentioned in the introduction, and possibly provide more details about them?\n- In order to prepare the paper for publication, I would strongly suggest addressing all of the presentation issues mentioned above.", "limitations": "Most of the limitations of the work are related to presentation, which the authors obviously have not addressed. The algorithm itself seems to be novel and the results suggest it could be a significant paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper presents a novel, interesting method that seems to do well empirically in some difficult settings. However, there are issues with presentation that need to be addressed before the paper is ready for publication.\n\n# Strengths\n\n - Originality: the paper appears to present a novel algorithm. It is most similar to [53], but learns an advantage weighted latent space that empirically performs better that [53] in the domains considered.\n - Significance: The paper considers the offline RL problem, which is highly relevant. A key issue for many offline learning algorithms is dealing with multi-modal data, and this paper seems to take a convincing step towards addressing this problem, making it significant.\n\n# Weaknesses\n - The paper's introduction was confusing. It seemed to be referencing experiments done in a toy environment that are not mentioned elsewhere in the paper. The authors seem to be using this to motivate their method, but the end result was unfocused.\n - The methods section's organization was confusing. The discussion around Equation 3 was fairly confusing without the explanations that came in the later sections, and looking at Algorithm 1 it appears that equation wasn’t actually used anywhere, with LAPO instead alternating between optimizing equations 4, 2/5 and 6.\n - Some of the claims could be better substantiated. Although the results are compelling, it would be helpful to see experiments done with more than 3 random seeds. Also, the fact that AWAC+GMM does not work well does not satisfactorily prove to me that GMMs are not as well fit for this problem as the learned latent embedding.\n - The paper suffers from some issues with presentation quality that distract from the purpose of the work and make it, in its current state, unpublishable. Namely:\n1. Figure 3 does not have any axis labels\n2. Figures 5 and 6 are scaled sloppily, and the text is warped\n3. There is some inconsistency in the results. Why are there different names for the walker2d task between figure 3 and table 1? For Walker2d-mix-jump-v1, why does AWAC win in the table while LAPO wins in the graph?\n4. In section 3, it seems as though $Q$ is referring to $Q_\\phi$ although that is not written\n5. In section 3, $\\lambda$ is used multiple times before being introduced.\n6. The legend for figure 6 does not agree with the caption\n7. The word 'tenerative' should be replaced with 'generative' on line 194", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept