Counterfactual Transport Flows for Offline Conservative Trajectory Refinement

2026-06-08

AI Summary

This work addresses the challenge in offline reinforcement learning of improving trajectory performance without exceeding the support of historical data. The authors propose a counterfactual transport flow framework that, for the first time, integrates counterfactual reasoning with optimal transport to retrieve high-return neighboring trajectories in latent trajectory space, forming local preference pairs that serve as weak supervision signals for conservative policy optimization. By introducing an adjustable intensity parameter, the method enables trajectory-level and instance-level policy improvements that are both interpretable and restrained from excessive extrapolation, thereby preserving the original behavioral characteristics. Empirical evaluations demonstrate that the approach significantly enhances trajectory performance on the D4RL benchmark, particularly in AntMaze and MuJoCo tasks, while generating interpretable optimization pathways.

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.

Problem

• offline reinforcement learning
• conservative policy improvement
• trajectory refinement
• counterfactual reasoning
• distributional shift

Innovation

• counterfactual transport flows
• offline reinforcement learning
• trajectory refinement
• conservative policy improvement
• latent trajectory space