Rethinking Optimal Transport in Offline Reinforcement Learning

1 Introduction

• Offline RL faces distribution shift between learned and behavior policies, which can destabilize critic evaluation and motivate behavior-cloning objectives [33] [15] [52].Source: §1 · p. 1
• Mixed expert demonstrations may be incorrect or effective only in parts of an environment, so cloning the full behavior can limit improvement; stitching selects and recombines strong state-specific actions [30].Source: §1 · p. 1
• The paper reframes offline RL as an optimal transport problem between state and action distributions, using the Q-function as transport cost and the policy as the transport map.The proposed policy identifies a partial distribution of the best actions for each state.Source: §1 · pp. 1–2
• The authors distinguish their approach from prior OT-based methods by not adding an OT regularizer or reward function [11] [40].Source: §1 · p. 1
• PPL achieves superior performance on D4RL compared with existing OT-based methods and state-of-the-art model-free offline RL techniques.The contribution reports evaluation across various environments in the D4RL benchmark suite.Source: §1 · p. 2

2.1 Optimal Transport

• Monge OT seeks a measurable map T that sends distribution μ to ν at minimum cost, while Kantorovich OT uses transport plans that allow mass splitting.An optimal transport plan is a coupling with marginals μ and ν.Source: §2.1 · p. 2
• Maximin formulations jointly recover an OT distance and map by optimizing over functions T, and neural OT methods have been used for strong, weak, and general OT problems [27, 43].The paper draws on these developments for offline RL.Source: §2.1 · p. 2
• Partial or extremal OT matches the full input distribution to only part of the target, with coefficient w controlling imbalance so cost favors closer target samples [17].For Euclidean costs such as ℓ2, the passage describes partial transport maps as selecting nearest target samples.Source: §2.1 · p. 3

2.2 Offline Reinforcement Learning

• In offline RL, policies are learned from fixed experience data; a critic estimates expected cumulative reward and can support policy recovery through methods such as DPG or TD3 [15].The background defines the MDP using state and action spaces, transition probabilities, rewards, and a discount factor.Source: §2.2 · p. 3
• Offline RL can suffer critic inefficiency and overestimation under policy–dataset distribution shift; behavior-cloning methods address this issue, while mixed-quality expert data makes cloning the whole dataset problematic [15].The background also notes that Wasserstein-1 computation requires a 1-Lipschitz potential and that the behavior-cloning coefficient α is task-dependent [25, 1, 6].Source: §2.2 · p. 3
• Datasets may combine near-optimal expert policies even when only one expert is optimal, and conventional cloning lacks a mechanism to weigh each action's importance.The passage also identifies task-dependent α and the 1-Lipschitz constraint as challenges for Wasserstein-1-based approaches [25, 1, 6].Source: §2.2 · p. 3

3.1 Considering RL as OT problem

• The Monge formulation uses negative Qπ as transport cost and the policy as a map from the state distribution to behavior-policy actions.Source: §3.1 · p. 4
• Equality constraints force this formulation to mimic the expert policy, making it inappropriate when the dataset contains sub-optimal paths.Source: §3.1 · p. 4

3.2 Stitching with Partial OT

• Partial OT replaces equality with inequality constraints so states map only to relevant parts of the action distribution, with Q defining action optimality.Source: §3.2 · p. 4
• With critic-based cost Qβ, a policy trained by (10) improves over the behavior policy β.Source: §3.2 · p. 4
• The maximin dual enables neural approximations, trained with stochastic gradient ascent-descent on batches sampled from the offline dataset.Source: §3.2 · pp. 4–5
• Increasing w shrinks the mapped action subset, and scaling it can hurt performance when actions come from a single expert.Source: §3.2 · p. 5
• Unlike prior OT approaches, the method treats the entire policy optimization as OT and does not require a 1-Lipschitz potential for Wasserstein-1.Source: §3.2 · p. 5

4 Toy Experiments

• The toy dataset contains three sub-optimal expert trajectories, with reward only at the final target; the task is to find the shortest path from S0 to ST.Source: Figure 1 · p. 5, §4 · p. 6
• PPL with w = 8 identified an optimal straight trajectory, outperforming the average expert solutions produced by the behavior-cloning comparison.Source: Figure 1 · p. 5, §4 · p. 6
• Both critic and policy used single-layer 32-neuron ReLU networks trained for 5,000 steps with Adam at learning rate 1e-3.Source: §4 · p. 6

5.1 Benchmarks and Baselines

• The evaluation uses D4RL tasks including Gym MuJoCo-v2, Antmaze-v2, and Android-v1, with comparisons against CQL, TD3+BC, IQL, ReBRAC, IQL+OTR, and behavior cloning.The reported score tables average final evaluations over five random seeds.Source: §5.1 · p. 6, Table 1 · p. 7, Table 2 · p. 7

5.2 Settings

• Experiments combine PPL with One-Step RL on MuJoCo, CQL on Antmaze, and ReBRAC on MuJoCo, Android, and Antmaze, using the respective method settings.The One-Step RL setup pre-trains β for 500k steps and Qβ for 2 million steps; the CQL and ReBRAC runs train for 1M steps without pre-trained models.Source: §5.2 · p. 6

5.3 Results

• PPL achieves state-of-the-art results on all Antmaze datasets, improving scores by +16 on large-play-v2 and +21 on large-play-diverse-v2 while consistently outperforming OTR+IQL.Table 1 reports averaged normalized scores from the final 100 evaluations and five random seeds.Source: §5.3 · p. 7, Table 1 · p. 7
• PPL is interpreted as balancing behavior cloning with Q maximization by selecting high-Q actions from the expert action distribution that remain close to critic-training actions.The authors contrast this with direct Q maximization, which can be sub-optimal due to overestimation bias, and staying too close to expert behavior, which prevents improvement.Source: §5.3 · p. 7

5.4 Parameter w

• The effect of w varies by task: its largest effect appears on Antmaze large-play and large-diverse, while some datasets do not need action-subspace reduction.Figure 2 compares normalized-score curves for w = 3, 5, 8, and 12; w controls the policy's action support range.Source: §5.4 · pp. 8–9, Figure 2 · p. 8

6 Related work

• BRAC regularizes policy optimization by minimizing a divergence between the learned policy and the expert policy, and TD3+BC combines TD3 with behavior-cloning loss.The divergence is described as a behavior-cloning loss.Source: §6 · p. 9
• DICE incorporates f-divergence regularization into the critic, while CQL and Adversarial Trained Actor-Critic use pessimistic critic terms on learned-policy and dataset actions..Source: §6 · p. 9
• One-step RL and IQL estimate better actions for each state, while prior OT methods use transport for imitation rewards or reward labeling rather than as the whole offline-RL formulation.One-step RL weights actions using an advantage function; IQL uses a state-conditional upper expectile.Source: §6 · p. 9

7 Conclusion and Discussion

• Partial optimal transport with the action-value function and policy selects the best expert actions for each state from noisy or sub-optimal datasets.The proposed method penalizes inefficient actions in the dataset.Source: §7 · p. 9

7.1 Reproducibility

• The parameter w is task dependent, and the authors provide source code in supplementary materials with hyperparameter details in the settings section.The cited reproducibility passage points to settings (M5.2) for hyperparameters.Source: §7 · p. 9, §7.1 · p. 9

7 Conclusion and Discussion

• The authors propose integrating other OT methods, including regularizations, general costs, and Weak Neural OT, into offline RL.They identify hierarchical RL and exploration with multimodal goals as potentially relevant settings.Source: §7 · p. 9

• The authors do not identify particular societal implications as needing specific emphasis.Source: §8 · p. 9

9.1 Proofs

• Proposition 9.1 states policy improvement for π obtained by solving Eq. 10 with the expert policy β’s Qβ cost, with the displayed return difference nonnegative.The proof compares the best actions in the expert policy’s support with its average actions.Source: §9.1 · p. 13
• As w approaches infinity, the solution of Eq. 10 converges to a problem whose support, in these settings, identifies expert actions maximizing Q.Source: §9.1 · p. 13
• The proof contrasts its maximization over actions in the expert policy’s support with classic policy iteration’s maximization over all actions.Source: §9.1 · p. 13

9.2 Additional Illustrations

• Figure 3 presents normalized Antmaze score curves during training for the proposed method and CQL; the caption identifies the methods by blue and red, respectively.Source: §9.2 · p. 13, Figure 3 · p. 14
• Table 4 reports averaged normalized MuJoCo scores from the final 10 evaluations and 5 random seeds, while Table 5 lists optimal w values by environment.For most tasks, higher w values are preferable, which the authors associate with suboptimal trajectories in the datasets.Source: Table 4 · p. 13, §9.2 · p. 13, Table 5 · p. 14

• The checklist directs readers to sections supporting the abstract and introduction’s contribution claims and states that experimental details are provided in Section 5.Source: NeurIPS Paper Checklist · p. 15
• The authors report using publicly available datasets and including code in supplementary materials, with the solver code to be made public after acceptance.Source: NeurIPS Paper Checklist · p. 15
• The checklist says experimental resources are discussed in the results section, while comparisons across environments are offered as sufficient for significance assessment in offline RL.Source: NeurIPS Paper Checklist · pp. 15–16
• The authors state that the work conforms to the NeurIPS Code of Ethics and that the paper discusses societal impact; they report no human-subject research.Source: NeurIPS Paper Checklist · p. 16