Design from Policies: Conservative Test-Time Adaptation for Offline Policy Optimization
Authors: Jinxin Liu
Hongyin Zhang
Zifeng Zhuang
Yachen Kang
Donglin Wang
Bin Wang
Venue: NeurIPS 2023 poster (2023)
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
Jinxin Liu
Hongyin Zhang
Zifeng Zhuang
Yachen Kang" 1.2
Donglin Wang**
Bin Wang
Westlake University
ZZhejinng University
"Huawei Noah's Ark Lab
Abstract
this work, we decouple the iterative bi-level offline RL (value estimation and policy extraction) from the offline training phase, forming non-iterative bi-level paradigm and avoiding the iterative error propagation over two levels. Specifi- cally, this non-iterative paradigm allows (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 the benefits concurrently conducting outer-level optimization during testing? Motivated by model-based optimization (MBO), 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 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 prior non-iterative offline RL counterparts, gains average improvement probabil- ity of more than 80%, and achieves comparable or better performance compared prior iterative baselines.
conduct inner-level optimization
Introduction
Offline reinforcement learning (RL. [37, 40] describes task of learning policy from static offline data. Due to the overestimation of values at our-tf-distributionn (OOD) state-actions, recent iterative offline RI. methods introduce various policy/value regularization avoid deviating from the offline data distribution 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 left), where the inner-level optimization refers trying eliminate the OOD issue, the outer-level optimization refers 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 eтrоr exploitation and propagation [6] over the two levels, performing such an iterative bi-level optimization completely training often struggles learn stable policy/value function.
"Corresponding author: Donglin Wang <wangdonglin westlake.edu.ccs.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Iterative bi-level offline RL optimization
Non-iterative bi-level offline RL optimization
Training
Testing/Deployment
Training
Testing/D eelooymenpllttttt
Outer-
Inner-level
Fixed nt" (aplsp)
Inner-level
It' (aalss)
Outer-level
(Online rollout
(Online rollout
Figure I: Framework of (iterative and non-iterative) bi-level offline RL optimization, where the inner-level optimization refers to regularizing the policy/value function (for OOD issues) and the outer-level refers to updating the policy for reward maximizing. The "I"" transferred from inner-level to outer-level depends on the specific choice of algorithm used (see different choices for ( in Table 1).
In this work, we thus advocate for non-iterative bi-level optimization (Figure 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, ie., 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 ("E"") 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- non-iterative framework have proposed to transfer ("IE" 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 new alternative method that transfers an embedding-onee coooeeddiinnnooiiiionneeeddd 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² [62]), which aims 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 "I""", 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 RI. task into multiple sub-tasks, each of which thus corresponds to behavior policy- return Dparameters-retumy pair. However, there are practical optimization difficulties when using high-dimensional policy's parameter space (as input for the score model) to learn such score model (inner-level) and conduct inference (outer-level). At inference, directly extrapolating the learned score model ("EE") 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' A1) learning low-dimensional embeddings for those sub-tasks decomposed in the MBO implementation, over which we estimate an booeeddinn-coddiionnee Q-vulue ("E"" 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 method DROP (Design fROm Policies). 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- thhthc-rr offline RL algorithms.
2 Pleease note that this MBO is different from the typical model-based RL. where the model MBO denotes a score model while that in the model-based RL. denotes a transition dynamics model. We use A1, A2, and A3 denote our answers the raised questions (Q1, Q2, and Q3) respectively.

---

2 Preliminaries
Reinforcement learning. We model the interaction between agent and environment as a Markov Decision Process (MDP) [61]. denoted by the tuple (S, P. R. po), where S is the state space, A is the action space, P: S x A x S + [0,1] is the transition kernel, R: S x A + R is the reward function, and Pu S + [0, is the initial state distribution. Let Il :- {- S A + [0, 1]} denotes a policy. In RL, we aim to find stationary policy that maximizes the expected discounted return [ ;= E, Ey~x E D ' R(s;, at) ) in the environment, where T (So. ag. To. $1, ai....). = R(s, at). is sample trajectory and y E (0, 1) is the discount factor. We also define the state-action value functiin Q" (s. a) == E (Er= R(st, az)]s, = 8, ap = a). which describes the expected discounted return starting from state and action and following æ afterwards, and the state value function V* (s) = (als(a)s) (0T a)]. To maximize J(r), the actor-critic algorithm alternates between policy evaluation and improvement. Given initial Qo and n², iterates
Qh+1 +- arg min (s. (a'l') [R((s, a) + (0 - Q(s. Q
n k+1 f- arg max (aw [ (s.
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) T (als) is updated to maximize the estimated (st (a(s)).
Offline reinforcement learning. In offline RL [40], the agent is provided with static data D (r) 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 can be dynamically updated, the agent in offline RL is not allowed interact with the environment to collect new experience data. As result, naively performing policy evaluation as in Equation may query the estimated Q* (s', a') on actions that lie far outside of the static offline data D, resulting in pathological value Q*+1(s, a) that incurs large error. Further, iterating policy evaluation and improvement will cause the inferred policy N ++1 (als) 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 imput x* with a given score function f X - y C I. i.e., x = arg maxx f (x). Typically, 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 static offline dataset {(xx y)} of labeled input-score pairs. To track such design from data question, we can parametric model f to the static offline data {(x, via the empirical
risk minimization (ERM): f + arg miny B(xx)) (J(xx y)² 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 == GradAscent, (x,, K), where
x + x + 17 (x))xxxx for k == O.... K 1.
For simplicity, we will omit subscript f in GradAscenty. 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 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 RI. tasks and answering the raised questions (Q1, Q2, and Q3). To begin with, first decompose the offline data D into N offline subsets
3

---

training phase (inner-level
2
Lesting/deploymmnt phase
1
Zz
outer-level
(using
$ (zlm)
3
2
m(x)
f.
11 as,
(a B(als,z))
B z)
environment
task embedding & conservative MBO
offline data (task)
sub-tasks
deployment adaptation
Figure 2: Overview of DROP. Given offline dataset D, we decompose the data into N (= in the diagram) subsets D D[N]- Over the decomposed sub-tasks, we learn task embedding (z2)))) and con- duct conservative MBO by learning score model (s, z) and N behavior policies (modeled by contextual policy ([a)s, z)). During test-time, at state s, we perform test-time adaptation by adapting the policy (contextual variable) with x² (a|s) 3 (als, z where z' arg max, f fs, ((a(s, 2), z).
Din] == (Di...., Dx). In other words, we decompose an offline task, learning with D, into multiple offline sub-tasks, learning with D, E DINj respectively. Then, for sub-task n E NJ. we can perform behavior cloning (BC) fit parametric behavior policy Bn S A [0. 1] model the corresponding offline subset D,:
Bn + arg max ((.))) log (a|s))
8x
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- aata-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 offine data exhibit hybrid behaviors. Thus, to encourage the emergence of diverse sub-tasks that capture distinct behavior modes, perform 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
3.2 Task Embedding and Conservative MBO
Naive MBO over behavior policies. Benefiting from the above offline task decomposition, we can conduct MBO over set of input-score (x, 3) pairs, where we model the (parameters of) behavior policies 3n as the design input x and the corresponding expected returns at initial states ((,,,) as the
score y. Note that, ideally, evaluating behavior policy 8n, i.e., calculating J(8,) ) = TE ( (so)]. with subset D, will never trigger the overestimation of values in the inner-level optimization. By introducing a score model f 11 + IR and setting it as the transfer information "["" in Q1, we can then perform outer-level policy inference with
(a s) f- arg max f (m), where f = arg min E,, ((((3))-(8B)))
(5)
However, directly performing optimization-bun based inference (outer-level optimization), max- f (T). will quickly find an invalid input for which the learned score model f outputs erroneously large values (Q2). Furthermore, 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.
particularly severe if we perform the inference directly over the
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 < min(N [N, dim())))) As long as the high- dimensional parameters of behavior policies can be inferred from the embedding E Z, we can use to represent the corresponding sub-task (or the behavior policy). Formally, we learn deterministic task embedding"
that replaces N + Z and a contextual behavior policy S Z + [0,
'See more data decomposition rules in Appendix A.2.
'We feed the one hot encoding of the sub-task specification (n = 1,..,, N) into the embedding network ø. We write the embedding function as ø(2 instead of ((n) emphasize that the output of o

---

N separate behavior policies in Equation 4:
B(als, z), (z) t arg max EDD ((,,)))D. [log ((als, p(z[n))] 8,0
(6)
Conservative MBO. principle, by substituting the learned task embedding d(z(n) and contextual behavior policy ((a)s, z) into Equation 5, we can then conduct MBO over the embedding space: first learning f: 2 + R with miny (x))())))) (f(z))
and then setting the optimal embedding with zt = arg maxx I (z) and the corresponding optimal policy with T (als) ([ass, However, we must deliberate 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 conser- vative score model objective, additionally regularizing the scores of OOD embeddings u(z):
f arg min E,, )),((x)))
(f(z)) - J(B.))"). s.t. Ex(x)) [f() Enja(ain) [f(z) S 7. (7)
f
where we set (z(z) to be the uniform distribution over Z-space. Intuitively, as long as the scores of OOD embeddings Ey((s) [f(z] is lower than that of in-distribution embeddings (( [f()] (up to threshold 7)), conducting embedding inference with x² = arg max, 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 (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 I) of RL tasks and answer the left Q3.
3.3 Test-time Adaptation
Recalling that we update fz) regress the value initial states Exg [ (so) in Equation 7, we then conduct outer-level inference with z' arg max, fz) and rollout the z"-conditioned policy x ra(als) == B[a|s, z²) until the end of testing/deployment rollout episode. In essence, such an inference can produce 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 ʃ : S x A x Z + R. taking state-action as an extra input. Then, we encourage the score model to regress the values of behavior policies over all state-action pairs, miny En, (ss.)) (J(s,
(s, a)) 2
For simplicity, instead of learning a separate value function Qº, for each behavior policy, we learn the score model directly with the TD-error used for learning the value function Q (s. a) as in Equation 1, together with the conservative regularization in Equation 7:
f + arg min ED..D (( (R(s,a) a) + xf((ss a', o(z(n)) f(s. a, (( (8)
s.t. En,u(e)) an (ajs.))
[f a, z)] Be..
[f(s, a, z)] S 1
where / denotes a target network and we update the target / with soft updates: j - (1 v)j + uv.
In test-time, we thus dynamically adapt the outer-level optimization, setting the inferred optimal policy n² (als) = B[ass, (s)), where = (s) arg max, [(ss, ((( is, 2). Specifically, at any state s in the test-time, we can perform gradient ascent find the optimal behavior embedding: z" (s) = x (s) == GradAscent(s, zj, K), where z is the starting point and for A = 0, 1....., , K 1.
z (s) t zi(s) + V 2 B(als, 2), 2)) 12=2g
3.4 Connection to Prior Non-iterative Offline Methods
In Table I, we summarize the comparison with prior representative non-iterative offline RL methods. Intuitively, our DROP (using returns decompose D) similar in spirit to F-BC" and RyS-R [14],
6F-BC performs BC over filtered trajectories with high returns.
5

---

Table 1: Comparison of non-iterative bi-level offline RL methods, where we use R(7) to represent the return of a sampling trajectory and use R(s, a) to represent the expected return of trajectories starting from (s, a). The checkmark in A2 indicates whether the exploitation to "["" is regularized for outer-level optimization and that in A3 indicates whether test-time adaptation supported.
Inner-level
Outer-level
"(F" in A1
A2
A3
F-BC
BC over filtered with high R(T)
T(lls)
x
x
RvS-R [14]
min, -E [log (alss, R(r))
handcraft Rurget
(ajs,-)
x
x
Onestep [6]
ming L(Q(s. a). a(s.a))
arg maxa Qa(s, B(ais))
@3(s.a)
X
r
COMs [62]
mins [(f(B,). R(T))
arg max f(3) arg maxy f(s, ((a)z),
f(3)
x
DROP (ours)
miny L(f(ss a(),
f().
both of which use return R(T) 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 target return (as the contextual variable for πT (a(ss,)) -)), which also probably triggers the potential distribution mismatch' 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 On- estep [6] and COMs [62], where Onestep answers Q3 and performs outer-level optimization in action space (arg max,), COMs answers Q2 and performs that in parameter space (arg max;), while our DROP answers both Q2 and Q3 and performs that in embedding space (arg max²). 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
Algorithm DROP (Training)
Algorithm 2 DROP (Testing/Deployment)
Require: Dataset of trajectories: == (T).
Require: Env, B(ass z), and f(ss a, 2).
1: Decompose D into N sub-sets DiNj-
1: Sy Env.Reset().
2: Initialize (z)nn), (als,, 2), and (ss
2: while not done do
3: while not converged do
3: Test-time adaptation:
Sample a sub-task: D,, Div)-
z (s.)) arg max, f( ( [s z),
5: Learn Φ, 3, and f with Equations 6 and 8,
4: Sample action: a ~ 3 |s, z' (s:)).
end while
5: Step Env: S141 ~ P(s at).
Return: B(als, z) and (s, a,
6: end while
We now summarize the DROP algorithm (see Algorithm for the training phase and Algorithm 2 for the testing phase). During training (inner-level optimization), we alternate between updating o(z/n). ([a)s, z). and fs, 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 2". Instead of simply sampling a single starting point zi, we choose N starting points corresponding to all the embeddings (z,,} of sub-tasks, and then choose the optimal z (s) from those updated embeddings for which the learned ʃ outputs the highest score, i.e., z (s) arg max, f(s. (is,B(a)s,, z), s.t. E [GradAscenttss Zn; K) I 1,.... N). Then, we sample action from x =(aas) = B(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, 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 be close to the behavior policy under a probabilistic
7 For more illustrative examples w.r.t. this mismatch, refer the reader to Figure 6 of the RvS-R paper [14].

---

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, 51J. Besides regularizing policy, 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, DROP only 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 RI. as a weighted or conditional imitation learning problem (Q1). Derived from the sehavior-regularizatio RI [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 "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 returns. Due to the simplicity and stability, RvS- byS-based methods advocate for learning a goal-conditioned or reward-conditionss 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 "I" in outer-level (Q2), or does not support test-time adaptation (Q3).
policy with supervised
Offiiee 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 offline- to-online fine-tuning results, ablation studies w.r.t. the decomposition rules and the conservative regularization, and training details on the
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 So, we choose the best embedding from those embeddings of behavior policies, zn (5)) arg max. [so, 3(an) so, z), z) s.t. E [[N] = (zi,..., ]. and keep this embedding fixed for the entire episode, i.e., setting nr <" (s.) = 3 (a, ss, (s)))
(2) DROP-Grad: At initial state Sp. we conduct inference (gradient ascent on starting point 2 (5p) with x (5o) = arg max; [So.
2), z). and keep z (5o) fixed throughout the test rollout.
(3) DROP-Best-Ada: We adapt the policy by setting π 7'(a; ss)) = (a (st, (s:)) where we choose the best embedding z (s:) directly from those embeddings of behavior policies for which the score model outputs the highest score, i.e., z (s.) = arg max, f(e, ((a,) (s;, 2), z) s.t. z Z(N]-
(4) DROP-Grad-ALe (as described in Section 3.5): We set T" [s)] 8(a, (st, 2: (s:)) and choose the best embedding x" (s:) from those updated embedddngg of behavior policies, iee. x (s.) = arg max, fs, ((a (s., z), z) s.t. E [GradAscent(ss, Zn, K)|n L.....N..
7

---

f
Grad. Ascent
2
Zo (5o)
z 2(50)
2ó (sc)
z z((s))
0.75
0.75
0.75
0.75
Normalired retum
100
0.50
0.50
0.50
0.50
|~0
1-50
50
0.25
0.25
0.25
0.25
0.00
0.00
0.00
0.00
0.0
(8)
0.8
0.0
(b)
0.8
0.0
(cl)
0.8
0.0
(c2)
0.8
(d)
Figure 3: Visualization of the embedding inference (a, b, cl, c2) and performance comparison (d),
[[N] denotes embeddings of all behavior policies. ( z" (So). 2 (st), and x (s²) denote the
lected embeddings DROP- Best, Grad, Best-Ada, and Grad- implementations respectively.
Median
IQM
Mean
DROP-GA
Probability impprremmmnnt
DROP-BA
Algerithm
Algorithm
Onestep
DROP-GA
Cneslep
RyS-R
DROP-DA
DROP BA
Onestep
COMs
DROP-RA
COMs
40
60 normalized return
50 nommalizdd return
40
60
normalized return
Figure 4: (Left) Aggregate metrics [2] on 12 D4RL tasks. Higher median, IQM, and mean scores are better. (Right) Each row shows the probability of improvement that the algorithm X on the left outperforms algorithm Y on the right. The CIs are estimated using the percentile bootstrap (95%) with stratified sampling. For all results of our method, we average the normalized returns across seeds; for each seed, we run 10 evaluation episodes. (GA: Grad-Ada. BA: Best-Ada.)
In Figure 3, we visualize the four different inference rules and report the corresponding performance in the wirichhettahmmmiumm----
task [18]. In Figure 3 (a), we set the starting point as the best embedding zūj (So) Z[N]. and perform gradient ascent to find the optimal (sa) for DROP- In Figure 3 (b), we can find that at different time steps, DROP-Best-Ade Best-Ada chooses different embeddings for ((ar/((,))))) (s.... Intuitively, performing such a dynamical embedding inference enables us combine different embeddings, thus stitching behavior policies al different states. Further, in Figure 3 (cl, c2), we find that performing additional inference (with gradient ascent) in DROP-Grad-Aee 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, provide the performance comparison in Figure 3 (d), where can find that performing gradient-ber based optimization r-Grad-*)) outperforms the natural selection among those embeddings of behavior policies/sub-task (*- (t-Best-*), and rollout with adaptive embedding inference (DROP-*-J) Ada)) outperforms that with fixed embeddings (DROP-Best and DROP- In subsequent experiments, if not explicitly stated, DROP takes the DROP-Grad-
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
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 basclines in 12 D4RL tasks (see individual comparison results in Appendix C). We can find our DROP (Best- Ada and Grad-Ada) consistently outperforms prior non- iterative 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 right. We can find that PIROP-Grad-AAee 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 D by conducting task decomposition in inner-level, we thus compare DROP 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 latent policy in the inner-level optimization, except that we adopt the non-iterative bi-level learning while baselines are instantiated
'Due the page limit, provide the comparison with prior iterative baselines in Appendix C.

---

Table 2: Comparison in AntMaze and Gym-MuJoCo domains (v2). Results are averaged over 5 seeds; for each seed, we run 10 evaluation episodes. u: umaze (antmaze). um: amaze-medium ul: larze p. play. d: diverse. wa: walker2d. ho: hopper. ha: halfcheetah. r: random. m: medium.
u-d
um-p
um-d
ul-p
ul-d
wa-r
ho-r
ha-r
wa-m
ho-m
ha-m
mean
CQL
74.0
84.0
61.2
53.7
15.8
14.9
-0.2
8.3
22.2
82.1
71.6
49.8
44.8
IQL
87.5
62.2
71.2
70.0
39.6
47.5
5,4
7.9
13.1
77.9
65.8
47.8
49.7
PLAS
70.7
45.3
16.0
0.7
0.7
0.3
9.2
6.7
26.5
75.5
51.0
44.5
28.9
LAPO
86.5
80.6
68.5
79.2
48.8
64,8
1.3
23.5
30.6
80.8
51.6
46.0
55.2
DROP
90.5
92.2
74.1
82.9
57.2
63.3
5.2
20.8
32.0
82.1
74,9
52.4
60.6
DROP std
22.4
117
#3.9
43.5
#5.5
2244
116
10.3
42.5
25.2
+2.8
+2.2
under the iterative paradigm. By answering Q3, DROP permits test-time adaptation, enabling to dynamically stitch "skills"(latent behaviorseemeedding 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 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 the embedding of behaviors, instead of conducting return-guides task decomposition. We provide implementation details (DROP+CVAE) in Appendix A.4 and the comparison results in Table (taking SIOO--Grad-Aaee implementation). We can observe that our DROP consistently achieves better performance than PLAS in all tasks, and performs better comparably 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 offiin iierativennon-iterative baselines.
Computation cost. One limitation of DROP that conducting test-time adaptation at each state is usually expensive for inference time. To reduce the computation cost, we find that resample (infer) the best embedding z' after a certain time interval, i.e., z does not necessarily need to be updated each testing state. In Figure 5, we plot the average inference time versus different models and DROP imple- mentations with varying inference time intervals, We see that compared to Diffuser (one offline RL method that also supports test-time adapta- tion, see quantitative performance comparison in Table 4) [26], DROP can achieve significantly lower computation costs. In Figure top-right, we also illustrate the trade-off between perfor- mance and runtime budget as we vary the infer- ence time intervals. We find that by choosing 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.
Average inference time
DROP' scores on AntMaee
10'
000
200
(50)
100
1100
(200)
10
Diffiss
Diffurer
Oautty
DROP l)
DROP (o)
DROPLE
DROPL
DROP(100)
DROPE10
DROP2RD
Figure 5: Computation cost across different models and DROP implementations with vary- ing inference time intervals. The number in parenthesis denotes the inference time intervals. In top-right, we also provide the average perfor- mance of DROP AntMaze tasks when using different inference time intervals. (Diffuser* de- notes the warm-start implementation of Diffuser.)
Comparison a "distilled" DROP imple- mentation. Our answer to Q3 is that we con- duct outer-level optimization phase. Similarly, we can design a "distilled" DROP implementation: we keep the same treatment as DROP for Qland Q2, but we per- form lever- 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
Table 3: Comparison to "distilled" DROP (D-
DROP), using Bra and CVAE implementation.
the testing
ho-r
ha-r
wa-m
ho-m
ha-m
CQL
-0.2
8.3
22.2
82,1
71.6
49.8
IQL
5,4
7.9
13.1
77.9
65.8
47.8
D-DROP
4.5
18.9
27.8
84.2
67.1
50.9
DROP
5.2
20.8
32.0
82.1
74.9
52.4

---

z" into fixed rollout policy. As shown in Table 3, we can see that such a "distilled" DROP im- plementation achieves competitive performance, approaching DROP's overall score and slightly outperforming CQL. and IQI,, which further supports the superiority afforded by the non-iterative offline RI paradigm versus the iterative one.
Comparison to adaptive baselines. One merit of DROP resides its test- time adaptation ability across rollout states. Thus, here we compare DROP with two offline RL methods that also support time adaptation: 1) APE- [23] learns an uncertainty-adaptive policy and dynamically updates a con-
Table 4: Comparison to adaptive baselines (APE-V and Diffuser), DROP uses Grad-Ada and CVAE implementation.
wa-mr
ho-mr
ha-mu
wa-me
ho-me
ha-me
APE-V
82.9
98.5
64.6
110.0
105.7
101.4
Diffuser
70.6
93.6
37.7
106.9
103.3
88.9
DROP
83.5
96.3
50.9
109.3
107.2
102.2
textual belief vector at test states, and 2) Diffuser [26] employs score/return function to guide the diffusion denoising process, i.e., ccore-guided action sampling. As shown in Table 4, we can see that in medium-replay (mr) and medium-exper (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 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 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 implemen- tation). 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 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
[1] Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller. Maximum a posteriori policy optimisation. arXiv preprint arX.v:1806.06920, 2018.
[2] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Belle- mare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 01229334--9322 2021.
[3] Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-basee offline reinforcement learning with diversified qeensebble. Advances Neural Information Processing Systems, 34, 2021.
10

---

[4] Marcin Andrychowiez, 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.
[5] Chenjia Bai, Lingxiao Wang, Zhuoran Yang, Zhihong Deng, Animesh Garg, Peng Liu, and Zhaoran Wang. Pessimistic bootstrapping for uncertainty-dee offline reinforcement learning. arXiv preprint arXiv:2202. 1566,, 2022.
[6] David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline r1 without off-policy evaluation. Advances in Neural Information Processing Systems, 14:4933-4946,
2021.
[7] 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.
[8] 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.091 19, 2021.
[9] Yevgen Chebotar, Karol Hausman, Yao Lu, Ted Xiao, Dmitry Kalashnikov, Jake Varley, Alex Irpan, Benjamin Eysenbach, Ryan Julian, Chelsea Finn, al. Actionable models: Unsupervised offline reinforcement learning of robotic skills. arXiv preprini arXiv:2 2021.
[10] 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.
[11] Xi Chen, Ali Ghadirzadeh, Tianhe Yu, Yuan Gao, Jianhao Wang. Wenzhe Li, Bin Liang, Chelsea Finn, and Chongjie Zhang. varible advantage-weightee policy optimization for offline rl. arXiv preprint wxxx:2203.089442 2022.
[12] 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 1rXiv:1910.121799 2019.
[13] Yiming Ding, Carlos Florensa, Pieter Abbeel, and Mariano Phielipp. Goal-conditioned imitation learning. Advances in neural information processing systems, 32, 2019.
[14] Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning? arXiv preprint arXiv:2112.1075/ 2021.
[15] Ben Eysenbach. Xinyang Geng, Sergey Levine, and Russ Salakhutdinov. Rewriting history with inverse rl: Hindsight inference for policy improvement. Advances in neural information processing systems, 33: 17795-47955 2020.
[16] Clara Fannjiang and Jennifer Listgarten. Autofocused oracles for model-based design. Advances in Neural Information Processing Systems, 13:455--122556 2020.
[17] Justin Fu and Sergey Levine. Offline madel-based optimization via normalized maximum likelihood estimation. arXiv preprint arXiv:2102.07970, 2021.
[18] 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..e orgg//..or//.ooo abs/a//s//aa/ /2044. 07219.
[19] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021.
[20] 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.
[21] Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. arXiv preprint 2rXiv:2111.10364, 2011.
11

---

[22] Matthicu 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, ICMI. 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 2160-2169. PMLR, 2019.
[23] 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.
[24] Caglar Gulcehre, Sergio Gómez Colmenarejo, Ziyu Wang, Jakub Sygnowski, Thomas Paine, Konrad Zolna, Yutian Chen, Matthew Hoffman, Razvan Pascanu, and Nando de Freitas. Regu- larized behavior value estimation. arXiv preprint Stti::2003.09573.. 2021.
[25] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34, 2021.
[26] Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205. 099911, 2022.
[27] Ying Jin, Zhuoran Yang. and Zhaoran Wang. pessimism provably efficient for offline ri? In International Conference on Machine Learning, pages 5084-5096 PMLR, 2021.
[28] Yachen Kang. Diyuan Shi, Jinxin Liu, Li He, and Donglin Wang. Beyond reward: Offline preference-guided policy optimization. arXiv preprint arXiv:2305.16217, 2023.
[29] Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. Advances in neural information processing systems, 33:21810-211 823, 2020.
[30] 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.
[31] Iyya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learnig, arXiv preprint 1rxiv:2110006/10 2021.
[32] Aviral Kumar and Sergey Levine. Model inversion networks for model-based optimization. Advances Neural Information Processing Systems, 33:5126-51377 2020.
[33] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- policy q-learming via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
[34] Aviral Kumar, Xue Bin Peng, and Sergey Levine. Rewarrd-condiioneed policies. arXiv preprint 1912. 2019.
[35] 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.
[36] 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:
/
[37] Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning, pages 45-73. Springer, 2012.
[38] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and F Huang. A tutorial on energy-based learning. Predicting structured data, 1(0), 2006.
[39] Jongmin Lee, Wonseok Jeon, Byungjun Lee, Joelle Pineau, and Kee-Eung Kim. Optidice: Offline policy optimization via stationary distribution correction estimation. In International Conference Machine Learning, pages 6120-6130. PMLR, 2021.
12

---

[40] 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.
[41] 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.
[42] 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.
[43] 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.
[44] 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.
[45] Jinxin Liu, Hongyin Zhang, and Donglin Wang. Dara: Dynamics-aware reward augmentation in offline reinforcement learning. arXiv preprint amXiv:2203.06660 2022.
[46] 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.
[47] Yao Liu, Adith Swaminathan, Alckh Agarwal, and Emma Brunskill, Off-policy policy gradient with state distribution correction. arXiv preprint 190v:1904008473 2019.
[48] 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.09996, 2021.
[49] Yecheng Ma, Dinesh Jayaraman, and Osbert Bastani. Conservative offline distributional reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021.
[50] Tatsuya Matsushima, Hiroki Furuta, Yutaka Matsuo, Ofir Nachum, and Shixiang Gu. Deploymenn-efffccient reinforcement learning via model-basee offline optimization. arXiv preprint arXiv:2006.036677 2020.
[51] Ofir Nachum, Bo Dai, Ilya Kostrikov, Yinlam Chow, Lihong Li, and Dale Schuurmans, Al- gaedice: Policy gradient from arbitrary experience. arXiv preprint arXiv: 19/2:0/9/2.0077: 2019.
[52] Ashvin Nair, Murtaza Dalal, Abhishek Gupta, and Sergey Levine. Accelerating online rein- forcement learning with offline datasets. arXiv preprint arXiv:2006.00335 2020.
[53] Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09355, 2020.
[54] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weighttd regression:: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00/177 2019.
[55] Vitchyr Pong, Shixiang Gu, Murtaza Dalal, and Sergey Levine. Temporal difference models: Model-free deep rl for model-based control. arXiv preprini 1xXi::800.00008 2018.
[56] Shideh Rezaeifar, Robert Dadashi, Nino Vieillard, Léonard Hussenot, Olivier Bachem, Olivier Pietquin, and Matthieu Geist. Offline reinforcement learning anti-exploration. arXiv preprint arXiv:2106.06431, 2021.
[57] Mare Rigter, Bruno I Lacerda, and Nick Hawes. Rambo-rl: Robust adversarial model-based offline reinforeement learning, arXiv preprint arXiv:2204.1258)) 2022.
13

---

[58] 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.08:99 08336, 2020.
[59] David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference machine learning, pages 387-395. PMLR, 2014.
[60] Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojeiech Jaškowski, and Jürgen Schmidhuber. Training agents using upside-down reinforcement learning. arXiv preprint 1rXiv:1912.00287, 2019.
[61] Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[62] 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.
[63] Brandon Trabucco, Xinyang Geng, Aviral Kumar, and Sergey Levine. Design-bench, Bench- marks for data-driven offline model-based optimization. arXiv preprint arX1v:220..08550, 2022.
[64] 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, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florine Balcan, and Hsuan-Tiee Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6- 12, 2020, virtual, 2020.
[65] Michael Wan, Jian Peng, and Tanmay Gangwani. Hindsight foresight relabeling for meta- reinforcement learning. arXiv preprini (XXiv:2109.0003) 2021.
[66] 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. 15 15134, 2020.
[67] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.1133 2019,
[68] 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 10X.v:2005.08140, 2021.
[69] Rui Yang, Yiming Lu, Wenzhe Li, Hao Sun, Meng Fang, Yali Du, Xiu Li, Lei Han, and Chongjie Zhang. Rethinking conditiodd supervised learning and its connection to offline rl. arXiv preprint arXiv:2202.04477, 2022.
[70] 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.
[71] Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Masel- offline policy optimization. Advances in Neural Information Processing Systems, 01:14129-14144 2020.
[72] Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances neural information processing systems, 34:28954-2888 2021.
[73] Chi Zhang, Sanmukh Kuppannagari, and Prasanna Viktor. Brac+: Improved behavior regular- ized actor critic for offline reinforcement learning. Asian Conference on Machine Learning, pages 204-219. PMLR, 2021.
14

---

[74] Wenxuan Zhou, Sujay Bajracharya, and David Held. Plas: Latent action space for offline reinforcement learning. arXiv preprint arXiv:2011.07213. 2020.
[75] Zifeng Zhuang, Kun Lei, Jinxin Liu, Donglin Wang, and Yilang Guo. Behavior proximal policy optimization. arXiv preprint urXiv22302.13112 2023.
15

---

Appendix
A
More Experiments
A.1
Single Behavior Policy v.s. Multiple Behavior Policies
antmazz- meannazz-medium-pple
untmaaz--medine diverse
antmmaze-arggeplay
antimazc-arggedivers
Training lass
Triinin Jass
Triinin lass
0.2
0.2
Training
0.2
0.2
0.1
0.1
0.1
100
50
100
50
100
Training steps (k)
Tiinnin steps (k)
Tmining steps
Training stops (k)
10
50
100
200
500
1000
Figure 6: Learning curves of behavior cloning on AntMaze suites (*-v2) in D4RL, where the x-axis denotes the training steps, and the y-axis denotes the training loss. The number N in the legend denotes the number of sub-tasks. If N we learn a single behavior policy for the whole of- fline dataset.
In Figure 6, we provide empirical evidence that learning a single behavior policy (using BC) is not sufficient characterize the whole offline dataset, and multiple behavior policies (conducting task decomposition) deliver better resilience to characterize the offline data than single behavior policy.
A.2 Decomposition Rules
In DROP algorithm, we explicitly decompose offline task into multiple sub-tasks, over which we then reframe the offline policy learning problem as one of offline model-based optimization. this section, we discuss three different designs for the task decomposition rule.
Random(N NV, .): We decomposition offline dataset D := {r} into N subsets, each of which contains at most M trajectories that are randomly sampled from the offline dataset.
Quantization NN 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 sub-task) from each bin. Specifically, in the i-th bin, the quantized trajectories {T,) satisfy Raiin A i < Return(T,) S
Rmio + * (i + 1), where A (Ra.s- Runn), Return(TT) denotes the return of trajectory T;, and Rmax and Rmin denote the maximum and minimum trajectory returns in the offline dataset respectively.
N
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 fundamental impact 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 suitable task decomposition rule. In our paper, we adopt the Rank decomposition rule, as it demonstrates more robust performance shown in Figure 7. In Appendix A.4, we adopt the conditional variational auto-encoder (CVAE) 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 RI. 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, 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.
16

---

halfcheetah random
halfcheetah medium
halfcheetah mwe--me-edddmeeeeeeeeee expert
hallichettah replay
24
100
37
Nomsitod
2.3
36-
35
22
35
Rand
Quan
Rank
Rand
Quan
Rank
Rand
Quan
Rank
Rand
Quaan
Rank
hopper-random
hopper-medium
hepger-mediun-cxper
hopper-medium-cplay
Nomalizzd retum
100
Rand
Quan
Rack
Rond
Quia
Rank
Rend
Quan
Rank
Rand
Quan
Rack
Normaliad tetur
walker2drrandom
walker2d-medium
wallker2d-meddum-xpper
walkerr2d-mcdiun--rplay
100
50
Quan
Rank
Rand
Rand
Rank
Quan
Normazzed sourn
aniinaze-umazee
antimaze antmmaze-mediinng mediusn--lay
cottmaze-mediurn-diin
25
75
40
30
20
25
0
Rand
Quan
Rank
Rand
Quan
Rank
Rand
Quin
Rank
Rand
Quan
Rank
Nomallzal ccam
antmaze-large-play
antmaae--argge lanmazzz-arge--iversse
30
Rand DROP Grod
Rank DROP Giaad
20
DROP-Bet-Aaa
25
Ruad DROPP Crod Ada
Rank DROP Grad-Ada
Quan DROP-ltestAdds DROP-Grul-daa
Quaan
Kank
Quan
Figure 7: Comparison of three different decomposition rules on D4RL MuJoCo-Gym suite ("-v0) and AntMaze suite (*-v2), where "Rand", "Quan" and "Rank" denote the Random, Quantization, and Rank decomposition rules respectively. We can find across 18 tasks (AntMaze and MuJoCo-Gym suites) and 3 embedding inference methods (DROP-Gres DROP-Best-ARe and DROP- Arad-Add-),,, Rank is more stable and yields better performance compared with the other two decomposition rules.
ha halfchectal:- -ncdium--cpiller
hopper- crnppe--edium--egie
m-eekeeddmedium-eee teplay
witmaze-largg-plag
Sgisook team
5000
5000
Peroonk remm
3000
Epiunk coum
3000
0.75
5000
2000
Untiode
0.50
2000
sund-1
1000
0.25
50
L00
50
100
50
100
Training
Training
Training
Troining
Figure 8: Learning curves of DROP, where the x-axis denotes the training steps (k), the y-axis denotes the evaluation return (using DROP-Best embedding inference rule). We only show two seeds for legi- bility.
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 We can find that in AntMaze tasks, the overall performance of DROP is higher than that of F-BC. For the MuJoCo- Gym suite, DROP-dased 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 more effective than simply performing imitation learning on subset of the dataset.
A.3 Online Fine-tuning
Online fine-tuning (checkpoint-eeve), In Figure we show the learning curves DROP-Best on four DR4L tasks. We can find that DROP exhibits high variance (in performance) across training
17

---

Table 5: Comparison between our DROP (using the Rank decomposition rule) and filtered behavior cloning (F-BC) on D4RL AntMaze and MuJoCo suites (*-v2). We take the baseline results of BC and B-BC from Emmons et al. [14], where F-BC is trained over the top 10% trajectories, ordered by the returns. Our DROP results are computed over seeds and 10 episodes for each seed.
Tasks
BC
F-BC
DROP-Grad
DROP-Bese Beee-Best-Aee Ada
DROP- Grad- Ada
antmaze-umaze
54.6
72 17.2
11.7
80
1 12.6
cemmaazz-umaz--diveee
48 22.3
66 =
anlmaze-medium-ppay
0
42.1
24 10.2
34 12
30 : 21
aatmaaze-mediun-diveeees
0
37.2
20 19
24 12
30 + 16.7
antmaze-large-plan
0
28
24 8
36 17.4
42 = 17.2
antmaze-lange-diverse
0
34.3
14 14:88 8
20 1 14.1
26 1 13.6
halfchectah
random
2.3
2.3
+ 0
2.3 +
2.3 +
hopper
random
4.8
5.1
+ 0.8
5.4 0.7
5.5 + 0.6
walker2d
random
1.7
1.7
1.7
1.6
3
halfcheetah
medium
42.6
42.5
42.4
0.7
42.9 0.4
43.1 0.4
hopper
medium
52.9
56,9
57.5
60.3 6.1
59.5 5.1
walker2d
medium
75.3
75
76.5 =
75.8 =
79.1 = 1.4
halfcheetah medium-replay
36.6
40.6
39.5
40.4 0.8
40.3 1.2
hopper
medium-repla
18.1
75.9
48 = 17.7
83.4 - 6.5
87.4 221
walker2d
medium-replay
26
62.5
37.4 13.5
60.9 = 7.4
61.9 = 2.3
steps". 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 LDROP-Grad, DROP-Beste Best-Ada and (see Algorithm in Section for more details): we evaluate each of the latest T checkpoint models and choose the best one that leads to the highest episode return.
we propose to conduct checkpoint-levé online fine-tuning
In Figure 9, show the total normalized returns across all the tasks in each suite (including Maze2d, AntMaze, and MuloCo-Gym) We can find that most tasks, fine-tuning (FT) can guarantee performance improvement. However, we also find such fine-tuning causes negative impacts on performance in AntMaze(*-vO)) 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, Ada) at the test time. Such finding also implies that the success of DROP's test-tine sst-time adaptation is not entirely dependent on the best embedding across sub-tasks 10 (i.e., the best embedding z (Su) in DROP Best),, but requires switching between some "suboptimal" embeddings (using -eest-Adaa or extrapolating new embeddings (using DROP-Grad-ARe -rPdrrad-Add..
PHOP-Best-Ada and DROP-Grade
Online fine-tuning (embedding-level). Beyond the above checkpoint-level fine-tuning procedure, can also conduct ombedding-lever online fine-tuning: aim to choose suitable gradient update step for the gradient-basee baseed embedding inference rules (including DROP-Grad and DROP-Grad-AMe Gra).Ada).. Similar to the checkpoint-level fine-tuning, we first conduct the test-time adaptation procedure (DROP-Gree and DROP- SROOP-GradAMee 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).
"In view of such instability, evaluate our methods over multiple checkpoints for each seed, instead choosing the final checkpoint models during the training loop (see detailed evaluation protocol in Appendix B).
"Conversely, if the performance of DROP depends on the best embedding PROP-Best),, then the checkpoint model we choose by fine-tuni mne-tuning DROP- Best should enable consistent performance improvement for rules that perform embedding inference with DROP-Best Best-Ada and DROP- Ada. However, we find performance drop in AntMaze suite, which means there is no explicit depen- dency between the best embedding z (So) and the inferred embedding using the adaptive inference rules.
sub-tasks (i.e., z (so) in
18

---

Total nomman wonnaaaaaddddaooo retum retum
Maze2d
AntMaze(-v0)
MiJoCo-Gyml-vd)
AntMazo(-v2)
MuJoCo-Gyml-2)
50
100
250
100
250
DROP- Gmal
DROP God FT
DROP Se--BBtt-hhhl hest-Ade
DROP Best Ada PT
DROP-Grdd
DROP-Gimt--id FT.
DROP Gind
DROP DROP-GGo -P---God FT
DROP Se---Btt--ll sest-Ade
DROP Best Ada
DROP
DROP-Giai--in
DROP Gidd
DROP DROP--Go CoP-Godd
DROP -e---et--a sut--Ade
DROP Best Ada
DROP-
PROOP-Gmml-Ain
DROP Cind
DROP DROP-God CPP-GGod
DROP Se--Bet--hal Best-Ade
DROP Best Ada
DROP Coo--iidd--d
DROP Godd dadd-da Ada Fl
DROP Cind
DROP PRoP-God
DRUP see---e-dll
DROP- MRoP-Ben-AAl
DROP -rad--rdd-d Add
DROP Godd adad-Ada FT
Figure 9: Total normalized returns across all the tasks in Maze2d, AntMaze, and MuJoCo- CalloC suites.
Table 6: Online fine-tuning results (initial performance > performance after online fine-tuning).
The baseline results AWAC, CQL, and IQL are taken from Kostrikov et al. [31], where they 1M online steps to fine-tune the learned policy. For our DROP method (DROP-Grad and DROP-Grad-Adie we run 0.3M (= Gcheckpoint X k0.... 1000steps per episode online steps to fine-tune (embedding- level) the policy, i.., aiming to find the optimal gradient ascent step that is used to infer the contextual embedding z" (so) or z" (se) for x (ar/s;) := 3(ay (s. (see Algorithm 4 for the details). Morcover, for medium- and large-* * tasks, conduct additional parametric-level fine-tuning, with 0.7M online steps to update the policy's parameters.
Task (+-v0)
CQL
IQI.
DROP-Grad
DROP- Grap-Grad-Ada
umaze
70.1 + 99.4
86.7
+ 96
70 + 96
= 1.2
76 + 98 0
umaze-diverse
31.1
+ 99.4
75
+
+ 88
8
66 + 94 4.9
medium-play
0 23
+
72
+
56
+ 8.9
30 + 50 6.3
- 94 2.9
medium-diverse
+ 32.3
68.3
+
12 + 44
4.9
22 38 4.9
-) 96 0.8
large-play
1
25.5
38
8.9
6.3
> 53 1.3
large-diverse
42.6
60.7
40
13.6
22
10.2 + 58 4.5
03M
In Table 6, we compare our DROP (DROP-Grad and GROP-Grad-Adoe to three offline RL methods (AWAC [53], CQL [35] and IQL I1,, reporting the initial performance and the performance after online fine-tuning. We can find that the embedding- improvement performance. The fine-tuned DROP- Caap--rrr--dd 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 1M online rollout steps, while DROP-based fine-tuning takes 0.3M steps). However, there is still big gap between the fine-tuned IQL and the embedding-level fine-tuned DROP (0.3M). Considering that there remain 0.7M online steps in the comparison, we further conduct "parametric-level" fine-tuning (updating the parameters of the policy network) for our DROP-Gradi Ada on medium- and large-* tasks, we can find which achieves competitive finettuning performance compared with IQL.
fine-tuning (0.3M) enables a significant
A.4 DROP + CVAE Implementation
CVAE-based embedding learning. Similar to LAPO [11] and PLAS [74], we adopt the conditional variational auto- anto-encoder (CVAE) to model offline data. Specifically, we learn the contextual policy and behavior embedding:
3 [|ass, 2), (z)8) +- arg max ((B.)))-BB()) 3,0
[log, ((as, (l((((2))))lp(())
(10)
Then, we learn the score model with the TD-error and the conservative regularization:
arg min ['...''''''''' a'D (R(s, a) + xJ(s's al, o(z)s)) f(s, a, a(sss (11)
s.t. Eg am [f a, z)] EB-D 1ol.p_ [f z)] S ?
where i denotes target network and u(z) denotes the uniform distribution over the Z-sacc.
19

---

In testing. we also dynamically adapt the outer-level optimization, setting policy inference with (a|s)
(( z" (s)). where 2" (s) arg max; f(s, ((a)s,)))
A.5
Ablation Study
Note that our embedding inference depends on the learned score model f. Without proper regu- larization, such inference will lead to out-of- distribution embeddings that are erroneously high-scored (Q2). Here we conduct an ablation study examine the impact of the conservative regularization used for learning the score model.
100
75
Grad
Grad-Ada
Grad wio Reg
25 Grad-Ada wio Reg
150 Figure 10: Ablation the conservative regular- ization. The y-axis represents the normalize re- turn, and the x-axis represents the number of gradient-ascent steps used for embedding infer- deployment. We plot each random seed as a transparent line; the solid line corresponds to the average across 5 seeds.
120
In Figure 10, compare DROP-Grad and DROP-Grad-Aee to their naive implementa- tion (w/o Reg) that ablates the regularization We can find that re- moving the conservative regularization leads to unstable performance when changing the update 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, embeeding--condttioned behavior policy can correct such deviation.
B Implementation Details
For the practical implementation of DROP, we parameterize the task embedding function d(z(n). the contextual behavior policy [[ass, z). and the score model f(s, a, z) with neural networks. For Equation the main paper, we construct a Lagrangian and solve the optimization through primal- dual 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-of- distribution extrapolation at inference.
To optimize the constrained objective in Equation in the main paper, we construct Lagrangian and solve the optimization through primal-dual gradient descent,
Lagrangian relaxation.
min max ASO Ep..- (' [ (R(s, a) - a', ((z(n)) f(s, a, d(z))))) +
A (Bm,tee) B BawD. ar [f(s, a, z)] - En, ((x (alss2)) Dr [f a, 2)] 11)
This unconstrained objective implies that if the expected difference in scores of disti--iist embeddings and in-distribution embeddings is less than a threshold 7, À going to be adjusted to 0, on the contrary, À is likely to take larger value, used to punish the over-estimated value function. This objective encourages that out-of-distributtin embeddings score lower than in-distribution embeddings, thus performing embedding inference will not lead these out-of-distributtion embeddings that falsely and over-optimistically scored by the learned score model.
In our experiments, 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 77 2.0.
Table 7, we provide the prrea-parameters of the task embedding d(z)s), the contextual behavior policy ((a ss, z). and the score function f(, z). For the gradient ascent update steps (used for embedding inference), we set = 100 for all the embedding inference rules in experiments.
Klyper-parameter
In Table 9, we provide the number of sub-tasks, the number trajectories in each sub-task, and the dimension of the embedding for each sub-task (behavior policy). The selection of hyperparameter N based on two evaluation metries: (1) the fitting loss of the decomposed behavioral policies the offline data, and (2) the testing performance of DROP. Specifically.
20

---

Table 7: Hyper-paramen of the task embedding function d(z)s), the contextual behavior policy
[(a)s, z). and the score function fss, a, z). The task embedding function (zz)s): z <- Enc_ [n__ The contextual behavior policy ((a)s, z): +- Enc 2(z, Enc Ene_1(s) (s)). The score function f a, 2):
+- Enc_ [nc__(z, Enc ann_3(s,a0 a)).
Enc 0
Enc
Enc 2
Enc 3
Enc 4
Optimizer
Adam
Adam
Adam
Adam
Adam
Hidden layer
2
3
Hidden dim
512
512
512
512
512
Activation function
ReLU
ReLU
ReLU
ReLU
ReLU
Learning rate
1.00E-03
1.00E-03
1.00E-033
1.00E-03
1.00E-03
Mini-batch size
1024
1024
1024
1024
1024
(Step1) Over hyperparameter (the number of sub-tasks) set, we conduct the hyperparam- eter 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 hyper- meters 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 (fil- tered) hyperparameters using seed, then evaluate the best hyperparameter by training an additional 4 seeds and finally report the results on the total seeds (see next "evaluation pro- tocol"). 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 use the above hyperparameter selection rules (Stepl and Step2) to choose the number.
Table Hyperparameter (the number of sub-tasks) set.
tasks
the number of sub-tasks
Antmaze
500 (v0), 150 (v2)
Gym-mujoco
10, 20, 50, 100, 200, 500, 800, 1000 10, 20, 50, 100, 200, 500, 800, 1000
Adroit
Baseline details. For the comparison of our method to prior iterative offline RL methods, we consider the v0 versions of the datasets in D4RL" 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), RyS-R, and Onestep are taken from Emmons et al. [14]. In our implementation of COMs, take the parameters (neural network weights) of behavior policies 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 seeds. For each seed, instead of taking the final checkpoint model produced by 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 5sced X 6checkpoint 10ppisede rollouts.
Online fine- tuning (checkpoint-level); Instead of re-training the learned (final) policy with online rollouts, we fine-tune our policy with enumerated
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.
We noticed that Maze2D-v6 in the D4RL dataset (https://rail used vI version instead in our experiment. For simplicity, still use vū in the paper exposition.
http://raalllee...... eecs. wer...eeeeee.e.ddaaattaae
is not available, so
21

---

Algorithm 3 DROP: Online fine-tuning (checkpoint-level)
Require: Env, last I checkpoint models: (a and fi(s, 1,... ,7).
1: RMAX-
2: Seen None.
3: fuent None.
4: while ,
5: so Env.Reset(
6: (50) < Conduct embedding inference with DROP-Best.
7: Return Evaluate 3. and Env, setting (a|s) B[ass, 20 (80))
RMaA < Retur then
9: Update the best checkpoint models: Breat + 3t. frea +- Sa.
10: Update the optimal return: RMAX +- Return.
11: end if
12:
end while
Return: Brest and fnest-
Algorithm 4 DROP: Online fine-tuning (embedding-level)
Require: Env, last T checkpoint models: (Ba(a)s, z) and fi 1,,
1: RMMA -00.
2: Bres + None.
3: fost < None.
4: Kbea 0.
5: while = 1, 1. max
do
while
7: So = Env.Reset().
Conduct embedding inference with DROP- Grad or DROP-Grad-AR
8:
Return +- Fvaluate 9c and fe Env, setting n (als) [Baass, (5o)) B(ass, x' (s)), where conduct gradient ascent steps obtain z" (8o) RMAX < Return
(s).
10:
Update the best checkpoint models: Shosa +- Br, Frea t St.
11:
Update the best gradient update step: knea <- ki Update the optimal return: RMAX +- Return.
12:
13:
end if
14:
end while
15:
end while
Return: Beest feom and Ribes.
Online fine-tuning (embedding-level): The embedding- fine-tuning aims to find suitable gradient ascent step that is used to conduct the embedding inference in DROP-Grad or DROP-Grad- Ada. Thus, enumerate list of gradient update steps and pick the best update step (according to the episode returns).
Codebase. Our code is based on d3ripy: https: //g thub. hot..//////akkkeee///rryy We provide our source code in the supplementary material.
Computational resources. The experiments were run 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-AMe 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 tasks and outperforming (or comparing) these value regularized algorithms (BEAR, BCQ, and CQL) 19 out of 33 tasks.
22

---

Table 9; The number (N) of sub-tasks, the number (M) of trajectories in each sub-task, and the dimension (z) the embedding for each sub-task.
Parameters (+-v0)
Parameters (*-v2)
Domain
Task Name
NV
M
dim(z)
N
M
dim(z)
umaze
500
5
5
Maze 2D
medium
150
50
5
large
100
umaze
500
150
50
umaze-diverse
150
50
5
Medium-play
500
50
5
150
50
5
Antmaze
Medium-diverse
500
50
5
150
50
5
Large-play
500
50
150
50
5
Large-diverse
150
50
5
random
1000
1000
5
medium
100
2
100
5
halfcheetah
medium-expert
1000
5
1000
5
medium-replay
50
10
5
50
10
5
random
100
2
5
100
5
medium
100
5
5
100
5
hopper
medium-expert
100
100
medium-replay
10
30
5
random
500
500
medium
50
50
5
walker2d
medium-expert
50
5
5
50
5
medium-replay
1000
5
5
10
50
5
cloned
1000
2
5
door
expert
500
5
human
cloned
1000
hammer
expert
500
5
human
20
3
5
cloned
500
5
5
expert
500
5
5
human
50
5
5
cloned
500
5
5
relocate
expert
500
5
human
4
Ablation studies. In Figure 11, we provide more results for the ablation of the conservative regularization Equation in the main paper. We can find that for the halfcheetah-mediun and hopper-medium tasks, the performance of DROP-Grad Grad-Ada w/o Reg depends on the choice of the gradient update steps, showing that too small too large number of gradient update step deteriorates the performance. Such result is also consistent with COMs [62], which also observes the sensitivity of naive gradient update (i.e., w/o Reg) the number of update steps used for design input inference. By comparison, the conservative score model learned with GROP-Grad-Ada exhibits more stable and robust performance to the gradient update steps.
Further, we also find that in malker2d-medinee and malker2d-mmdiune medium-expert expert tasks, the naive gradient update (w/ Reg) does affect performance significantly across 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-saffegard" is special merit in the offline RL domain 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.
23

---

halfchectall-meddee mediimm
hopper-medium
Polker2d-meddinee
walker2d-medium- expert
40
50
26
Normaized
30
24
80
20
30
60
90
120
150
30
60
90
120
150
01
30
601
90
1.00
1500
30
60
901
120
150
Dro--Giad-Ada
DROP-Grad-Ada Reg
Figure 11: The performance comparison DROP-Grade PRRO-Grad-Adke and DROP-Grad-Are w/o Reg, where ablate the conservative regularization for the w/o Reg implementation. The y-axis denotes the normalized return, the x-axis denotes the number of gradient-ascent steps used for embedding inference deployment.
Table 10: Comparison of our method to prior offline methods that perform iterative (regularized) RL paradigm on D4RL. We take the baseline results of BEAR, BCQ, CQL, and BRAC-p from Fu et al. [18], and the results of TD3-BC from Fujimoto and Gu [19]. For all results of our method (DROP), average the normalized returns across 5 seeds; for each seed, For proper comparison, we use and to denote DROP (*-Ada) achieves comparable or better performance compared with value and policy regularized offline RL methods respectively.
10 evaluation episodes.
Value Reg.
Policy Reg.
DROP-
Task Name
BEAR
BCQ
CQL
BRAC-p
TD3+BC
Grad
Best-Ada
Grad-Ada
umaze
73.0
78.9
74.0
50.0
72.0
78.0
80.0
umaze-diverse
61.0
55.0
84.0
40.0
48.0
62.0
66.0
antmaze
medium-play
0.0
61.2
0.0
34.0
30.0
medium-diverse
8.0
0.0
53.7
0.0
20.0
24.0
A
30.0
large-play
0.0
6.7
15.8
0.0
24.0
36.0
42.0
large-diverse
0.0
2.2
14.9
0.0
14.0
20.0
26.0
halicheetah
random
25.1
2.2
35.4
24.1
10.2
2.3
2.3
2.3
medium
41.7
40.7
44.4
43.8
42.8
42.4
42.9
43.1
medium-expert
53.4
64.7
44.2
97.9
86.6
88.5
88.9 A
medium-replay
38.6
38.2
46.2
45.4
43.3
39.5
40.4
40.3
random
11.4
10.6
10.8
11.0
11.0
5.1
5.4
5.5
hopper
medium
52.1
54.5
58.0
32.7
99.5
57.5
60.3
59.5
medium-expert
96.3
110.9
98.7
1.9
112.2
103.5
102.5
105.9 A
medium-replay
33.7
33.1
48.6
0.6
31.4
48.0
83.4
87.4
random
7.3
4.9
7.0
-0,2
1.4
2.8
3.0
3,0
walker2d
medium
59.1
53.1
79.2
77.5
79.7
75.8
medium-expert
40.1
57.5
111.0
76.9
101.1
107.5
106.8
106.9
medium-replay
19.2
15.0
26.7
-0.3
25.2
37.4
60.9
61.9
cloned
-0.1
0.0
0.4
-0.1
0.5
2,5
2.7 A
door
expert
103.4
99.0
101.5
-0.3
98.6
102.2
102.6
human
-0.3
0.0
9.9
40.3
3.3
1.9
3.0
hammer
cloned
0.3
0.4
2.1
0.3
0.3
0.3
0.3
expert
127.3
107.2
86.7
0.3
65.7
73.3
77.7
human
0.3
4.4
0.3
1.1
0.3
2.1
cloned
26.5
44.0
39.2
1.6
76.7
82.4
expert
105.9
114,9
107.0
-3.5
113.1
118.6
116.7
human
-1.0
68.9
37.5
8.1
71.1
85.2
81.5
relocate
cloned
-0.3
-0.3
-0.1
-0.3
0.1
0.5
0.2 A
expert
98.6
41.6
95.0
-0.3
2.5
6.2
5.4
human
-0.3
-0.1
0.2
-0.3
0.0
0.0
0.0
24


Reviews:
[{"review_id": "JBq8HEwo1V", "reviewer": "Reviewer_YuS3", "summary": "This 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.", "questions": "(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", "limitations": "The authors discuss several limitations of their work, and the reviewer personally agrees with that.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "(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", "weaknesses": "(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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**About Weaknesses:**\n\n**(a): concerns about the ideas.**\n+ *the difference between CQL and DROP:* The main difference is that DROP performs non-iterative bi-level offline RL optimization, while CQL performs iterative bi-level offline RL optimization (see Figure 1 in the main paper).\n+ *comparison to [1, 2, 3] and more advanced algorithms:*  Please refer to the global response for the comparison. \n+ *goal-conditioned methods ought to be included numerically:* Thank you for the suggestion. We will include it in our revision. \n+ *stronger methods are expected:* See CCVL in the  global response. \n+ *criteria in Table 1:* First, in Appendix B.5, we have conducted an ablation study to examine the impact of the conservative regularization used to learn the score model (Q2A2), and find that removing such regularization leads to unstable performance, thus demonstrating the effectiveness of Q2A2. Second, in Figure 3, we also demonstrate the effectiveness of answering Q3A3, i.e. that performing adaptive policy inference can produce better results. \n+ *concerns on the task decomposition rule:* First, we point out that the purpose of performing task decomposition is to establish a link between offline RL and MBO, thereby analyzing different offline RL methods from a unified perspective. Second, the purpose of performing task decomposition (using returns) is to provide a fair comparison with previous approaches (i.e., RvS) in the non-iterative paradigm, thereby demonstrating the superiority of DROP's design (by answering Q1Q2Q3). Third, DROP can treat each trajectory as a single task, as in our DROP+CVAE implementation (lines 296-306), which does not require complex task decomposition. Our experiments also show the effectiveness of DROP. Thus, the algorithm designer does not need to worry much about the choice of decomposition rule in deployment, just take the DROP+CVAE implementation (treating each trajectory as a single subtask). \n\n\n\n\n**(b) whether the score function will overestimate even with the introduced regularization** \n\nIf we assume that the problem (overestimation, raised by the verification) occurs, there is actually an implicit advantage of DROP: the overestimation occurs on top of the inference to z, not on top of the actions. That is, even if the inference yields an OOD z, the contextual behaviour policy still has the ability to produce in-distribution actions. We have also done an ablation study on this in Appendix D (lines 383-428) and find that the learned behaviour policy can protect against the shift of the inferred embeddings. Similarly, the same benefit (implicit OOD regularisation) is also observed in the PLAS paper. Thus, we believe that there is no need to be overly concerned about DROP's overestimation.\n\n\n**(c) concerns on the experiment section.**\n+ *CQL results:* Our CQL results in Antmaze domain are taken from the RvS paper and results in Gym domain are taken from the LAPO paper. \n+ *Environment version in Table 2:* All results are on the \"v2\" version datasets, except the results of LAPO on antmaze, which are on the \"v1\" dataset. To unify the dataset versions used, we re-implement the LAPO algorithm (not released by the authors) and report the results. We refer the reviewer to question (7) in the authors' response to reviewer SsTF for our reproduced results. We will include them in our revision of the paper. \n+ *which variant of DROP in Table 2:* DROP-Grad-Ada.\n+ *unify the versions of the dataset used:* Thank you for the suggestion. We will unify it in our paper revision. \n+ *searching the best number of subsets:* In our experiments, we take the number of subsets as a hyperparameter and choose it with typical hyperparameter tuning strategies. In response to the review's concern about algorithmic complexity, I don't think it's worth worrying about. Because in the implementation we can treat each trajectory as a separate task, similar to the DROP+CVAE implementation in the paper. This simple implementation can give better performance, and at the same time, the algorithm designer does not have to worry much about the decomposition rule.\n+ *large batch size:* We respectfully point out that the issue raised by the review in paper [4] does not make sense. Paper [4] does emphasize the importance of large batch size, but this paper is still based on Q-ensemble (eliminates OOD problems). Our approach to OOD problems has nothing to do with this paper (and Q-ensemble ideas). Large batch size can improve the robustness of model training, and we respectfully don't think we need to force all offline RL work (and even future works) to set 256 sizes for the sake of \"fairness\". \n\n**(d) concerns on real-world applications.**\n\nWe note that the reviewer's main concern is that real-world deployment requires real-time inference and that this is time consuming. We think the reviewer is overly concerned about this. First, even if our z is fixed (updated only at s_0), we have found experimentally that DROP still outperforms the RvS and decision-transform baselines. Second, we do not actually need to run inference every time. We found that we can still get good results by running inference every 10 steps. We tested the inference time on a very early RTX 2080 Ti, which only takes about 0.01 seconds to infer. In particular, in many of today's robotic arm tasks, the speed of the arm's movement is quite slow, so 0.01 seconds will basically have no effect. Thirdly, our model only uses the MLP network, and the inference speed is much slower than that of transformer/diffusion based methods. Therefore, we do not believe that inference time will be an obstacle for DROP in real applications. \n\n\n\n**About Questions:**\n\n**(a) Why do you choose yes to theory?**\n\nWe apologize for any confusion about this. We will correct it.\n\n**(b) train the contextual behavior policy in an IQL-style manner**\n\nWe thank the reviewer for pointing out this option. We refer the reviewer to the global response for the comparison.\n\n"}, {"review_id": "YSyQOC75sS", "reviewer": "Reviewer_7FJj", "summary": "The 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.", "questions": "1. 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.", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. 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", "weaknesses": "1. 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the positive comments. \n\n**Q1: which of the three decomposition rules from B.2 is used?**\n\nWe use Rank(N, M) in the main paper. See Line 035 in the appendix. We will clarify it in our main paper. Thank you.\n\n**Q2: Rank(N, M) rule: My guess is first M trajectories (with the highest returns) become task 1, the next M trajectories become task 2, and so on.**\n\nYes, you are right!\n\n**Q3: 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?**\n\nYes, you are right. It is true that some subtasks may contain only low-return trajectories. The main motivation is that we expect to build *a diverse set of task distributions*. Then trajectories of different qualities will have different latent embeddings, which will benefit contextual policy learning and optimal embedding inference (test-time adaptation). \n\n**Q4: an ablation to see the effect of N on the performance.**\n\nGood suggestion. We report the ablation results in the following table. \nWe can see that when N is very small, DROP basically performs worse than Onestep, proving that when the number of subtasks is small, updating z when doing outer-level optimization does not bring as much benefit as updating the policy directly (Onestep). A general trend can be observed that DROP becomes progressively better as N increases and outperforms Onestep for some values, suggesting that optimizing for low-dimensional z (out-level inference) brings more benefits at these points. However, we also note that larger N is not always better, and that performance degradation can occur when N is too large. We speculate that this is due to the fact that when N is too large, the learning (of the corresponding behavioral policies and Q-values) is underfitting and instead leads to worse performance in the end.\n\n|    | 1  | 2  | 5  | 8  | 10  | 20  | 50  | 100  | Onestep  |\n|  ----  | ----  | ----  | ----  | ----  | ----  | ----  | ----  | ----  | ----  |\n|  hopper-medium-replay-v2  | 65.9| 67.6  | 82.0| 79.5  | **87.4**  | 87.0  | 85.9  | 77.8  | 77.3 |\n|  halfcheetah-medium-replay-v2  | 33.8| 35.7| 34.8| 35.8| 34.1| 39.5  | **40.3**  | 38.6  | 38.3 |\n|  walker2d-medium-replay-v2  | 53.8| 57.9| 60.6| 59.7  | 61.9  | 59.9| 60.6| 56.9  | **66.4** |\n\n\n\n**Q5: Is the computation time affected by N?**\n\nIn fact, N has a negligible effect on the inference time. \n\n**Q6: What is the advantage of using task embeddings instead of directly fitting an MBO on the collected data?**\n\nThe main advantage of explicitly using task embedding is that we can perform test-time adaptation by exploiting the sequential structure of RL tasks, rather than simply performing inference at the beginning of the test rollout (which is what we do when we fit a simple MBO to the collected data). Empirically, we also find that test-time adaptation can yield better results than fitting a simple MBO (i.e., baseline COMs).  \n\n**Q7: the version of the environments in Table 2.**\n\nAll results are on the \"v2\" version of the datasets, except for the results of LAPO on antmaze, which are on the \"v1\" dataset. To unify the versions of the dataset used, we re-implement the LAPO algorithm (not released by the authors) and report the results. We refer the reviewer to question (7) in the authors' response to reviewer SsTF for our reproduced results. We will include the new results in our revision of the paper.\n\n"}, {"review_id": "JVOhXXzlaU", "reviewer": "Reviewer_hhjh", "summary": "The 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.", "questions": "How 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?", "limitations": "Limitations 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", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 4, "strengths": "I 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.", "weaknesses": "While 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nThank you very much for the insightful comments! \n\n**(1) description of the three questions and answers.**\n\nThank you for the suggestion, we will clarify this in our revision of the paper. \n\n**(2) term: \"non-iterative, bi-level optimization\". I think what you mean is that there is no iteratively learned value function that bootstraps off of itself and thus propagates errors further and further.**\n\nYes, you are right. We will clarify this in the introduction. \n\n**(3) comparison to a \"distilled\" DROP implementation, i.e., moving the outer-loop into the optimization and distilling the learned policy + conditions into a single fixed model.**\n\nThank you for such a valuable suggestion! As you suggested, we have performed new experimental comparisons under the DROP+CVAE implementation (performing inference with DROP-Grad-Ada). We can see that such a \"distilled\" DROP implementation achieves competitive performance, approaching DROP's overall score (though still slightly below DROP). It really is a good proposal! \n\n|    | DROP  | \"distilled\" DROP   |\n|  ----  | ----  | ----  |\n|  walker2d-random-v2  | **5.2** $\\pm$1.6  | 4.5 $\\pm$0.8  |\n|  hopper-random-v2  | **20.8** $\\pm$0.3  | 18.9 $\\pm$0.2  |\n|  halfcheetah-random-v2  | **32.0** $\\pm$2.5 | 27.8 $\\pm$1.54  |\n|  walker2d-medium-v2  | 82.1 $\\pm$5.2  | **84.2** $\\pm$2.4  |\n|  hopper-medium-v2  | **74.9** $\\pm$2.8  | 67.1 $\\pm$2.1  |\n|  halfcheetah-medium-v2  | **52.4** $\\pm$2.2  | 50.9 $\\pm$1.7  |\n|  sum  | **267.4**  | 255.4  |\n\n\n**(4) The authors mention and qualitatively compare prior \"non-iterative\" algorithms (RvS, Onestep, COMs, F-BC), yet do not include them in their empirical evaluation.**\n\nWe refer the reviewer to Figure 4, where we provide the empirical comparison to these baselines. \n\n**(5) comparison to recent works [1,2,3].**\n\nWe thank the reviewer for raising these related works. We refer the reviewer to the global response for the comparison. \n\n\n[1] Confidence-Conditioned Value Functions for Offline Reinforcement Learning.\n\n[2] User-Interactive Offline Reinforcement Learning. \n\n[3] Offline rl policies should be trained to be adaptive.\n\n**(6) Why did you not use all of the available locomotion ones?**\n\nThis is mainly due to page limitations, so we present only some of the comparative experiments in the main text, and include more results in the appendix.\n\n**(7) standard deviations for the other algorithms?**\n\nThank you for pointing this out. The mean of most of the baselines is outside of DROP's 95% CI. We will include this in our revision. Thank you very much. \n\n**(8) What do you mean exactly with non-iterative?**\n\nWe know that offline RL 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 propose to 1) introduce a policy/value regularization in the iterative loop or 2) try to eliminate the iterative loop itself. The term “non-iterative” refers to the paradigm of eliminating the iteration loop. \n\n**(9) Why are the non-iterative prior works not part of the empirical evaluation?**\n\nThe main starting point is that we have already compared DROP and non-iterative methods in Figure 4, so we did not list the non-iterative results in Table 2. Thank you for raising this point. We will include it in our revision. \n\n\n"}, {"review_id": "4ay19rrduA", "reviewer": "Reviewer_SsTF", "summary": "This 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.", "questions": "In 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.", "limitations": "The limitations of the number of subtasks and how to split subtasks are adequately discussed in the paper.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This 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. ", "weaknesses": "The 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)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the insightful comments.\n\n**(1) the terms \"inner level\" and \"outer level\", and the raised three questions.**\n\nThank you for your valuable suggestion. We will include it in our revision. Thank you very much! \n\n**(2) Comparison to related works in offline RL with test time adaptation.**\n\nWe thank the reviewer for raising the related work. We refer the reviewer to the global response for the comparison. \n\n\n**(3) In Figure 5, is it possible to also include the inference time of other baselines (COMs, RvS-R, etc)?**\n\nWe believe that it is not very necessary to include the inference time of COMs or RvS-R because they only perform inference at s_0 (the initial state of a test rollout). The subsequent test time steps do not require any additional inference about the context variables. Therefore, we did not include their inference time in Figure 5.\n\n**(4) How do Onestep and DROP compare when they use the same amount of computation?**\n\nThank you for this suggestion. We implement this comparison by searching for a time interval that makes DROP's inference time consistent with Onestep, and then we compare Onestep and DROP+CVAE (with the searched inference interval, DROP-Grad-Ada implementation). We show the experimental results in the following table (D4RL *-v2). We can see that overall, with the same amount of computation, DROP still performs better than Onestep, except for the tasks walker2d-medium-expert and walker2d-medium-replay. \n\n|     | Onestep  | DROP  |\n|  ----  |  ----  |  ----  |\n|  umaze  | 64.3  | **75.0**$\\pm$2.3  |\n|  umaze-diverse  | 60.7  | **66.2**$\\pm$1.9  |\n|  antmaze-medium-play  | 0.3  | **33.2**$\\pm$3.6  |\n|  antmaze-medium-diverse  | 0  | **38.1**$\\pm$2.5  |\n|  antmaze-large-play  | 0.3  | **21.5**$\\pm$1.6  |\n|  antmaze-large-diverse  | 0  | **28.9**$\\pm$3.1  |\n|  walker2d-medium-replay  | **66.4**  | 63.5$\\pm$1.5  |\n|  hopper-medium-replay  | 77.3 | **78.8**$\\pm$1.7 |\n|  halfcheetah-medium-replay  | 38.4  | **42.0**$\\pm$1.9  |\n|  walker2d-medium-expert  | **111.8**  | 102.7$\\pm$1.7  |\n|  hopper-medium-expert  | 81.4  | **100.2**$\\pm$2.3  |\n|  halfcheetah-medium-expert  | 77.0  | **89.1**$\\pm$1.8  |\n\n\n\n**(5) procedure of determining hyperparameters (the number of subtasks).**\n\nIn our experiments, we take the number of subsets as a hyperparameter and choose it with typical hyperparameter tuning strategies. Thank you for pointing this out. We will clarify this in our revisions. \n\nWe also point out that our DROP+CVAE treats each trajectory as a single subtask and can achieve better performance. Thus, in the real world, the algorithm designer does not need to worry much about the number of subtasks, just take each trajectory as a single subtask (and use the DROP+CVAE implementation). \n\n\n**(6) In Table 2, which variant is the listed DROP?**\n\nWe use DROP-Grad-Ada in Table 2.\n\n\n**(7) When comparing DROP with LAPO in Table 2, why are the results of LAPO on u, um-p, and ul-p missing?**\n\nThe main reason is that LAPO doesn't release their source code, and our private reproduction was not as good as the results in their paper (there may be some details we didn't notice). So we reported their results directly and then left it open if LAPO did not report their results in the paper.\n\nTo provide a reference for the reviewer, we provide our reproduction results (antmaze *-v2) in the following table.\n|  |  umaze  |  umaze-diverse  |  medium-play  |  medium-diverse  |  large-play |  large-diverse  |\n|  ----  |  ----  |  ----  |  ----  |  ----  |  ----  |  ----  |\nLAPO (in its original paper, *-v1) |– |91.3 |– |**85.7** |– |61.7 |\nLAPO (our implementation, *-v2) | 86.5 $\\pm$ 1.5 |80.6$\\pm$2.1 |68.5$\\pm$1.8 |79.2 $\\pm$ 1.8| 48.8$\\pm$ 2.6 |**64.8**$\\pm$ 3.0 |\nDROP (*-v2) |**90.5**$\\pm$2.4 |**92.2**$\\pm$1.7 |**74.1**$\\pm$3.9 |82.9$\\pm$3.5 |**57.2**$\\pm$5.5 |63.3$\\pm$2.4 |\n\n\n**(8) results of Diffuser.**\n\nThank you for pointing this out. We will include the Diffuser results in table 2 in our revision. \n\n**(9) why the iterative process can cause error propagation for potential overestimation of values**\n\nThe main reason is that naively performing policy evaluation (Equation 1 in the main paper) may query the estimated Q^k(s', a') for actions that are far outside the static offline data, resulting in pathological value Q^{k+1}(s, a) with large error. Such an iterative process (iterative policy evaluation and improvement) will further cause the inferred policy \\pi_{k+1}(a|s) to be biased towards OOD actions with erroneously overestimated values. Thus, performing policy evaluation and improvement iteratively leads to the (potential) policy/value errors being propagated iteratively, which in turn leads to collapsed performance.\n"}]

Author rebuttal:



**(1) comparison to prior adaptive baselines.**

1) The following table shows a comparison with APE-V [1] in the D4RL suite. The APE-V results are taken from the original paper, which uses SAC-N as the base offline RL implementation. Despite the comprehensive approach we have taken, it is important to acknowledge that the performance of our DROP does not generally outperform the specific baseline APE-V [1]. Nevertheless, we believe that this result provides valuable insights for future research directions on MBO solutions for offline RL problems. 

|        | APE-V [1]  | DROP |
|  ----  | ----  | ----  |
|  halfcheetah-random  | 29.9 $\pm$ 1.1  | 32.0 $\pm$ 2.5  |
|  halfcheetah-medium | 69.1 $\pm$ 0.4 | 52.4 $\pm$ 2.2  |
|  halfcheetah-medium-expert  | 101.4 $\pm$ 1.4  | 102.2 $\pm$ 1.5  |
|  halfcheetah-medium-replay  | 64.6 $\pm$ 0.9  | 50.9 $\pm$ 1.6  |
|  hopper-random  | 31.3 $\pm$ 0.2 | 20.8 $\pm$ 0.3  |
|  hopper-medium |   | 61.5 $\pm$ 3.7  |
|  hopper-medium-expert  | 105.7 $\pm$ 3.7 | 107.2 $\pm$ 1.5  |
|  hopper-medium-replay  | 98.5 $\pm$ 0.5 | 96.3 $\pm$ 2.5  |
|  walker2d-random  | 15.5 $\pm$ 8.5 | 5.2 $\pm$ 1.6  |
|  walker2d-medium | 90.3 $\pm$ 1.6 | 82.1 $\pm$ 5.2  |
|  walker2d-medium-expert  | 110.0 $\pm$ 1.5  | 109.3 $\pm$ 0.4  |
|  walker2d-medium-replay  | 82.9 $\pm$ 0.4 | 83.5 $\pm$ 1.2  |

2) As CCVL [2] only experiments on the discrete action Atari games and does not test its performance on the common continuous control tasks of D4RL. For comparison, we also deploy our DROP method on 4 offline Atari games. We present the results in the following table (using an initial 10\% of the replay dataset after 12.5M gradient steps). We can see that our DROP outperforms CCVL in 2 out of 4 tasks. 

|          | CCVL [2]  | DROP  |
|  ----  | ----  | ----  |
|  Asterix | 7576.0 $\pm$ 360.2  | 6517.9 $\pm$ 564.4  |
|  Breakout  | 121.4 $\pm$ 10.3  | 139.5 $\pm$ 25.0  |
|  Pong  | 13.4 $\pm$ 6.1  | 10.7 $\pm$ 8.2  |
|  Seaquest  | 1211.4 $\pm$ 437.2  | 1358.1 $\pm$ 352.0  |



3) Regarding the paper [3] mentioned by the reviewer, we think it is difficult to make a fair experimental comparison. The main reason is that paper [3] asks the user to adjust the policy behaviour after training, which is too subjective and therefore we cannot make a fair comparison. We will cite this paper in our paper and explain the connection between us. Thank you again for bringing up this valuable related work. 



[1] Ghosh, D., Ajay, A., Agrawal, P., & Levine, S. (2022, June). Offline rl policies should be trained to be adaptive. ICML 2022

[2] Hong, J., Kumar, A., & Levine, S. (2022). Confidence-Conditioned Value Functions for Offline Reinforcement Learning. ICLR 2023

[3] Swazinna, P., Udluft, S., & Runkler, T. (2022). User-Interactive Offline Reinforcement Learning. ICLR 2023



**(2) comparison to RoMA.**

As suggested by reviewer YuS3, we also make a comparison with RoMA [4]. In the implementation, we take the parameters (neural network weights) of the behavioural policies as the design input and the behavioural scores as the outputs. In the pre-training phase, we also use Gaussian noise for input perturbation; in the inference phase, we perform parameter inference (outer-level optimisation) with 200-step updates. The comparison results are shown in the following table. We can see that RoMA improves the performance of COMs and can outperform our DROP in 1 out of 3 tasks. 


|        | COMs  | RoMA | DROP |
|  ----  | ----  | ----  | ----  |
|  halfcheetah-medium-replay  | 41.4  | 56.8 $\pm$ 2.9  | 50.9 $\pm$ 1.6  |
|  hopper-medium-replay  | 49.7| 61.2 $\pm$ 2.0  | 96.3 $\pm$ 2.5  |
|  walker2d-medium-replay  | 33.9 | 67.4 $\pm$ 4.1  | 83.5 $\pm$ 1.2  |


[4] Yu, S., Ahn, S., Song, L., & Shin, J. (2021). Roma: Robust model adaptation for offline model-based optimization. NeurIPS 2021.



**(3) train the contextual behavior policy in an IQL-style manner.**

We note that IQL essentially weights the behavior policy with Q-values, and also note that its policy is not a contextual policy, but rather a simple policy that inputs states and outputs actions. In this paradigm, we believe that a reasonable way to understand the review's proposal (IQL style) is to first perform a policy improvement for each contextual policy, and then distill multiple improved contextual policies into a single fixed policy. In the implementation, we first use the score function to optimize z (implicit policy improvement), and then distill the improved contextual policy into a fixed policy, in the same spirit as the Q-weighted BC in IQL. The following table shows the comparison results. We can see that this method ("distilled" DROP) is able to match DROP on walker2d-medium-v2 and halfcheetah-medium-v2, but it still performs worse than our non-iterative DROP implementation overall. 


|    | DROP  | "distilled" DROP   |
|  ----  | ----  | ----  |
|  walker2d-random-v2  | **5.2** $\pm$1.6  | 4.5 $\pm$0.8  |
|  hopper-random-v2  | **20.8** $\pm$0.3  | 18.9 $\pm$0.2  |
|  halfcheetah-random-v2  | **32.0** $\pm$2.5 | 27.8 $\pm$1.54  |
|  walker2d-medium-v2  | 82.1 $\pm$5.2  | **84.2** $\pm$2.4  |
|  hopper-medium-v2  | **74.9** $\pm$2.8  | 67.1 $\pm$2.1  |
|  halfcheetah-medium-v2  | 52.4 $\pm$2.2  | **53.9** $\pm$1.7  |
|  sum  | **267.4**  | 258.4  |







Meta-review: The paper introduces an approach designed for adapting policies during inference in offline reinforcement learning. It introduces a bi-level optimization process, which enhances policy adaptability by segregating inner and outer loop optimizations. Reviewers commend the paper for its well-structured presentation and insightful figures. The approach is recognized for its potential in transferring part of the optimization process to test-time. Furthermore, the paper's claims are substantiated by experiments conducted on D4RL benchmarks, which take into account both performance metrics and computational overhead.

Several concerns regarding hyperparameters and comparisons were raised. It is also identified that there are some inconsistency between the results shown in figures and tables. 

Overall, the meta-reviewer recommends acceptance, contingent on that the authors carefully addressing the raised concerns and integrating reviewers' feedback..
Decision: Accept (poster)