Counterfactual Transport Flows for Offline Conservative Trajectory Refinement
2026-06-08 • Machine Learning
Machine Learning
AI summaryⓘ
The authors introduced a method called counterfactual transport flows to improve decision-making using only previously collected data. Their approach gently adjusts existing behavior by comparing and refining trajectories based on better-performing examples nearby in a learned space. This method allows controlling how much the behavior changes, balancing safety and improvement. They tested it on standard benchmark tasks and showed it can improve performance while offering clear paths of how behaviors are refined.
offline reinforcement learningtrajectory refinementcounterfactual reasoninglatent spacepolicy improvementD4RL benchmarksAntMazeMuJoCo
Authors
Lena Krieger, Xuan Zhao, Zhuo Cao, Qin Wang, Hanno Scharr, Ira Assent
Abstract
Offline reinforcement learning (RL) offers a path to policy improvement from logged data alone, using historical returns or other measurable outcomes as world feedback. A key difficulty is improving observed behavior without extrapolating beyond what the offline data supports. We propose \emph{counterfactual transport flows}, a source-conditioned trajectory refinement framework for offline decision-making guided by world feedback. Given a low-feedback candidate trajectory, we construct local preference pairs from offline data by retrieving nearby trajectories in latent trajectory space with higher task-specific feedback, and use them as weak supervision for conservative refinement. The framework learns instance-specific refinement directions: at inference time, a refinement strength parameter controls how far the candidate trajectory is transported, enabling a trade-off between preserving the original behavior and applying stronger improvement. Experiments on D4RL benchmarks, including AntMaze and MuJoCo tasks, show that our method improves behavior from historical returns as world feedback, while providing interpretable trajectory-level refinement paths.