Rethinking Optimal Transport in Offline Reinforcement Learning
Authors: Arip Asadulaev
Rostislav Korst
Alexander Korotin
Vage Egiazarian
Andrey Filchenkov
Evgeny Burnaev
Venue: NeurIPS 2024 poster (2024)
Keywords: Optimal Transport
Reinforcement Learning
Primary area: reinforcement_learning

Abstract: We propose a novel algorithm for offline reinforcement learning using optimal transport. Typically, in offline reinforcement learning, the data is provided by various experts and some of them can be sub-optimal. To extract an efficient policy, it is necessary to \emph{stitch} the best behaviors from the dataset. To address this problem, we rethink offline reinforcement learning as an optimal transportation problem. And based on this, we present an algorithm that aims to find a policy that maps states to a \emph{partial} distribution of the best expert actions for each given state. We evaluate the performance of our algorithm on continuous control problems from the D4RL suite and demonstrate improvements over existing methods.

Rethinking Optimal Transport in Offline Reinforcement Learning
Arip Asadulaey Alexander Korotin
Rostislav Korst
Vage Egiazarian
Andrey Filchenkov
Evgeny Burnaey4,1
AIRI
2TTMO
MIPT
4Skottech Yandex
HSS University
asadulaev@airi..nt
Abstract
We propose novel algorithm for offline reinforcement learning using optimal transport. Typically, in offline reinforcement learning. the data is provided by various experts and some of them can sub-optimal. To extract an efficient policy, it is necessary to stitch the best behaviors from the dataset. To address this problem, we rethink offline reinforcement learning as an optimal transport problem. And based on this, we present an algorithm that aims to find policy that maps states a partial distribution the best expert actions for each given state. We evaluate the performance of our algorithm on continuous control problems from the D4RL suite and demonstrate improvements over existing methods.
Introduction
Deep reinforcement learning (RL) has shown remarkable progress in complex tasks such as strategy games [5I], robotics [37], and dialogue systems (39]. In online RL, agents operate in the environment and receive rewards that are used to update policies. However, in some domains, such healtheare. industrial control, and robotics interactions with the enviromment can be costly and risky. For these reasons, offline RL. algorithms that use historical data of expert interactions with the environment are more applicable. Given only a dataset of experiences, offline RL. algorithms learn a policy without any online actions
33
Due to the off- nature of offline RL, naive training of online RL algorithms in offline settings does provide effective solutions K3 This is primarily due distribution shift between the learned and the behavior policy [15), which causes the instability of the critie function during off- policy evaluation. To solve this problem and make actions of the learned policy more similar the one on which critic function was trained, a Behavior Cloning (BC) objectives was proposed [52] 33
As example of BC loss, Optimal Transport (OT) [50] distance has been introduced for RL. This application OT' is particularly relevant because by minimizing the OT distance between policies, we can efficiently clone the behavior [52
34 20 g (3). But there is another problem, in offline RL, the dataset often consists of various experts demonstrations. Some of these demonstrations may be incorrect, or efficient only for certain parts of the environment. Consequently, cloning the behavior policy in such scenarios may limit policy improvement (30] To solve these types of problems, RL algorithms need apply stitching, which means that the algorithm selects the best action for each state and recombines different trajectories provided by multiple experts [30]
[ 40
In our paper, rethink OT framework for offline reinforcement learning that aims to stitch the best trajectories, rather than clone the policy. Unlike previous OT-baased approaches in RL, we do not introduce new OT-based regularization reward function []]
40 Instead, propose
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

novel perspective that views the entire offline RL problem as an optimal transport problem between state and actions distribution ($3). By utilizing the Q-function as a transport cost and the policy as optimal transport map. we formalize offline reinforcement learning as a Maximin OT optimization problem. This opens doors for applying recent advantages OT methods to RI. directly. In particular, we propose an algorithm that trains policy to identify the partial distribution of the best actions for each given state.
We proposed novel optimal transport-based policy extraction method and provided an analysis of its performance on various RL-based cost functions. The core of our method is the Partial Policy Learning (PPL) algorithm, which efficiently identifies the best action for each state in the dataset. We evaluated our method across various environments using the D4RL benchmark suite, achieving superior performance compared to state-of ofut-oof-the-ar model- free offline RL techniques.
Contribution:
2
Background and Related Work
2.1 Optimal Transport
Primal form of the optimal transport problem was first proposed by Monge [50]. Suppose there are two probability distributions fs and y over measurable spaces X' and y respectively, where X, C RD. We want to find a measurable map T : X - y such that the mapped distribution is equal to the target v, for a cost function e: X' x y - IR, the OT problem between u,
min Ex [em=[c((x; T()))] Thpzv
(1)
where T: is the push forward operator and respectively Tiu = v represents the mass preserving condition. Informally, we can say that the cost is measure of how hard it is to move mass piece between points x E X' and y E y from distributions fr and v correspondingly. That is, an OT map 7 shows how optimally move the mass of μ V, i.e., with the minimal effort. A widely recognized alternative formulation for optimal transport was introduced by Kantorovich [22]. Unlike the Monge's OT problem formulation, this alternative allows for mass splitting. The Kantorovich OT problem can be written as:
min B(ra)) [((1, y)].
(2)
gEll(w.))
In this case, the minimum is obtained over the transport plans g, which refers to the couplings II with the respective marginals being and v. The optimal g belonging to [(u, v) is referred to as the optimal transport plan.
Dual form of optimal transport, following the Kantorovich-Rubbnstein duality 49), can be obtained from 2 and written as:
max Eaow (x)) + Ey-- [ [f(w]]
f
where f x X + R and (x)) min [c(rr f is referred to as the c-transform of potential
f(y]
yEu
[50] For cost function c(x, y) = | y 1l22 the resulted OT cost called the Wasserstein- I distance, see [50 $1] or [44] 11, 2]. was shown that to compute Wasserstein- distance, instead of computing conjugate function f we can simply consider to be from the set of 1 1-Lipschitz functions (50
Maximin formulation for simultaneously computing the OT distance and recovering the OT map T was recently proposed. According to (27), by applying the Rockafellar interchange theorem [3] Theorem 3AJ, we can replace the optimization over points y E y with an optimization over functions
T X + ),, which reformulates the problem 3 as saddle-point optimization problem for the potential f and map T:
max min Raz [c(r, T(x)) - S(T())) + EEy-r
Using this formulation, significant progress has been achieved in utilizing neural networks for computing OT maps address strong [36]] 24 [8), weak [27], and general [2] OT problems. Our work is inspired by these developments, as we leverage neural optimal transport methods to improve offline reinforcement learning.
Partial OT: cases where is necessary to ignore some data and map the input to part of the target
distribution, techniques like unbalanced or partial optimal transport are commonly employed 12

---

17]. For example, partial optimal transport is useful in resource allocation problems where only a subset of resources needs to be optimally allocated to match subset of demands [2]
In our paper. we consider the partial OT formulation proposed by 17 This framework, named extremal OT, solves the partial alignment between the full input distribution and part of the target distribution. Essentially, for Euclidean cost functions such as l² partial transport maps can be seen as a tool for finding nearest neighbors (maximally close) to the input samples from the target according to the cost function. Please see Figure 1 in 17.. Formally, this method can be written as: min TipSav
Ex ((z,T(x))]
(5)
We can call a coefficient of unbalance. The intuition behind this formulation is follows: when
w = 1. the constraint Tip S wD is equivalent to Tiu = v. This equivalence holds because there is only one probability measure that is less than or equal to v, and is v itself. However, with values of
w 2 I, the mass of the second measure is scaled. Consequently. it is sufficient to match f only to part of the second measure v to transfer the full mass contained in jz. Since the cost is our objective, T will tend to map to the samples from v that are closest to the input distribution. It was shown that this formulation can also be considered in maximin settings which is the core of our method ($3.2).
2.2 Offline Reinforcement Learning
Reinforcement Learning is a well-established framework for decision-making processes in environ- ments modeled as Markov Decision Processes (MDPs). The MDP is defined as M = (S, A, P.r,y) by the state space S, action space A, transition probability P(s' s, a), reward function r(s, a), and discount factor y. The goal in RL is to find a policy x (als) that maximizes the expected cumulative
reward over time t: J(T) daf E, E ('r(st, a.)]. Also we define the distribution over the state space following policy
d= (s). To estimate of the expected cumulative reward following a given policy π the critic function Q is used:
Q" (s) r (s,a)) + IQ" (s',a'))]
In deep RI. neural networks are used to parameterize critic Q7 a) and policy no- Critic can be learned by minimizing the mean squared Bellman error over the experience replay dataset D = [(si, ai, sy,, r:)}. which consist of trajectories implied by policy n:
min s") (r(s,a) a) - (lEe'= (Q ad)] - Q ())) 23
Using trained critic function we can recover near-optimal policy, for example, using Deterministic Policy Gradient (DPG) (45) or its improved version named Twin Delayed Deep Deterministic Policy Gradient (TD3) 15).
Offinee RL is fully data-driven approach for decision-making problems. Conventional offline RL algorithms use the 7 formula to recover the Q function, using data D collected according to the expert policy 3. However, due to distribution shift between actions from DD and those induced by policy π, the Q function may suffer from inefficiency and overestimation bias during evaluation [[5]. To mitigate this issue, various offline RL algorithms with behavior cloning objectives was proposed. See the related work section (36 for more details.
OT in offline RL has been proposed primarily for efficient behavior cloning. Compared to BC/KL, OT offers more flexibility, while including many building blocks to properly consider the geometry of the problem. The simplest method of integrating OT into offline RL is by utilizing the Wasserstein-1 distance the measure of dissimilarity between the learned and expert policies. This method is known as the W-BRAC algorithm [52]. Given the critic function Q(s, a), the Wasserstein- distance defined by f, and the behavior cloning coefficient a, we have
min= max (f $1 Ex-Davx(())
(s. + a (((Ea))D [f(s, a)] - Ba-Da [fss (8)
Critic
Wassersteie distance
In real-world applications, collecting large and high-quality datasets may be too costly or impractical. Due to that, the provided data it is often collected by many expert playing several near-optimal policies and only a single expert can optimal dataset D. The formulation (8) has no mechanism to infer the importance of each action, its just clone the entire data. Moreover, the problem is that computing Wasserstein-I / distance is typically complicuted 25). because the potential f is required to satisfy the 1 Lipschhitz condition [1. It is also important to note that the coefficient is task-dependent(6.
3

---

Algorithm Partial Policy Learning
Input: Dataset a,, Initialize Qo. TT, fw. / for in I...N do
(s. a, r.) t- D: sample a batch of transitions from the dataset.
Qh+1 &- Update cost function Q using the Bellman update in (2).
fk+t
Update fw using outputs of To and samples from dataset:
arg min, - (awn" (e) (s,a)) + wB..a-Dl (8,(
n nh+1 f- Update policy wo as transport map: arg min, E (a - a) - f (s,
end for
3 Rethinking OT in RL
3.1 Considering RL as OT problem
To extend connection and apply OT methods in RL, we can view the entire offline RL prob- lem as an optimal transport problem. For this purpose, let's consider (T). By replacing the cost function c(2, y) with the critic Q* (s,a) and treating our policy as map that moves mass from the state distribution (s). (d(s) for shorthand), to the corresponding distribution of ac- tions given by the behavior policy (-|ss, we obtain a primal state-conditional Monge OT problem.
EgwD,a> ((( (= a)] which is equal to:
min
EA-Da (s.
( (rdd())=)((11)
The objective is minimize the expectation of the negative critic function Q*, (we can also consider Q9) while mapping exclusively to the distribution of actions given by the expert policy B.
Remark 1: In this formulation, the imposed equality constraints force the policy to mimic the behavior of the expert policy. However, since the provided dataset often contains sub-optimal paths, this formulation remains inappropriate for avoiding inefficient actions provided by the expert policy.
In the following subsection, introduce algorithm that integrates Partial Optimal Transport [17 for deep RL.
3.2 Stitching with Partial OT
In many real-world applications, collecting high-quality datasets may be costly or impractical. Therefore, offline RL. algorithms need to select the best actions provided by the experts and ignore sub-optimal ones. Our idea is that partial OT can help us deal with sub-optimal trajectories by mapping states only to the most relevant parts of the action distribution with respect the critic function Q To obtain partial formulation 3f instead of equality constraints, we need to consider optimization problem with the following inequality constraints [17]
min
Eg-D.awa(s) Q (s.
(10)
}))}
As discussed in ($2.1). the parameter w scales the mass contained in the target distribution. For values of w 2 1, the mass of the second measure is increased. Consequently, it is not necessary to cover the full distribution 3(1(s) to transport the Dirac mass given by d(s). As a result, a learned policy (map) matches the state only part of the action distribution B(-ls). This is particularly relevant for offline RL, as the resulting policy will not clone all actions but will select only the most optimal ones, where optimality is defined by the critic function Q.
Proposition 1: (Informal) for critic-based cost function Q", an offline Policy trained by (10) improves over the behavior policy 3. Proofs are given an Appendix ($9.1).
To find practical solution to the given problem using neural networks, we can also consider dual and maximin formulation for (10). The dual form for the partial transport problem (17 can be ex expressed as
max E. a~a (s, a)] + WIE. a [f(s. f20
(11)

---

To obtain the maximin formulation, we expand the c-transform using the function Q as the cost This operation can be represented as EE-D,a> a)] E min, { (-Q* a) f(s, a)} Subsequently, by applying the Rockafellar interchange theorem 44 Theorem 3AJ, we replace the
optimization over points with an equivalent optimization over functions क SS - A. The solution can be computed using neural approximations, which is useful for RL tasks
max min En (aaa((())) /20
-Q,(s,a) a() + wel (s))))))
(12)
Cost
Constraints
Remark 2: In the 17 Proposition 4] was shown that f a) vanishes on the outliers i.e on the samples outside the scaled target distribution. The policy update in (22) can be written as min, Ex- a) f(s, a)}. This reveals that if f(s, a) exclusively on the outlier actions, it then applies extra weight to correct actions, which reduces the value of -(s, [(s a).
Remark 3: The value of w controls the size of the action space in which we map; the higher the value of w. the smaller the subset([7). When the actions are provided by a single expert, scaling can be detrimental to performance.
The primary distinction between our method as follows:
12 and previous OT for RL approaches, particularly
OT is not treated as a component of the problem of o. Instead, the entire policy optimization is considered as OT.
regularization term with coefficient
Our formulation for computing OT does not necessitate 1- I-Lipschitz function constraints on f. as it does not compute the Wasserstein-I distance. Rather, it addresses a maxmin OT problem with critic-based cost function, where f can be any arbitrary scalar function satisfying f 2 O.
Unlike traditional OT approaches that fully aligns two distributions, our method maps only to the part of the target distribution with the best actions with respect to the critic.
In practice,
R parameterize æ and respectively. These neural networks serve as approximators that can capture complex mappings of states. We train these neural networks using stochastic gradient ascent-descent (SGAD) by sampling random batches from the dataset D. At each training step, we sample batch of transitions from the offline dataset D and then adjust the value function Q, and the potential f. Then for the l update step, given the last policy value function Q* and f we update the policy, see Algorithm
neural networks Ty S A and fw S x
Toy Experiments
Eo Bortor
Calisaa
Teceve Kitaiie Pilyy?
(a)
(b)
Toy experiments. (a) Left point So denoting start and Sr is the only rewarded, target location. Black curves visualize behavior trajectories 3. (b) Best behavior policy 3 according to the data, and the optimal policy x² that provides the optimal (shortest path) solution. (c) Results of the policy 33r trained by minimizing - -Q
Figure 1:
((, (s))+ and policy æ trained by our algorithm
To demonstrate the ability of our algorithm to extract the best policy from the sub-optimal behavior trajectories, we constructed a simplified offline RL dataset. The goal is to find the shortest path from a starting location So to a reward point S7, navigating through 50 intermediate states S, uniformly spaced along the æraxxi from -1.3 to 0. The action space is continuous, with the agent's 's task being to determine the optimal y- values for the each r-coordinate. For the each intermediate state-action pair, the reward is defined negative Euclidean distance to the final state-action action pair Sr, Ar.
5

---

The offline RL dataset consisted of three sub-optimal expert trajectories, as shown in Figure ((a The best policy that can be extracted by recombining the expert trajectories and an actual optimal policy x¹ that provides the shortest path are highlighted in Figure ((b)
We compared our method with the offline RI. algorithm based behavior cloning. For this, trained the critic Q* by (2), using actions drawn from the observed behavior trajectories & presented B(s)) + []8(s) - all²] to obtain the policy g Next, using the given critic function Qº we trained the policy (8) using our algorithm with parameter vo = 8. Compared to offiine RL Tbb] which produced average expert solutions, our method demonstrated superior performance by extracting and exploiting the most strategic insights from the data, maximizing the reward and ignoring all inefficient actions, see Figure [igure](c).
in the dataset D. Given Q we minimize aa
In all experiments, single-layer neural network with 32 neurons and ReLU activation the configuration for both the critic and policy networks. For our method, we used a Lagrangian network [s, a) with similar parameters. There are 5000 steps was done for each network, with Adam 23 optimizer with a learning rate (Ir) of 1e-3.
used as
5 D4RL Experiments
5.1
Benchmarks and Baselines
We evaluate our proposed method using the Datasets for Deep Data-Driven Reinforcement Learning (D4RL. 13 benchmark suite, which a collection of diverse datasets designed for training and evaluating deep RI. agents in a variety of settings. It includes tasks in continuous control, discrete control, and multi-agent settings with a variety of reward structures and observation spaces. First, we tested our method on the Gym's 's MuJoCo-v2 v2 environments, such as Walker, Hopper, and HalfCheetah. We also tested our method on the complex Antmaze-v2 and Android-v1 environments.
To compare the performance of our proposed method, we selected four state- of-the-art offline RL algorithms. These include Conservative Q-Learning (CQL) [32]. Twin Delayed Deep Deterministic Policy Gradient with behavior Cloning (TD3+BC) [14], Implicit Q-Learning IQL (28). ReBRAC [7], and IQL. with Optimal Transport Reward Labeling (IQL+OTR) (35]. More details on the methods are given in ($6). The results of the supervised behavior cloning are also included.
5.2 Settings
We tested our method in conjunction with various algorithms that avoid overestimation of Q-function. These overesiination-avoii method, showing that our method can work efficiently with any of them.
methods are used to obtain different types of cost functions in our
PPL. First, to be consistent with Proposition 1, we considered our method in the One-Step RL settings. Following the original setting provided by One-Step RI. (). we pre-trained / for 500k steps. Next, we pre-trained Qa for 2 million steps. Then, to avoid overestimation bias, we applied simplified conservative update Q3 (see (32] Eq.1]. IJ. Finally, to improve beyond the behavior policy 3, we trained the policy for the 100k steps using our method For these experiments, a two-layer feed-forward network with hidden layer size of 1024 and a learning rate of 0.001 was used with the Adam optimizer (23). We tested this approach for the Mujoco environment see Table for the comparison with the One RL method.
PPL CQL. Second, we tested our method in conjunction with the CQL method the Antmaze problems. this setting, no pre-trained models were used; we used the code provided by the CORL library 48). with all hyperparameters and architectures set identical to those originally used in CORL for the CQL method. We trained the algorithm for IM steps, with w set to 8 for all experiments. See the comparison with CQL in T
R Finally, we coupled our method with the improved version of TD3+BC called ReBRAC [4 In these experiments, for a cost function, the TD3+BC objective was used, which combines the Q function with supervised loss. The same hyperparameters proposed by the authors for ReBRAC were kept in our experiments. No pre-trained models were used, we trained the agent from scratch for
'e also tested direct minimization of (s, a). However, even in toy settings this method suffered from overestimation bias, causing the resulting policy to consistently predict values equal 1.

---

1M steps. We have tested our method on the Mujoco, Android, and Antmaze datasets. While the ReBRAC framework recommends task specific hyperparameters [47], the results for the Antmaze are also reported for the best w hyperparameter. The results can be found in the Tables and
5.3 Results
In the Table I we consider side-by side-b comparison, and the color purple indicates the best results between columns. In the Tables and 3 the color indicates the top results. We compared the
results with those reported in the ReBRAC [47] and OTR+IQL papers [35]] For the Antmaze datasets, the main comparison is with the oracle method which our method coupled. For comparison, we reproduced the CQL and ReBRAC results using the code provided in CORL. The CQL score curves are provided in Appendix 3
For completeness, IQL and OTR+IQL results are also included.
Table I: Averaged normalized scores Antmaze-v2 tasks. Reported scores are the results of the final 100 evaluations and 5 random seeds.
Datasct
IQL
OTR+IQL
CQL
(Ours))
ReBRAC
PPL"(Ours)
umaze
87.5 a 2.6
83.4 + 3.3
86.3 43.7
9042.6
97.8 # 1.0
98.0 +14
umaze-diverse
62.2 + 13.8
68.9 13.6
34.6 +20.9
40+2,6
88.3 13.0
93.6#6.1
medium-play
71.2
70.5 + 6.6
63.0 49.8
67.3+10.1
84.0
90.2 #3.1
medium-diverse
70.0 # 10.9
70.4 4.8
59.6 #3.5
65.3+8.0
76.3 13.5
84.8 #14.7
large-play
3.8
45.3 6.9
20.0 +10.8
25.6+3.7
60.4 26.1
76.8 +4.0
large-diverse
47.5 9.5
45.5 + 6.2
20.0 45.1
23.6 #11.0
54.4 25.1
76.6 = 7.4
Total
378
384
283.5
311.8
461.2
520
Table Averaged normalized scores on MuJoCo tasks. Reported scores are the results of the final 10 evaluations and 5 random seeds.
Dataset
BC
One-RI.
CQL
IQL.
OTR+IQI
TD3+BC
ReBRAC
PPI."
Half.
42.6
48.4
44.0
47.4
43.3
48.3
65.6
64.95+0.2
M
Hopper
52.9
59.6
58.5
66.3
78.7
59.3
102.0
93.49+7.2
Walker
75.3
81.1
72.5
78.3
79.4
65.5
82.5
85.66+0.6
Half.
36.6
38.1
45.5
44.2
41.3
44.6
51.0
51.1
MR
Hopper
18.1
97.5
95.0
94.7
84.8
60.9
98.1
100.0+2
Walker
26.0
49.5
77.3
73.9
66.0
81.8
77.3
78. 78.66422.0
Half.
55.2
93.4
91.6
86.7
89.6
90.7
101.1
104.8540.
ME
Hopper
52.5
103.3
105.4
91.5
93.2
98.0
107.0
109.0*1..
Walker
107.5
113.0
108.8
109.6
109.3
110.1
112.3
111.74+1.1
Total
467.7
684.9
698.6
692.6
685.6
659.2
796.9
799.45
In offline RL, the problems noisy, complex, and diverse, and task-specific hyperparameter search required. Despite this fact, our method shows promising results and consistently outperforms methods on which it is based. Most importantly achieved scores for the Antmaze problems (Table []. In particular, our method provides state-of- the-art results for all datasets this task, and gives significant improvement of (+16) and (+21) points for the complicated large-play-v2 and prrge--lay-diveree-ee environments. Importantly, we consistently outperform the previous best OT offline RI algorithm, OTR+IQL.
method gives significant improvement in the
We can interpret that our method lies between behavior cloning and direct maximization of the Q function. Recent studies have shown that direct maximization
lead to sub-optimal results due to overestimation bias [6]. Conversely, being too close to the expert's 's policy prevents improvement (30). Intuitively, our extremal formulation allows to strike balance, maximizing Q by actions from the expert action distribution that are not radically different from those on which the Q was trained.
Our formulation aligns the state space with only a portion of the action space. This property is partieularly relevant for offline datasets that contain a mixture of expert demonstrations, Our method can help to select the best possible action from set of expert trajectories [3 In other words, optimal f is to encourage the discovery new actions that can yield high rewards. To illustrate the task-specific dependence of the parameter w, we performed ablation for different values ($5.4)
7

---

Table 3: Averaged normalized scores on Android tasks. Reported scores are the results of the final 10 evaluations and random seeds.
Dataset
BC
TD3+BC
CQL
IQL
OTR+IQL
ReBRAC
PPL (Ours)
Human 34.4
81.8
37.5
81.5
66.82
103.5
108.43
Pen
Cloned
56.9
61.4
39.2
46.8
78.7
91.8
113.97
Expert
85.1
146.0
107.0
133.6
154.1
152.33
Human
0.5
-0.1
9.9
3.1
5.9
0.0
1.0
Door
Cloned
-0.1
0.1
0.4
0.8
0.01
1.1
0.44
Expert
34.9
101.5
105.3
104.6
104.66
Human
1.5
4.4
1.79
Hammer Cloned
0.8
2.1
0.8
4.74
Expert
125.6
117.0
86.7
106.5
133.8
130.2
Human
0.0
0.2
-0.2
0.23
Relocate
Cloned
-0.1
-0.1
-0.1
0.2
0.1
0.45
Expert 101.3
107.3
95.0
106.5
106.6
92.99
Total
340.9
600.5
484.8
588.5
703.2
711.44
Run time: The code is implemented in the PyTorch [4] and JAX frameworks and will be publicly available along with the trained networks. Our method converges within 2-3 hours on Nvidia 1080 (12 GB) GPU. We used WanDB (5] for babysitting training process. The code is available in supplementary materials.
5.4
Parameter w
80
40
20
20
Stop
Stop
Stop
200
400
600
800
1k
200
400
900
800
1k
200
400
600
800
(a) umaze-v2
(b) umaze-diverse-v2
medium-play-e
80
20
20
Stop
Stop
Stop
200
400
500
800
1k
200
400
800
800
1k
200
400
800
800
1k
(d) medium-diverse-v2
(c) large-play-vo
(f) large-diverse-v2
Figure 2: Exponential moving average (coef. 0.3) curves of the normalized score curves for the Antmaze. Different colors the curves represent results for w 3, w 5, w = 8, w = 12.
The parameter w influences the action selection process by determining the support range over which the policy operates. With the parameter equal to 1, our method considers all actions provided in the dataset for each state. By increasing this value, we consider the smaller subspace of possible actions, which can be usefuul when it is necessary to select the best one from the data. In perfect scenario, want to find as few actions as possible that maximize the score, so we favor the higher values of w.
For some datasets, reducing the action subspace is unnecessary. We studied different parameters w=[,,,,8, 12) for the range tasks, results for the Antmaze are presented Figure We can see that the biggest effect of the parameter was accrued for the large-play and large-diverse tasks 2 (e). This means that selecting the subspace of actions important for this environments, and our method allows for this nuanced control depending the task.

---

Related work
There is rich family of offline RL algorithms. To solve the distribution shift problem, several methods have been proposed [6 29 53
BIDE 28]. The simplest approach the behavior divergence D, between the learned and the expert policy 3. the BRAC framework can be written as distance regularized policy optimization: 3)) is BC loss. Based on BRAC framework [14] proposed TD3+BC approach that combines standard TD3 [5] with BC loss minimization between policy and expert actions, and showed strong performance.
regularized actor-criti (BRAC) 152 Minimizing
Where (=)(s,),, B (s)
Policy Gradient from Arbitrary Experience via DICE 38
proposes to incorporate f f-divergence
regularization into the critic function. Based on a similar concept, Conservative Q-Learning (32) and Adversarial Trained Actor-Critic [8] extends the critic loss with pessimistic terms that minimize the Q values on the samples of a learned policy and maximize the values of the dataset actions.
Methods such as one-step RL 6 solves stitching by approximating the behavior Qa function and extract the corresponding policy by weighting actions using the advantage function A (s. a) to find the best action for the given state [4). Based on this concept, Implicit Q-Learning (IQL) [28 approximates the policy improvement step by treating the state value function V (s) as random
variable and taking condicoontiional upper expectile to estimate the value of the best actions.
Application of optimal transport in offline RL was wiedly explored for behavioral cloning. For
example it used to construct pseudo-reward function [1] 40 34 20 9. Primal Wasserstein
Imitation Learning [1I and Sinkhorn Imitation Learning [40] use the OT distance between the expert and imitator create reward function. Recently, the Optimal Transport Reward Labeling (OTR) method 35 was proposed to generate offline RL algorithms. However, like W-BRAC, these methods primarily use OT as additional loss.
dataset that can then be used by various
7 Conclusion and Discussion
In this paper, we have established the novel algorithm for offline reinforcement learning. Our work introduces the concept using the action-value function and the policy as components of the OT problem between states and actions. To demonstrate the potential our formulation, we propose the practical algorithm that penalize the policy to avoid inefficient actions provided in the dataset. By applying partial optimal transport, our algorithm effectively selects and maps the best expert actions for each given state, ensuring efficient policy extraction from noisy or sub-optimal datasets. The limitation of our method is that parameter w is task dependent.
This work is a step towards our larger vision towards deepening the connection between OT and offline RL RL. Using our formulation other OT methods also can be integrated into RL. For example, various regularizations (26) [0] and general costs [42 2 can be used to incorporate task-specific information into the learned map, which can be particularly relevant in hierarchical RL problems ( Additionally Weak Neural OT (27) can be relevant in RL where stochastic behavior is preferred for exploration in the presence of multimodal goals (19].
7.1
Reproducibility
To reproduce our experiment we provide source code in htttppe//ggaab.cc/.aaaaeeeee/e// PPL/ Details on used hyperparameters are presented in settings (85.2).
7.2
Societal Impact
This paper presents research aimed at advancing the field RL. While our work may have various societal implications, do not believe any particular ones need to be specifically emphasized here.
Acknowledgment
The work was supported by the Analytical center under the Government (subsidy agreement 000000DD733321S500002 Grant No. 70-2021-00145 02. 11.202))

---

References
[1] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference machine learning, pages 214-223. PMLR, 2017.
[2] Arip Asadulaev, Alexander Korotin, Vage Egiazarian, and Evgeny Burnaev, Neural optimal transport with general cost functionals. arXiv preprint arXiv:2205.15403 2022.
[3] Pie Bacon, Jean Harb, and Doina Precup. The option-critic architecture. In Proceedings of the AAAI conference on artificial intelligence, volume 31, 2017.
[4] Yikun Bai, Bernhard Schmitzer, Matthew Thorpe, and Soheil Kolouri. Sliced optimal partial transport. In Proceedings of the IEEE/CVF Conference on Computer Viston and Pattern Recognition, pages 13681-13690, 2023.
[5] Lukas Biewald. Experiment tracking with weights and biases, 2020. Software available from wandb.com.
[6] David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline rl without off-policy evaluation. Advances neural information processing systems, 34:4933 4946, 2021.
[7] Xi Chen, Ali Ghadirzadeh, Tianhe Yu, Yuan Gao, Jianhao Wang, Wenzhe Li, Bin Liang, Chelsea Finn, and Chongjie Zhang. varibl- ad advantage-weighted policy optimization for offline rl. arXiv preprint arXiv:2203. 003.08949, 2022.
[8] Ching-An Cheng, Tengyang Xie, Nan Jiang, and Alekh Agarwal. Adversarially trained actor critic for offline reinforcement learning. In International Conference on Machine Learning, pages 3852-3878. PMLR, 2022.
[9] Samuel Cohen, Brandon Amos, Mare Peter Deisenroth, Mikacl Henaff, Eugene Vinitsky, and Denis Yarats. Imitation learning from pixel observations for continuous control. 2021.
[10] Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamoniy. Optimal transport for domain adaptation. IEEE transactions pattern analysis and machine intelligence, 39(9): 1853- 1865, 2016.
[11] Robert Dadashi, Léonard Hussenot, Matthieu Geist, and Olivier Pietquin. Primal wasserstein imitation learning. arXiv preprint arXiv:2006.04657 046778 2020.
[12] Alessio Figalli. The optimal partial transport problem. Archive for rational mechanics and analysis, 195())533--560, 2010.
[13] 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.
[14] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 34:20132-201 2021.
[15] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. International conference on machine learning, pages 1587-1596. PMLR, 2018
[16] 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.
[17] Milena Gazdieva, Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev, Extremal domain translation with neural optimal transport. arXiv preprint arXiv:2301.12874 2023.
[18] Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Improved training of Wasserstein GANs. In Advances in Neural Information Processing Systems, pages 5767-5777, 2017.
[19] Tuomas Haarnjja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International conference on machine learning, pages 1352-1361. PMLR, 2017.
10

---

[20] Siddhant Haldar, Vaibhav Mathur, Denis Yarats, and Lerrel Pinto. Watch and match: Su- percharging imitation with regularized optimal transport. arXiv preprint AXiv:2206.15699 2022.
[21] Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. In Proceedings of the Nineteenth International Conference on Machine Learning, pages 267- 274, 2002.
[22] Leonid Kantorovitch. On the translocation of masses. Management Science, 5(1):1-4, 1958.
[23] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[24] Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations, 2021.
[25] Alexander Korotin, Alexander Kolesov, and Evgeny Burnaev. Kantorovich strikes back!
wasserstein gans are not optimal transport? arXiv preprint arXiv:2206.07767, 2022.
[26] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Kernel neural optimal transport. arXiv preprint arXiv:2205. 15269, 2022.
[27] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. arXiv preprint arXiv:2201.12220, 2022.
[28] lya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofrr Nachum. Offline reinforcement learning with fisher divergence critic regularization. In International Conference Machine Learning, pages 5774-5783. PMLR, 2021.
[29] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- policy q-learning via bootstrapping reduction. Advances in Neural Information Processing Systems, 32, 2019.
[30] Aviral Kumar, Joey Hong, Anikait Singh, and Sergey Levine. When should we prefer offline reinforcement learning over behavioral cloning? arXiv preprint arXiv:2204.0566 2022.
[31] Aviral Kumar, Sergey Levine, Yevgen Chebotar, et al. Dasco: Dual-generator adversarial support constrained offline reinforcement learning. In Advances in Neural Information Processing Systems, 2022.
[32] 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.
[33] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint wrXiv:2005.01643 2020.
[34] Zhuo Li, Xu Zhou, Taixin Li, and Yang Liu. An -raimal-t based reinforcement learning approach for computation offloading. In 2021 IEEE Wireless Communications and Networking Conference (WCNC), pages 1-6. IEBE, 2021.
[35] Yicheng Luo, Zhengyao Jiang, Samuel Cohen, Edward Grefenstette, and Mare Peter Deisenroth. Optimal transport for offline imitation learning. arXiv preprint arXiv:2303. /3971, 2023.
[36] Ashok Vardhan Makkuva, Amirhossein Taghvaci, Sewoong Oh, and Jason D Lee. Optimal transport mapping via input convex neural networks. arXiv preprint arXiv: 10962, 2019.
[37] Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540:559-533,, 2015.
[38] Ofir Nachum, Bo Dai, Ilya Kostrikov, Yinlam Chow, Lihong Li, and Dale Schuurmans. Al- gaedice: Policy gradient from arbitrary experience. arXiv preprint 192.:2111..02074, 2019. 2019.

---

[39] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida. Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35:27730-27744, 2022.
[40] Georgios Papagiannis and Yunpeng Li. Imitation learning with sinkhorn distances. arXiv preprint 2001092008.09166, 2020.
[41] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[42] François-Pierre Paty and Marco Cuturi. Regularized optimal transport is ground cost adversarial. In International Conference on Machine Learning, pages 3332-7542. PMLR, 2020.
[43] 'Tyrrlll Rockafellar. Integral functionals, normal integrands and measurable selections. In Nonlinear operators and the calculus of variations, pages 157-207. Springer, 1976.
[44] Filippo Santambrogio. Optimal transport for applied mathematicians, Birkäuser, NY, 55(58- 63):94, 2015.
[45] David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference on machine learning, pages 387-395. PMLR, 2014.
[46] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[47] Denis Tarasov, Vladislav Kurenkov, Alexander Nikulin, and Sergey Kolesnikov. Revisiting the minimalist approach to offline reinforcement learning, arXiv preprint wrXi::2305.09836, 2023.
[48] Denis Tarasov, Alexander Nikulin, Dmitry Akimov, Vladislav Kurenkov, and Sergey Kolesnikov. CORL: Research-oriented deep offline reinforcement learning library. In 3rd Offline RL. Workshop: Offline RL "Launchpad", 2022.
[49] Cédric Villani. Topics in optimal transportation. Number 58. American Mathematical Soc., 2003.
[50] Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
[51] Oriol Vinyals, Igor Babuschkin, Junyoung Chung, Michael Mathieu, Max Jaderberg, Wo- jeiech M Czarnecki, Andrew Dudzik, Aja Huang, Petko Georgiev, Richard Powell, et al. Alphastar: Mastering the real-time strategy game starcraft ii. DeepMind blog, 2, 2019.
[52] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning, arXiv preprint 1rXiv:1911.11361, 2019.
[53] Wenxuan Zhou, Sujay Bajracharya, and David Held. Plas: Latent action space for offline reinforcement learning. In Conference on Robot Learning, pages 1719-1735. PMLR, 2021.
12

---

Appendix
9.1 Proofs
Proposition 9.1 (Policy Improvement with Partial Policy Learning). For any policy n, let's define
its performance as ((r def B 7'r(8, at) over the trajectories obtained by following the policy (sg S. at ((.)), Sc-1 P(. Sc, ot)). Let be the policy of an expert and is the solution to Eq. 10 with the Q cost function. Then it holds that:
( 2 J(3).
Proof. According to (1I, to compare the performance of any two policies и and we can use the Performance difference lemma:
1
((x) - J(B)
s-ud (s) m]]
(13)
By applying [17 Theorem 3]. the solution of 100) as the parameter w oo, converges to the solution of the following problem:
(((-))) Ea-Daarl--((s..
min
(14)
Here support (Supp) is the of points where the probability mass of the distribution lives. In our settings (upp(3((1s)) indicates the best action from 3 which maximizes Q.
Now, show the improvement over the behavior policy, let's consider the difference between the behavior policy 3 and the new policy æ obtained after the update of B3 using our Eq 14 then we have:
J(T) - J(3)
E snide
[Q (s, x) (²)
(15)
1 E
max oCSupp(P((Is))
(s, Ba) 12 (s, a)ll 2
(16)
1-7
Where first we used the definition of æ, and then noted that the maximum over the actions given by the expert is always greater than the average over some actions from the experts policy 33. D
This result is analogous the classic policy iteration improvements proof. The difference is that we takes max over the Supp(3(-Is) rather than max over the all possible actions A.
9.2 Additional Illustrations
For completeness, we provide an additional illustration of the results of the proposed method. Specifically, we show the scores for the one-step settings Table this subsection we show the normalized score curves of the CQL and our method during training Figure 3
Table 4: Averaged normalized scores on MuJoCo tasks. Results of the final 10 evaluations and random seeds.
Dataset
OneStep-RI.
PPL.
HalfCheetah
48.4
51.4+0.2
Medium
Hopper
59.6
80.4+7.4
Walker
81.1
84.320.6
HalfCheetah
38.1
44.8=0.5
Medium- Replay
Hopper
97.5
92.1
Walker
49.5
86.624,8
HalfCheetah
93.4
74.4
Medium-
Expert
Hopper
103.3
110.2+1.9 .9
Walker
113.0
111.2+0.7
Parameters settings As mentioned, we considered different values of w for the ReBRAC-basee method. The parameters can be seen the Table For most tasks, higher values of are preferable, which means that some sort of suboptimal trajectories present in the datasets.
13

---

Antmaze-umaze-divers.
Antmaze-medium-play
Antmare-medium-dliers.
50
40
40
20
10
Step
Step
Stop
200k
400k
600K
800K
1M
200k
400k
600k
800K
1M
200k
400k
600k
BO0K
1M
(a) Umaze-diverts
(b) Medium-play
(c) Medium-diverse
Antmaze-large play
Antmazze-aad entraaze-largeddivree
45
30
25
25
20
15
10
5
Step
Step
200k
400K
600k
800K
200k
400k
600k
800K
(d) Large-play
(e) Large-diverse
Figure Normalized score curves on the Antmaze tasks, IPL IPL CQL algorithm blue, CQL is red
w
Environments
hopper-medium, pen-cloned, relocate-cloned,
relocate-expert, antmaze-large-divers.
halfcheetah-meddimm, walker-medium,
haffcheerah-mediun--epaaa hopper-medium-replly
walker-meddim--rpplyy meikkermediiu--xpper.
umaze,, door-cloned, door-expert, hammer-cloned
hopper-medium, pen-cloned, relocate-cloned,
relocate-expert, antmaze-large-divee
antmaze-large-ple animaze- antmaze-medium medium-play antmaze-mediun-divvrs.. pen-expert, hammer-expert umaze-diverse, pen-human, relocate-huma
Table 5: Optimal w values for different environments
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: For each contribution listed in abstract and introduction we provide the links to the sections about them.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss the limitations of our solver in Appenix 7
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We provide the proofs and the assumptions in Appendix 9.1
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide full list of the experimental details. We use publicly available datasets. The code for the experiments is provided in supplementary material.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: The code for our solver included in supplementary material and will be made public after acceptance of our paper. We use publicly available datasets.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We provide the experimental details in Experiments section
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: We have provided comparisons in tables on the range different environments, which usually sufficient for comparing algorithms in the offline RL domain.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time exccution) needed to reproduce the experiments?
Answer: [Yes]
15

---

Justification:
The computational resources are discussed in results section.
9. Code Of Ethies
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of http://tt/..eappp.ccc////a///aaa
Answer: [Yes]
Justification: The research conforms with NeurIPS Code of Ethics. We discuss the societal impact of our paper in Appendix 7.2
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We discuss the societal impact of our paper in 7.2
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The research does release data models that have high risk for misuse and does not need safeguards.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cited the creators all of the used assets.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The code is included in supplementary material. The license for the code will be provided after the paper acceptance.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not include crowdsourcing experiments or research with human subjects.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not include crowdsourcing experiments or research with human subjects.
16


Reviews:
[{"review_id": "2VIAocKjSg", "reviewer": "Reviewer_qCSm", "summary": "The authors address the problem of offline RL. They rethink offline RL using optimal transport. In offline RL, often the datasets consist of several sub-optimal trajectories that are needed to be stitched together. The authors use partial OT to incorporate stitching and a maxmin formulation of this partial OT. They perform experiments majorly on D4RL.", "questions": "(1) What is the motivation for using OT here?\n\n(2) what is d^\\beta(s). Is it the visitation of the behavioral policy \\beta? Then why do you say that you want to learn a policy that transfers mass from d^\\beta(s) to the corresponding distribution given by the behavioral policy?", "limitations": "Limitations are addressed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "(1) While a lot of methods that use OT, use Wasserstein distance and that requires optimizing a function constraint to be Lipschitz. This is often hard. The authors have used a maxmin formulation which does not need the function to be Lipschitz. \n\n(2) They treat offline RL as an OT problem rather than using OT as a regularizer.", "weaknesses": "(1) There should have been comparisons to W-BRAC.\n\n(2) The results show that PPL^{CQL} produces some marginal improvement over PPL and PPL^{R} with ReBRAC. \n\n(3) From Equation 12, you must not need \\beta. But in experiments you constantly talk about being in conjugation with something or training a \\beta. Maybe this wasn't clear or I misunderstood, why do you need to be in conjugation with CQL or ReBRAC or one-step RL? Why can't you simply train the maxmin objective in equation 12.\n\nThe work is interesting but I would like the authors to clarify why at all there is a need for conjugation? IQL performs better than CQL, so why would I use PPL^{CQL} and not train IQL directly? I request the authors to clarify the contribution of this paper in context to this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful questions. We have managed to address and improve the paper based on them. Below, we address each of your concerns: we include the W-BRAC comparison, provide an explanation of the conjunction, clarify the OT motivation, and provided details on the behavior policy visitation.\n\n**Q1: There should have been comparisons to W-BRAC.**\n\nWe will include W-BRAC's results in Table 2 of our paper. Initially we compared our method with other methods that have already shown significant improvements over W-BRAC, suggesting a clear performance improvement hierarchy.\n\n\n\n**Q2: ... Why can't you simply train the maxmin objective in equation 12.? ... Why at all there is a need for conjugation  with CQL or ReBRAC or one-step RL? IQL performs better than CQL, so why would I use $\\text{PPL}^{CQL}$ and not train IQL directly? I request the authors to clarify the contribution of this paper in context to this.**\n\nTo address the comment regarding the conjunction, we would like to clarify several reasons for that:\n\n- We don't simply train Eq. (12), because the Q-function trained via Eq. (7) and then used in Eq. (12) can suffer from overestimation bias (lines 107-113).\n    \n- Consequently, we tested our method in conjunction with various methods that avoid overestimation of Q-function. $\\beta$ is actually necessary for the methods used. Please note that the contribution of our method is policy extraction, not solving the overestimation bias. These overestimation-avoiding methods are used to obtain different types of cost functions in our method, showing that our method can work efficiently with any of them.\n\n- We do not use IQL as a backbone because this method does not align with the Optimal Transport framework. If we examine the optimization problems (Eq. 1, 4, 12), we can see that the map (policy) outputs are used as inputs for the cost function (Q-function in our case) during optimization. However, the IQL method is a weighted behavior cloning approach, which does not use policy outputs as inputs for the critic function. Instead, only actions from the dataset are weighted by the action-value function.\n \n\nIn summary, we clarify our contribution: We proposed a novel optimal transport-based policy extraction method and provided an analysis of its performance on various RL-based cost functions. We will add this very explicitly to the final version of the paper.\n\n\n**Q3: The results show that $\\text{PPL}^{CQL}$ produces some marginal improvement over PPL and $\\text{PPL}^{R}$ with ReBRAC.**\n\nThe goal of these experiments was to show a side-by-side comparison between the basic method and its improved version via our method. In Table 1, we compare $\\text{PPL}^{CQL}$ to CQL and $\\text{PPL}^{R}$ to ReBRAC, showing that regardless of the backbone used, our method allows improvement.  \n\n\n**Q4: What is the motivation for using OT here?**\n\nThis motivation follows from the need to have a *partial* policy that maps only to the best action distribution when dealing with suboptimal data in offline learning. In OT, *partial* alignment methods have been extensively developed. To integrate OT methods into RL in the simplest way, we considered offline RL as an *max-min* OT problem, meanwhile avoiding the limitations of existing OT in RL methods. \n\n**Q5: What is $d^\\beta(s)$. Is it the visitation of the behavioral policy $\\beta$? Then why do you say that you want to learn a policy that transfers mass from $d^\\beta(s)$ to the corresponding distribution given by the behavioral policy?**\n\nYes, you are right. Indeed, it is the visitations of the behavioral policy. Please note that our method is completely offline. Thus, the probability over the states, $d^\\beta(s)$, is the only distribution over the state space that we have. We have no ability to get our distribution, $d^\\pi(s)$, as we cannot interact with the environment. This is the standard way for offline RL, not something we came up with. For each state visited by the expert, we aim to map it to the most efficient part of the distribution over the actions for this state, provided by $\\beta(\\cdot|s)$.\n\n\n**Concluding remarks:** We truly value your reviews. Please respond to us, to let us know if the clarifications above, suitably address your concerns. If you finds the responses above are sufficient, we kindly ask that you consider raising score.", "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": "5ssIzWSaT6", "reviewer": "Reviewer_iRBG", "summary": "The authors propose a novel perspective for offline reinforcement learning by formulating offline reinforcement learning as a partial optimal transport problem. They view the policy $\\pi$ as a transport map from the state distribution $d^\\beta$ to $\\beta(\\cdot\\mid s)$ and show that the dual form of the partial optimal transport problem can be expressed as a maximin optimization problem. The resulting maximin problem can be easily optimized, unlike other optimal transport problems, due to the absence of 1-Lipschitz constraints. Experimental analyses on various offline RL benchmarks manifest the effectiveness of the proposed algorithm, especially in the antmaze environments.", "questions": "Have you tried using a pre-trained Q-function learned by IQL or any other in-sample value learning algorithms instead of training the Q-function in parallel?", "limitations": "The authors adequately addressed the limitations of their work.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 4, "strengths": "The paper introduces an interesting perspective of formulating offline RL as a partial optimal transport problem. Also, the proposed algorithm outperforms existing baselines on antmaze tasks by a huge margin.", "weaknesses": "1. The authors formulate offline RL as a partial optimal transport problem between $d^\\beta$ and $\\beta(\\cdot\\mid s)$ by regarding the policy $\\pi$ as a transport map. Since the problem depends on the choice of $s$ in $\\beta(\\cdot\\mid s)$, the resulting policy $\\pi$ will also depend on it. This does not make much sense.\n\n2. The explanation of the environment used in the toy example is unclear. Line 193 states that the final state yields a reward of 1, while other states have a zero reward. Then, what would necessitate the agent to seek the shortest path? The reward seems independent of the agent's action.\n\n3. It is difficult to understand how PPL, PPL-CQL, and PPL-R work. A pseudocode for each variant would be helpful.\n\n4. The paper has a huge room for improvement in terms of formatting. Excessive use of underlines, underfull hbox on Line 6 of Algorithm 1, and misaligned $\\pm$s and ragged purple boxes in the tables hurt the paper's readability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper and provide useful feedback. Your questions will help us improve the manuscript. Below are the answers to your questions.  Please let us know if any issues remain!\n\n**Q1: The authors formulate offline RL as a partial optimal transport problem between  and  by regarding the policy  as a transport map. Since the problem depends on the choice of $s$ in $\\beta(\\cdot|s)$, the resulting policy will also depend on it. This does not make much sense.**\n\nPlease note that our method is completely offline, with no access to the online environment. The distribution over the states provided by the expert policy is everything we have in offline learning. All we have to learn from are the states visited by the expert policy $\\beta$. The goal of our method is to extract the best policy using the given distribution of expert states and actions $\\beta(\\cdot|s)$. The dependence on $s$ in $\\beta(\\cdot|s)$ is the standard approach for offline RL, not something we invented (lines 16-17, 107-112).\nPlease let us know if this answer addresses your concern.\n\n\n**Q2: The explanation of the environment used in the toy example is unclear. Line 193 states that the final state yields a reward of 1, while other states have a zero reward. Then, what would necessitate the agent to seek the shortest path? The reward seems independent of the agent's action.?**\n\nOur method works in the most common MDP settings with a discount factor of $\\gamma < 1$ (lines 94-106), which prioritize immediate rewards over distant future rewards. By taking the shortest path to a reward, the agent ensures that it receives a higher discounted reward compared to taking a longer path where the same reward would be less valuable due to discounting. The Q-function trained by Eq. 7 also tends to favor the shortest path. In these experiments, we simply show that BC limits the agent's performance by closely mimicking the provided suboptimal dataset. In contrast, our method allows the extraction of the best actions according to the Q-function, ignoring the suboptimal actions. We will clarify this section in the final version. \n\n**Q3. It is difficult to understand how PPL, PPL-CQL, and PPL-R work. A pseudocode for each variant would be helpful.**\n\nAll of these methods mainly differ in the way they learn the Q-function (cost), not the policy extraction. We agree with the reviewer that the description of the variants is an important component that needs to be discussed. Below is the high-level pseudocode for these methods. We will add the complete pseudocode in the appendix and refer to them in main text. Please note that for ReBrac $-Q^k(s, a) +BC$ cost is using.  \n\n```\nInput: Dataset D(s,a,r,s')  \nInitialize:  Q, π, f, β, α, w, ℛ \n\n----------Update the Q-Function (cost)---------------\nfor k in 1...N do\n\n    (s, a, r, s') ← sample a batch of transitions from D\n\n    if One-Step RL then pre-train Q by:\n        Q^{k+1} ← argmin_{Q} E_{(s, a, s') ~ D} [(r(s, a) + γ E_{a' ~ β(s')} [Q^k(s', a')] - Q^k(s, a))^2]\n\n    else if CQL then\n        Q^{k+1} ← argmin_{Q} E_{s ~ D, a ~ π(s)} [Q^k(s, a)] - E_{s ~ D, a ~ β(s)} [Q^k(s, a)] \n                  + (1/2) E_{(s, a, s') ~ D} [(r(s, a) + γ E_{a' ~ π(s')} [Q^k(s', a')] - Q^k(s, a))^2] + ℛ(π)\n\n    else if ReBrac then \n        Q^{k+1} ← argmin_{Q} E_{(s, a, s') ~ D} [(r(s, a) + γ E_{a' ~ π(s')} [Q^k(s', a') - α(π(s') - a')] - Q^k(s, a))^2]\n    end if\n\n----------Update OT---------------\n\n    f^{k+1} ← argmin_f E_{s ~ D, a ~ π^k(s)} [f^k(s, a)] + w E_{s, a ~ D} [f^k(s, a)]\n    π^{k+1} ← argmin_π E_{s~D, a~π^k(s)} [-Q^k(s, a) - f^k(s, a)]\nend for\n```\n\n**Q4. The paper has a huge room for improvement in terms of formatting. Excessive use of underlines, underfull hbox on Line 6 of Algorithm 1, and misaligned s and ragged purple boxes in the tables hurt the paper's readability.**\n\nWe will incorporate your suggestions to improve the formatting. All misalignments will be corrected and the purple boxes will be replaced with underlining. In the main text, underlining will be minimized. If you have any other suggestions for improving the presentation of the paper, we would appreciate it.\n\n**Q5: Have you tried using a pre-trained Q-function learned by IQL or any other in-sample value learning algorithms instead of training the Q-function in parallel?.**\n\nNo. We do not use IQL because this method does not align with the Optimal Transport framework. If we look at the OT optimization problems (Eq. 1, 4, 12), we can see that the map (policy) outputs are used as inputs for the cost function (Q-function in our case) during optimization. However, the IQL method is a weighted behavior cloning approach and not use learned policy outputs as inputs for the action-value function. \n\n**Concluding remarks:** In conclusion, we truly value your reviews. We hope that the revisions and clarifications will influence and improve your overall opinion. If we've managed to resolve your principal concerns and questions, we'd be thankful for your endorsement through an elevated score of our submission. On the other hand, if there are remaining issues or questions on your mind, we're more than willing to address them.", "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": "FZjVaJyV8W", "reviewer": "Reviewer_7ePM", "summary": "This paper provided a novel view of offline RL, which roughly is maximizing return while keeping close to the data, as OT problem. The key contribution is demonstrating that partial OT can effectively address the challenge of stitching—a fundamental issue in offline RL—both theoretically and empirically.", "questions": "Could you explain the theoretical/intuitive benefit of OT in offline RL compared with other technique to make policy close to the data such as behavior cloning or KL regularization?", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Formulation of offline RL as OT from stationary dist for state to datasets dist by some policy is promising. \n2. The formalisation as OT is made convincing b y presenting the capability of problem of stitching by partial OT in principled way (hyperparameter w controls the degree of the algorithm between BC and naive value maximization)\n3. The main claim—that partial OT effectively handles stitching—is supported by experimental results, particularly in maze environments known to require stitching.", "weaknesses": "In your method, stitching is done by the dynamic programing, application of bellman operator and your contribution regarding stitching is introducing partial OT to allow it work by alleviating the matching restriction. And, in high-level, it is what existing regularization methods like TD3+BC(BC regularization) or ReBRAC(KL regularization) are doing. While this is not a weakness, the paper could be strengthened by providing a clearer explanation of the advantages of OT compared to BC or KL regularization, which is partially addressed in the experiment section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and for a positive assessment of the paper! Your valuable feedback will help us improve the manuscript! Please find below the answers to your questions.\n\n**Q1: In your method, stitching is done by the dynamic programing, application of bellman operator and your contribution regarding stitching is introducing partial OT to allow it work by alleviating the matching restriction. And, in high-level, it is what existing regularization methods like TD3+BC(BC regularization) or ReBRAC(KL regularization) are doing. While this is not a weakness, the paper could be strengthened by providing a clearer explanation of the advantages of OT compared to BC or KL regularization, which is partially addressed in the experiment section. Could you explain the theoretical/intuitive benefit of OT in offline RL compared with other technique to make policy close to the data such as behavior cloning or KL regularization?**\n\nCompared to BC/KL, OT offers more flexibility, while including many building blocks to properly consider the geometry of the problem. For example\n\n\n- BC/KL regularization match the policy to the data distribution at a pointwise level, which can lead to suboptimal policy learning in scenarios where exact matching is not feasible or necessary. In OT, however, different constraints and regularizations can be used. The partial alignment proposed in our paper is a clear example.\n- OT methods allow the choice of arbitrary cost functions that take into account the domain of the problem. For example, we proposed to use Q-function-based cost function for transformation in RL domain. In comparison, BC/KL do not inherently consider the transformation costs between distributions. \n\nFrom a more general perspective, considering RL as OT bridges the gap between these two areas, allowing tools from OT to be applied in RL for better efficiency.  We will explicitly include this discussion in the final version of the paper.\n\n\n**Concluding remarks:** We truly value your reviews. We hope that clarifications are helpful. If there are remaining issues or questions on your mind, we're more than willing to address them.", "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:

Meta-review: The paper presents an approach to offline reinforcement learning based on multiple different (potentially suboptimal) expert data. The problem is then to "stitch" together a policy from the best actions across experts. The authors formalize the problem as partial Optimal Transport. This is a novel, interestingly different role for Optimal Transport than has appeared in previous work on reinforcement learning. The results show promising performance against various baselines. 

The authors are in consensus that this is interesting and sound work. I have read the reviews and the paper and agree. This is novel work that will likely be of broad interest to the community.
Decision: Accept (poster)