Compositional Conservatism: A Transductive Approach in Offline Reinforcement Learning
Authors: Yeda Song
Dongwook Lee
Gunhee Kim
Venue: ICLR 2024 poster (2024)
Keywords: offline reinforcement learning
compositional generalization
conservatism
transduction
Primary area: reinforcement learning
TL;DR: We encourage conservatism in the compositional input space of the policy and Q-function, independently to the prevalent behavioral conservatism.

Abstract: Offline reinforcement learning (RL) is a compelling framework for learning optimal policies from past experiences without additional interaction with the environment. Nevertheless, offline RL inevitably faces the problem of distributional shifts, where the states and actions encountered during policy execution may not be in the training dataset distribution. A common solution involves incorporating conservatism into the policy or the value function to safeguard against uncertainties and unknowns. In this work, we focus on achieving the same objectives of conservatism but from a different perspective. We propose COmpositional COnservatism with Anchor-seeking (COCOA) for offline RL, an approach that pursues conservatism in a _compositional_ manner on top of the transductive reparameterization (Netanyahu et al., 2023), which decomposes the input variable (the state in our case) into an anchor and its difference from the original input. Our COCOA seeks both in-distribution anchors and differences by utilizing the learned reverse dynamics model, encouraging conservatism in the compositional input space for the policy or value function. Such compositional conservatism is independent of and agnostic to the prevalent _behavioral_ conservatism in offline RL. We apply COCOA to four state-of-the-art offline RL algorithms and evaluate them on the D4RL benchmark, where COCOA generally improves the performance of each algorithm. The code is available at https://github.com/runamu/compositional-conservatism.

Published as a conference paper at ICLR 2024
COMPOSITIONAL CONSERVATISM: A TRANSDUCTIVE APPROACH IN OFFLINE REINFORCEMENT LEARNING
Yeda Song", Dongwook Lee" & Gunhee Kim Seoul National University
yeda songevii sion con.ssuuuae anu.ac. kr, (dwsma rt32, gunhee)esnu.aco ac.
ABSTRACT
Offline reinforcement learning (RL) is a compelling framework for learning op- timal policies from past experiences without additional interaction with the envi- ronment. Nevertheless, offline RI. inevitably faces the problem of distributional shifts, where the states and actions encountered during policy execution may not in the training dataset distribution. A common solution involves incorporating conservatism into the policy the value function safeguard against uncertain- ties and unknowns. In this work, we focus on achieving the same objectives of conservatism but from different perspective. We propose COmpositional COn- servatism with Anchor- mechor-seeking (COCOA) for offline RL, an approach that pursues conservatism in a compositional manner on top of the transductive reparameteri- zation (Netanyahu et al., 2023), which decomposes the input variable (the state in our case) into an anchor and its difference from the original input. Our COCOA seeks both in-distribution anchors and differences by utilizing the learned reverse dynamics model, encouraging conservatism in the compositional input space for the policy or value function. Such compositional conservatism is independent of and agnostic to the prevalent behavioral conservatism in offline RL. We apply COCOA to four state-of-the-art offline RL algorithms and evaluate them on the D4RL benchmark, where COCOA generally improves the performance of each algorithm. The code is available at https://github.. https://gthh...ommmammm/mmumm com/ruuammmmmmoootoonnntto conservatism.
INTRODUCTION
Reinforcement learning (RL) has achieved notable successes across various domains, from guiding robotic movements (Dasari et al., 2020) and optimizing game strategies (Mnih et al., 2015) to re- cently promising training of language models (Rajpurkar et al., 2016). Despite these achievements, the challenges posed by real- peal-time interaction in complex and sensitive environments have prompted the development of offline RL as viable direction. Offline RL (Wiering & Van Otterlo, 2012; Levine et al., 2020) or batch RL (Lange et al., 2012) learns policies solely from pre-existing data, without any direct interaction with the environment. Offline RL is becoming increasingly popular in weal-world applications such as autonomous driving (Yu et al., 2020a) or healthcare (Gottesman et al., 2019) where prior data are abundant.
By its nature of learning from prior datasets, offline RL is often susceptible distributional shifts. This issue arises when the distribution of states and actions encountered during policy execution differs from that of the training dataset, situation particularly challenging in machine learning (Levine et al., 2020). Numerous existing offline RL algorithms tackle this by reducing distributional shifts through conservative approaches, including constraining the policy estimating uncertainty to measure distributional deviations (Kim & Oh, 2023; Ran et al., 2023; Kostrikov et al., 2022; Kumar al.. 2020; Wu et al., 2019; Kumar et al., 2019; Fujimoto et al., 2019; Sun et al., 2023; Rigter et al., 2022; Wang et al., 2021; Yu et al., 2021; 2020b; Kidambi et al., 2020). These strategies aim keep the agent within known distributions, mitigating risks of unexpected behaviors. this work, we also pursue the same goal of conservatism, focusing on aligning the test data distribution with the seen distribution, but from different perspective.
Authors equally contributed.

---

Published as a conference paper at ICLR 2024
We begin by recognizing that the state distributional shift problem is closely related to addressing how to deal with the out-of-support input points of the function approximators. We explore the pos- sibility of transforming the out-of-suppot learning problem into cotoo--combinaiionn problem by injecting inductive biases into function approximators of the policy or the Q-value function. Such transformation has been previously proposed by Netanyahu et al. (2023), where a transductive approach named bilinear transduction makes predictions through a bilinear architecture after repa- rameterizing the target function. This reparameterization decomposes the input variable into two components, namely an anchor and delta, where the anchor is variable in the input space and the delta the difference between the input variable and the anchor. If the reparameterized training and test data distribution satisfy certain assumptions, and if the target function has certain properties, the bilinear transduction can address the out-of-combination problem, which potentially resolves the out-of-support problem with the original target function.
We propose COmpositional COnservatism with Anchor Anchor-seeking (COCOA) for offline RI., a frame- work that adopts compositional approach to conservatism, building upon the transductive repa- rameterization (Netanyahu et al., 2023). Our approach transforms the distributional shift problem into an cnttof-combination problem. This shifts the key factors for generalizability from data to decomposed components and the interrelations between them, demanding the anchor and delta to be selected close to the training dataset distribution.
We suggest a new anchor-seeking approach with an additional policy, named anchor-seeking policy, which enforces the agent find anchors within the seen area of the state space. Hence, COCOA encourages anchors to be close to the offline dataset while confining the deltas within a narrow range by identifying anchors among neighboring states. This approach can reduce the input space and guide it toward the space predominantly explored during the training phase. In summary, by learning a policy to seek in-distribution anchors and differences from the learned dynamics, we can encourage conservatism in the compositional input space of the function approximator for the Q-function and policy. This approach is independent of and agnostic to the prevalent behavioral conservatism in offline RI.
We empirically find that our method improves the performance of four representative offline RL. methods, including CQL (Kumar et al., 2020), IQL (Kostrikov et al., 2022), MOPO (Yu et al., 2020b), and MOBILE (Sun et al., 2023) on the D4RL benchmark (Fu et al., 2020). We also show through an ablation study that learning anchor-seeking policy is effective in improving the perfor- mance of our method. Our main contributions can be summarized as follows:
We pursue conservatism in the compositional inpul space for the function approximators of the Q-function and policy, independently and agnostically to the prevalent behavioral conservatism in offline RL.
We introduce COmpositional COnservatism with Anchor-seeking (COCOA) that finds in- distribution anchors and deltas with the learned dynamics model, which is crucial for com- positional generalization.
We empirically show that COCOA improves the performance of four state-of- -hee---ff-hh--tt-t offline RL algorithms on the D4RL benchmark. Additionally, our ablation study shows the efficacy of anchor-seeking policy compared to heuristic anchor selection.
PRELIMINARIES
2.1 OFFLINE RL
We assume an MDP problem AS,A,T, R) with continuous state space S, continuous action
space A, a transition function T : x A + S, and a reward function R : S x A- R. The goal is to find policy æ S + A that maximizes the expected return J(x) E, Eee (s,, where 7 E [0, 1) is a discount factor.
In offline RL, also known as batch RL, we are given a dataset Dom = {( ag $141, ri)] N ated with a behavior policy. The goal in offline RL is find a policy æ that maximizes the expected return ((T) using only the fixed dataset Denv Like most model-baed offline RL algorithms, we learn dynamics model 7(s++1l 8i,a,) az) that predicts next state 8 given current state 8; and action a;-
i= gener-
2

---

Published as a conference paper at ICLR 2024
In addition to the forward dynamics model, also learn reverse dynamics model T(s:(5s41, that predicts a current state 8, given a next state 8,41 and action aN.
2.2 BILINEAR TRANSDUCTION
We follow the formulation of Netanyahu et al. (2023) about the generalization problem. With no assumption on the train and test distribution, the generalization performance of a function approx- imator is limited. This problem occurs especially when the test distribution is not contained in the train distribution, also known as out-of-support (OOS) learning problem. As special case of OOS, an out-of-combination (OOC) problem occurs when the input space is decomposed into two components, and the marginal of the train distribution of each component includes that of the test distribution while the joint train distribution does not necessarily contain the joint test distribu- tion. Under certain assumptions, Netanyahu et al. (2023) propose transductive reparameterization method called bilinear transduction to convert an OOS problem into an OOC problem.
Bilinear transduction. It solves extrapolation under certain assumptions. First, the target function fz) is reparameterized as
f(x) = f(x- - I,
where á is termed as an anchor, selected from the training dataset. The difference (x - -) between the input variable and the anchor ī is termed as delta. The reparameterized target function f is approximated as a bilinear function of the embeddings 401 and 422
fo(x) 4p (rr a) (p)(.).
Intuitively, it facilitates the low-rank property of the embeddings (1 and y02, enabling the function approximator to generalize to OOC points.
Sufficient conditions for bilinear transduction. Netanyahu al. (2023) introduce sufficient con- ditions for bilinear transduction to be applicable. The assumptions are about both the dataset and the target function f. The first assumption is about combinatorial coverage of the dataset. The test dataset has to have a bounded combinatorial density ratio with respect to the training dataset. It im- plies that the support of the joint distribution the training distributions the components should include the support of the joint distribution of the test distributions of the components. Second, the target function f should be bilinearly transducible. i.e., there exists a deterministic function f such that f(x) f(x ă, .) for all x, E X. Lastly, the training distribution of anchors should not degenerate (Shah et al., 2020). Under these three conditions, it is possible to generalize the target function to O0C points with theoretically guaranteed risk bound.
Connection to compositional generalization. In light of the literature on compositional general- ization (Wiedemer et al., 2023), we interpret bilinear transduction special case of compositional generalization, where the (P1, 92 models serves as component functions, extracting low-rank fea- tures of the input, and the inner product serves as composition function.
COMPOSITIONAL CONSERVATISM WITH ENCHOR-SEEKING (COCOA)
3.1 OFFLINE RL WITH BILINEAR TRANSDUCTION
The base algorithms in offline RL, such Deep N-Networks (DQN) (Mnih et al., 2015) and Actor- Critic methods (Mnih et al., 2016; Haarnoja et al., 2018), frequently use deep neural networks as function approximators. Hence, we employ bilinear transduction ($ 2.2) to the function approxima- tors of policy and O-function. In both train and test phases, we decompose the current state into an anchor and delta As 5 5, where Denv Then, the policy and the Q-function will be
no P0,1 (As) (8),
Qo(s,a) = P (As,a) (),,2($,a)
The architecture of the policy and Q-function is illustrated in Figure 1c. The policy (a]s) is trained to maximize the expected return (TT), and the Q-function Q(s, a) is trained to minimize its loss function fg defined in the base offline RL algorithm.
Different state decompositions can result in different compositional input spaces, resulting in differ- generalization capabilities. In order satisfy the assumptions of bilinear transduction in $ 2.2,
3

---

Published as a conference paper at ICLR 2024
S
S
S
state decomposition
As
As
π
As
transductive predictor
(dx) (P((F)
sechir--eeking rollow: intermediate state reverse madel polloct soin area
tesusilione Das IN ai ((-la) sem
state Dew
(a)
(b)
(c)
Figure 1: (a) An illustration of anchor-secking rollouts that find anchors close the seen area of the state space S. Given the current state s, the anchor-seeking policy गे gives actions to reach the anchor Its behavior is derived by utilizing reverse model rollouts, which diverge from the offline dataset. (b) An illustration of the current state and an anchor s. Ideally, the anchor 3 has been observed during the training phase when served as an anchor for another state. Similarly, the difference, delta, had also been encountered previously in the ideal case but in combination with different anchor. (c) The architecture of our policy π m(ass) that aims to generalize to an unfamiliar state by decomposing the state into familiar components (seen anchor 8 and seen delta As) and applying a transductive predictor. The architecture the Q-function is similar to that the policy.
the ideal approach is find the decomposition that fulfills these two criteria: in-distribution anchor and in-distribution delta. This ideal case is illustrated in Figure 1c. However, in contrast to the prior works (Netanyahu al., 2023; Pinneri al., 2023) that only focus on the transducible property of the goal state, we try to handle each state in every step, and it is computationally infeasible to enforce these constraints using brute-force methods like comparing the current states to all other points. Hence, we introduce new anchor-seeking policy that learns to find in-distribution anchors and deltas and exploit the learned dynamics model that prevents arbitrary decomposition, to exploit the power of bilinear transduction. We also enforce each delta to be within distance of few steps of the dynamics model to confine the distribution delta in both the train and test phase to a similar range. This approach reduces the input space and guides it toward the space that is predominantly explored during the training phase, thereby further enhancing generalizability.
3.2 LEARNING TO SEEK IN- DISTRRIBUUTIO DECOMPOSITION
We describe the reverse dynamics model, the anchor-seeking trajectory, and the random divergent reverse policy, which are required components before training the anchor-seeking policy.
3.2.1 ANCHOR-SEEKING TRAJECTORY
Training a reverse dynamics model. Given a transition (s,0, 8') sampled from the dataset De train reverse transition dynamics model (s)s', (Wang al., 2021; Lai et al., 2020; Goyal et al., 2018; Edwards et al., 2018; Holyoak & Simon, 1999) to predict the state given the next state s' and action In other words, the reverse dynamics model T(s', a) predicts "From which state do we come if we arrive at s' by taking action a?". This is done by minimizing the loss function with the dataset's state defined as
C, = E(s.,,,
(r(5',a) $
(4)
Random divergent reverse policy. We do not use a trained reverse policy but instead use a heuristic reverse policy that randomly selects action from the dataset Danv The subsequent actions in reverse rollouts, after the initial action, follow the same direction as the initial action but are slightly scaled down with small added Gaussian noise. This ensures that the reverse rollout diverges away from the dataset. Since we use random actions and maintain a consistent direction throughout the reverse rollout, it is likely venture into unexplored regions beyond the offline dataset. In
4

---

Published as a conference paper at ICLR 2024
Algorithm Generation of Anchor-Seeking Trajectory
Require: Offline dataset Denv. reverse dynamics T,, anchor-secking horizon h, rollout epoch e. for in 1...e do
2: Sample anchor state 8t Denv.
3: Generate reverse model rollout f = : S1+1 from by using the
reverse dynamics 7, and random divergent actions {at-il
4: Add model rollouts to replay buffer, Dreverse Preverse U a1 - st41-i)
5: end for
6: return Deverse
Algorithm Anchor- Secking Rollout
Require: State s, anchor-seeking policy म, forward dynamics model T, anchor-seeking horizon h.
1: Set the initial state for anchor-seeking: 8.
2: for i 1 to h do
3: Compute the anchor-seeking action 17 + (().
4: Update anchor á T(5, n).
5: end for
6: return
summary, the reverse policy gives an action a, at each rollout step j as follows:
a = 6o Ej, where a Denv. Eg NN0,, j 1.2,...,
h is horizon length, ó is scale coefficient, and is noise coefficient. We set ø 0.8 and 0.1 when the maximum action value 1.0.
Anchor-seeking trajectory. We rollouts of the reverse model make anchor-seeking trajec- tories for training the anchor-seeking policy. First, we sample an anchor state from the dataset and generate reverse transition Drevess = {(S+41, a,, r.)) from the anchor state using the reverse dynamics model and the random divergent reverse policy. Note that the direction of anchor-seeinng trajectory is reverse to that of reverse transition, Dreverse- By doing so, we can effectively generate anchor-seeking trajectories for training of the anchor-seeking policy. Utilizing reverse model roll- outs to address the OOD problem is first proposed by Wang al. (2021), who augment the offline dataset with reverse transition, train a policy using this augmented dataset,, and demonstrate the ef- ficacy of such approach in the offline RL setting. Finally, the detail of generating anchor-seeking trajectory is summarized in Algorithm 1.
3.2.2 TRAINING OF DYNAMICS-AWar ANCHOR- ANCHOR-SSEKING POLICY
We train the anchor-secking policy 7 (a]s) before training the main policy. We utilize anchor-seeking trajectories in $ 3.2.1, which are the reverse direction to the dataset Drevense By following the path of an anchor-seeking trajectory, the anchor-seeking policy is trained to select actions that guide the agent direction moving from the external boundary towards the seen area as illustrated in Figure la. Given the anchor-seeking rollout is generated by the anchor-seeking policy मे (aas) and
the dynamics model T'ss, a), the reverse transition Dreverse relies T, r|s', a). Since the reverse transition was designed to diverge from the offline dataset, the anchor-seeking trajectory, with its reversed direction, ensures the transition converges back the dataset from unfamiliar states. As result, we train the anchor-seeking policy to minimize the MSE loss between the predicted action and the action the dataset Dreverse. The loss function defined as
Lanchor (0) = EE (ee)" (o [(7 - a)?l
(6)
In this way, the anchor-seeking policy ((ass) can provide a proper action to move toward an in- distribution anchor. This action is given to the dynamics model f(s. a) to predict the next state. Thereby, as shown Figure 1, the anchor-seeking rollout a sequence of transitions that starts from the current state and ends the anchor state 5.
5

---

Published as a conference paper at ICLR 2024
Algorithm 3 Bilinear Transduction with Anchor-Secking (Actor)
Require: Input state s, anchor-secking policy म, forward dynamics model T, anchor-seeking hori- zon h; embedding layers {411. )-2)
1: Start from 8p + 8 and perform anchor-seeking rollout F = {{*E, ai si41 ) h-1 by using the anchor-seeking policy +({a]s:) and forward dynamics T(s++1 (8;; ai).
2: Decompose the state $ by using the final state as an anchor: $ f- Sp, As + 8/.
3: Bilinear transduction of the target function: f + 40 (8
8) (p2(8).
4: Process with an additional MLP layer: f <- MLP(/).
5: return
3.3 SUMMARY OF THE METHOD
We illustrate the integration of the anchor-seeking model into the bilinear transduction framework. We choose the Soft Actor-Critic (SAC) algorithm (Haarnoja et al., 2018) as representative RI. algorithm that employs function approximators. We utilize the bilinear transduction supplemented with anchor-secking for the SAC's actor and critic networks.
Given an input state 8n, use rollout to obtain an action from the anchor-seeking policy. The anchor & is derived by this action through the forward step of dynamics and then updated to be the next state Sn+1- After rolling out a few times can pinpoint the final anchor and use it to decompose the state into the anchor and delta. The difference between the initial state s, and this anchor 8, delta, is computed as As 5 Sn-
We then carry out the bilinear transduction as described in E2,()). We respectively embed As and 3 4 (As) and (5)) and compute the inner product between them. The output is then fed into a small MLP layer to enhance the flexibility of the function approximator. This step introduces nonlinearity, as the policy Q-function may not be linear their inputs.
Algorithm summarizes the whole process in the actor module. In the critic module, we concatenate the action with both the anchor and delta in the forward operation before executing bilinear trans- duction. Subsequently, the values of action and Q-value, derived as facoor and feritic respectively, are employed to update the actor and critic networks of the SAC policy.
4 EXPERIMENTS
In our experiments, we aim to empirically answer the following two questions: (i) how much does method improve the performance of prior model-free and model-based algorithms? and (ii) what the effect the anchor-seeking on performance?
4.1 RESULTS OF D4RL BENCHMARK TASKS
We evaluate our method on the Gym-MuJoCo tasks in D4RL benchmark (Fu et al., 2020), which consists of 12 tasks from the OpenAI Gym (Brockman et al., 2016) and MuJoCo (Todorov al., 2012) environments. Refer to A.I I for the details of the tasks.
Baselines. We apply COCOA several prior offline RL algorithms, both model-based and model- free methods. They include (i) CQL (Kumar et al., 2020) that penalizes Q-values on out-of- distribution samples for safety, (ii) IQL (Kostrikov et al., 2022) that leverages the generalization capability of the function approximator by viewing the state value function (iii) MOPO (Yu al. 2020b) as model-hased approach that penalizes rewards based on uncertainty from predicting subsequent states, and (iv) MOBILE (Sun et al., 2023) that quantifies uncertainty through the inconsistency of the Bellman estimation using an ensemble of dynamics models. We also report the results of (v) Behavior Cloning (BC), which learns tasks by imitating expert data. To stabilize training, all baseline algorithms are reproduced with layer normalization applied.
random variable,
Results. Table summarizes the results of our experiments. The baseline algorithms denoted as "Alone", and our method is denoted as "+COCOA". We report the average return of the last 10 training epochs across 4 seeds, with the standard deviation. For all algorithms, reproduce the

---

Published as a conference paper at ICLR 2024
Table 1: Results of the D4RL benchmark. We report the normalized average return of the last 10 training epochs across 4 seeds on the D4RL benchmark tasks.
CQL
IQL
MOPO
MOBILE
Task
Alone
+COCOA
Alone
+COCOA
Alone
+COCOA
Alone
+COCOA
halfcheetah--arn random
2.2
31.3
23.8 a 0.8
37.3
33.9
40.9
29.5 a
hopper-random
3.7
5.3
8.8 = 1.6
31.7
32.3
17.4
walker2d-randon
8.5
23.7
2.1
halfchectal-mmedimm
43.2
49.0 0.2
47.4
4.1
72.4
3.7
73.5
69.2
hopper-medium walker2d medium
54.1
61.9
66.0 = 1,4
66.3
65.5 1.5
62.8
2.3
92.9
103.7 = 6.1
70.9
79.5
83.1 + 0.3
78.3
81,9
84.1
89.4
80.8
84.4 +: 1.3
halfcheetah-
seffchhetaa-mmdin-eeee medimm--edmmm-- replay
37.6
45.3
46.4 a 0.4
44.2
40.3 =
72.1
67.0
66.8
68.6 0.9
hopper-medium--rpllay
16.6
86.3
96.5
94.7
8.9
92.8
7
87.9
105.3 1.7
mwdimme-ihmmpppeeee replay
20.3
76.8
83.9 + 2.9
73,9
66.9 + 10.0
85.2
93.1 * 5.1
81.1
83,5 1.7
hallcheetah-medium. expert
44.0
95.0
9.11
86.7
91.4 1.4
83.6
109.1 = 1.4
86.75
104,1 + 2.0
hopper- medeer--mddiiieeeeeee
53.9
96.9
103.3
91.5
105.3
74.9
101.4 13,9
102.3
109.1 + 4.7
walker2ddmediim--
90.1
109.1
109.4
109.6
109.1
105.3
109.1 0.4
107.0
107.5 0.9
Averuge
36.5
61.6
64.0
77.0
77.3
67.2
70.1
70.6
74.8
Fultchoatal-mooim---a-
hallchootah-modiumce eoppptrt
App. Rours
Aug. Rouum
Narm Aug.
40
CILL COCOA COL
IZL
COCOA
MOPO COCOA
MOBILE COCOA
IQL
MOPO
MOBILE
Training Epochs
Training Epocrs
Traning Epochs
Traning Epochs
(a) CQL
(b) IQL
(c) MOPO
(d) MOBILE
Figure 2: Performance of CQL, IQL, MOPO, and MOBILE with and without COCOA on the halfcheetal-mediun-eeperer task of D4RL.
results with the codebase described in Appendix A.3, Our method enhances the performance of all baseline algorithms, as measured by the average improvement across tasks excluding the random tasks of IQL as the original IQL paper does. In summary, our COCOA improves the performance of the original methods in 10 out of 12 tasks for CQL, out of 9 tasks for IQI., out of 12 tasks for MOPO, and out 12 tasks for MOBILE.
4.2 ABLATION STUDY: THE EFFECT OF ANCHOR-SEEKING
To examine the impact of anchor selection on performance, we experiment with variant of our method that does not anchor-seeking. For this ablation study, we CQL (Kumar et al., 2020) as the base algorithm due to its computational efficiency as model-free algorithm and its completeness in supporting all task types, including "random" tasks.
Baseline. The baseline for this ablation study is denoted as "+COCOA (w/o A.S.)" in Table 2. In this baseline, adapt the heuristic anchor selection procedure from Netanyahu al. (2023), incorporating key modifications for our context. Unlike the original method, which selects anchors based on goal states, our baseline selects anchors based on the current state, addressing the absence of goal state in our setting. To mitigate the computational demands of this method, we limit our selection to a subset candidate anchors, randomly sampled from the dataset.
The heuristic anchor selection operates as follows. We initially draw N candidate anchors 8; from the dataset and calculate the difference As between the candidates and the current state, defined by
Asn
Sns
N}. [L,...,N), Sn E Dew-
Subsequently, we assess every pairwise difference among another N sampled states from Deny as
Asij 8, 89z i,j E (1, (I,...,N,,
NN, i i 81,8 Dew-
(8)

---

Published as a conference paper at ICLR 2024
Finally, we select the candidate anchor that minimizes the distance to the current state:
8 = S with ñ = arg min
min ||4s, - A8ts f.5
(9)
This baseline enforces the results of state de- composition to be close to in-distribution data through direct distance calculation. While it can be effective and feasible if the dataset is small and N is sufficiently large, its scalabil- ity is limited as the required amount of com- putation scales quadratically with the data size. Since the computation cost escalates cubically with the sample size, we set N to 30, matching our computation budget with "+COCOA".
Table 2: Ablation study for anchor-seeking. We report the normalized average return of the last 10 training epochs across seeds on the D4RL benchmark tasks.
CQL
Task
+COCOA
Alone
(wlo A.S.)
+COCOA
mulfcheetah-randon
31.3
22.5
hopper-random walker2d-randon
5.3
25.8 7.4
88
Results. We examine whether this variant im- proves the performance of CQL. The results are summarized Table 2. We report the average return of the last 10 training epochs across 4 seeds, with the standard deviation. The baseline "+COCOA (w/o A.S.)" achieves higher perfor- mance in only two tasks, "hopper-random" and 2anakkeo2d-randoo" and comparable lower performance in the other tasks compared to the original "Alone" baseline. In contrast, our method "+COCOA" improves the performance of the CQL models in 10 out of 12 tasks. This result suggests that the anchor-seeking is a crucial component for the success of our method.
5.4
8.3
5.5
halfcheetah medium hopper- medium walker2d-mediun
46.9
49.0 4.2
61.9
2.4
66.0 = 1.4 83.1 + 0.3
79.5
80.3 0.9
hulfcheetah- mullihhetah-mecdiu--repp. medium-replly
45.3
4.3
46.4 + 0.4
nopper-medium-reple
86.3
8.97
96.5 = 1.9
76.8
78.7 2.2
83.9 2.9
halfcheetah medium expert
95.0
92.4 1.2
hopper-medium wwweemm-dddddeeeeeeee walker2d 2el-ddd-ei----- exppert
96.9
30.4 9.8
103.3 + 1.4
109.1
86.8 + 2.3
109.4 + 0.4
Average
61.6
48.4
64.0
5 RELATED WORK
Offliie RL. In offline RL, agents use a predefined dataset without additional interactions with the environment, typically following either the model-based or model-free strategy. Model-free RL algorithms (Kim & Oh, 2023; Ran et al., 2023; Kostrikov et al., 2022; Kumar et al., 2020; Wu et al., 2019; Kumar et al., 2019; Fujimoto et al., 2019) optimize policy directly using prior experiences in the replay buffer, applying conservatism to the value function or policy. In contrast, model- based offline RL methods (Sun et al., 2023; Rigter et al., 2022; Wang et al., 2021; Yu et al., 2021; 2020b; Kidambi et al., 2020) use model trained in the environment to create additional data that are employed for policy learning. Through such synthesized data, this approach becomes stronger in generalization and robust even to unseen states.
onttrf-Disttibution Generalization in Offline RL. Much study has been done to improve the out-of-distribution (OOD) generalization of offline RL algorithms. Lou et al. (2022) tackles the action distributional shift problem by introducing a mutual information action embedding model. In a similar pursuit, Gu et al. (2022) propose pseudometric action rep- resentation learning method that measures both behavioral and data-distributional relations between actions. Bai et al. (2022) develop an uncertainty-drrven method that uses the disagreement of boot- strapped O-functions. It augments the dataset with OOD data on which more refined penalty is imposed. Pitis et al. (2022) propose local factorization of transition dynamies and state augmenta- tion for improved generalization of offline RL algorithms. They also provide theoretical proofs for sample complexity and generalization ability. Our method is similar to theirs in that we also employ the factorized architecture of the policy and Q-function. Unlike them, however, we do not use factorized dynamics model and instead leverage a bilinear transduction framework.
approach to learn an
Compositional Generalization and Extrapolation. Compositional generalization. which strives to generalize to unseen combinations components, is explored by various studies. Wiedemer et al. (2023) highlight a two-step generative procedure essential for tackling wide range of composi- tional problems. This procedure involves the complex generation of individual components and their straightforward combination into single output. They provide set of sufficient conditions under
8

---

Published as a conference paper at ICLR 2024
which models trained the data can generalize compositionally. On related note, Shah et al. (2020) presents a sample-efficiem RL algorithm that exploits the low-rank structure the optimal Q-function, which is bilinear function of states and actions. They prove quantitative sample com- plexity improvement for RI. with continuous state and action spaces via low-rank structure. Dong & Ma (2023) explore the extrapolation of nonlinear models for structured domain shift. They prove that a specific family of nonlinear models can successfully extrapolate to unseen distributions, pro- vided the feature covariance is well-conditioned. (Netanyahu al., 2023) propose an extrapolation strategy based on bilinear embeddings to enable combinatorial generalization, thereby addressing the out-or-support problem under certain conditions.
6 CONCLUSION
We explored a new perspective of conservatism for offline RL that does not regard the behavior space of the agent but the compositional input space of the policy and Q-function. We proposed a practical framework, COCOA, for finding better decomposition of states to encourage such conservatism. COCOA is a simple yet effective approach that can be applied to any offline RL algorithm that utilizes a function approximator. We empirically found that our method generally enhanced the performance of offline RI algorithms through our experiments across various tasks in the Gym- MuJoCo environment of the D4RL benchmark.
As study primarily engages in empirical exploration, further investigation may be demanded for a more comprehensive understanding of the mechanism behind the performance improvement or the properties of the compositional input space. Moreover, since experiments were limited to control-based robotics environments with continuous state and action spaces, it could be a valu- able extension of this work to apply compositional conservatism framework to other domains, including environments with discrete action spaces, image-based observations, or highly complex dynamics.
7 ACKNOWLEDGMENTS
We thank Jackyeom Kim, Soochan Lee, Seohong Park, Aviv Netanyahu, and the anonymous review- ers for their valuable discussions and feedback. This work was supported by Institute of Information & Communications Technology Planning & Evaluation (ITTP) grant funded by the Korean govern- ment (MSIT) (No. 2019-0-01082, SW StarLab), Institute of Information & communications Tech- nology Planning Evaluation (IITP) grant funded by the Korean government (MSIT) (No. 2022-0- 00156, Fundamental research on continual meta-lcarning for quality enhancement of casual videos and their 3D metaverse transformation), Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korean govemmeen((MSIT) (NO.2021-0-01 Artificial Intelligence Graduate School Program (Seoul National University)), and Center for Ap- plied Research in Artificial IntelligenneeCARRAI] grant funded by Defense Acquisition Program Ad- ministrationnAAAA) and Agency for Defense Development(ADD) (UD19003 (RD)). Gunhee Kim is the corresponding author.
8 REPRODUCIBILITY STATEMENT
For reproducibility, we provide the code of our method at conservatism. For the codebase for the baseline algorithms,, please refer to Appendix A.3. The hyperparameters and the model architecture are described in Appendix A.4 and Appendix A.2, re- spectively.
REFERENCES
Chenjia Bai, Lingxiao Wang, Zhuoran Yang. Zhi-Hong Deng, Animesh Garg, Peng Liu, and Zhao- Wang. Pessimistic bootstrapping for uncertainty-driven offline reinforcement learning. In International Conference on Learning Representations, 2022. URL https: /pppnreview. net/ http:/fooummmixxcslllls

---

Published as a conference paper at ICLR 2024
Greg Brockman, Vicki Cheung. Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.
Sudeep Dasari, Frederik Ebert, Stephen Tian, Suraj Nair, Bermadette Bucher, Karl Schmeckpeper, Siddharth Singh, Sergey Levine, and Chelsea Finn. Robonet: Large-scale multi-robot learning. 2020.
Kefan Dong and Tengyu Ma. First steps toward understanding the extrapolation of nonlinear models to unseen domains. In The Eleventh International Conference on Learning Representations, 2023. URL https: http::://pppeeeeee
net forum? add=wwaqvvccbe
Ashley Edwards, Laura Downs, and James Davidson. warkwdd reinforcement learn- ing. arXiv preprint 18iv::803.10222 2018.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, arXiv preprint arXiv:2004.07219, 2020.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pp. 2052 2062. PMLR, 2019.
Omer Gottesman, Fredrik Johansson, Matthieu Komorowski, Aldo Faisal, David Sontag, Finale Doshi-Velez, and Leo Anthony Celi. Guidelines for reinforcement learning in healthcare. Nature medicine, 25(1):16-18, 2019.
Anirudh Goyal, Philemon Brakel, William Fedus, Soumye Singhal, Timothy Lillicrap, Sergey Levine, Hugo Larochelle, and Yoshua Bengio. Recall traces: Backtracking models for efficient reinforcement learning. arXiv preprint /804.00370 2018.
Pengjie Gu, Mengchen Zhao, Chen Chen, Dong Li, Jianye Hao, and Bo An. Learning pseudometric- based action representations for offline reinforcement learning. Proceedings of the 39th Inter- national Conference on Machine Learning. volume 162, pp. 7902-7918. PMLR, 2022.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning. pp. 1861-1870. PMLR, 2018.
Keith J Holyoak and Dan Simon. Bidirectional reasoning in decision making by constraint satisfac- tion. Journal of Experimental Psychology: General, 128(1):3, 1999.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. Advances in neural information processing systems, 33: 21810-21823, 2020.
Byeongchan Kim and Min-hwan Oh. Model-based offline reinforcement learning with count-based conservatism. 2023.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. In International Conference on Learning Representations, 2022. URL https:/
openreview. httt://orrmmmdd-nn/..mm
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- policy q- learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1179: 179-1199, 2020.
Hang Lai, Jian Shen, Weinan Zhang, and Yong Yu. Bidirectional model-based policy optimization. In International Conference on Machine Learning, pp. 5618-5627. PMLR, 2020.
Sascha Lange. Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforce- ment learning: State-of-the-ant pp. 45-73. Springer, 2012.
10

---

Published as a conference paper at ICLR 2024
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems, 2020.
Xingzhou Lou, Qiyue Yin, Junge Zhang, Chao Yu, Zhaofeng He, Nengjee Cheng. and Kaiqi Huang. Offline reinforcement learning with representations for actions. Information Sciences, 610:746- 758. 2022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 5 518(7540):529-555 -3322 2015.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforeement learning. In International conference machine learning, pp. 1928-1937. PMLR, 2016.
Aviv Netanyahu, Abhishek Gupta, Max Simchowitz, Kaiqing Zhang, and Pulkit Agrawal. Learning to extrapolate: A transductive approach. In The Eleventh International Conference on Learning Representations, 2023. URL https: : /ooporreevveeeeeeeee net. /tt//forut/tteeeeee
forum? perrumiid=11 1d 4UKLPd4.
Cristina Pinneri, Sarah Bechtle, Markus Wulfmeier, Arunkumar Byravan, Jingwei Zhang, William F. Whitney, and Martin Riedmiller. Equivariant data augmentation for generalization in offline rein- forcement learning, 2023.
Silviu Pitis, Elliot Creager, Ajay Mandlekar, and Animesh Garg. Mocoda: Model-based counterfac- tual data augmentation. Advances in Neural Information Processing Systems, 35:18143-181555 2022.
Rong-Jun Qin, Xingyuan Zhang, Songyi Gao, Xiong- Hui Chen, Zewen Li, Weinan Zhang, and Yang Yu. NeoRL: near real-world benchmark for offline reinforcement learning. In Thirty- sixth Conference Neural Information Processing Systems Datasets and Benchmarks Track, 2022. URL https:://///... htt.:tttss//peeeeeeee...ee.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arxiv: 2016.
Yuhang Ran, Yi-Chen Li, Fuxiang Zhang, Zongzhang Zhang, and Yang Yu. Policy regularization with dataset constraint for offline reinforcement learning. In International Conference on Machine Learming. 2023.
Mare Rigter, Bruno Lacerda, and Nick Hawes. Rambo-rl: Robust adversarial model-based offline reinforcement learning. Advances in neural information processing systems, 35:16082-16097, 2022.
Devavrat Shah, Dogyoon Song. Zhi Xu, and Yuzhe Yang. Sample efficient reinforcement learning via low-rank matrix estimation. Advances in Neural Information Processing Systems, 33:12092- 12103, 2020.
Yihao Sun. Offlinerl-kit: An elegant pytorch offline reinforcement learning library. https:/ github. com/ www.hshunliliaa 121///22///f/e////
2023.
Yihao Sun, Jiaji Zhang, Chengxing Jia, Haoxin Lin, Junyin Ye, and Yang Yu. Model-Bellman incon- sistency for model-based offline reinforcement learning. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 33177-33194, PMLR, 23-29 Jul 2023.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference Intelligent Robots and Systems, pp. 5026-5033. IEEE, 2012. doi: 10. 1090000O. 201 2.6386109.
Jianhao Wang. Wenzhe Li, Haozhe Jiang, Guangxiang Zhu, Siyuan Li, and Chongjie Zhang. Offline reinforcement learning with reverse model-based imagination. Advances in Neural Information Processing Systems, 1--2442--29433 2021.

---

Published as a conference paper at ICLR 2024
Thaddäus Wiedemer, Prasanna Mayilvahanan, Matthias Bethge, and Wieland Brendel. Composi- tional generalization from first principles. arXiv preprint arXiv:3307.05596 2023.
Marco A Wiering and Martijn Van Otterlo. Reinforcement learning. Adaptation, learning, and
optimization, 12(3):729, 2012.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning-
arXiv preprint arXiv: / 1361, 2019.
Fisher Yu, Haofeng Chen, Xin Wang, Wenqi Xian, Yingying Chen, Fangchen Liu, Vashisht Madha- van, and Trevor Darrell. Bdd100k: diverse driving dataset for heterogeneous multitask learn- ing. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp- 2636-2645, 2020a.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 33:14129-14142, 2020b.
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, 14-29944-28967 2021.
12

---

Published as a conference paper at ICLR 2024
A EXPERIMENT SETTINGS AND IMPLEMENTATION DETAILS
A.1 D4RL BENCHMARK TASKS
The half-cheetah is a two-dimensional bipedal robot composed of 8 solid links, encompassing two legs and torso, coupled with motorized joints. The state space is 17- dimensional, encompassing both joint angles and velocities. An adversary destabilizes it by exerting a 6-dimensional action with 2-dimensional forces the torso and each foot.
HalfCheetah:
The hopper is a planar monopod robot, assembled with 4 solid links that represent the torso, upper leg. lower leg. and foot, and includes motorized joints. It has an I1-dimensional state space, including joint angles and velocities. An adversary employs a 2-dimensional force on the foot to disrupt its stability.
Hopper:
The walker operates as two-dimensional bipedal robot with structure of 7 links, representing two legs and torso, along with 6 actuated joints. Within its 17-dimensional state space, joint angles and velocities are included. An adversary employs a 4-dimensional action with 2-dimensional forces on both feet to disrupt its equilibrium.
Walker2D:
Adroit: Adroit is a complex task where a simulated robotic hand with 24 Degrees of Freedom is used for tasks such as hammering a nail, opening a door, spinning a pen, or moving a ball. We use two kinds of datasets for this: the "human" dataset, which includes 25 human demonstration trajectories, and the "cloned" dataset, which is an equal mix of demonstration data and behavior cloned from the demonstration policy.
NeoRL: NeoRL(Qin et al., 2022) is a benchmark designed to mirror real-world conditions by gath- ering datasets using more cautious policy, aligning closely with realistic data collection methods. The scarcity and specificity of the data present a significant challenge for offline RL algorithms. Our research examines nine datasets, encompassing three different environments (HalfCheetah-v3,, Hopper-v3, Walker2d-v3) and three levels of dataset quality (L, M, H), indicating low, medium, and high quality, respectively. Notably, NeoRL offers varying quantities of training data trajectories (100, 1000, 10000) for each environment. For our experiments, we uniformly chose 1000 trajecto- ries.
A.2 MODEL ARCHITECTURE
Dynamics Model Architecture: with previous works, we used a neural network as the back- bone for dynamics model, which outputs Gaussian distribution for the next state and reward. By ensembling these networks, we achieved greater stability and enhanced performance. From an ensemble of seven, we selected the top five models based validation error. The backbone of the dynamics model comprises four layers, each with a hidden dimension of 200.
Actor & Critic Architecture: The actor-critic framework like SAC (Haarnjj et al., 2018) com- prise actor and critic modules. Typically, an actor possesses a backbone constructed from a neural network. Features embedded within this backbone relayed through last layer that outputs Gaussian distribution, yielding a non-deterministic result. Although MOPO, MOBILE, CQL, and IQL (Yu al., 2020b; Sun et al., 2023; Kostrikov et al., 2022; Kumar et al., 2020), traditionally use 2, 2, 3, and 2 backbone layers with dimension of 256 respectively, upon integrating COCOA, we standardized the use of two backbone layers with 100 hidden dimensions.
Anchor-seeking Policy Architecture: The anchor-secking policy acts an add-on module shared between the actor and critic. The input data, consisting of delta and anchor, is embedded through a neural network and subsequently processed by bilinear architecture. Initially, inputs are embedded to the dimension of 4 with two neural networks with 64 channels, and the bilinear architecture produces output with a dimension of 64 using those embedded features. Then, the outputs of bilinear architecture are passed through the actor and critic backbone architectures, resulting in the determination of the action and Q value, respectively.
Parameter Size: The anchor- seeking policy is built upon a compact neural network. For model- based algorithms like MOPO and MOBILE, the dynamics parameter size is approximately 1.9M, similar to that of COCOA. However, the parameter size needed to train the actor and critic for MOPO and MOBILE E is equivalent to 0.2 However, when COCOA is added to these algorithms,
13

---

Published as a conference paper at ICLR 2024
the parameter size decreases to 0.19M. Given the significant magnitude of the dynamics parameters, the cumulative parameter requirement for training across COCOA-added model-based algorithms consistently stands at 2.2M. In contrast, IQL+COCOA and COL+COCOA, which operate without dynamics model, each have parameter size of 2.0M.
A.3 CODE IMPLEMENTATION
Our method is designed as add-on enhancement to existing offline RL algorithms. Consequently, rather than developing a new implementation, we adapted the established codebases of base algo- rithms. For consistent and reliable code adaptation, all base algorithms, including CQL (Kumar et al., 2020), IQL (Kostrikov et al., 2022), MOPO (Yu al., 2020b) and MOBILE (Sun al., 2023). The reliability of this codebase supported by de- tailed training logs and results that align with those in the original papers. Additionally, Sun (2023) offers results for the Gym-MuJoCo-v2 datasets not present the original CQL and MOPO papers, satisfying our needs. Note that an author of MOBILE (Sun et al., 2023) provides this codebase. Our adaptations to the code are shared as a demo in the supplementary material.
relied on Sun (2023) the foundation for
A.4 HYPERPARAMETER FOR EACH ALGORITHM
CQL. For both CQL and COL+COCOA, we use a = 5.0 for all D4RL-Gym tasks because the reproduced codebase (Sun, 2023) which provides the results for MuJoCo-v2 tasks, which are not included in the original paper (Kumar et al., 2020), uses this value. For COCOA, the anchor seeking for most tasks, except for "halfcheetah-medee medium-eepppetetv""""" "hopper- medium-expert-v22, and "walker2d-mediim-exppee
horizon length h was
where h was set to 3.
IQL. For both IQI., we use the same hyperparameters described in the original paper (Kostrikov et al., 2022), т 0.7 and A = 3.0, which is also used in the reproduced codebase (Sun, 2023). For IQL+COCOA, we used r = 0.6 and 3 = 3.0. For COCOA, we set the anchor seeking horizon length h to I for all tasks. We reproduce the random value for halfcheetah, hopper, walker2d, which are 6.62 to 6, 8.1 to 7, 6.1 to 6.5 respectively.
the hyperparameters used in the reproduced codebase (Sun, 2023), which provides the results for MuJoCo-v2 tasks not included the original paper (Yu et al., 2020b). As in the original paper, we use aleatoric uncertainty for MOPO and MOPO+COCOA. For MOPO+COCOA, we search for the best penalty coefficient À and rollout length h, for each task in the following ranges: À E {0.1, 0.5, 1.0, 5.0, 10.0), h, {1,5, 7, 10) except for the case of wiwichhhttah-mdiuu--eee
MОPО. For MOPO,
The best hyperparameters are described in Table 3. For COCOA, we the anchor seeking horizon length h to 1 for all tasks.
MOBILE. We use the same hyperparameters described in the original paper (Sun et al., 2023) for MOBILE. For MOBILE+COCOA, we search for the best penalty coefficient À and rollout length h, for each task in the following ranges: À e {0.1, 1.5, 2.0}. h, E {1. 5, 10} except for the case of walker-mediuu--rplay The best hyperparameters are described Table 3. For COCOA, we set the anchor seeking horizon length h to for all tasks. Additionally, after verifying convergence, we limited our training to a maximum of 2000 epochs and obtained results from this specific epoch range.
A.5 ADDITIONAL EXPERIMENTAL RESULTS
We experimented with two more benchmarks D4RL Adroit and NeoRL. The outcomes of these ex- periments are summarized in the table4 and 5. This broader analysis reveals that COCOA enhances the performance of IQL. and MOBILE most tasks. All experiments on additional benchmarks were conducted without applying layer normalization to allow for direct comparison with the per- formance reported in their original papers.
Our method demonstrated consistent performance enhancements across six D4RL Adroit tasks we tested, showcasing its robustness and adaptability. While COCOA encountered challenges in com- plex tasks like the door and hammer, akin to its original algorithm, this reflects the inherent difficulty of these tasks due to sparse rewards. Notably, in tasks like the pen, our method achieved noticeable performance improvements.
14

---

Published as a conference paper at ICLR 2024
Table 3: Hyperparameters for MOPO+COCOA and MOBILE&COCOA on D4RL benchmark.
MOPO+COCOA
MOBILE+COCOA
Task
X
14
X
hr
halfcho
0.1
7
1.0
10
hopper-random
10.0
5
walker2d-random
0.5
5
1.0
5
balfcheetaa--mddimm
0.5
7
1.0
5
hopper-medium
5.0
10
1.5
10
walker2d-medium
1.0
7
1.0
10
halfcheetah senddium-rpplay
1.0
10
1.0
10
hopper-medium- replay
0.1
5
1.0
10
walker2d-mediuum-ree mepllm replay
0.5
10
1.0
mallihhetta-mmciu--exppee
5
2.5
10
hopper-medium -epperrreeeert
5.0
10
2.0
10
malker2d-mediunee medium-expert expert
1.0
10
1.0
10
Table 4: Adroit benchmark results. We report the normalized average return of the last 10 training epochs across 2 seeds on the Adroit benchmark tasks.
IQL
MOBILE
Task
Alone
+COCOA
Alone
+COCOA
door-cloned-vl
2.11
2.81
-0.27
0.92
door-human-v1
5.27
6.81
-0.28
-0.05
hammer-cloned-vv
0.53
0.57
0.23
0.29
Kammer-human-w
1.33
2.21
0.25
1.11
pen-cloned-vl
72.55
73.98
54.95
48.9
pen-human-vl
74.4
78.15
14.76
41.65
Average
26.03
27.22
11.61
15,47
In Adroit, the training epoch limited to 200 as described in (Sun et al., 2023). Plus, we used the same hyperparameters for MOBILE&COCOA Adroit described in the paper. Hyperparameters for Adroit and NeoRI. benchmark are described in Table 6.
A.6 COMPARISON WITH COMBO ALGORITHM
COL+COCOA and COMBO exhibit some similarities, notably in their use of dynamics and a less conservative approach towards state-action space. However, their methodologies in pursuing con- servatism differ significantly: COCOA focuses conservatism in the compositional input space, whereas COMBO emphasizes regularizing the values for unfamiliar actions. Thus, COCOA and COMBO are orthogonal, it would be insightful comparison to integrate COCOA with COMBO, as COCOA can be an add-on to any algorithm.
Similar to COMBO, MBPO-based methods like MOPO and RAMBO also show a tendency to outperform model-free methods in random and medium settings. It seems that data augmentation through MBPO is particularly beneficial in these tasks. It would be interesting theoretically or empirically compare the state-specific value functions among CQL, COL+COCOA, and COMBO for further analysis.
15

---

Published as a conference paper at ICLR 2024
Table NeoRL benchmark results. We report the normalized average return of the last 10 training epochs across 4 seeds on the NeoRL benchmark tasks.
Task
MOPO
MOPO+COCOA
HalfCheetah-I.
40.1
47,47
Hopper-L
6.2
23.02
Walker2d-L
11.6
14.23
HalfCheetah-M
62.3
78.26
Hopper-M
61.65
Walker2d-M
39.9
49.80
HalfCheetah-H
65.9
42.44
Hopper-H
11.5
29.31
Walker2d-H
18
53.38
Average
28.5
44.40
Table 6: Hyperparameters of MOBILE&COCOA for Adroit benchmark and MOPO+COCOА for NeoRL benchmark.
Task
hp
door-cloned-vl
0.5
7
door-human-vl
3
hammer-cloned-vl
hammer-human-vl
pen-cloned-vl
0.5
1
pen-human-vl
10
HallChectah-vv-llw
0.5
5
Hopper-v3-low
2.5
5
Walker2d-v3-low
2.5
HalfCheettab-x3mmmdiim
0.5
5
Hopper-x3-medium
1.5
5
Walker2d--33meddimm
2.5
HalfCheetab-vv3-high
1.5
5
Hopper-v3-high
2.5
5
Walker2d-v3-hige
2.5
B DETAILED RESULTS
B. THEORETICAL ANALYSIS
Here, we simply show how our transductive predictor can be approximated in bilinear operation with the Stone-Weierstras theorem.
Theorem (Stone- Weierstrass Theorem for Locally Compact Spaces) Suppose a locally compact Hausdorff space and A is subalgebra of Cu R). Then A is dense in Co(X, R) with spect to the topology of uniform convergence if and only if separates points and vanishes nowhere.
Let X' and ) be locally compact Hausdorff (LCH) spaces. Plus, let F c. C(X; and G c C(); RR) be dense subvector spaces in the topology uniform convergence on compacta. Then the Theorem tells us that
Σ (4 fi...., EF.g....... E EG CC(X));
which forms an algebra, is dense in the topology of uniform convergence on compacta. other words, if we have joint embedding fo X Rdd and go: y Rd then ha, y) = fo(x)) go(
16

---

Published as a conference paper at ICLR 2024
universal approximator, such that (d, width) - (oo, oo) and fe(z), go(y) have depth 2 2. Considering we utilize a parameterized network to approximate a transductive predictor and our input space (s, a) is subset of T x R", where and n denote their respective dimensions, 4 and 40,2 which are described in Section 3.1, can correspond to fo and Jø. respectively.
B.2 PERFORMANCE GRAPHS OF D4RL BENCHMARK TASKS
In this section, we provide the performance graphs of each algorithm on D4RL benchmark tasks. We include only the 9 tasks that are not "random" tasks because the checkpoints of the baseline methods for the "random" tasks are not provided.
modum-v2
happer-mallm-22
valee22--medu--d
8
CQL COCDA
COL COCOA
COL COCDA
COL
CQL
COCOA (Mto ASS.)
CQL COCOK AS)
CQL COCOA (WlC AS).
60
60
92
10
24
50
10
15
Training Epocts
Training Epechs
Training Epocts
hallchhotahh v2
hoopar-redi.mmppppy wap.aa
waker26-madum vuplayy
Nor Aug. Reaums
Norm. Retums
Norm
OQL COCOA
GQLL + COCOA
CQL + COCDA
COL
COCOA (Mto AS)
COL COCOA (wic AS.
COL COCCA (ulc AS).
6.1 60
20
24
Training Epochs
Tranning Epochs
Training Epochs
halichootah madium-edmecc experte22
wakker26-mosimmooppett
Darm Avg. Returns
Narm Avp. Retums
Narm Avg. Returss
CQL COCOA
CQe COCOA GEL CQL pato
OQL COCOA (win A.S.)
0OL CUCDA
COL
COL COCOA (wlc A.S.)
Training
Training Epochs
Training Epochs
Figure 3: Performance comparison of CQL, CQL+COCOA and CQL+COCOA without anchor-
seeking across D4RL tasks except for "random" tasks.
17

---

Published as a conference paper at ICLR 2024
maltchersh--mmeue
kerzd medium
Norm Aug. Retums
Norm. Aug. Retuma
Norm Aug. Retuma
COCOA
COCDA
K2L
IQL
IQL
Trairing EernggEEpoo Epoctts
Triining Epochs
Triiiing Epocch
hallchestah medium reaayy V2
rooper-medum reslay
walker23 medium replay
Darm Avg. Retumns
Norm Returs
Meg
KIL CDCOA
IQL GOCOA
COCOA
Training Eppcchs
training Eppccs
Training Eppcchi
halichestah modium expertv22
hoger-inedum-ooe un-oxpert-22
walker25 medium expen-v2
Norm. Avg. Reums
Aug. Retarts
Aog.
IUL CDCOA
IQL COCOA
KQL + COCOA
Epochs
Epocts
Figure 4: Performance comparison of IQL and IQL+COCOA across all D4RL tasks except for "random" tasks.
Irafchhetthh frafihhetallmeddmmee medum-22
hopper-mecium-g
unlee26--medumm-c
MOPO COCDA
Norm Aug. Rasums
MOPO
Retumss
MOPO COCOA
MCPO COCOA
MOPO
20
Training Epochs
Training Epochs
Training Epochs
helichestah medium replay
hopper medium reploy
wash Bellerz mediumm replay
Narm Aup. Feaums
Narm (ug. Retuins
Aug.
MOPO COCOA
MOPO COCOA
MOPO COCOA
Training
Traniing
Trainin Epochs
hulichertah medium expert v2
hogpermmeuum experivv
walker2d medium expertu22
MOnO
Narm Avg. Reeums
Nar. Aug. Returts
MOPO
Narm Avg. Rewms
MOPO COCOA
MOP3 + COCOA MOPO
MOPO
Training Epoche
Training Epochs
Triinin Epochs
Figure 5: Performance comparison of MOPO and MOPO+COCOA across all D4RL tasks except for "random" tasks.
18

---

Published as a conference paper at ICLR 2024
haltheetth medum-v2
hoppermedium.ee x2
walkers mediumie
Norm AND. Renumms
Norm Reaum
Ju AWiN
MOBILE COCOA
MOBILE COCOA MOBILE
MOBILE COCOA MOBILE
MOBILE
10
24
20
20
Lo
20
Training Epochs
Training Epochs
Training Epocte
hallcheetah medium replay-
hoppor mredi.m repay-2
walker2o mediume caplay
Norm Avp. Reauss
Norm Avg. Reaums
Aug.
MOBILE COCOA
MOBILE COCOA
MOBILE COCOA
Training
Training Epochs
Traiing Epocts
halicheatah medium expertv22
hooper-meduum-uueae
walker26-ma mwd.mmmmuuue www.m--mpe---tteeee
Norm. Avg. seerums
Nor Aug. Renunts
Norm Avg. Retums
MOBILE COCOA MOBILE
MOBILE COCOA MOBILE
MORILE COCOA MOBILE
Training Excchs
Training Epochs
Trainin Epocts
Figure 6: Performance comparison of MOBILE and MOBILE+COCOA across all D4RL tasks
except for "random" tasks.
19


Reviews:
[{"review_id": "aPKlYQKrl4", "reviewer": "Reviewer_HX7o", "summary": "The paper presents an approach to enhance offline reinforcement learning by reparameterizing out-of-distribution states into two simpler components: an anchor and a difference from the anchor, leveraging bilinear transduction for better generalization. The proposed anchor-seeking policy, which utilizes reverse model rollouts, aims to identify these components from within the distribution seen during training, ensuring relevance and improving computational efficiency. The policy is complemented by a reverse dynamics model that diversifies the training data, allowing the model to generalize to new environments effectively. This novel method has been empirically shown to enhance the performance of several state-of-the-art offline RL algorithms.", "questions": "please see weakness", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Pros:\n1. The paper is well-written and easy to follow.  The reproducibility is good as it provided code.\n2. Introduce a novel perspective in finding conservatism in the compositional input space, different from the previous works in finding conservatism in the behavioral space.\n3. COCOA can be used in-combination with other offline RL algorithms to improve the performance of previous model-free and model-based algorithms.\n4. Good experimental results with SOTA performance on the D4RL benchmarks.\n5. Ablation study in comparing with a baseline anchor selection process demonstrate that their anchor-seeking methods is crucial for their performance.", "weaknesses": "Cons:\n1. Why not train a reverse policy but use a random divergent reverse policy? Is there an ablation study on this?\n2. The evaluation is only limited to D4rl benchmarks and why in IQL, adding COCOA does not improve much among all the d4rl tasks, could you provide some insights into this?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9DfdBaUwtc", "reviewer": "Reviewer_vdDy", "summary": "This paper proposes a new form of conservatism in offline RL called COCOA, specifically with \"anchor states\" defined by a learned dynamics model. The use of anchor states is motivated by work transforming the out-of-support generalization problem to an out-of-combination problem in the offline RL context.\n\nExperiments show that the addition of COCOA to common offline RL algorithms results in improved performance across the D4RL benchmark suite, and ablations show that the proposed anchor-seeking method is crucial for good performance across said domains.", "questions": "No real big questions here actually -- very solid!", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall, the paper is very well written, and generally easy to follow. While I haven't worked in this area in RL research a lot, it is interesting to see the extension of a supervised learning approach for use in offline RL, and for it to work as well as it does seem to work here.\n\nThe performance results are very interesting, and show that anchor-seeking methods can improve performance especially in the medium-expert domains, where maybe excess conservatism leads to suboptimal performance for most offline RL algorithms. The experiments and ablations are generally solid and thorough, across many relevant benchmarks.", "weaknesses": "It seems that this model-based approach here is similar to those in forward model-based offline RL algorithms such as COMBO [1], where the goal is to use a larger coverage state-action distribution for offline RL through the use of a forward dynamics model. It would be interesting to see the comparison between this approach (combined with CQL) and COMBO (if COMBO results can be reproduced), as COMBO also shows significantly less conservatism when used with CQL compared to CQL on its own. I wouldn't say that this is a big weakness, but I feel like from a high level it would be worth it to do this comparison.\n\nThere has also been work showing that even synthetic experience replay [2] from strong generative models is useful for RL -- might be worth citing or comparing against as well, but I may be wrong here. Again, not a huge weakness in my eyes but would be interesting to see the comparison.\n\n[1] Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, Chelsea Finn; COMBO: Conservative Offline Model-based Policy Optimization (NeurIPS 2021)\n\n[2] Cong Lu, Philip J. Ball, Yee Whye Teh, Jack Parker-Holder; Synthetic Experience Replay (NeurIPS 2023)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eKE8emwY9z", "reviewer": "Reviewer_7X6Q", "summary": "Authors propose method COCOA for adapting bilinear transduction to offline RL setup with training environment dynamics and anchor-seeking policy.", "questions": "* What is the training time for algorithms with and without COCOA?\n\n* What is hyperparameters sensitivity for MOPO and MOBILE when COCOA is added compared to the same algorithms without COCOA? EOP (https://arxiv.org/abs/2110.04156) can be used for this purpose.\n\n* What are IQL scores on random datasets?\n\n* Are there any thoughts why IQL suffer from COCOA?\n\n* How would COCOA behave with offline RL algorithms which regularize policy? E.g. ReBRAC (https://arxiv.org/pdf/2305.09836.pdf) which is the state-of-the-art algorithm from this family.\n\n* Could you please run experiments on AntMaze or Adroit domains or test your approach in offline-to-online setup? At least for IQL and MOBILE.\n\n* Please add the average scores for each of the approaches in the Table 1 so it is clear what is actual impact of COCOA on algorithms?\n\nMinor text issues: the second paragraph of section 2.1 duplicates information from the first one. Results section has  \"**FFor** IQL,\".", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Proposed approach may boost the performance of the offline RL algorithms and applicable to any actor-critic algorithm. Approach is tested on 4 different algorithms.", "weaknesses": "* Training world models might be very time-consuming.\n\n* Proposed approach lead to the improvement only in approximately 58% of cases. And performance might be decreased dramatically. I've also calculated the differences between \"+ COCOA\" and \"Alone\" in Table 1 and it appeared that COCOA decreases performance down by 0.8 points. This indicates that only specific algorithms are expected to benefit on average by COCOA. And for MOPO and MOBILE hyperparameters were heavily tuned for each dataset which might be the cause why they benefited.\n\n* Evaluation is performed only on the Gym MuJoCo datasets which I think is not enough now and evaluation on D4RL AntMaze or Adroit is essential. As an alternative, offline-to-online setup might be tested.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)