Transport or Discard: Robust Unbalanced Optimal Transport for Cross-Domain Policy Adaptation
Abstract
Cross-domain offline reinforcement learning leverages a source dataset to improve policy learning in a data-scarce target domain, but dynamics mismatch makes many source transitions kinematically infeasible and can cause negative transfer. Recent non-parametric geometric methods (e.g., standard optimal transport and k-nearest neighbors) avoid overfitting yet often yield only relative rankings under an implicit matching or retrieval budget, making performance sensitive to hand-tuned thresholds when the true cross-domain overlap is unknown. We formulate availability estimation as soft subset selection by learning a source reweighting that geometrically aligns with the target. We propose Robust Offline unbalanced Optimal Transport (ROOT): (i) a robust ambiguity set for uncertainty under limited target samples, and (ii) an unbalanced transport objective that penalizes mass deviation, enabling a principled transport-or-discard mechanism. ROOT thus down-weights or discards high-cost source samples rather than forcing them onto the target support. Moreover, the induced weights decay exponentially with transport cost, guaranteeing outlier suppression. On D4RL dynamics-shift benchmarks, ROOT improves downstream offline RL and outperforms strong baselines on most tasks without task-specific threshold tuning.
Lay Summary
Cross-domain offline reinforcement learning aims to train decision-making agents for a target environment by reusing data collected from related source environments. This is useful when collecting target data is expensive or limited, such as in robotics or physical control tasks. However, source data can be misleading when the source and target environments have different dynamics: some source transitions may be physically incompatible with the target environment and can harm policy learning if used directly. This paper proposes ROOT, a method that automatically estimates how useful each source transition is for the target task. Instead of forcing all source data to be matched with the limited target data, ROOT allows incompatible source samples to be down-weighted or discarded. It also accounts for uncertainty caused by having only a small amount of target data. As a result, the learned policy can benefit from relevant source data while avoiding negative transfer from mismatched data. Experiments on standard offline reinforcement learning benchmarks show that ROOT improves policy performance across diverse dynamics shifts without requiring task-specific filtering thresholds.