Title: Design from Policies: Conservative Test-Time Adaptation for Offline Policy Optimization
Venue: NeurIPS 2023
Accepted: yes
Keywords: offline reinforcement learning

test-time adaptation
Abstract: In this work, we decouple the iterative bi-level offline RL (value estimation and policy extraction) from the offline training phase, forming a non-iterative bi-level paradigm and avoiding the iterative error propagation over two levels. Specifically, this non-iterative paradigm allows us to conduct inner-level optimization (value estimation) in training, while performing outer-level optimization (policy extraction) in testing. Naturally, such a paradigm raises three core questions that are not fully answered by prior non-iterative offline RL counterparts like reward-conditioned policy: (q1) What information should we transfer from the inner-level to the outer-level? (q2) What should we pay attention to when exploiting the transferred information for safe/confident outer-level optimization? (q3) What are the benefits of concurrently conducting outer-level optimization during testing? Motivated by model-based optimization (MBO), we propose DROP (design from policies), which fully answers the above questions. Specifically, in the inner-level, DROP decomposes offline data into multiple subsets, and learns an MBO score model (a1). To keep safe exploitation to the score model in the outer-level, we explicitly learn a behavior embedding and introduce a conservative regularization (a2). During testing, we show that DROP permits deployment adaptation, enabling an adaptive inference across states (a3). Empirically, we evaluate DROP on various tasks, showing that DROP gains comparable or better performance compared to prior methods.

Design from Policies: Conservative Test-Time Adaptation for Offline Policy Optimization

Abstract

In this work, we decouple the iterative bi-level offline RL (value estimation and policy extraction) from the offline training phase, forming a non-iterative bi-level paradigm and avoiding the iterative error propagation over two levels. Specifically, this non-iterative paradigm allows us to conduct inner-level optimization (value estimation) in training, while performing outer-level optimization (policy extraction) in testing. Naturally, such a paradigm raises three core questions that are not fully answered by prior non-iterative offline RL counterparts like rewardconditioned policy: Q1 ) What information should we transfer from the inner-level to the outer-level? Q2 ) What should we pay attention to when exploiting the transferred information for safe/confident outer-level optimization? Q3 ) What are the benefits of concurrently conducting outer-level optimization during testing? Motivated by model-based optimization (MBO), we propose DROP ( D esign f RO m P olicies), which fully answers the above questions. Specifically, in the inner-level, DROP decomposes offline data into multiple subsets and learns an MBO score model ( A1 ). To keep safe exploitation to the score model in the outer-level, we explicitly learn a behavior embedding and introduce a conservative regularization ( A2 ). During testing, we show that DROP permits test-time adaptation, enabling an adaptive inference across states ( A3 ). Empirically, we find that DROP, compared to prior non-iterative offline RL counterparts, gains an average improvement probability of more than 80%, and achieves comparable or better performance compared to prior iterative baselines.

1 Introduction

Offline reinforcement learning (RL) [37, 40] describes a task of learning a policy from static offline data. Due to the overestimation of values at out-of-distribution (OOD) state-actions, recent iterative offline RL methods introduce various policy/value regularization to avoid deviating from the offline data distribution in the training phase. Then, these methods directly deploy the learned policy in an online environment to test its performance. To unfold our following analysis, we term this kind of learning procedure as iterative bi-level offline RL (Figure 1 left ), where the inner-level optimization refers to trying to eliminate the OOD issue, the outer-level optimization refers to trying to infer an optimal policy that will be employed at testing. Here, we use the "iterative" term to emphasize that the inner-level and outer-level are iteratively optimized in the training phase. However, due to the iterative error exploitation and propagation [6] over the two levels, performing such an iterative bi-level optimization completely in training often struggles to learn a stable policy/value function.

Iterative bi-level offline RL optimization

Non-iterative bi-level offline RL optimization

𝜋 ∗

(Online rollout       )

Inner-level

Outer-level

Inner-level

Outer-level

𝜋 ∗

(𝑎𝑡 |𝑠

𝑡

)

Fixed

(Online rollout       )

Training

Testing/Deployment

𝜋 ∗

(𝑎𝑡 |𝑠

𝑡

)

Training

Testing/Deployment

In this work, we thus advocate for non-iterative bi-level optimization (Figure 1 right ) that decouples the bi-level optimization from the training phase, i.e. , performing inner-level optimization (eliminating OOD) in training and performing outer-level optimization (maximizing return) in testing. Intuitively, incorporating the outer-level optimization into the testing phase can eliminate the iterative error propagation over the two levels, i.e. , there is no iteratively learned value function that bootstraps off of itself and thus propagates errors further. Then, three core questions are: Q1 ) What information (" ") should we transfer from the inner-level to the outer-level? Q2 ) What should we pay special attention to when exploiting " " for safe/confident outer-level optimization? Q3 ) Notice that the outer-level optimization and the testing rollout form a new loop, what new benefit does this give us?

Intriguingly, prior works under such a non-iterative framework have proposed to transfer (" " in Q1 ) filtered trajectories [10], a reward-conditioned policy [14, 34], and the Q-value estimation of the behavior policy [6, 24], all of which, however, partially address the above questions (we will discuss these works in Section 3.4). In this work, we thus propose a new alternative method that transfers an embedding-conditioned score model (Q-value) and we will show that this method sufficiently answers the raised questions and benefits most from the non-iterative bi-level framework.

Before introducing our method, we introduce a conceptually similar task (to the non-iterative bi-level optimization) - offline model-based optimization (MBO 2 [62]), which aims to discover, from static input-score pairs, a new design input that will lead to the highest score. Typically, offline MBO first learns a score model that maps the input to its score via supervised regression ( corresponding to inner-level optimization ), and then conducts inference with the learned score model (as " "), for instance, by optimizing the input against the learned score model via gradient ascent ( corresponding to the outer-level ). To enable this MBO implementation in offline RL, we are required to decompose an offline RL task into multiple sub-tasks , each of which thus corresponds to a behavior policyreturn (parameters-return) pair. However, there are practical optimization difficulties when using high-dimensional policy's parameter space (as input for the score model) to learn such a score model ( inner-level ) and conduct inference ( outer-level ). At inference, directly extrapolating the learned score model (" ") also tends to drive the high-dimensional candidate policy parameters towards OOD, invalid, and low-scoring parameters [32], as these are falsely and over-optimistically scored by the learned score model in inner-level optimization.

To tackle these problems, we suggest 3 A1 ) learning low-dimensional embeddings for those sub-tasks decomposed in the MBO implementation, over which we estimate an embedding-conditioned Q-value (" " in Q1 ), and A2 ) introducing a conservative regularization, which pushes down the predicted scores on OOD embeddings. Intuitively, this conservation aims to avoid over-optimistic exploitation and protect against producing unconfident embeddings when conducting outer-level optimization ( Q2 ). Meanwhile, we suggest A3 ) conducting test-time adaptation, which means we can dynamically adjust the inferred embeddings across testing states (aka deployment adaptation). We name our method DROP ( D esign f RO m P olicies). Compared with standard offline MBO for parameter design [62], test-time adaptation in DROP leverages the sequential structure of RL tasks, rather than simply conducting inference at the beginning of testing rollout. Empirically, we demonstrate that DROP can effectively extrapolate a better policy that benefits from the non-iterative framework by answering the raised questions, and achieves better performance compared to state-of-the-art offline RL algorithms.

2 Preliminaries

Reinforcement learning. We model the interaction between agent and environment as a Markov Decision Process (MDP) [61], denoted by the tuple ( S , A , P, R, p 0 ) , where S is the state space, A is the action space, P : S × A × S → [0 , 1] is the transition kernel, R : S × A → R is the reward function, and p 0 : S → [0 , 1] is the initial state distribution. Let π ∈ Π := { π : S × A → [0 , 1] } denotes a policy. In RL, we aim to find a stationary policy that maximizes the expected discounted return J ( π ) := E τ ∼ π [ ∑ ∞ t =0 γ t R ( s t , a t )] in the environment, where τ = ( s 0 , a 0 , r 0 , s 1 , a 1 , . . . ) , r t = R ( s t , a t ) , is a sample trajectory and γ ∈ (0 , 1) is the discount factor. We also define the state-action value function Q π ( s , a ) := E τ ∼ π [ ∑ ∞ t =0 γ t R ( s t , a t ) | s 0 = s , a 0 = a ] , which describes the expected discounted return starting from state s and action a and following π afterwards, and the state value function V π ( s ) = E a ∼ π ( a | s ) [ Q π ( s , a )] . To maximize J ( π ) , the actor-critic algorithm alternates between policy evaluation and improvement. Given initial Q 0 and π 0 , it iterates

where the value function (critic) Q ( s , a ) is updated by minimizing the mean squared Bellman error with an experience replay D + and, following the deterministic policy gradient theorem [59], the policy (actor) π ( a | s ) is updated to maximize the estimated Q k +1 ( s , π ( a | s )) .

Offline reinforcement learning. In offline RL [40], the agent is provided with a static data D = { τ } which consists of trajectories collected by running some data-generating policies. Note that here we denote static offline data D , distinguishing from the experience replay D + in the online setting. Unlike the online RL problem where the experience D + in Equation 1 can be dynamically updated, the agent in offline RL is not allowed to interact with the environment to collect new experience data. As a result, naively performing policy evaluation as in Equation 1 may query the estimated Q k ( s ′ , a ′ ) on actions that lie far outside of the static offline data D , resulting in pathological value Q k +1 ( s , a ) that incurs large error. Further, iterating policy evaluation and improvement will cause the inferred policy π k +1 ( a | s ) to be biased towards OOD actions with erroneously overestimated values.

Offline model-based optimization. Model-based optimization (MBO) [63] aims to find an optimal design input x ∗ with a given score function f ∗ : X → Y ⊂ R , i.e. , x ∗ = arg max x f ∗ ( x ) . Typically, we can repeatedly query the oracle score model f ∗ for new candidate design, until it produces the best design. However, we often do not have the oracle score function f ∗ but are provided with a static offline dataset { ( x , y ) } of labeled input-score pairs. To track such a design from data question, we can fit a parametric model f to the static offline data { ( x , y ) } via the empirical risk minimization (ERM): f ← arg min f E ( x ,y ) [ ( f ( x ) -y ) 2 ] . Then, starting from the best point in the dataset, we can perform gradient ascent on the design input and set the inferred optimal design x ∗ = x ◦ K := GradAscent f ( x ◦ 0 , K ) , where

For simplicity, we will omit subscript f in GradAscent f . Since the aim is to find a better design beyond all the designs in the dataset while directly optimizing score model f with ERM can not ensure new candidates (OOD designs/inputs) receive confident scores, one crucial requirement of offline MBO is thus to conduct confident extrapolation.

3 DROP: Design from Policies

We present our framework in Figure 2. In Sections 3.1 and 3.2, we will answer questions Q1 and Q2 , setting a learned MBO score model as " " ( A1 ) and introducing a conservative regularization over the score model ( A2 ). In Section 3.3, we will answer Q3 , where we show that conducting outer-level optimization during testing can produce an adaptive embedding inference across states ( A3 ).

3.1 Task Decomposition

Our core idea is to explore MBO in the non-iterative bi-level offline RL framework (Figure 1 right ), while capturing the sequential characteristics of RL tasks and answering the raised questions ( Q1 , Q2 , and Q3 ). To begin with, we first decompose the offline data D into N offline subsets

environment

𝛽(a|𝑠, 𝑧)

𝑓(𝑠, a, 𝑧)

v

𝜋 ∗

a 𝑠

= 𝛽(a|𝑠, 𝑧

∗

)

𝑧1

𝑧2

𝑧3

𝑧4

𝜇(𝑧)

offline data (task)              sub-tasks

deployment adaptation

𝒟

𝑠

a

𝒟1

𝒟2

𝒟3

𝒟4

𝜙(𝑧|𝑛)

⇌

↦

training phase (inner-level     )

testing/deployment phase

𝛽, 𝑓

outer-level       (using

𝑓

)

task embedding &

conservative MBO

𝑧 ∗

D [ N ] := {D 1 , . . . , D N } . In other words, we decompose an offline task, learning with D , into multiple offline sub-tasks, learning with D n ∈ D [ N ] respectively. Then, for sub-task n ∈ [1 , N ] , we can perform behavior cloning (BC) to fit a parametric behavior policy β n : S × A → [0 , 1] to model the corresponding offline subset D n :

Additionally, such a decomposition also comes with the benefit that it provides an avenue to exploit the hybrid modes in offline data D , because that D is often collected using hybrid data-generating behavior policies [18], which suggests that fitting a single behavior policy may not be optimal to model the multiple modes of the offline data distribution (see Appendix A for empirical evidence). In real-world tasks, it is also often the case that offline data exhibit hybrid behaviors. Thus, to encourage the emergence of diverse sub-tasks that capture distinct behavior modes, we perform a simple task decomposition according to the returns of trajectories in D , heuristically ensuring that trajectories in the same sub-task share similar returns and trajectories from different sub-tasks have distinct returns 4 .

3.2 Task Embedding and Conservative MBO

Naive MBO over behavior policies. Benefiting from the above offline task decomposition, we can conduct MBO over a set of input-score ( x , y ) pairs, where we model the (parameters of) behavior policies β n as the design input x and the corresponding expected returns at initial states J ( β n ) as the score y . Note that, ideally, evaluating behavior policy β n , i.e. , calculating J ( β n ) ≡ E s 0 [ V β n ( s 0 ) ] , with subset D n will never trigger the overestimation of values in the inner-level optimization. By introducing a score model f : Π → R and setting it as the transfer information " " in Q1 , we can then perform outer-level policy inference with

However, directly performing optimization-based inference (outer-level optimization), max π f ( π ) , will quickly find an invalid input for which the learned score model f outputs erroneously large values ( Q2 ). Furthermore, it is particularly severe if we perform the inference directly over the parameters of policy networks, accounting for the fact that the parameters of behavior policies lie on a narrow manifold in a high-dimensional parametric space.

Task embedding. To enable feasible out-level policy inference, we propose to decouple MBO techniques from high-dimensional policy parameters. We achieve this by learning a latent embedding space Z with an information bottleneck ( dim( Z ) ≪ min( N, dim(Π)) ). As long as the highdimensional parameters of behavior policies can be inferred from the embedding z ∈ Z , we can use z to represent the corresponding sub-task (or the behavior policy). Formally, we learn a deterministic task embedding 5 ϕ : R N →Z and a contextual behavior policy β : S ×Z ×A → [0 , 1] that replaces

N separate behavior policies in Equation 4:

Conservative MBO. In principle, by substituting the learned task embedding ϕ ( z | n ) and contextual behavior policy β ( a | s , z ) into Equation 5, we can then conduct MBO over the embedding space: first learning f : Z → R with min f E n,ϕ ( z | n ) [ ( f ( z ) -J ( β n )) 2 ] , and then setting the optimal embedding with z ∗ = arg max z f ( z ) and the corresponding optimal policy with π ∗ ( a | s ) = β ( a | s , z ∗ ) . However, we must deliberate a new OOD issue in the Z -space, stemming from the original distribution shift in the parametric space when directly optimizing Equation 5.

Motivated by the energy model [38] and the conservative regularization [35], we introduce a conservative score model objective, additionally regularizing the scores of OOD embeddings µ ( z ) :

where we set µ ( z ) to be the uniform distribution over Z -space. Intuitively, as long as the scores of OOD embeddings E µ ( z ) [ f ( z )] is lower than that of in-distribution embeddings E n,ϕ ( z | n ) [ f ( z )] (up to a threshold η ), conducting embedding inference with z ∗ = arg max z f ( z ) would produce the best and confident solution, avoiding towards OOD embeddings that are far away from the training set.

Now we have reframed the non-iterative bi-level offline RL problem as one of offline MBO: in the inner-level optimization ( Q1 ), we set the practical choice for " " as the learned score model f ( A1 ); in the outer-level optimization ( Q2 ), we introduce task embedding and conservative regularization to avoid over-optimistic exploitation when exploiting f for policy/embedding inference ( A2 ). In the next section, we will show how to slightly change the form of the score model f , so as to leverage the sequential characteristic (loop in Figure 1) of RL tasks and answer the left Q3 .

3.3 Test-time Adaptation

Recalling that we update f ( z ) to regress the value at initial states E s 0 [ V β n ( s 0 ) ] in Equation 7, we then conduct outer-level inference with z ∗ = arg max z f ( z ) and rollout the z ∗ -conditioned policy π ∗ ( a | s ) := β ( a | s , z ∗ ) until the end of testing/deployment rollout episode. In essence, such an inference can produce a confident extrapolation over the distribution of behavior policy embeddings. Going beyond the outer-level policy/embedding inference only at the initial states, we propose that we can benefit from performing inference at any rollout state in deployment ( A3 ).

To enable test-time adaptation, we model the score model with f : S × A × Z → R , taking a state-action as an extra input. Then, we encourage the score model to regress the values of behavior policies over all state-action pairs, min f E n,ϕ ( z | n ) E ( s , a ) ∼D n [ ( f ( s , a , z ) -Q β n ( s , a ) ) 2 ]

For simplicity, instead of learning a separate value function Q β n for each behavior policy, we learn the score model directly with the TD-error used for learning the value function Q β n ( s , a ) as in Equation 1, together with the conservative regularization in Equation 7:

.

In test-time, we thus can dynamically adapt the outer-level optimization, setting the inferred optimal policy π ∗ ( a | s ) = β ( a | s , z ∗ ( s )) , where z ∗ ( s ) = arg max z f ( s , β ( a | s , z ) , z ) . Specifically, at any state s in the test-time, we can perform gradient ascent to find the optimal behavior embedding: z ∗ ( s ) = z ◦ K ( s ) := GradAscent( s , z ◦ 0 , K ) , where z ◦ 0 is the starting point and for k = 0 , 1 . . . , K -1 ,

3.4 Connection to Prior Non-iterative Offline Methods

In Table 1, we summarize the comparison with prior representative non-iterative offline RL methods. Intuitively, our DROP (using returns to decompose D ) is similar in spirit to F-BC 6 and RvS-R [14],

both of which use return R ( τ ) to guide the inner-level optimization. However, both F-BC and RvS-R leave Q2 unanswered. In outer-level, F-BC can not enable policy extrapolation, which heavily relies on the data quality in offline tasks. RvS-R needs to handcraft a target return (as the contextual variable for π ( a | s , · ) ), which also probably triggers the potential distribution mismatch 7 between the hand-crafted contextual variable ( i.e. , the desired return) and the actual rollout return induced by the learned policy when conditioning on the given contextual variable.

Diving deeper into the bi-level optimization, we can also find DROP combines the advantages of Onestep [6] and COMs [62], where Onestep answers Q3 and performs outer-level optimization in action space ( arg max a ), COMs answers Q2 and performs that in parameter space ( arg max β ), while our DROPanswers both Q2 and Q3 and performs that in embedding space ( arg max z ). Specifically, DROP thus allows us to ( A2 ) conduct safe extrapolation in outer-level and ( A3 ) perform test-time adaptation in testing rather than simply conducting outer-level optimization at initial states as in COMs.

3.5 Practical Implementation

We now summarize the DROP algorithm (see Algorithm 1 for the training phase and Algorithm 2 for the testing phase). During training (inner-level optimization), we alternate between updating ϕ ( z | n ) , β ( a | s , z ) , and f ( s , a , z ) , wherein we update ϕ with both maximum likelihood loss and TD-error loss in Equations 6 and 8. During testing (outer-level optimization), we use the gradient ascent as specified in Equation 9 to infer the optimal embedding z ∗ . Instead of simply sampling a single starting point z ◦ 0 , we choose N starting points corresponding to all the embeddings { z n } of sub-tasks, and then choose the optimal z ∗ ( s ) from those updated embeddings for which the learned f outputs the highest score, i.e. , z ∗ ( s ) = arg max z f ( s , β ( a | s , z ) , z ) s.t. z ∈ { GradAscent( s , z n , K ) | n = 1 , . . . , N } . Then, we sample action from π ∗ ( a | s ) := β ( a | s , z ∗ ( s )) .

4 Related Work

Offline RL. In offline RL, learning with offline data is prone to exploiting OOD state-actions and producing over-estimated values, which makes vanilla iterative policy/value optimization challenging. To eliminate the problem, a number of methods have been explored, in essence, by either introducing a policy/value regularization in the iterative loop or trying to eliminate the iterative loop itself .

Iterative methods: Sticking with the normal iterative updates in RL, offline policy regularization methods aim to keep the learning policy to be close to the behavior policy under a probabilistic

distance [8, 19, 30, 33, 45, 46, 52, 54, 58, 67, 73]. Some works also conduct implicit policy regularization with variants of importance sampling [39, 47, 51]. Besides regularizing policy, it is also feasible to constrain the substitute value function in the iterative loop. Methods constraining the value function aim at mitigating the over-estimation, which typically introduces pessimism to the prediction of the Q-values [9, 27, 35, 41, 48, 49] or penalizes the value with an uncertainty quantification [3, 5, 56, 68], making the value for OOD state-actions more conservative. Similarly, another branch of model-based methods [29, 71, 72, 57] also perform iterative bi-level updates, alternating between regularized evaluation and improvement. Different from this iterative paradigm, DROPonly evaluates values of behavior policies in the inner-level optimization, avoiding the potential overestimation for values of learning policies and eliminating error propagation between two levels.

Non-iterative methods: Another complementary line of work studies how to eliminate the iterative updates, which simply casts RL as a weighted or conditional imitation learning problem ( Q1 ). Derived from the behavior-regularization RL [22, 64], the former conducts weighted behavior cloning: first learning a value function for the behavior policy, then weighing the state-action pairs with the learned values or advantages [1, 12, 66, 54]. Besides, some works propose an implicit behavior policy regularization that similarly does not estimate values of new learning policies: initializing the learning policy with a behavior policy [50] or performing only a "one-step" update (policy improvement) over the behavior policy [20, 24, 75]. For the latter , this branch method typically builds upon the hindsight information matching [4, 15, 28, 36, 44, 43, 55, 65], assuming that the future trajectory information can be useful to infer the middle decision that leads to the future and thus relabeling the trajectory with the reached states or returns. Due to the simplicity and stability, RvS-based methods advocate for learning a goal-conditioned or reward-conditioned policy with supervised regression [10, 13, 14, 21, 25, 42, 60, 69]. However, these works do not fully exploit the non-iterative framework and fail to answer the raised questions, which either does not regularize the inner-level optimization when exploiting " " in outer-level ( Q2 ), or does not support test-time adaptation ( Q3 ).

Offline MBO. Similar to offline RL, the main challenge of MBO is to reason about uncertainty and OOD values [7, 16], since a direct gradient-ascent against the learned score model can easily produce invalid inputs that are falsely and highly scored. To counteract the effect of model exploitation, prior works introduce various techniques, including normalized maximum likelihood estimation [17], model inversion networks [32], local smoothness prior [70], and conservative objective models (COMs) [62]. Compared to COMs, DROP shares similarities with the conservative model, but instantiates on the embedding space instead of the parameter space. Such difference is nontrivial, not only because DROP avoids the adversarial optimization employed in COMs, but also because DROP allows test-time adaptation, enabling dynamical inference across states at testing.

5 Experiments

In this section, we present our empirical results. We first give examples to illustrate the test-time adaptation. Then we evaluate DROP against prior offline RL algorithms on the D4RL benchmark. Finally, we provide the computation cost regarding the test-time adaptation protocol. For more offlineto-online fine-tuning results, ablation studies w.r.t. the decomposition rules and the conservative regularization, and training details on the hyper-parameters, we refer the readers to the appendix.

Illustration of test-time adaptation. To better understand the test-time adaptation of DROP, we include four comparisons that exhibit different embedding inference rules at testing/deployment:

(1) DROP-Best: At initial state s 0 , we choose the best embedding from those embeddings of behavior policies, z ∗ 0 ( s 0 ) = arg max z f ( s 0 , β ( a 0 | s 0 , z ) , z ) s.t. z ∈ z [ N ] := { z 1 , . . . , z N } , and keep this embedding fixed for the entire episode, i.e. , setting π ∗ ( a t | s t ) = β ( a t | s t , z ∗ 0 ( s 0 )) .

DROP-Grad: At initial state s 0 , we conduct inference (gradient ascent on starting point z ∗ 0 ( s 0 ) ) with z ∗ ( s 0 ) = arg max z f ( s 0 , β ( a 0 | s 0 , z ) , z ) , and keep z ∗ ( s 0 ) fixed throughout the test rollout.

(3) DROP-Best-Ada: We adapt the policy by setting π ∗ ( a t | s t ) = β ( a t | s t , z ∗ 0 ( s t )) , where we choose the best embedding z ∗ 0 ( s t ) directly from those embeddings of behavior policies for which the score model outputs the highest score, i.e. , z ∗ 0 ( s t ) = arg max z f ( s t , β ( a t | s t , z ) , z ) s.t. z ∈ z [ N ] .

(4) DROP-Grad-Ada (as described in Section 3.5): We set π ∗ ( a t | s t )= β ( a t | s t , z ∗ ( s t )) and choose the best embedding z ∗ ( s t ) from those updated embeddings of behavior policies, i.e. , z ∗ ( s t ) = arg max z f ( s t , β ( a t | s t , z ) , z ) s.t. z ∈ { GradAscent( s t , z n , K ) | n = 1 , . . . , N } .

Grad. Ascent

z

[

N

]

z

*

0

(

s

0)

z

*

(

s

0)

z

*

0

(

st

)

z

*

(

st

)

0.0

0.8

0.00

0.25

0.50

0.75

(a)

0.0

0.8

0.00

0.25

0.50

0.75

(b)

t=0

t=50

0.0

0.8

0.00

0.25

0.50

0.75

(c1)

t=17

0.0

0.8

0.00

0.25

0.50

0.75

(c2)

t=50

0

50

100

Normalized return

(d)

40

50

60

normalized return

COMs

DT

RvS-R

Onestep

DROP-BA

DROP-GA

Median

40

50

60

normalized return

IQM

40

50

60

normalized return

Mean

0.8

0.6

P(X > Y)

DROP-BA

DROP-BA

DROP-GA

DROP-GA

COMs

Onestep

COMs

Onestep

Probability of improvement

In Figure 3, we visualize the four different inference rules and report the corresponding performance in the halfcheetah-medium-expert task [18]. In Figure 3 (a), we set the starting point as the best embedding z ∗ 0 ( s 0 ) in z [ N ] , and perform gradient ascent to find the optimal z ∗ 0 ( s 0 ) for DROP-Grad. In Figure 3 (b), we can find that at different time steps, DROP-Best-Ada chooses different embeddings for β ( a t | s t , · ) . Intuitively, performing such a dynamical embedding inference enables us to combine different embeddings, thus stitching behavior policies at different states . Further, in Figure 3 (c1, c2), we find that performing additional inference (with gradient ascent) in DROP-Grad-Ada allows to extrapolate beyond the embeddings of behavior policies, and thus results in sequential composition of new embeddings (policies) across different states . For practical impacts of these different inference rules, we provide the performance comparison in Figure 3 (d), where we can find that performing gradient-based optimization (*-Grad-*) outperforms the natural selection among those embeddings of behavior policies/sub-tasks (*-Best-*), and rollout with adaptive embedding inference (DROP-*-Ada) outperforms that with fixed embeddings (DROP-Best and DROP-Grad). In subsequent experiments, if not explicitly stated, DROP takes the DROP-Grad-Ada implementation by default.

Empirical performance on benchmark tasks. We evaluate DROP on a number of tasks from the D4RL dataset and make comparisons with prior non-iterative offline RL counterparts 8 .

As suggested by Agarwal et al. [2], we provide the aggregate comparison results to account for the statistical uncertainty. In Figure 4, we report the median, IQM [2], and mean scores of DROP and prior non-iterative offline baselines in 12 D4RL tasks (see individual comparison results in Appendix C). We can find our DROP (Best-Ada and Grad-Ada) consistently outperforms prior noniterative baselines in both median, IQM, and mean metrics. Compared to the most related baseline Onestep and COMs (see Table 1), we provide the probability of performance improvement in Figure 4 right . We can find that DROP-Grad-Ada shows a robust performance improvement over Onestep and COMs, with an average improvement probability of more than 80% .

Comparison with latent policy methods. Note that one additional merit of DROP is that it naturally accounts for hybrid modes in D by conducting task decomposition in inner-level, we thus compare DROP to latent policy methods (PLAS [74] and LAPO [11]) that use conditional variational autoencoder (CVAE) to model offline data and also account for multi-modes in offline data. Essentially, both our DROP and baselines (PLAS and LAPO) learn a latent policy in the inner-level optimization, except that we adopt the non-iterative bi-level learning while baselines are instantiated

under the iterative paradigm. By answering Q3 , DROP permits test-time adaptation, enabling us to dynamically stitch "skills"(latent behaviors/embeddings as shown in Figure 3) and thus encouraging high-level abstract exploitation in testing. However, the aim of introducing the latent policy in PLAS and LAPO is to regularize the inner-level optimization, which fairly answers Q2 in the iterative offline counterpart but can not provide the potential benefit (test-time adaptation) by answering Q3 .

Note in our naive DROP implementation, we heuristically use the return to conduct task decomposition (motivated by RvS-R), while PLAS and LAPO take each trajectory as a sub-task and use CVAE to learn the latent policy. For a fair comparison, here we also adopt CVAE to model the offline data and afterward take the learned latent embedding in CVAE as the embedding of behaviors, instead of conducting return-guided task decomposition. We provide implementation details (DROP+CVAE) in Appendix A.4 and the comparison results in Table 2 (taking DROP-Grad-Ada implementation). We can observe that our DROP consistently achieves better performance than PLAS in all tasks, and performs better or comparably to LAPO in 10 out of 12 tasks. Additionally, we also provide the results of CQL and IQL [31]. We can find that DROP also leads to significant improvement in performance, consistently demonstrating the competitive performance of DROP against state-of-the-art offline iterative/non-iterative baselines.

Computation cost. One limitation of DROP is that conducting test-time adaptation at each state is usually expensive for inference time. To reduce the computation cost, we find that we can resample (infer) the best embedding z ∗ after a certain time interval, i.e. , z ∗ does not necessarily need to be updated at each testing state. In Figure 5, we plot the average inference time versus different models and DROP implementations with varying inference time intervals. We see that compared to Diffuser (one offline RL method that also supports test-time adaptation, see quantitative performance comparison in Table 4) [26], DROP can achieve significantly lower computation costs. In Figure 5 top-right , we also illustrate the trade-off between performance and runtime budget as we vary the inference time intervals. We find that by choosing a

Diffuser

Diffuser*

Onestep

DROP(1)

DROP(10)

DROP(20)

DROP(50)

DROP(100)

DROP(150)

DROP(200)

10

2

10

0

Average inference time

(1)

(10)

(20)

(50)

(100)

(150)

(200)

40

45

DROP's scores on AntMaze

suitable time interval, we can reduce the inference cost with only a modest drop in performance: DROP(50) brings less computational burden while providing stable performance.

Comparison to a "distilled" DROP implementation. Our answer to Q3 is that we conduct outer-level optimization in the testing phase. Similarly, we can design a "distilled" DROP implementation: we keep the same treatment as DROP for Q1 and Q2 , but we perform outer-level optimization on the existing offline data instead of on the testing states, i.e. , conducting optimization in z -space for

all states in the dataset. Then, we can "distill" the resulting contextual policy and the inferred

z ∗ into a fixed rollout policy. As shown in Table 3, we can see that such a "distilled" DROP implementation achieves competitive performance, approaching DROP's overall score and slightly outperforming CQL and IQL, which further supports the superiority afforded by the non-iterative offline RL paradigm versus the iterative one.

Comparison to adaptive baselines. One merit of DROP resides in its testtime adaptation ability across rollout states. Thus, here we compare DROP with two offline RL methods that also support test-time adaptation: 1) APEV [23] learns an uncertainty-adaptive policy and dynamically updates a con-

textual belief vector at test states, and 2) Diffuser [26] employs a score/return function to guide the diffusion denoising process, i.e. , score-guided action sampling. As shown in Table 4, we can see that in medium-replay (mr) and medium-expert (me) domains (v2), DROP can achieve comparable results to APE-V, with a clear improvement over Diffuser.

6 Discussion

In this work, we introduce non-iterative bi-level offline RL, and based on this paradigm, we raise three questions ( Q1 , Q2 , and Q3 ). To answer that, we reframe the offline RL problem as one of MBO and learn a score model ( A1 ), introduce embedding learning and conservative regularization ( A2 ), and propose test-time adaptation in testing ( A3 ). We evaluate DROP on various tasks, showing that DROP gains comparable or better performance compared to prior methods.

Limitations. DROP also has several limitations. First, the offline data decomposition dominates the following bi-level optimization, and thus choosing a suitable decomposition rule is a crucial requirement for policy inference (see experimental analysis in Appendix A.2). An exciting direction for future work is to study generalized task decomposition rules ( e.g. , our DROP+CVAE implementation). Second, we find that when the number of sub-tasks is too large, the inference is unstable, where adjacent checkpoint models exhibit larger variance in performance (such instability also exists in prior offline RL methods, discovered by Fujimoto and Gu [19]). One natural approach to this instability is conducting online fine-tuning (see Appendix A.3 for our empirical studies).

Going forward, we believe our work suggests a feasible alternative for generalizable offline robotic learning: by decomposing a single robotic dataset into multiple subsets, offline policy inference can benefit from performing MBO and the test-time adaptation protocol.

Acknowledgments and Disclosure of Funding

We sincerely thank the anonymous reviewers for their insightful suggestions. This work was supported by the National Science and Technology Innovation 2030 - Major Project (Grant No. 2022ZD0208800), and NSFC General Program (Grant No. 62176215).

References

Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller. Maximum a posteriori policy optimisation. arXiv preprint arXiv:1806.06920 , 2018.

Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems , 34:29304-29320, 2021.

Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-based offline reinforcement learning with diversified q-ensemble. Advances in Neural Information Processing Systems , 34, 2021.

Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAI Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay. Advances in neural information processing systems , 30, 2017.

Chenjia Bai, Lingxiao Wang, Zhuoran Yang, Zhihong Deng, Animesh Garg, Peng Liu, and Zhaoran Wang. Pessimistic bootstrapping for uncertainty-driven offline reinforcement learning. arXiv preprint arXiv:2202.11566 , 2022.

David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline rl without off-policy evaluation. Advances in Neural Information Processing Systems , 34:4933-4946, 2021.

David Brookes, Hahnbeom Park, and Jennifer Listgarten. Conditioning by adaptive sampling for robust design. In International conference on machine learning , pages 773-782. PMLR, 2019.

Catherine Cang, Aravind Rajeswaran, Pieter Abbeel, and Michael Laskin. Behavioral priors and dynamics models: Improving performance and domain transfer in offline rl. arXiv preprint arXiv:2106.09119 , 2021.

Yevgen Chebotar, Karol Hausman, Yao Lu, Ted Xiao, Dmitry Kalashnikov, Jake Varley, Alex Irpan, Benjamin Eysenbach, Ryan Julian, Chelsea Finn, et al. Actionable models: Unsupervised offline reinforcement learning of robotic skills. arXiv preprint arXiv:2104.07749 , 2021.

Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems , 34, 2021.

Xi Chen, Ali Ghadirzadeh, Tianhe Yu, Yuan Gao, Jianhao Wang, Wenzhe Li, Bin Liang, Chelsea Finn, and Chongjie Zhang. Latent-variable advantage-weighted policy optimization for offline rl. arXiv preprint arXiv:2203.08949 , 2022.

Xinyue Chen, Zijian Zhou, Zheng Wang, Che Wang, Yanqiu Wu, and Keith Ross. Bail: Best-action imitation learning for batch deep reinforcement learning. arXiv preprint arXiv:1910.12179 , 2019.

Yiming Ding, Carlos Florensa, Pieter Abbeel, and Mariano Phielipp. Goal-conditioned imitation learning. Advances in neural information processing systems , 32, 2019.

Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning? arXiv preprint arXiv:2112.10751 , 2021.

Ben Eysenbach, Xinyang Geng, Sergey Levine, and Russ R Salakhutdinov. Rewriting history with inverse rl: Hindsight inference for policy improvement. Advances in neural information processing systems , 33:14783-14795, 2020.

Clara Fannjiang and Jennifer Listgarten. Autofocused oracles for model-based design. Advances in Neural Information Processing Systems , 33:12945-12956, 2020.

Justin Fu and Sergey Levine. Offline model-based optimization via normalized maximum likelihood estimation. arXiv preprint arXiv:2102.07970 , 2021.

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR , abs/2004.07219, 2020. URL https: //arxiv.org/abs/2004.07219 .

Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in Neural Information Processing Systems , 34, 2021.

Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning , pages 2052-2062. PMLR, 2019.

Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. arXiv preprint arXiv:2111.10364 , 2021.

Matthieu Geist, Bruno Scherrer, and Olivier Pietquin. A theory of regularized markov decision processes. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA , volume 97 of Proceedings of Machine Learning Research , pages 2160-2169. PMLR, 2019.

Dibya Ghosh, Anurag Ajay, Pulkit Agrawal, and Sergey Levine. Offline rl policies should be trained to be adaptive. In International Conference on Machine Learning , pages 7513-7530. PMLR, 2022.

Caglar Gulcehre, Sergio Gómez Colmenarejo, Ziyu Wang, Jakub Sygnowski, Thomas Paine, Konrad Zolna, Yutian Chen, Matthew Hoffman, Razvan Pascanu, and Nando de Freitas. Regularized behavior value estimation. arXiv preprint arXiv:2103.09575 , 2021.

Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems , 34, 2021.

Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991 , 2022.

Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rl? In International Conference on Machine Learning , pages 5084-5096. PMLR, 2021.

Yachen Kang, Diyuan Shi, Jinxin Liu, Li He, and Donglin Wang. Beyond reward: Offline preference-guided policy optimization. arXiv preprint arXiv:2305.16217 , 2023.

Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. Advances in neural information processing systems , 33:21810-21823, 2020.

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.

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06169 , 2021.

Aviral Kumar and Sergey Levine. Model inversion networks for model-based optimization. Advances in Neural Information Processing Systems , 33:5126-5137, 2020.

Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing offpolicy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems , 32, 2019.

Aviral Kumar, Xue Bin Peng, and Sergey Levine. Reward-conditioned policies. arXiv preprint arXiv:1912.13465 , 2019.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 33: 1179-1191, 2020.

Yao Lai, Jinxin Liu, Zhentao Tang, Bin Wang, HAO Jianye, and Ping Luo. Chipformer: Transferable chip placement via offline decision transformer. ICML , 2023. URL https: //openreview.net/pdf?id=j0miEWtw87 .

Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning , pages 45-73. Springer, 2012.

Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and F Huang. A tutorial on energy-based learning. Predicting structured data , 1(0), 2006.

Jongmin Lee, Wonseok Jeon, Byungjun Lee, Joelle Pineau, and Kee-Eung Kim. Optidice: Offline policy optimization via stationary distribution correction estimation. In International Conference on Machine Learning , pages 6120-6130. PMLR, 2021.

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

Jinning Li, Chen Tang, Masayoshi Tomizuka, and Wei Zhan. Dealing with the unknown: Pessimistic offline reinforcement learning. In Conference on Robot Learning , pages 1455-1464. PMLR, 2022.

Qinjie Lin, Han Liu, and Biswa Sengupta. Switch trajectory transformer with distributional value approximation for multi-task reinforcement learning. arXiv preprint arXiv:2203.07413 , 2022.

Jinxin Liu, Hao Shen, Donglin Wang, Yachen Kang, and Qiangxing Tian. Unsupervised domain adaptation with dynamics-aware rewards in reinforcement learning. Advances in Neural Information Processing Systems , 34:28784-28797, 2021.

Jinxin Liu, Donglin Wang, Qiangxing Tian, and Zhengyu Chen. Learn goal-conditioned policy with intrinsic motivation for deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI) , volume 36, pages 7558-7566, 2022.

Jinxin Liu, Hongyin Zhang, and Donglin Wang. Dara: Dynamics-aware reward augmentation in offline reinforcement learning. arXiv preprint arXiv:2203.06662 , 2022.

Jinxin Liu, Ziqi Zhang, Zhenyu Wei, Zifeng Zhuang, Yachen Kang, Sibo Gai, and Donglin Wang. Beyond ood state actions: Supported cross-domain offline reinforcement learning. arXiv preprint arXiv:2306.12755 , 2023.

Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Off-policy policy gradient with state distribution correction. arXiv preprint arXiv:1904.08473 , 2019.

Xiaoteng Ma, Yiqin Yang, Hao Hu, Qihan Liu, Jun Yang, Chongjie Zhang, Qianchuan Zhao, and Bin Liang. Offline reinforcement learning with value-based episodic memory. arXiv preprint arXiv:2110.09796 , 2021.

Yecheng Ma, Dinesh Jayaraman, and Osbert Bastani. Conservative offline distributional reinforcement learning. Advances in Neural Information Processing Systems , 34, 2021.

Tatsuya Matsushima, Hiroki Furuta, Yutaka Matsuo, Ofir Nachum, and Shixiang Gu. Deployment-efficient reinforcement learning via model-based offline optimization. arXiv preprint arXiv:2006.03647 , 2020.

Ofir Nachum, Bo Dai, Ilya Kostrikov, Yinlam Chow, Lihong Li, and Dale Schuurmans. Algaedice: Policy gradient from arbitrary experience. arXiv preprint arXiv:1912.02074 , 2019.

Ashvin Nair, Murtaza Dalal, Abhishek Gupta, and Sergey Levine. Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09359 , 2020.

Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09359 , 2020.

Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177 , 2019.

Vitchyr Pong, Shixiang Gu, Murtaza Dalal, and Sergey Levine. Temporal difference models: Model-free deep rl for model-based control. arXiv preprint arXiv:1802.09081 , 2018.

Shideh Rezaeifar, Robert Dadashi, Nino Vieillard, Léonard Hussenot, Olivier Bachem, Olivier Pietquin, and Matthieu Geist. Offline reinforcement learning as anti-exploration. arXiv preprint arXiv:2106.06431 , 2021.

Marc Rigter, Bruno Lacerda, and Nick Hawes. Rambo-rl: Robust adversarial model-based offline reinforcement learning. arXiv preprint arXiv:2204.12581 , 2022.

Noah Y Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, Michael Neunert, Thomas Lampe, Roland Hafner, Nicolas Heess, and Martin Riedmiller. Keep doing what worked: Behavioral modelling priors for offline reinforcement learning. arXiv preprint arXiv:2002.08396 , 2020.

David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference on machine learning , pages 387-395. PMLR, 2014.

Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojciech Ja´ skowski, and Jürgen Schmidhuber. Training agents using upside-down reinforcement learning. arXiv preprint arXiv:1912.02877 , 2019.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Brandon Trabucco, Aviral Kumar, Xinyang Geng, and Sergey Levine. Conservative objective models for effective offline model-based optimization. In International Conference on Machine Learning , pages 10358-10368. PMLR, 2021.

Brandon Trabucco, Xinyang Geng, Aviral Kumar, and Sergey Levine. Design-bench: Benchmarks for data-driven offline model-based optimization. arXiv preprint arXiv:2202.08450 , 2022.

Nino Vieillard, Tadashi Kozuno, Bruno Scherrer, Olivier Pietquin, Rémi Munos, and Matthieu Geist. Leverage the average: an analysis of KL regularization in reinforcement learning. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020.

Michael Wan, Jian Peng, and Tanmay Gangwani. Hindsight foresight relabeling for metareinforcement learning. arXiv preprint arXiv:2109.09031 , 2021.

Ziyu Wang, Alexander Novikov, Konrad Zolna, Jost Tobias Springenberg, Scott Reed, Bobak Shahriari, Noah Siegel, Josh Merel, Caglar Gulcehre, Nicolas Heess, et al. Critic regularized regression. arXiv preprint arXiv:2006.15134 , 2020.

Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361 , 2019.

Yue Wu, Shuangfei Zhai, Nitish Srivastava, Joshua Susskind, Jian Zhang, Ruslan Salakhutdinov, and Hanlin Goh. Uncertainty weighted actor-critic for offline reinforcement learning. arXiv preprint arXiv:2105.08140 , 2021.

Rui Yang, Yiming Lu, Wenzhe Li, Hao Sun, Meng Fang, Yali Du, Xiu Li, Lei Han, and Chongjie Zhang. Rethinking goal-conditioned supervised learning and its connection to offline rl. arXiv preprint arXiv:2202.04478 , 2022.

Sihyun Yu, Sungsoo Ahn, Le Song, and Jinwoo Shin. Roma: Robust model adaptation for offline model-based optimization. Advances in Neural Information Processing Systems , 34, 2021.

Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems , 33:14129-14142, 2020.

Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in neural information processing systems , 34:28954-28967, 2021.

Chi Zhang, Sanmukh Kuppannagari, and Prasanna Viktor. Brac+: Improved behavior regularized actor critic for offline reinforcement learning. In Asian Conference on Machine Learning , pages 204-219. PMLR, 2021.

Wenxuan Zhou, Sujay Bajracharya, and David Held. Plas: Latent action space for offline reinforcement learning. arXiv preprint arXiv:2011.07213 , 2020.

Zifeng Zhuang, Kun Lei, Jinxin Liu, Donglin Wang, and Yilang Guo. Behavior proximal policy optimization. arXiv preprint arXiv:2302.11312 , 2023.

Appendix

A More Experiments

A.1 Single Behavior Policy v.s. Multiple Behavior Policies

0

50

100

Training steps (k)

0.1

0.2

0.3

Training loss

0

50

100

Training steps (k)

0.1

0.2

0.3

Training loss

0

50

100

Training steps (k)

0.1

0.2

0.3

Training loss

antmaze-large-play

0

50

100

Training steps (k)

0.1

0.2

0.3

Training loss

antmaze-large-diverse

1

5

10

50

100

200

500

1000

In Figure 6, we provide empirical evidence that learning a single behavior policy (using BC) is not sufficient to characterize the whole offline dataset, and multiple behavior policies (conducting task decomposition) deliver better resilience to characterize the offline data than a single behavior policy.

A.2 Decomposition Rules

In DROP algorithm, we explicitly decompose an offline task into multiple sub-tasks, over which we then reframe the offline policy learning problem as one of offline model-based optimization. In this section, we discuss three different designs for the task decomposition rule.

Random ( N,M ) : Wedecomposition offline dataset D := { τ } into N subsets, each of which contains at most M trajectories that are randomly sampled from the offline dataset.

Quantization ( N,M ) : Leveraging the returns of trajectories in offline data, we first quantize offline trajectories into N bins, and then randomly sample at most M trajectories (as a sub-task) from each bin. Specifically, in the i -th bin, the quantized trajectories { τ i } satisfy R min +∆ ∗ i < Return( τ i ) ≤ R min +∆ ∗ ( i +1) , where ∆ = (R max -R min ) N , Return( τ i ) denotes the return of trajectory τ i , and R max and R min denote the maximum and minimum trajectory returns in the offline dataset respectively.

Rank ( N,M ) : We first rank the offline trajectories descendingly based on their returns, and then sequentially sample M trajectories for each subset. ( We adopt this decomposition rule in main paper. )

In Figure 7, we provide the comparison of the above three decomposition rules (see the selected number of sub-tasks and the number of trajectories in each sub-task in Table 9). We can find that across a variety of tasks, the decomposition rule has a fundamental impact on the subsequent model-based optimization. Across different tasks and different embedding inference rules, Random and Quantization decomposition rules tend to exhibit large performance fluctuations, which reveals the importance of choosing a suitable task decomposition rule. In our paper, we adopt the Rank decomposition rule, as it demonstrates a more robust performance shown in Figure 7. In Appendix A.4, we adopt the conditional variational auto-encoder (CV AE) to conduct automatic task decomposition (treating each trajectory in offline dataset as an individual task) and we find such implementation (DROP+CVAE) can further improve DROP's performance. In future work, we also encourage better decomposition rules to decompose offline tasks so as to enable more effective model-based optimization for offline RL tasks.

Comparison with filtered behavior cloning. We also note that the Rank decomposition rule leverages more high-quality (high-return) trajectories than the other two decomposition rules (Random and Quantization). Thus, a natural question to ask is, is the performance of Rank better than that of Random and Quantization due to the presence of more high-quality trajectories in the decomposed sub-tasks? That is, whether DROP (using the Rank decomposition rule) only conducts behavioral cloning over those high-quality trajectories, thus leading to better performance.

Rand

Quan

Rank

2.2

2.3

2.4

Normalized return

halfcheetah-random

Rand

Quan

Rank

35

36

37

Rand

Quan

Rank

50

100

Rand

Quan

Rank

35

40

Rand

Quan

Rank

9

10

11

Normalized return

hopper-random

Rand

Quan

Rank

0

50

hopper-medium

Rand

Quan

Rank

50

100

hopper-medium-expert

Rand

Quan

Rank

20

40

hopper-medium-replay

Rand

Quan

Rank

0.0

2.5

5.0

Normalized return

walker2d-random

Rand

Quan

Rank

0

25

50

walker2d-medium

Rand

Quan

Rank

0

50

100

walker2d-medium-expert

Rand

Quan

Rank

10

20

walker2d-medium-replay

Rand

Quan

Rank

25

50

75

Normalized return

antmaze-umaze

Rand

Quan

Rank

25

50

75

antmaze-umaze-diverse

Rand

Quan

Rank

0

25

50

Rand

Quan

Rank

0

20

40

Rand

Quan

Rank

0

25

50

Normalized return

antmaze-large-play

Rand

Quan

Rank

0

20

antmaze-large-diverse

Quan DROP-Best-Ada

Rank DROP-Grad

Rank DROP-Best-Ada

Rank DROP-Grad-Ada

0

50

100

Training steps

2000

3000

4000

5000

Episode return

seed-1

0

50

100

Training steps

1000

2000

3000

Episode return

hopper-medium-replay

0

50

100

Training steps

0

1000

2000

3000

Episode return

walker2d-medium-replay

0

50

100

Training steps

0.00

0.25

0.50

0.75

Episode return

antmaze-large-play

To answer the above question, we compare DROP (using the Rank decomposition rule) with filtered behavior cloning (F-BC), where the latter (F-BC) performs behavior cloning after filtering for trajectories with high returns. We provide the comparison results in Table 5. We can find that in AntMaze tasks, the overall performance of DROP is higher than that of F-BC. For the MuJoCoGym suite, DROP-based methods outperform F-BC on these offline tasks that contain plenty of sub-optimal trajectories, including the random, medium, and medium-replay domains. This result indicates that DROP can leverage embedding inference (extrapolation) to find a better policy beyond all the behavior policies in sub-tasks, which is more effective than simply performing imitation learning on a subset of the dataset.

A.3 Online Fine-tuning

Online fine-tuning (checkpoint-level). In Figure 8, we show the learning curves of DROP-Best on four DR4L tasks. We can find that DROP exhibits a high variance (in performance) across training

steps 9 , which means the performance of the agent may be dependent on the specific stopping point chosen for evaluation (such instability also exists in prior offline RL methods, Fujimoto and Gu [19]).

To choose a suitable stopping checkpoint over which we perform the DROP inference (DROP-Grad, DROP-Best-Ada and DROP-Grad-Ada), we propose to conduct checkpoint-level online fine-tuning (see Algorithm 3 in Section B for more details): we evaluate each of the latest T checkpoint models and choose the best one that leads to the highest episode return.

In Figure 9, we show the total normalized returns across all the tasks in each suite (including Maze2d, AntMaze, and MuJoCo-Gym). We can find that in most tasks, fine-tuning (FT) can guarantee performance improvement. However, we also find such fine-tuning causes negative impacts on performance in AntMaze(*-v0) suite. The main reason is that, in this checkpoint-level fine-tuning, we choose the "suitable" checkpoint model using the DROP-Best embedding inference rule, while we adopt the other three embedding inference rules (DROP-Grad, DROP-Best-Ada and DROP-GradAda) at the test time. Such a finding also implies that the success of DROP's test-time adaptation is not entirely dependent on the best embedding across sub-tasks 10 ( i.e. , the best embedding z ∗ 0 ( s 0 ) in DROP-Best), but requires switching between some "suboptimal" embeddings (using DROP-Best-Ada) or extrapolating new embeddings (using DROP-Grad-Ada).

Online fine-tuning (embedding-level). Beyond the above checkpoint-level fine-tuning procedure, we can also conduct embedding-level online fine-tuning: we aim to choose a suitable gradient update step for the gradient-based embedding inference rules (including DROP-Grad and DROP-Grad-Ada). Similar to the checkpoint-level fine-tuning, we first conduct the test-time adaptation procedure (DROP-Grad and DROP-Grad-Ada) over a set of gradient update steps, and then choose the best step that leads to the highest episode return (see Algorithm 4 in Section B for more details).

DROP-Grad

DROP-Grad FT

DROP-Best-Ada

DROP-Best-Ada FT

DROP-Grad-Ada

DROP-Grad-Ada FT

0

50

Total normalized return

Maze2d

DROP-Grad

DROP-Grad FT

DROP-Best-Ada

DROP-Best-Ada FT

DROP-Grad-Ada

DROP-Grad-Ada FT

0

100

200

AntMaze(-v0)

DROP-Grad

DROP-Grad FT

DROP-Best-Ada

DROP-Best-Ada FT

DROP-Grad-Ada

DROP-Grad-Ada FT

0

250

500

MuJoCo-Gym(-v0)

DROP-Grad

DROP-Grad FT

DROP-Best-Ada

DROP-Best-Ada FT

DROP-Grad-Ada

DROP-Grad-Ada FT

0

100

200

AntMaze(-v2)

DROP-Grad

DROP-Grad FT

DROP-Best-Ada

DROP-Best-Ada FT

DROP-Grad-Ada

DROP-Grad-Ada FT

0

250

500

MuJoCo-Gym(-v2)

Table 6: Online fine-tuning results (initial performance → performance after online fine-tuning). The baseline results of AWAC, CQL, and IQL are taken from Kostrikov et al. [31], where they run 1 M online steps to fine-tune the learned policy. For our DROP method (DROP-Grad and DROP-Grad-Ada), we run 0 . 3 M( = 6 checkpoint × 50 K max × 1000 steps per episode) online steps to fine-tune (embedding-level) the policy, i.e. , aiming to find the optimal gradient ascent step that is used to infer the contextual embedding z ∗ ( s 0 ) or z ∗ ( s t ) for π ∗ ( a t | s t ) := β ( a t | s t , · ) (see Algorithm 4 for the details). Moreover, for medium-* and large-* tasks, we conduct additional parametric-level fine-tuning, with 0 . 7 Monline steps to update the policy's parameters.

In Table 6, we compare our DROP (DROP-Grad and DROP-Grad-Ada) to three offline RL methods (AWAC [53], CQL [35] and IQL [31]), reporting the initial performance and the performance after online fine-tuning. We can find that the embedding-level fine-tuning (0.3M) enables a significant improvement in performance. The fine-tuned DROP-Grad-Ada (0.3M) outperforms the AWAC and CQL counterparts in most tasks, even though we take fewer rollout steps to conduct the online fine-tuning (baselines take 1 Monline rollout steps, while DROP-based fine-tuning takes 0 . 3 Msteps). However, there is still a big gap between the fine-tuned IQL and the embedding-level fine-tuned DROP (0.3M). Considering that there remain 0 . 7 Monline steps in the comparison, we further conduct "parametric-level" fine-tuning (updating the parameters of the policy network) for our DROP-GradAda on medium-* and large-* tasks, we can find which achieves competitive fine-tuning performance even compared with IQL.

A.4 DROP + CVAE Implementation

CVAE-based embedding learning. Similar to LAPO [11] and PLAS [74], we adopt the conditional variational auto-encoder (CVAE) to model offline data. Specifically, we learn the contextual policy and behavior embedding:

Then, we learn the score model f with the TD-error and the conservative regularization:

where ¯ f denotes a target network and µ ( z ) denotes the uniform distribution over the Z -space.

In testing, we also dynamically adapt the outer-level optimization, setting policy inference with π ∗ ( a | s ) = β ( a | s , z ∗ ( s )) , where z ∗ ( s ) = arg max z f ( s , β ( a | s , z ) , z ) .

A.5 Ablation Study

Note that our embedding inference depends on the learned score model f . Without proper regularization, such inference will lead to out-ofdistribution embeddings that are erroneously high-scored ( Q2 ). Here we conduct an ablation study to examine the impact of the conservative regularization used for learning the score model.

In Figure 10, we compare DROP-Grad and DROP-Grad-Ada to their naive implementation ( w/o Reg ) that ablates the regularization on halfcheetah-medium-expert. We can find that removing the conservative regularization leads to unstable performance when changing the update

0

30

60

90

120

150

25

50

75

100

Grad

Grad w/o Reg

0

30

60

90

120

150

25

50

75

100

Grad-Ada

Grad-Ada w/o Reg

Figure 10: Ablation on the conservative regularization. The y-axis represents the normalize return, and the x-axis represents the number of gradient-ascent steps used for embedding inference at deployment. We plot each random seed as a transparent line; the solid line corresponds to the average across 5 seeds.

steps of gradient-based optimization. However, we empirically find that in some tasks such a naive implementation ( w/o Reg ) does not necessarily bring unstable inference (Appendix C). Although improper gradient update step leads to faraway embeddings, to some extent, embedding-conditioned behavior policy can correct such deviation.

B Implementation Details

For the practical implementation of DROP, we parameterize the task embedding function ϕ ( z | n ) , the contextual behavior policy β ( a | s , z ) , and the score model f ( s , a , z ) with neural networks. For Equation 8 in the main paper, we construct a Lagrangian and solve the optimization through primaldual gradient descent. For the choice of µ ( z ) , we simply set µ ( z ) to be the uniform distribution over the Z -space and empirically find that such uniform sampling can effectively avoid the out-ofdistribution extrapolation at inference.

Lagrangian relaxation. To optimize the constrained objective in Equation 8 in the main paper, we construct a Lagrangian and solve the optimization through primal-dual gradient descent,

This unconstrained objective implies that if the expected difference in scores of out-of-distribution embeddings and in-distribution embeddings is less than a threshold η , λ is going to be adjusted to 0 , on the contrary, λ is likely to take a larger value, used to punish the over-estimated value function. This objective encourages that out-of-distribution embeddings score lower than in-distribution embeddings, thus performing embedding inference will not lead to these out-of-distribution embeddings that are falsely and over-optimistically scored by the learned score model.

In our experiments, we tried five different values for the Lagrange threshold η ( 1 . 0 , 2 . 0 , 3 . 0 , 4 . 0 , and 5 . 0 ). We did not observe a significant difference in performance across these values. Therefore, we simply set η = 2 . 0 .

Hyper-parameters. In Table 7, we provide the hyper-parameters of the task embedding ϕ ( z | s ) , the contextual behavior policy β ( a | s , z ) , and the score function f ( s , a , z ) . For the gradient ascent update steps (used for embedding inference), we set K = 100 for all the embedding inference rules in experiments.

In Table 9, we provide the number of sub-tasks, the number of trajectories in each sub-task, and the dimension of the embedding for each sub-task (behavior policy). The selection of hyperparameter N is based on two evaluation metrics: (1) the fitting loss of the decomposed behavioral policies to the offline data, and (2) the testing performance of DROP. Specifically,

(Step1) Over a hyperparameter (the number of sub-tasks) set, we conduct the hyperparameter search using the fitting loss of behavior policies, then we choose/filter the four best hyperparameters;

(Step2) We follow the normal practice of hyperparameter selection and tune the four hypermeters selected in Step1 by interacting with the simulator to estimate the performance of DROP under each hyperparameter setting.

We provide the hyperparameter sets in Table 8. In Step2, we tune the (filtered) hyperparameters using 1 seed, then evaluate the best hyperparameter by training on an additional 4 seeds and finally report the results on the 5 total seeds (see next "evaluation protocol"). In Antmaze domain, a single fixed N works well for many tasks;

while in Gym-mujoco and Adroit domains, we did not find a fixed N that provides good results for all tasks in the corresponding domain in D4RL, thus we use the above hyperparameter selection rules (Step1 and Step2) to choose the number.

Baseline details. For the comparison of our method to prior iterative offline RL methods, we consider the v0 versions of the datasets in D4RL 11 . We take the baseline results of BEAR, BCQ, CQL, and BRAC-p from the D4RL paper [18], and take the results of TD3+BC from their origin paper [19]. For the comparison of our method to prior non-iterative offline RL method, we use the v2 versions of the dataset in D4RL. All the baseline results of behavior cloning (BC), Decision Transform (DT), RvS-R, and Onestep are taken from Emmons et al. [14]. In our implementation of COMs, we take the parameters (neural network weights) of behavior policies as the design input for the score model; and during testing, we conduct parameters inference (outer-level optimization) with 200 steps gradient ascent over the learned score function, then the rollout policy is initialized with the inferred parameters. For the specific architecture, we instantiate the policy network with dim( S ) input units, two layers with 64 hidden units, and a final output layer with dim( A ) .

Evaluation protocol. We evaluate our results over 5 seeds. For each seed, instead of taking the final checkpoint model produced by a training loop, we take the last T ( T = 6 in our experiments) checkpoint models, and evaluate them over 10 episodes for each checkpoint. That is to say, we report the average of the evaluation scores over 5 seed × 6 checkpoint × 10 episode rollouts.

Online fine-tuning (checkpoint-level): Instead of re-training the learned (final) policy with online rollouts, we fine-tune our policy with enumerated trial-and-error over the last T checkpoint models (Algorithm 3). Specifically, for each seed, we run the last T checkpoint models in the environment over one episode for each checkpoint. The checkpoint model which achieves the maximum episode return is returned. In essence, this fine-tuning procedure imitates the online RL evaluation protocol: if the current policy is unsatisfactory, we can use checkpoints of previous iterations of the policy.

Online fine-tuning (embedding-level): The embedding-level fine-tuning aims to find a suitable gradient ascent step that is used to conduct the embedding inference in DROP-Grad or DROP-Grad-Ada. Thus, we enumerate a list of gradient update steps and pick the best update step (according to the episode returns).

Codebase. Our code is based on d3rlpy: https://github.com/takuseno/d3rlpy . We provide our source code in the supplementary material.

Computational resources. The experiments were run on a computational cluster with 22x GeForce RTX 2080 Ti, and 4x NVIDIA Tesla V100 32GB for 20 days.

C Additional Results

Comparison with iterative offline RL baselines. Here, we compare the performance of DROP (Grad, Best-Ada, and Grad-Ada ) to iterative offline RL baselines (BEAR [33], BCQ [20], CQL [35], BRAC-p [67], and TD3+BC [19]) that perform iterative bi-level offline RL paradigm with (explicit or implicit) value/policy regularization in inner-level. In Table 10, we present the results for AntMaze, Gym-MuJoCo, and Adroit suites in standard D4RL benchmark (*-v0), where we can find that DROP-Grad-Ada performs comparably or surpasses prior iterative bi-level works on most tasks: outperforming (or comparing) these policy regularized methods (BRAC-p and TD3+BC) on 25 out of 33 tasks and outperforming (or comparing) these value regularized algorithms (BEAR, BCQ, and CQL) on 19 out of 33 tasks.

Ablation studies. In Figure 11, we provide more results for the ablation of the conservative regularization term in Equation 8 in the main paper. We can find that for the halfcheetah-medium and hopper-medium tasks, the performance of DROP-Grad-Ada w/o Reg depends on the choice of the gradient update steps, showing that too small or too large number of gradient update step deteriorates the performance. Such a result is also consistent with COMs [62], which also observes the sensitivity of naive gradient update ( i.e. , w/o Reg ) to the number of update steps used for design input inference. By comparison, the conservative score model learned with DROP-Grad-Ada exhibits more stable and robust performance to the gradient update steps.

Further, we also find that in walker2d-medium and walker2d-medium-expert tasks, the naive gradient update ( w/o Reg ) does not affect performance significantly across a wide range of gradient update steps. The main reason is that although the excessive gradient updates lead to faraway embeddings, conditioned on the inferred embeddings, the learned contextual behavior policy can safeguard against the embeddings distribution shift. Compared to prior model-based optimization that conducts direct gradient optimization (inference) over the design input itself, such "self-safeguard" is a special merit in the offline RL domain as long as we reframe the offline RL problem as one of model-based optimization and conduct inference over the embedding space. Thus, we encourage the research community to pursue further into this model-based optimization view for the offline RL problem.

0

30

60

90

120

150

20

30

40

Normalized returns

0

30

60

90

120

150

30

40

50

hopper-medium

0

30

60

90

120

150

20

22

24

26

walker2d-medium

0

30

60

90

120

150

60

80

walker2d-medium-expert

DROP-Grad-Ada

DROP-Grad-Ada w/o Reg

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an iterative bi-level offline RL algorithm that separates an \"inner level\" and \"outer level\" optimization for reinforcement learning. A novel variant of model-based optimization (including task decomposition and task embedding) is proposed to enable flexible test-time adaptation. Experiments on popular d4rl benchmarks show state-of-the-art performances over other recent baselines.\n\nSTRENGTHS:\nThis paper proposes an approach that incorporates model-based optimization into offline rl with test-time adaptation. Each component in the methodology is well justified. This reformulation of offline rl as model-based optimization seems novel and can be of value to the community.\n\nAdequate comparisons have been made to prior iterative and non-iterative offline rl methods, and the differences between the proposed method and related works are well justified.\n\nSystematic evaluations on popular D4RL benchmarks are included, and statistically significant improvement over other state-of-the-art methods is observed.\n\nThe weakness of increased computation cost is discussed.\n\nWEAKNESSES:\nThe writing of the paper (especially that of the abstract and introduction) is hard to follow. Specifically, the terms \"inner level\" and \"outer level\" are very confusing. Using other terms such as \"value estimation\" and \"policy extraction\" or the like will make reading much easier. The three questions Q1, Q2, and Q3 raised in the abstract and introduction are also very vague and do not make sense until I have carefully read the methodology and experiments. In my opinion, the main contribution of this paper is reformulating offline rl as a model-based optimization problem and it makes flexible adaptation to offline states possible. \n\nI am missing a comparison to related works in offline rl with test time adaptation. In particular, a comparison with [x] should be discussed.\n\nIn Figure 5, is it possible to also include the inference time of other baselines (COMs, RvS-R, etc)? How do Onestep and DROP compare when they use the same amount of computation?\n\nAs discussed in the limitations section, the number of subtasks can be important for this method. However, in offline rl when online data is difficult to collect, we are not allowed to tune hyperparameters using online interactions [y]. The procedure of determining hyperparameters should be included in the paper. \n\n\n[x] Offline RL Policies Should be Trained to be Adaptive (https://arxiv.org/abs/2207.02200)\n\n[y] Conservative Objective Models for Effective Offline Model-Based Optimization (https://arxiv.org/pdf/2107.06882)\n\nQUESTIONS:\nIn Table 2, which variant is the listed DROP? DROP Best, Grad, Best-Ada, or Grad-Ada?\n\nWhen comparing DROP with LAPO in Table 2, why are the results of LAPO on u, um-p, and ul-p missing? How is the mean calculated with missing entry?\n\nIn Figure 5, the inference time of Diffuser is reported, is it possible to also report its performance in Table 2 or Figure 4?\n\nIn line 226-228, \"Different from this iterative paradigm, DROP only evaluates values of behavior policies in the inner-level optimization, avoiding the potential overestimation for values of new learning policies and eliminating the error propagation between the two levels.\" I do not quite understand why the iterative process can cause error propagation for potential overestimation of values, if you consider the iterative process as a min max optimization process? And there is an additional overestimation of MBO introduced by this paper.\n\nI am willing to increase the score if the weaknesses and questions can be adequately addressed.\n\nLIMITATIONS:\nThe limitations of the number of subtasks and how to split subtasks are adequately discussed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a non-iterative offline RL algorithm, in which the policy is trained in a bi-level optimisation process. As opposed to commonly known iterative algorithms, the authors propose to split the optimisation into an inner loop training to mitigate any OOD issues, and an outer loop optimisation for reward maximisation. The paper compares the proposed algorithm on a subset of the D4RL benchmark datasets with prior iterative offline RL algorithms and demonstrates new state-of-the-art performance.\n\nSTRENGTHS:\nI think the proposed way of thinking about offline RL is still relatively new and a logical, very promising way forward. It seems wasteful to not make use of the information presented to the policy online (which is what many offline algorithms do since after training the policy stays fixed no matter what happens) - moving a part of the optimisation into this outer loop thus makes a lot of sense to me. A particularly new element to me is that the value that the final policy conditions on is not something handcrafted, but a learned representation of the potential behaviour policies. The combination with MBO is also a very interesting concept that I am unaware of being used in prior offline RL methods. I thus consider the method to score well in terms of novelty. I also believe that in the future we will see a lot more algorithms operating in a similar fashion, at least in the regard that part of the optimisation is moved to the online / outer loop.\n\nThe paper provides meaningful experiments and statistical analysis thereof: The algorithm is evaluated on maze as well as robotic locomotion tasks from the D4RL suite and the authors report measures beyond mean performance, such as probability of improvement and IQM performance. They also perform an analysis of the computational burden.\n\nWEAKNESSES:\nWhile I very much like the idea, I personally found the framing with the three questions and answers rather confusing. E.g. the questions are framed differently in abstract and introduction and leave room for lots of interpretation. For example, What should we pay attention to when exploiting the transferred information in the outer-level optimization? -> The question lacks a goal (i.e. what should we pay attention to ... in order to achieve WHAT?, otherwise the answer could be anything).\n\nAlso, the term \"non-iterative, bi-level optimisation\" is used very often. It is obvious what is meant by bi-level (inner and outer are described), but it is at least a little ambiguous what is meant by \"non-iterative\". I'm sure DROP also has iterations in some sense (i.e. mini batches during the learning of the conditional behaviour policy) - I think what you mean is that there is no iteratively learned value function which bootstraps off of itself and thus propagates errors further and further. Is that right? Please clarify.\n\nThe answer to question 3 is that the outer loop allows using the online data to adapt the policy. This is intuitive, but I would say an experiment to really show this is lacking. DROP is compared with some iterative baselines and shows better performance, but this could be due to other differences as well. What I would have liked to see would be something like:\n- do everything like DROP, but try to put the outer-loop directly into the optimisation\n- e.g.: perform optimisation in z-space for all states in the dataset and distill the resulting global policy model + conditioning values into a fixed model\n- then you could isolate the difference that the outer-loop actually makes\n\nThe authors mention and qualitatively compare prior \"non-iterative\" algorithms (RvS, onestep, COMs, F-BC), yet do not include them in their empirical evaluation, which seems a little odd since one could assume they have in a way similar qualities and it should be interesting what benefit the newly proposed non-iterative method has over prior methods of this type.\nFurther, it seems to me that some prior works are missing in this context, i.e. there exist already offline RL algorithms that condition on something which can then adaptively change the policy behaviour by altering the value that it conditions on: In [1], the policy conditions on a level of confidence, in [2] on a balance between conservatism and more liberal behaviour, and in [3] it conditions on a belief over the MDP dynamics, which is updated based on each online state-action pair that is processed. I realise these are all rather recent, but I think they could at least be part of the qualitative comparison.\n\n[1] Hong, J., Kumar, A., & Levine, S. (2022). Confidence-Conditioned Value Functions for Offline Reinforcement Learning. ICLR 2023\n\n[2] Swazinna, P., Udluft, S., & Runkler, T. (2022). User-Interactive Offline Reinforcement Learning. ICLR 2023\n\n[3] Ghosh, D., Ajay, A., Agrawal, P., & Levine, S. (2022, June). Offline rl policies should be trained to be adaptive. ICML 2022\n\nMost weaknesses should be rather easy to address. What I am most concerned about is more relevant baselines (prior non-iterative algorithms, the missing ones that I mentioned, or the experiment proposed to actually show the merit of the outer-loop).\n\nQUESTIONS:\nHow is the subset of D4RL datasets chosen? Why did you not use all of the available locomotion ones, e.g. medium-expert and medium-replay?\n\nWhy are there standard deviations only for DROP, but not for the other algorithms? How is the probability of improvement then calculated - is it whether the other methods mean is outside DROP's 95% CI or is it wither both 95% CI's don't overlap?\n\nWhat do you mean exactly with non-iterative?\n\nWhy are the non-iterative prior works not part of the empirical evaluation?\n\nLIMITATIONS:\nLimitations are discussed & I enjoyed reading them since I also thought about these points while reading the paper.\n\npotential negative societal impacts are not discussed", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper tackles offline reinforcement learning, and considers two classes of algorithms — iterative (at each step, policy evaluation and policy improvement is done sequentially) and non-iterative. The main contribution of the paper is to move the part of policy improvement from offline training to during test-time. This can increase the computational time of deploying the model, but seems to improve performance. The paper also make clever use of task decomposition and learning a low-dimensional task embedding to improve their results.\n\nSTRENGTHS:\n1. The idea is novel and very interesting.\n2. The experimentations of the paper is thorough and organized, even though potentially some ablations can be added.\n3. The paper is also written well, has high clarity and readability.\n\nWEAKNESSES:\n1. The final algorithm has many moving parts and components, and can feel to be over-engineered.\n2. Ablation studies are not sufficient. One can potentially add ablations for effects of the hyper-parameters, and for example, it would be very important to see how the number of subtasks affect the performance of the algorithm.\n\nQUESTIONS:\n1. How is the task decomposition done for appendix section B.1? I.e., which of the three decomposition rules from B.2 is used? In general, more information about the task decomposition in the main paper would be appreciated.\n2. How is the sequential sampling done for Rank(N, M) task decomposition? My guess is first M trajectories (with highest returns) become task 1, the next M trajectories become task 2 and so on.\n3. Appendix line 048, “rank decomposition rule leverages more high quality trajectories” → what defines the quality of a trajectory? If quality = return of the trajectory, then wouldn’t some subtasks have lower quality in rank decomposition due to having low return trajectories only?\n4. Could we have an ablation to see the effect of N on the performance of the algorithm? For example, the onestop algorithm [1] does not have a sub-task level decomposition. Maybe checking if there is a certain threshold for N, below which the onestop algorithm performs better, would be interesting. In other words, I am curious about how much improvement is coming from task divisions, if that is one of the main driving force for the given algorithm’s effectiveness. \n5. Also is the computation time affected by N?\n6. What is the advantage of using task embeddings $z$ instead of directly fitting a MBO on the collected data?\n7. In table 2 (main paper), the umaze environment is antmaze-umaze or maze2d-umaze? Also the version of the environments should be included in the main paper.\n\n\n[1] David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline RL without off-policy evaluation. Advances in Neural Information Processing Systems, 34:4933–4946, 2021.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a method, namely DROP, to adapt the policy during the inference time. The authors achieve this by dividing the optimization process into two phases. In the first phase, the authors train a contextual behavior policy, a score model, and a deterministic task embedding model. At the second phase, the authors utilzie the score model for adapting the policy execution, i.e., follow the optimal embedding. The authors belives their proposed method can better achieve \"stitching\" across the states in the offline data. They provide some experiments on D4RL datasets to show the benefits of their proposed method.\n\nSTRENGTHS:\n(a) this paper is well-structured and easy to follow\n\n(b) the figures are nice and helpful for the readers to understand the claims from the authors\n\n(c) it is good to see the authors consider statistical uncertainty and significance in the paper. I believe this is quite important to RL fields\n\nWEAKNESSES:\n(a) I have several concerns on the ideas presented in the main text\n\n - the adopted method shares many similarities with CQL, i.e., the authors utilize a CQL-style optimization objective in Equation 8. The regularization term is hence not novel. The differences are, CQL penalizes actions, while DROP penalizes the embedding $z$.\n\n - though the author claim that they learn a *score model*, the learnt score function $f$ is actually ***action-value function***, but conditioned on the task embedding. This can be observed in Eqaution 8, where the score model is updated via bellman error. This is highly similar to [1, 2]. Both [1] learns a value function conditioned on the confidence $\\delta$, while [2] learns the value function conditioned on the evolving belief. These two papers are relevant to this paper, while they are not cited in this paper. I believe they are important baselines to compare against. Specially, [1,2] can also achieve the claimed *test-time adaption* by adjusting the confidence or the belief. Actually, this has already been achieved by [2]. To realize this in [1], perhaps another objective function is needed (or simply by maximizing the confidence conditioned value function). \n\n - the authors cite RoMA [3], but do not involve it as a baseline. I can tell that RoMA is also strongly related to the topics covered in this paper, and RoMa seems to exhibit better performance than COMs\n\n - the generated embeddings seem to serve as *goal* for the learnt policy. I believe some goal-conditioned and return-conditioned algorithms ought to be included **numerically** as baselines. The authors compare their DROP against IQL, CQL, PLAS, LAPO while none of them are goal-conditioned or return-conditioned algorithms. The authors provide IQM comparison against decision transformer, RvS-R, while more advanced and stronger methods are expected\n\n - Meanwhile, the comparison in Table 1 is misleading. These criterias are manually proposed by the authors. The success of DROP in answering A1-A3 do not necessiarily indicate that DROP is better than prior methods\n\n - the proposed method, DROP, seems quite redundant. It requires to first split the offline dataset into some subsets, and then learn the embedding network and the contextual behavior policy upon it. Additionally, one needs to learn a score model, and finally utilizes this score model for querying the optimal embedding during inference. It is unclear whether there is a need for dividing the datasets into some subsets and learn embeddings on them. The authors provide a design choice comparison concerning on which decomposition rule is better, while the corresponding analysis on whether decomposition itself is needed is missing. The authors write that \"suggests that fitting a single behavior policy may not be optimal to model the multiple modes of the offline data distribution\" (line 125-126), but the authors themselves choose to learn one single contextual behavior policy instead of multiple policies conditioned on the task embedding. Since the behavior policy is unique, why do we still need to split the datasets? For example, why not the authors directly receive $s,a$ as inputs to train $\\phi$. I believe this eliminates the need of dividing the datasets. Also, the reviewer does not think that training loss is a good indicator to show that multiple policies can better characterize the offline datasets (as presented in appendix B.1).\n\n[1] Confidence-conditioned value functions for offline reinforcement learning\n\n[2] Offline rl policies should be trained to be adaptive\n\n[3] Roma: Robust model adaptation for offline model-based optimization\n\n(b) no theoretical understandings are provided either in the main text or in the appendix. I do not want to blame the authors too much on this point. However, it is unclear whether the socre function will overestimate even with the introduced regularization, especially during deployment (i.e., whether poor embedding is provided even if we optimize the score function at inference time). It is also unclear how large should we penalize the embedding. One of the advantages of CQL is that it theoretically answers that one ought to use a large $\\alpha$ to enforce conservatism. The authors comment that the inference can be unstable if the number of subtasks is large. I believe the instability is due to the wrong output embedding. It will be better if the authors can theoretically characterize the relationship between the embedding and the number of subsets.\n\n(c) I also have the following concerns on the experiment section\n\n - in Table 2 of main text, some baseline results are lower than their typically reported ones in prior work, e.g., CQL\n\n - in Table 2 of main text, the results of LAPO on antmaze are on \"v1\" version datasets, and I can tell that the results of IQL on antmaze are from \"v2\" datasets. It is unclear whether the results of baselines and DROP are also conducted on \"v2\" datasets, since the authors do not specify the version of datasets in the main text. This is huge problematic. \n\n - in Table 2 of main text, it is unclear which variant of DROP is used to report the results, given the fact that the authors introduce four variants of DROP in section 5\n\n - I can see that the authors use MuJoCo \"v0\" datasets in Figure 2 while \"v2\" in Table 1 in the appendix. You should unify the versions of dataset used.\n\n - in line 72, the authors write that \"achieves better performance compared to state-of-the-art offline RL algorithms\". However, the compared baselines are NOT state-of-the-arts, especially on MuJoCo datasets. The authors ought to use this term carefully.\n\n - the superior performance of DROP are acquired via a very careful tuning of hyperparameters. Based on Table 5 in the appendix, the authors carefully search for the optimal number of subsets $N$ for each dataset. This value varies across different datasets, and may lies in a large range (e.g., 1000 subsets on halfcheetah-medium-expert, but 50 subsets on halfcheetah-medium-replay). The choice of the number of subsets is quite weird. For example, the halfcheetah-medium-expert dataset only contains two modes, while DROP requires 1000 subsets to learn a contextual behavior policy. Meanwhile, it is unclear how this value affects the performance of DROP practically. The authors say that \"we find that when the number of sub-tasks is too large, the inference is unstable\", but how unstable is it? I believe this is a quite important hyperparamter, and its parameter study ought to be included. The need of searching the best number of subsets decreases the practicality of DROP\n\n - based on the table 3 in the appendix, the authors use a large batch size 1024 for DROP. As far as the reviewer can tell, the baseline algorithms all use a batch size of 256. A recent work show that a larger batch size is helpful for offline RL learning [4], hence I believe the comparison of DROP against other methods is not fair\n\n - it is difficult for the user to find the best time interval when deploying DROP in practice, and the best values may differ across different scenarios\n\n[4] Q-Ensemble for Offline RL: Don't Scale the Ensemble, Scale the Batch Size\n\n(d) It is unclear whether the proposed method can be applied in the real-world applications, where we may require a fast response from the policy. Since DROP requries test-time adaption, I think at the current stage it is not suitable to be deployed in the real-world robots. Moreover, in real-world scenarios, we actually expect the robots to be able to adapt their policies during deployment, as it is highly possible that the robots encounter the unseen scenarios (e.g., different landforms) and need to adapt the policies. I do not see the potential of using DROP to mitigate these challenges. The generalization capability of DROP to unseen scenarios is limited.\n\n(e) Minor issues\n\n - the format of appendix seems to be another venue instead of NeurIPS\n\n - appendix A contains nothing\n\nQUESTIONS:\n(a) this paper does not seem to involve any theoretical analysis, why do you choose `yes` to `theory`?\n\n(b) it seems the performance of DROP can be largely affected by the performance of the behavior policy, can it be improved if we train the contextual behavior policy in a IQL-style manner, i.e., weighted behavior cloning\n\nLIMITATIONS:\nThe authors discuss several limitations of their work, and the reviewer personally agrees with that.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}