Technical Summary: Counterfactual Transport Flows for Offline Conservative Trajectory Refinement

1. Problem Formulation

The paper addresses a fundamental limitation in offline reinforcement learning (RL): the difficulty of improving observed behavior without extrapolating beyond the support of logged data. While standard RL learns reactive state-to-action mappings, it often fails to explicitly model how a specific low-outcome trajectory could be locally revised to succeed. Existing offline methods typically either constrain policies via value penalties or use generative models (diffusion/flow) to create new behaviors, often operating at the transition level or learning global policies rather than instance-specific improvements.

The authors formulate the problem as source-conditioned trajectory refinement. Given an initial candidate trajectory $\tau^-$ with low world feedback (e.g., low return, safety violations), the goal is to learn a local revision toward a nearby alternative $\tau^*$ that offers improved feedback while remaining behaviorally plausible within the offline data distribution. This is conceptualized as finding a feasible trajectory that minimizes deviation from the source while satisfying an improvement margin $\delta$:
$$\tau^* = \arg \min_{\tau' \in \mathcal{T}_{feas}} d_\tau(\tau', \tau^-) \quad \text{s.t.} \quad R(\tau') \ge R(\tau^-) + \delta$$
where $R(\tau)$ represents measurable world feedback and $d_\tau$ measures trajectory-level deviation.

2. Methodology: Counterfactual Transport Flows

The proposed framework, Counterfactual Transport Flows, learns to transport a source trajectory toward a better alternative in a latent trajectory space using conditional flow matching. The method operates in three main stages:

A. Local Preference Pair Construction

Instead of learning a global policy, the method constructs weak local preference pairs from offline data. For a given source trajectory $\tau^-$, the algorithm:

1. Encodes trajectories into a latent space using an encoder $\phi$.
2. Identifies a local latent neighborhood $\mathcal{N}(\tau^-)$ of similar trajectories.
3. Selects a target trajectory $\tau^+$ from this neighborhood that satisfies the improvement criterion ($R(\tau^+) \ge R(\tau^-) + \delta$).
4. Forms a pair $(\tau^-, \tau^+)$ representing a "counterfactual" revision: what the agent could have done differently under similar conditions to achieve better feedback.

B. Source-Conditioned Flow Matching

The refinement is modeled as a continuous transport process rather than a one-shot prediction.

• Latent Interpolation: For each pair, a linear interpolation path is defined: $z_s = (1-s)z^- + s z^+$, where $z^- = \phi(\tau^-)$ and $z^+ = \phi(\tau^+)$, with $s \sim U(0, 1)$.
• Conditional Vector Field: The model learns a vector field $v_\theta(z_s, s \mid z^-)$ conditioned specifically on the source trajectory $z^-$. This ensures the refinement direction is instance-specific, learning how this particular trajectory can be revised, rather than learning a global direction toward high-feedback regions.
• Training Objective: The model is trained using the flow matching objective to minimize the difference between the predicted velocity and the target velocity ($z^+ - z^-$):
  $$\mathcal{L}_{FM} = \mathbb{E}_{z^-, z^+, s} \left[ \| v_\theta(z_s, s \mid z^-) - (z^+ - z^-) \|_2^2 \right]$$

C. Inference with Refinement Strength

At inference time, the method takes an initial candidate trajectory $\tau^-$, encodes it to $z^-$, and integrates the learned vector field. The integration is stopped at a refinement strength parameter $\alpha \in [0, 1]$:
$$\tilde{z} = z_\alpha$$
The resulting latent state is decoded to produce the refined trajectory $\tau^* = \psi(\tilde{z})$.

• $\alpha = 0$: Recovers the original trajectory (conservative).
• $\alpha = 1$: Fully transports the trajectory toward the local high-feedback target.
• $0 < \alpha < 1$: Allows a continuous trade-off between preserving the original behavior and applying stronger improvement.

3. Key Contributions

The authors highlight three primary contributions:

1. Formulation: They reframe offline decision improvement as source-conditioned trajectory refinement from world feedback, distinct from global policy optimization or unconstrained trajectory generation.
2. Methodology: They propose constructing local preference pairs from observed outcomes and utilizing conditional flow matching to learn instance-specific refinement directions in latent trajectory space.
3. Validation: They validate the approach on D4RL benchmarks (AntMaze and MuJoCo), demonstrating improved feedback with conservative trajectory deviations, requiring no human preference labels or separate reward models.

4. Experimental Results

Experiments were conducted on D4RL datasets, including AntMaze and MuJoCo tasks (HalfCheetah), using a trajectory autoencoder for latent space construction.

• Baselines: The method was compared against:
  • Nearest Improved Neighbor: Directly replacing the source with the nearest higher-feedback trajectory.
  • Random Improved Trajectory: Replacing the source with a random higher-feedback trajectory.
  • Non-local Flow Matching: Training the flow model without locality constraints on target selection.
• Performance:
  • AntMaze: The proposed method achieved a feedback improvement of +69.44 with the lowest action deviation (1.41) and latent deviation (27.00). In contrast, the "Random Improved Trajectory" achieved higher raw gains (+117.50) but with significantly larger deviations (Action: 2.18, Latent: 62.50), indicating a lack of conservatism. The "Non-local Flow Matching" baseline actually degraded feedback (-11.48).
  • MuJoCo: Similar trends were observed. The proposed method achieved +1180.20 feedback improvement with deviations of 1.58 (action) and 31.40 (latent). The "Nearest Improved Neighbor" achieved slightly lower raw improvement (+1050.30) but with higher deviations (2.14/42.50).
• Analysis: The results indicate that local target construction and source conditioning are both necessary for conservative refinement. Global or non-local approaches fail to learn stable refinement directions, while direct replacement strategies sacrifice too much behavioral fidelity.

5. Significance and Claims

The paper claims that Counterfactual Transport Flows offer a promising direction for reliable offline improvement by:

• Grounding Improvement in Data: The method improves behavior using only scalar world feedback (e.g., returns) from logged data, eliminating the need for human preference labels, separate reward models, or online environment interaction.
• Interpretability: The transport perspective provides trajectory-level interpretability. The learned vector field traces explicit refinement paths, offering "counterfactual-style" explanations of how low-feedback behaviors can be revised.
• Controllability: The refinement strength parameter $\alpha$ enables a tunable trade-off between conservatism (preserving the original proposal) and aggressive improvement.

The authors maintain a modest stance, noting that the term "counterfactual" is used in the trajectory-refinement sense (grounded in observed data) rather than formal causal inference. They identify future work directions, including adaptive neighborhood selection, theoretical guarantees for decoded trajectories, and integration with learned world models for iterative hypothetical planning.