Rethinking Optimal Transport in Offline Reinforcement Learning

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 22:27:09

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.

Main paper: 
 introduction: 
 Recently, deep reinforcement learning (RL) has shown remarkable progress in complex tasks such as strategy games (Vinyals et al., 2019) , robotics (Mnih et al., 2015) , and dialogue systems (Ouyang et al., 2022) . Standard RL algorithms interact with the environment and receive rewards that are used to update the policy. However, in some areas, such as healthcare, industrial control, and robotics, implementing RL algorithms that interact with the environment can be costly and risky. To avoid these risks, offline RL algorithms utilize historical data of expert interactions with the environment. In offline settings, only a dataset of experiences is provided, and the objective is to learn an effective policy without any online interactions (Levine et al., 2020) .Standard actor-critic methods (Sutton & Barto, 2018) have demonstrated poor performance in offline settings due to the instability of the critic function when dealing with out-of-distribution actions sampled by the learned actor (Fujimoto et al., 2018) . To avoid this issue, a different behavior cloning objective is used, which keeps the actor close to the actions on which the critic was trained. However, the provided offline data is often sub-optimal, and offline reinforcement learning algorithms are used to find a policy more efficient than that provided in the data (Kumar et al., 2022a) . In offline learning, the goal is to find a trade-off between staying close to the behavior policy for stable out-of-distribution evaluation and improving upon the behavior policy.To enable RL algorithms to be trained on offline datasets, various methods incorporating Optimal Transport (OT) (Villani, 2008) into the training, have been proposed (Wu et al., 2019; Dadashi et al., 2020; Papagiannis & Li, 2020; Li et al., 2021; Haldar et al., 2022; Cohen et al., 2021; Luo et al., 2023) . OT is particularly relevant for addressing distribution shift problems by minimizing deviation from the behavior policy ( 2.4).In our paper, we rethink OT as a framework for offline reinforcement learning that aims at policy improvement, rather than avoiding distribution shift. Unlike previous OT-based approaches in RL, we do not introduce a new OT-based regularization or reward function (Dadashi et al., 2020; Papagiannis & Li, 2020) . Instead, we propose a novel perspective that views the entire offline RL problem as an optimal transport problem between state and action distributions ( 3.1). By utilizing the action-value function as a transport cost function and the policy as an optimal transport map, we formalize offline reinforcement learning as a maximin OT optimization problem.Our formulation naturally allows to apply rich OT theory to the reinforcement learning. As example, we integrate recently proposed Extremal Optimal Transport (Gazdieva et al., 2023) into offline reinforcement learning ( 3.2). In context of the unpaired image translation, extremal optimal transport offers the theoretically best possible translation with respect to a given cost function (Gazdieva et al., 2023) . By employing an RL-based cost function, extremal transport implicitly trains the map (policy) to identify the best possible action for a given state with respect to the cost function, and to improve the policy beyond the behavior policy, see proposition 3.1).Contribution: We propose a novel OT-based formulation for offline reinforcement learning ( 3.1). Building on this formulation, we introduce a new algorithm that employs extremal optimal transport techniques to derive the most effective policy from the dataset, which outperforms the behavior policy ( 3.2). To evaluate our methodology, we conducted various experiments using the D4RL benchmark suite (Fu et al., 2020) . Our results demonstrate that the proposed approach not only outperforms other OT-based offline RL methods but also consistently improves the performance of current state-of-the-art offline RL techniques ( 4). 
reinforcement learning: 
 Reinforcement learning is a well-established framework for decision-making processes in environments modeled as Markov Decision Processes (MDPs). An MDP is defined by the state space S, action space A, and their respective distributions S and A. It also includes the reward function r(s, a), transition distribution between states T (s ′ | s, a), and discount factor γ. The goal in RL is to find a policy π(a|s) that maximizes the expected cumulative reward over time t: E π H t=0 γ t r(s t , a t ) , where H can be infinity. The action-value function Q π is a crucial component in the reinforcement learning framework, as it provides an estimate of the expected cumulative reward that can be obtained by following a given policy π. This can be written mathematically as:EQUATIONwhere s ′ and a ′ are the next state and the next action respectively. This recursive formulation allows the action-value function Q to be estimated by iteratively updating its estimate until convergence is reached. In continuous state spaces, finding the exact optimal policy is intractable due to the curse of dimensionalityy (Sutton & Barto, 2018) . To address this, the popular Actor-Critic algorithm can be used (Sutton & Barto, 2018) . Actor-Critic includes two components: an action-value function Q(s, a) and a policy function π. In practice, Q is updated by minimizing the mean squared Bellman error over the experience replay dataset (s, a, s ′ , r):EQUATION)Accurate estimation of the action-value function allows a near-optimal policy to be recovered. For continuous control tasks, Deterministic Policy Gradient (DPG) (Silver et al., 2014) method can be employed. In this method, the actor is simply updated in the direction of the gradient of the action-value function with respect to the action:EQUATIONIn combination with deep neural networks, the algorithm based on (3) is called Deep Deterministic Policy Gradient (DDPG) (Lillicrap et al., 2015) . Its improved version, called Twin Delayed Deep Deterministic Policy Gradient (TD3) (Fujimoto et al., 2018) , finds application in complex continuous control problems.Offline RL and behavior cloning (BC) approaches provide a fully dataset-driven RL that does not require interactions with the environment. The dataset D frequently contains expert as well as potentially irrelevant behaviors, and offline RL algorithms combine useful behavior segments spread across suboptimal trajectories (Levine et al., 2020) . In contrast to standard RL, offline algorithms may face challenges due to distribution shift during the evaluation step. Specifically, distribution shift refers to the difference between the distribution of state-action pairs given by the expert β and the distribution of actions induced by the policy π (Fujimoto et al., 2018) . In general, there are two classes of algorithms designed to address these issues:Policy penalty: To estimate the maximum Q-value over actions while constraining the learned policy to stay within the behavior data distribution's and avoid distribution shift, many methods have been proposed (Brandfonbrener et al., 2021; Kumar et al., 2019; Zhou et al., 2021; Chen et al., 2022; Kumar et al., 2022b; Fujimoto & Gu, 2021; Kostrikov et al., 2021) . For instance, Fujimoto & Gu (2021) proposed a minimalist approach, TD3+BC, that combines TD3 (Fujimoto et al., 2018) with minimizing the supervised loss between policy and expert actions, and showed strong performance.Another policy penalty method is one-step RL (Brandfonbrener et al., 2021) . This approach involves maximizing the policy log probability for the actions provided in the dataset without off-policy evaluation. The method approximates the behavior Q β -function and extracts the corresponding policy π by weighting actions using the advantage function A β (s, a) (Sutton & Barto, 2018).Implicit Q-Learning (IQL) (Kostrikov et al., 2021) is another method that does not employ off-policy evaluation. IQL approximates the policy improvement step by treating the state value function V β (s) as a random variable and taking a state-conditional upper expectile to estimate the value of the best actions in the state.Critic penalty: An alternative approach to address challenges in offline RL is to penalize the Q function. The conventional Q learning technique for offline RL uses (2), but with data D collected according to an expert policy β. To avoid overestimation bias (Fujimoto et al., 2018) during critic learning in offline settings, Policy Gradient from Arbitrary Experience via DICE (Nachum et al., 2019) proposes to incorporate f -divergence regularization into the critic function.Based on a similar concept, Conservative Q-Learning (CQL) (Kumar et al., 2020) and ATAC Cheng et al. (2022) apply an adversarial training of the citric by extending the loss with additional terms that minimize the Q values on the samples of a learned policy and maximize the values of the dataset actions. Another notable offline reinforcement learning method is Fisher Behavior Regularized Critic (F-BRC) (Kostrikov et al., 2021) . The F-BRC technique augments the standard Bellman critic loss with a regularization based on Fisher divergence.Monge's problem (Villani, 2008) was the first example of the OT problem. Monge's formulation of OT with a cost function c aims at finding a mapping T of the two probability distributions µ and ν, where T ♯µ = ν is the mass-preserving condition and T ♯ is the push-forward operator [ 1] (Villani, 2008) :EQUATIONInformally, we can say that the cost is a measure of how hard it is to move a mass piece between points x ∈ X and y ∈ Y from distributions µ and ν correspondingly. That is, an OT map T shows how to optimally move the mass of µ to ν, i.e., with the minimal effort.A widely recognized alternative formulation for optimal transport was introduced by Kantorovich (Kantorovitch, 1958) . Unlike the Monge's OT problem formulation, this alternative allows for mass splitting. It is important because for certain values of µ and ν, there may not be a map T satisfying T # µ = ν in the Monge's OT (4). The Kantorovich OT problem can be written as:EQUATIONIn this case, the minimum is obtained over the transport plans γ, which refers to the couplings Π with the respective marginals being µ and ν. The optimal γ * belonging to Π(µ, ν) is referred to as the optimal transport plan. The dual form of ( 5) following (Villani, 2003) can be written as:EQUATIONwhere max is taken over all f ∈ L 1 (ν), and , 2008) . For cost functions c(x, y) = ∥x -y∥ 2 , the solution of the optimization problems (4, 5) is called the Wasserstein-1 distance, see (Villani, 2008, 1) or (Santambrogio, 2015, 1, 2) . In recent years, significant progress has been achieved in utilizing neural networks for computing OT maps to address strong (Makkuva et al., 2019; Korotin et al., 2021; Gulrajani et al., 2017) , weak (Korotin et al., 2022c) , and general (Asadulaev et al., 2022) OT problems. Our work is inspired by these developments, as we leverage neural optimal transport methods to improve offline RL.f c (x) = min y c(x, y)-f (y) is c-transform of f (VillaniThe simplest OT-based approach to avoid distribution shift is the behavior regularized actor-critic (BRAC) (Wu et al., 2019) . This method is based on a behavior-regularized actor-critic that includes a different divergence D and design choices for offline RL approaches. In general, the BRAC framework can be written as distance regularized policy optimization:EQUATIONIf we put the Wasserstein-1 as a divergence D, we obtain W-BRAC algorithm. The core idea is to use the OT distance to measure the difference between the policy π and the expert policy β. Given the action-value function Q(s, a), the discriminator f , and the behavior cloning impact coefficient α, we have in ( 7):min π max ∥f ∥ L ≤1 E s∼D,a∼π(s) [-Q(s, a)] Critic function +α E s∼D,a∼β(•|s) [f (a)] -E s∼D,a∼π(•|s) [f (a)] Wasserstein-1 distance (8)The problem is that finding the analytical solution to the dual Wasserstein-1 formulation typically is intractable, and the potential f is required to satisfy the 1-Lipschitz condition. It was shown that this optimal transport method has drawbacks in computing Wasserstein-1 distance in practice (Korotin et al., 2022a) . Also, it is important to note that the coefficient α should be carefully chosen.Another widely explored way to apply optimal transport in offline RL is to use it to construct a pseudo-reward function (Dadashi et al., 2020; Papagiannis & Li, 2020; Li et al., 2021; Haldar et al., 2022; Cohen et al., 2021) . Primal Wasserstein Imitation Learning (Dadashi et al., 2020) and Sinkhorn Imitation Learning (Papagiannis & Li, 2020) utilizing OT distance between the expert and imitator to create a reward function. Recently, the Optimal Transport Reward Labeling (OTR) method (Luo et al., 2023) was proposed to generate a reward-annotated dataset, which can then be employed by various offline RL algorithms. However, as the W-BRAC these methods use OT primarily as an additional loss function. In summary, we can say that the main aspects of the existing OT in RL methods are:• 1) Wasserstein regularized actor-critic requires 1-Lipschitz function constraints on f , which complicates the optimization. • 2) Optimal transport incorporated in RL with a coefficient α, and it was shown that the choice of α value is crucial (Brandfonbrener et al., 2021 ). • 3) Previously proposed OT approaches focuses on behavior cloning and does not address policy improvement beyond the behavior policy.In our paper, we tackle the challenge of developing an OT-based approach that avoids these limitations.In the following section, we present our maximin OT as RL formulation and formally compare it to the W-BRAC framework. 
monge reinforcement learning: 
 To rethink application of OT in RL, lets consider the (4). If we replace cost c(x, y) with the optimal Q(s, a) function for the given data, distribution µ with the state distribution given by the expert S, and distribution ν with the expert actions distribution A, we obtain a following problem:EQUATIONIf policy π is deterministic, we can view this formulation as the DPG algorithm (3) with distributional constraints on the policy π. To convert the problem into a one that can be applied to deep RL and solved via neural network approximations (Fan et al., 2021) , we can consider a dual form of (5).To do this, we expand the c-transform of the dual form in (6) using the function Q as the cost c, which can be represented asE s∼S [f Q (s)] = E s∼S [min a {-Q(s, a) -f (a)}].Subsequently, by applying the Rockafellar interchange theorem (Rockafellar, 1976, Theorem 3A), we replace the optimization over points a ∈ A with an equivalent optimization over functions π : S → A:EQUATIONAfter substituting ( 10) with ( 6), we get:max f min π E s∼S,a∼π(s) [-Q(s, a)] Cost function -E a∼π(s) [f (a)] + E a∼A [f (a)]Constraints preserving (11)Let's denote the expression in the large square brackets by L(f, π). Then, following (Korotin et al., 2022c) for every optimal potential f * ∈ arg max f min π:S→A L(f, π) , we can solve a saddle point problem in (11) and extract an OT map π * between S, A from the optimal pair (f * , π * ).We call this method Monge Reinforcement Learning (MRL). This formulation provides a maximin optimal transport problem, where the optimal policy is defined as an optimal mapping from states to actions with an action-value cost function. This optimization problem instantly bypass the 1-Lipschitz function constraints on f . The specific role of the potential f in the objective of ( 11) is different from its function in (8). Here, the f is used to verify that the constraints on the optimization have been satisfied, rather than to compute the Wasserstein-1 distance itself.More formally, the potential f can be seen as a Lagrangian multiplier (Peyre, 2021) that ensures that the generated distribution π♯S matches the distribution of actions given by the dataset D, and penalizes the former if it does not. This formulation provides a foundation for OT in RL and compared to the previously proposed W-BRAC method, our method has the following advantages• 1) It does not require 1-Lipschitz function constraints on f . Instead, the function f plays the role of a Lagrangian multiplier. • 2) There is no dependence on the hyperparameter α, which controls the impact of optimal transport loss during policy optimization. The whole OT problem is aimed at finding an optimal policy.However, in (11) the policy maps the distribution of states to the complete distribution of actions. In offline reinfrocement learning, the provided dataset often contains suboptimal paths, and the policy should apply outlier detection in the action space to avoid inefficient actions for each state in the dataset. In the following subsection, we introduce an algorithm that avoids inefficient actions using optimal transport.An efficient offline RL algorithm should be able to stitch between trajectories for each state provided by the dataset (Kumar et al., 2022a) . To do this, the policy should implement outlier detection in the action space and select only the effective actions from the mixture of expert demonstrations. More formally, we can say that the goal is to identify a policy that maps into the support of the action space, π(s) ∈ supp(A), that maximizes Q(s, a) for each state. Here support is the set of points where the probability mass of the distribution lives.Recently, the OT formulation with similar constraints on the transport map was considered in the Extremal Optimal Transport (ET) method (Gazdieva et al., 2023) . The ET formulation represents the theoretically best possible unpaired domain translation based on a given cost function. Essentially, for the Euclidean cost functions as ℓ 2 , the extremal transport maps can be seen as the tool for finding nearest neighbors (maximally close) to the input samples from the target, according to the cost function, see Figure 1 in (Gazdieva et al., 2023) .In the context of RL, we can say that the extremal map will align the state space with the actions that maximize the Q cost function. Formally, for RL, we can express the optimization problem as: In this case, the min is taken over a function π : S → A that maps from S to the support of A, denoted as Supp(A). The goal is to find such a support, Supp(A), that minimizes the expectation E s∼S [-Q s, π(s) ]. The main difference between this and the optimal transport problem discussed earlier in (9), lies in the constraints used. The previous method used equality constraints, while this technique imposes inclusion constraints, allowing only partial alignment between state and action spaces. See, visual representation of the Extremal Monge RL in Figure 1 .EQUATIONTo show that (12) allows to improve the policy π beyond the behavior policy β, let's consider a onestep RL algorithm. Starting from the expert policy β, we solve the extremal transport problem ( 12) with respect to the given cost function Q β . Proposition 3.1 (Policy Improvement with Extremal Transport). For any policy π, let's define its performance as J(π)def = E τ ∼π [ ∞ t=0 γ t r(s t , a t , s t+1 )]where τ ∼ π indicates the random trajectories obtained by following the policy π, (s 0 ∼ S, a t ∼ π(s t ), s t+1 ∼ P (• | s t , a t )). Let β be the policy of an expert and π is the solution to (12) with the Q β cost function. Then it holds that: J(π) ≥ J(β).The proof is given in Appendix ( A). With this proposition, we show that we avoid a problem 3 mentioned in ( 2.4) by proposing an OT for policy improvement beyond behavioral policy.To obtain a saddle-point reformulation that can be used to compute the solution to (12) in practice, the steps analogous to those presented in ( 3.1), can be considered (Gazdieva et al., 2023) . This reformulation replaces set inclusion constraints with inequality constraints, and if the parameter w → +∞, it provides an extremal transport. This leads to the following optimization problem:EQUATIONThe changes compared to (11) are emphasized by the green color. In the (13), the w coefficient is used to control the extreme of the chosen action, and the potential f is always greater or equal to 0. Here f is checking that the constraints in the (12) are met. According to the (Gazdieva et al., 2023, Proposition 4) the optimal potential f * vanishes on the outliers, i.e. f * (a) = 0 holds for all a ∈ Supp(A) \ Supp(π * ♯S). It is important that the algorithm defines outliers by the critic function Q β (s, a), because the optimization goal is to find a Supp(A) for which the value -Q is minimum.The policy update in (13), in short can be written asmin π E s∼D,a∼π(s) [-Q(s, a) -f (a)].Here we can see that if f (a) = 0 only on the out-of-support actions, then it adds extra weight to the correct actions, indicating policy to choose actions in A that minimize the -Q function. As a result, a map (policy) learned with extremal transport will ignore inefficient actions provided in dataset.We can interpret that our method lies between behavior cloning and direct maximization of the Q function. Recent studies have shown that direct maximization can lead to suboptimal results due to overestimation bias (Fujimoto et al., 2019) . Conversely, being too close to the expert's policy prevents improvement (Kumar et al., 2022a) . Intuitively, our extremal formulation allows us to strike a balance that maximizes Q by taking actions from the expert action distribution that are not radically different from those on which the Q function was trained. In the following section, we propose the optimization procedure for the Extremal Monge-RL (XMRL) algorithm .In this section, we provide a practical solution for learning policies from offline data. While the optimal Q(s, a) for the given data is unknown, we can approximate it by a neural network Q ϕ : R S × R A → R, trained on the given expert data D via (2). If the Q function is trained only on actions sampled from the behavior policy β, we get Q β and Q π when the next actions a ′ are sampled from the learned policy π.Unfortunately, in practice, the Q function learned by (2) suffers from overestimation bias (Fujimoto et al., 2018; 2019; Fujimoto & Gu, 2021; Kumar et al., 2020; 2019; 2022a; Levine et al., 2020) . In OT, Algorithm 1 Extremal Monge Reinforcement Learning (XMLR) Input: Dataset D(s, a, r, s ′ ) Initialize Q ϕ , π θ , f ω , β for k in 1...N do (s, a, r, s ′ ) ← D: sample a batch of transitions from the dataset. Q k+1 ← Update cost function Q ϕ using the Bellman update in (2) or ( 14) π k+1 ← Update policy π θ as a transport map using the objective in (15) f k+1 ← Update f ω using outputs of π θ and samples from dataset, using ( 16) end for the cost function that does not match the given problem can lead to suboptimal results (Liu et al., 2020; Asadulaev et al., 2022) . To avoid this issue, we tested our method in pair with the two well-known approach to obtain not overestimated critic: CQL (Kumar et al., 2020) and TD3+BC (Fujimoto & Gu, 2021) . For the k update steps, the conservative Q function can be obtained by:EQUATION)Where R(π) is a regularization, usually entropy based. For the TD3+BC based cost, we used the Q function learned by ( 2), but coupled with the BC loss, this cost can be written as c(s, π(s)) = -λQ(s, π(s)) + (π(s) -a) 2 . The algorithm with the CQL-based cost (critic) we call XMRL CQL and the algorithm with a BC-based cost we call XMRL BC . In general, any RL-based cost function can be used. The core of our method is the policy penalty, which can help to avoid suboptimal actions provided by the expert and not the critic learning.To obtain a mapping π from the distribution of states S to the distribution of actions A, we use neural networks π θ : R S → R A and f ω : R D → R to parameterize π and f respectively. These neural networks serve as function approximators that can capture complex mappings of states. For the k update step, given the last policy π k , the value function Q k , and f k , we have:EQUATION)EQUATIONBy assigning a weight w, we can impose a more significant increase in the values of the highly rewarded actions. 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 a batch of transitions from the offline dataset D, and then adjust the value function Q, the policy π, and the potential f , see Algorithm 1. Our algorithm provides a scalable approach to learning policies from offline data, see the next section for practical evaluation. 
benchmarks and baselines: 
 We evaluate our proposed method using the Datasets for Deep Data-Driven Reinforcement Learning (D4RL) (Fu et al., 2020) benchmark suite, which is a collection of diverse datasets designed for training and evaluating deep RL 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 MuJoCo environments, such as Walker, Hopper, and HalfCheetah (HalfC) . Importantly, we also tested our method on the complex Antmaze 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) (Kumar et al., 2020) , Twin Delayed Deep Deterministic Policy Gradient with behavior Cloning (TD3+BC) (Fujimoto & Gu, 2021) , Implicit Q-Learning IQL (Kostrikov et al., 2021) , and IQL with Optimal Transport Reward Labeling (IQL+OTR) (Luo et al., 2023) . The results of the supervised behavior cloning are also included.MuJoCo. To be consistent with the theory given in ( 3.2) For this environment, we tested how our proposed algorithm improves beyond the expert policy β. For this, following the setting provided by One-Step RL (Brandfonbrener et al., 2021) , we pretrained β for 500k steps using the provided data. Then we pre-trained Q β for 2 million steps. A two-layer feed-forward network with a hidden layer size of 1024 and a learning rate of 1e-4 was used with the Adam optimizer (Kingma & Ba, 2014) . Then, we applied a simplified conservative update to Q β (see (Kumar et al., 2020)[Eq.1] ).Using the next actions by β, this update provides a lower bound on the real Q β (Kumar et al., 2020)[Theorem.1] . The simplified version does not include the penalty term E s∼D,a∼β(s) Q(s, a) and the regularization R(π). Then to improve beyond behavior policy β we trained policy for the 100k steps using Algorithm 1 with a learning rate of 1e-5. The results are shown as XMRL CQL in Table 1 .For XMRL BC , no pre-trained models were used, we trained TD3+BC from scratch for 1m steps in conjunction with our method. All hyperparameters were the same as the original (Fujimoto & Gu, 2021) . See results in Table 1 . For all experiments we set w equal to 8. The ablation study on the parameter w is given in the Appendix ( A.2).Antmaze. While this environments are more complicated than MuJoCo, for the XMRL CQL we used full CQL update ( 14). No pre-trained models were used; we used the code provided by the CORL library (Tarasov et al., 2022) , with all hyperparameters and architectures set identical to those originally proposed by CORL for the CQL method. We trained algorithm for 1m steps, with set w equal to 8 for all experiments. See the results in Table 2 .For XMRL BC , as a backbone we used the improved version of TD3+BC called ReBRAC (Tarasov et al., 2023) , keeping the same hyperparameters proposed by the authors. The ReBRAC framework recommends task-specific hyperparameters, to achieve high results, thus we tested different values of w for this algorithm. For the large-play, medium-diverse and umaze, it was equal to 5. For the medium-play, large-diverse, it was equal to 1, and for umaze-diverse, w was set equal to 8. The results are shown in Table 2 .We present the qualitative results of our XMRL algorithms in Tables 1 2. For MuJoCo, the blue color in the tables highlights the top 3 scores, and bold indicates the best. We compared the results with those reported in the OTR+IQL paper (Luo et al., 2023) . For the Antmaze, the main comparison is with the oracle to which our method is coupled, bold and blue indicates if the our algorithm is better than the oracle. In these experiments we reproduced CQL and ReBRAC results the code provided in CORL. The CQL score curves are provided in Appendix 4. For completeness, IQL and OTR+IQL results are also included in Table 2 .Our method shows promise for improving the field of offline RL using optimal transport. Importantly we show improvements over the previously best OT-based offline RL algorithm OTR+IQL. By introducing a novel mathematical task formulation, we provide effective solutions to complex RL problems. In offline RL, the problems are noisy, complex, and diverse, and task-specific hyperparameter search is required. But, we used the same set of parameters for each MuJoCo task, and our algorithm XMRL CQL is consistently in the top performing algorithms.For the Antmaze, our method consistently outperforms the CQL on which it is based. For the XMRL BC , our method provides state-of-the-art results for 4/6 environments.Our formulation aligns the state space with only a portion of the action space. This property is particularly relevant for offline datasets that contain a mixture of expert demonstrations. Our method can help to select the best possible action from a set of expert trajectories (Levine et al., 2020) . In other words, optimal f is to encourage the discovery of new extreme actions that are atypical for the behavior policy at each state, but can yield high rewards. However, our algorithm has a limitation in the setting of the hyperparameters w, as it was said, the value w controls the size of the action space in which we are mapping, the higher value w the more extreme actions will be chosen. However, in some environments, if the actions are provided only by a single expert, reducing this space can be detrimental to performance. The task specific value of w is required. For ablation study over w please see Appendix (A.2). Nevertheless, our approach represents a significant step toward improving the performance and stability of offline RL algorithms, and we believe it can serve as a foundation for future research in this area.The code is implemented in the PyTorch (Paszke et al., 2019) and JAX frameworks and will be publicly available along with the trained networks. Our method converges within 2-3 hours on an Nvidia 1080 (12 GB) GPU. We used WanDB (Biewald, 2020) for babysitting training process.In comparison to many offline RL methods that use ensemble-based training methods, our approach provides a simpler framework with a few modifications built upon the oracle methods . 
conclusion and discussion: 
 In this paper, we have established a formal connection between neural optimal transport and offline reinforcement learning. Our work introduces the concept of optimal transport to deep reinforcement learning by using the action-value function and the policy as components of the OT problem between states and actions. To demonstrate the potential of our formulation, we propose an XMLR algorithm that uses the OT potential f to penalize the policy and avoid inefficient actions provided in the dataset.According to the theory, our potential f (a) considers the actions from the entire expert distribution of actions. It might be reasonable to choose actions based on the state-conditional action distributions. With respect to (13), this implies that we need to modify our potential f (a) by adding another input parameter to make it state-conditional: f (s, a). However, our experiments show that even without state conditioning and without selecting actions from the expert's full set of actions, our method still achieves robust results. For more details please see Appendix (A.3).Using our formulation other OT methods also can be integrated into RL. For example, various regularizations (Korotin et al., 2022b; Courty et al., 2016) and general costs (Paty & Cuturi, 2020; Asadulaev et al., 2022) can be used to incorporate task-specific information into the learned map, which can be particularly relevant in hierarchical RL problems (Bacon et al., 2017) . Also, the Weak Neural OT (Korotin et al., 2022c ) can be relevant in RL where stochastic behavior is preferred for exploration in the presence of multimodal goals (Haarnoja et al., 2017) .To reproduce our experiment we provide source code in supplementary materials. Details on used hyperparameters are presented in settings ( 4.2). 
 
 Appendix: A.1 PROOFS Proposition A.1 (Policy Improvement with Extremal Transport). For any policy π, let's define its performance as J(π)def = E τ ∼π [ ∞ t=0 γ t r(s t , a t , s t+1 )]where τ ∼ π indicates the random trajectories obtained by following the policy π, (s 0 ∼ S, a t ∼ π(s t ), s t+1 ∼ P (• | s t , a t )). Let β be the policy of an expert and π is the solution to Eq. 12 with the Q β cost function. Then it holds that:J(π) ≥ J(β).Proof. According to Kakade & Langford (2002) , to compare the performance of any two policies π and β we can use the Performance difference lemma:EQUATIONTo show the improvement over the behavior policy, let's consider the difference between the behavior policy β and the new policy π obtained after the update of β using our Eq.12, then we have:EQUATIONWhere in first we used the definition of π, and then noted that the maximum over the all actions given by the expert is always greater than the average over some actions from the experts policy β. □This result is analogous to the classic policy iteration improvements proof. The difference is that we takes max over the supp(A) rather than max over the all possible actions A.A.2 ABLATION ON THE PARAMETER w The MRL formulation is hyperparameter free, applying extremal transport requires a hyperparameter w. But in XMRL, w is a parameter that is used. This parameter influences the action selection process by determining the support range over which the policy operates. With the parameter w equal to 1, our method considers all actions provided in the dataset; by increasing this value, we consider the smaller subspace of possible actions, which can be useful when the dataset is provided by a number of experts and it is necessary to select the best one from the data.Under review as a conference paper at ICLR 2024 More formally, according to the (Gazdieva et al., 2023, Proposition 4) , the optimal potential f * vanishes on the outliers, i.e. f * (a) = 0 holds for all a ∈ Supp(A) \ Supp(π * ♯S) and the coefficient w is used to control the size of this support in practice. In a perfect scenario, we want to find as few actions as possible that maximize the score, so we usually favor the higher values of w.But in practice, results will vary for different datasets; for some datasets, reducing the action subspace is unnecessary. Thus, we can say that our parameter allows for nuanced control depending on the task. See the results on the MuJoCo in Figure 2 and Table 3 . For the Antmaze, Pen and Door environments, see Figure 3 . With respect to (13), this implies that we need to modify our potential f (a) by adding another input parameter to make it state-conditional: f (s, a). However, our experiments show that even without state conditioning and without selecting actions from the expert's full set of actions, our method still achieves robust results. Below, we provide a side-by-side comparison of the proposed XMRL and its conditioned version. As can be seen in Table ( 4), the state-conditioned version XMRL CQL s does not provide an improvement in performance. We assume that this is because the state-conditioned distributions provided by the expert are not rich enough, and expanding the action space allows to act more efficiently. We can say that our marginal-based method lies between behavioral cloning and direct maximization of the function. Recent studies have shown that direct maximization can lead to suboptimal results due to overestimation bias. Conversely, adhering too closely to the expert's policy prevents improvement. Intuitively, our extremal formulation allows us to strike a balance, maximizing by taking actions that are not radically different from those on which the function was trained.A.4 ADDITIONAL ILLUSTRATIONS For completeness, we provide an additional illustration of the results of the proposed method. Specifically, in this subsection we show the normalized score curves of the CQL and our method during training.  


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: (1) There should have been comparisons to W-BRAC.

(2) The results show that PPL^{CQL} produces some marginal improvement over PPL and PPL^{R} with ReBRAC. 

(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.

The 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.
questions: questions: (1) What is the motivation for using OT here?

(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: Limitations are addressed.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

2. 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.

3. It is difficult to understand how PPL, PPL-CQL, and PPL-R work. A pseudocode for each variant would be helpful.

4. 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.
questions: 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: limitations: The authors adequately addressed the limitations of their work.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
questions: 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: limitations: Yes

Comments: title: Thank you for your detailed response
Please respond to the authors
Please respond to authors
comment: Now I get it. So by the *state distribution* $d(s)$ in line 129, you meant the Dirac mass $\delta(s)$. Please explicitly mention it in the paper. *State distribution*, especially when written as $d^\beta$, often means the state occupancy measure under policy $\beta$, so I highly recommend you use another term to prevent misunderstandings like mine. As all of my concerns were resolved, I have updated my score accordingly.
**Q1:** Dear reviewer. As we already noted, we consider the conditional optimal transport setup. This means that we want to simultaneously learn a family of transport plans (indexed by some condition $c$; each plan is between some distributions $p_c, q_c$). Following the standard practices of neural OT, this would means that we have to learn a map $T(c, x, z)$, where $c$ is a condition, $x$ is an input point (from $p_c$) and $z$ is a random noise to be able to learn stochastic plans.

In our case, we need to learn a set of OT plans, each plan is conditioned on the given state $c=s$. Following our problem, each such plan should be a plan between distribution $p_c=\delta_s$ (this is our choice) and $q_c=\beta(\cdot |s)$. Hence, we would have to learn a function $T(c,x,z)=T(s,s',z)$ with 3 arguments. However, since the input $s'$ comes from $\delta_s$, it always coincides with $s$ with probability 1. Hence, one may merge the arguments together to simply get a function of the form $T(s,z)$. In turn, we also remove the random noise component $z$, which is standard practice in neural OT methods, unless additional regularization (variance, entropy, etc.) is used[1]. Hence, our final function is $T(s)$. We will add this details into the final version. 

**Q2:** The transition function is $P(s'|s,a)=1$. As we said in the previous answer, we have redefined the environment to avoid any confusion. Now the reward is given for each state-action pair $r(s_t,a_t) = -\ell^2((s_t,a_t), (S_T, A_T))$, where $S_T=0$ and $A_T=0.2$ is the target state-action pair. In this scenario, trajectories deviating from the straight line (actions define y-axis coordinates) from $S_0$ to $S_T$ will have the lower cumulative reward $\sum^{T}_0\gamma^t r(s_t,a_t)$.

**References**
[1] Neural Optimal Transport: https://openreview.net/pdf?id=d8CBRlWNkqH
**Q1.** I'm confused now. Let's consider the simplest case where $\mathcal{D}$ has two elements $(s_1, a_1)$ and $(s_2, a_2)$. Then the state distribution $d(s)$ is $\frac{1}{2}\delta(s=s_1)+\frac{1}{2}\delta(s=s_2)$ and the behavior policies are $\beta(a\mid s_1)=\delta(a=a_1)$ and $\beta(a\mid s_2)=\delta(a=a_2)$. The optimal transport map $T_{s_1}$ from $d$ to $\beta(\cdot\mid s_1)$ should be $T_{s_1}\equiv a_1$ and the map $T_{s_2}$ from $d$ to $\beta(\cdot\mid s_2)$ should be $T_{s_2}\equiv a_2$. If you try to use the same $T$ in both cases, conflict would occur since $a_1\neq a_2$, wouldn't it? Am I missing something?

**Q2.** What is the transition function? From the description in the paper, it looks like there are states $s_1, s_2, \cdot, s_T$ uniformly spaced across the x-axis and the agent will transition from $s_i$ to $s_{i+1}$ regardless of the action it takes. After 50 steps, the agent will arrive at state $s_T$ no matter what and get a reward of 1. What would encourage the agents to follow the shortest path?
**Q1.** Thank you for the clarification. Optimal transport does not cause conflicts between different $s$. Our final objective (Eq.12) for policy is similar to the BC example that your gave: $\min_\pi\mathbb{E}_{s \sim \mathcal{D}} [-Q(s, \pi(s))-f(s,\pi(s))]$. For each $s$ the potential $f$ is affected by the value of $\pi(s)$ and nothing else. 
In the optimal transport literature, several methods with a similar objective have been considered to map into the conditional distributions [1, Eq. 8c][2, Sec. 6.2]. Please note that [1] considered the case when even $T$ is conditional, and then, use the single $T$ for all distributions, see Sec. 2.4. 

**Q2.** Yes, intermediate steps are not rewarded, but the $Q$-function is particularly given to estimate these intermediate steps. We added different lengths to the dataset, 20, 30, 50. Consequently, trajectories that lead to the reward faster have a higher $Q$- value. 

For a simpler illustration, we also made the reward for each step equal to the Euclidean distance between the current state-action pair and the final rewarded state-action pair. This simplifies the learning problem, but gives a clear intuition why the shortest path is optimal. If you find this more relevant, we will include such an experiment in the final version. 


**Reference:**

[1] Nonlinear Filtering with Brenier Optimal Transport Maps: https://openreview.net/pdf/70633c38d3ce64c9b3b29dd7abd18c2f6b6e1dc6.pdf

[2] Neural Monge Map estimation and its applications: https://openreview.net/pdf?id=2mZSlQscj3
**Q1.** Say we are going to use the log-likelihood as our regularizer, which means the objective function would be something like $\mathbb{E}_{s\sim \mathcal{D}}\left[Q(s, \pi(s))-\alpha\log\beta(\pi(s)\mid s)\right]$. For each $s\in S$, the regularizer $\log\beta(\pi(s)\mid s)$ is only affected by the value of $\pi(s)$ and nothing else. If we ignore the expressivity of a neural network, the resulting $\pi$ is straightforward to interpret: $\pi(s)=\text{arg}\max_a Q(s, a)-\alpha\log\beta(a\mid s)$. However, in the case of OT, the value of $T_s(s')$ for all $s'\in S$ matters, where $T_s$ is the transport map between the state distribution $d(\cdot)$ and $\beta(\cdot\mid s)$. I'm concerned that trying to force the same $T$ for all $\beta(\cdot\mid s)$ would cause conflicts between different $s$.

**Q2.** Does it mean that an agent's actions in the intermediate 49 steps wouldn't affect the return?
**Q1.** We learn a single neural network $\pi_\theta$ that approximates all policies $\pi_\theta(s) \leq w(\beta(\cdot|s))$ for each state $s$. But the same is true for any BC method in offline RL when we learn $\pi_\theta(s)=\beta(\cdot|s)$ using $\ell^2$, KL, or any other discrepancy. The neural approximator for the policy $\pi$ is able to generalize and provide a solution for the measure-zero set of the $d^\beta(s)$. Did we understand your question correctly?

**Q2.**  Yes, you are right, this was an oversight on our part. Only this toy experiment was affected. We thank the reviewer for pointing this out. We fixed this by providing the expert trajectories with a different length in the dataset. This gave us the same visual results. 

P.s. Although there were no explicit signals with fixed episode length, the model still learned the near shortest path. We think this is because the $Q$ function approximation implicitly learned to value actions similar to the final rewarded action more highly.

**Q5.** Yes, it was really interesting results. We will make these experiments public along with the rest of the code.
Thank you for your response.

**Q1.** I think I've failed to deliver my point. A transport map between the state distribution $d$ and the behaviour policy $\beta(\cdot\mid s)$ for some $s$ depends on the choice of $s$. This means the optimal transport framework will produce multiple policies $\pi_s$, one for each $s\in S$. How are you going to combine them into one policy $\pi$? It seems like you are defining $\pi(s)=\pi_s(s)$ for each $s$ (am I correct?), but this does not make much sense because any change on a measure-zero set {$s$} will not affect the result. Therefore, any arbitrary function $\pi$ would be a solution, which is definitely not what we want.

**Q2.** If the episode length is fixed to 50, the discounted return will always be $\gamma^{50}$ regardless of the agent's policy.

**Q5.** Thank you for running the experiments. It is interesting that using IQL Q-functions as the cost function fails so miserably.
**Q1. If I understood the paper correctly, you are viewing  as a transport map from the state distribution d(s) to the behavior policy $\beta(\cdot|s')$ for some s'. The transport map will depend on the choice of s'. How do you choose this s'?** 

Sorry for the misunderstanding. We are randomly sampling $s$ from the given dataset $\mathcal{D}$, please see derivations at lines 131 and Eq.9. For the sampled $s$ we choose the conditional distribution of the action space provided by the dataset (expert) $\beta(\cdot|s)$ as the target. We will make this explicit in the final version. 

**Q2. Lines 190 and 191 state that the agent has to "navigate through 50 intermediate steps." Doesn't this mean that the episode length is fixed to 51 regardless of the agent's path?**

Yes, the length of the episode is fixed. This length of 50, corresponds to the number of steps provided by each expert policy in the dataset.

**Q5. The value learning part of IQL can be completely separated from the policy learning part, which uses AWR, as you mentioned. IQL Q function can be trained without weighted behaviour cloning and thus can be used together with the proposed algorithm.**

Thanks for the clarification, this experiment is really interesting. Following your question, we conducted a series of experiments on MuJoCo using the CORL[1] IQL implementation. In these tests, we found that the value function trained using expectile regression from the IQL method is inappropriate for off-policy evaluation.

- First, we conducted an experiment *without* using our method. We trained a $Q$-function using  expectile regression loss and attempted to extract a policy through direct optimization: $\min_{\pi}\mathbb{E}_{s \sim \mathcal{D}, a\sim\pi(s)} \big[-Q^{\text{IQL}}(s, a)\big]$. This approach resulted in zero rewards.
- Second, we added a BC objective to avoid distribution shift: $\min_{\pi}\mathbb{E}_{s \sim \mathcal{D}, a\sim\pi(s)} \big[-Q^{\text{IQL}}(s, a)\big]+(a-\beta(s))^2$. The results were the same — $Q$-functions trained via expectile regression dramatically overestimate actions sampled by the learned policy $\pi$. 
- Finally, we tested our method: $\min_{\pi}\mathbb{E}_{s \sim \mathcal{D}, a\sim\pi(s)} \big[-Q^{\text{IQL}}(s, a)-f(s,a)\big]$ and observed improvements in the scores! Even with the ill-suited  cost function, policy optimization with respect to the potential $-f(s,a)$ yielded the highest scores.

We also tested the advantage $A(s,a) $and exponential advantage functions exp$A(s,a)$ from IQL, but did not observe any improvements. The table below summarizes our analyses.

---

**Table: Averaged normalized scores on MuJoCo tasks. Reported scores are the results of the final 10 evaluations and 3 random seeds.**

| Dataset                 | $-Q^{IQL}(s, a)$ | $-Q^{IQL}(s, a)$ + BC | $-Q^{IQL}(s, a) - f(s,a)$ (Ours) |
|-------------------------|------------------|-----------------------|---------------------------------|
| HalfCheetah-medium             | -2.53 ± 0.1      | -2.54 ± 0.1           | 48.7 ± 0.3                     |
| HalfCheetah-medium-expert      | -2.53 ± 0.2      | -2.54 ± 0.2           | 39.9 ± 1.4                     |
| Hopper-medium           | 0.6 ± 0.1        | 0.6 ± 0.1             | 27.9 ± 15.4                    |
| Hopper-medium-expert    | 0.7 ± 0.1        | 0.6 ± 0.1             | 8.7 ± 4.6                      |
| Walker-medium           | -0.16 ± 0.1      | -0.16 ± 0.0           | 37.6 ± 3.6                     |
| Walker-medium-expert    | -0.23 ± 0.2      | -0.21 ± 0.1           | 17.5 ± 14.3                    |
---

While the IQL method can be formally decoupled, we see that both components, expectile-based in-sample value learning and weighted behaviour cloning, are important parts of each other to achieve strong results. Adapting IQL value learning for better off-policy evaluation is a promising direction, but beyond the scope of our contribution. We believe that our strong performance on most tasks, combined with different types of a cost function, justifies our formulation and makes it valuable to the community. We will include these analyses in the final version of the paper. 

**Reference:**

[1] JAX-CORL: https://github.com/nissymori/JAX-CORL
Thank you for your replay, the claim that by formulating as OT, we can leverage well-developed tool in OT in RL sounds nice. Looking forward to your future work showing their effective usage in RL.
Thank you for clarifying my doubts. I have increased my score.
Hello reviewer 7ePM: The authors have responded to your comments. I would expect you to respond in kind.
Hello reviewer qCSm: The authors have responded to your comments. I would expect you to respond in kind.
Thank you for your response. I noticed some misunderstandings and wanted to clarify a few.

**Q1.**  If I understood the paper correctly, you are viewing $\pi\colon S\to A$ as a transport map from the state distribution $d(\cdot)$ to the behavior policy $\beta(\cdot\mid s')$ for some $s'\in S$. The transport map will depend on the choice of $s'$. How do you choose this $s'$?

**Q2.** Lines 190 and 191 state that the agent has to "navigate through 50 intermediate steps." Doesn't this mean that the episode length is fixed to 51 regardless of the agent's path?

**Q5.** The value learning part of IQL can be completely separated from the policy learning part, which uses AWR, as you mentioned. IQL Q function can be trained without weighted behaviour cloning and thus can be used together with the proposed algorithm.