🤖 machine learning

STITCH-OPE: Trajectory Stitching with Guided Diffusion for Off-Policy Evaluation

STITCH-OPE is a model-based generative framework that leverages guided diffusion and trajectory stitching to overcome the limitations of existing off-policy evaluation methods in high-dimensional, long-horizon settings, achieving significantly reduced variance and improved accuracy.

Original authors: Hossein Goli, Michael Gimelfarb, Nathan Samuel de Lara, Haruki Nishimura, Masha Itkina, Florian Shkurti

Published 2026-07-13

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Hossein Goli, Michael Gimelfarb, Nathan Samuel de Lara, Haruki Nishimura, Masha Itkina, Florian Shkurti

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to play a video game, but you can't let the robot actually play the game in the real world. Maybe the robot is too expensive, or the game is too dangerous. Instead, you have a giant library of old video recordings showing a different robot (let's call it "Old Bot") playing the game. Your goal is to figure out how well a new robot ("New Bot") would do, just by looking at those old recordings. This is called Off-Policy Evaluation (OPE).

The problem is, Old Bot and New Bot might play very differently. If New Bot tries to do something Old Bot never did, the old recordings are useless. If you try to guess New Bot's score by just stretching the old data, the math gets messy and explodes into chaos, especially for long, complicated games.

Enter STITCH-OPE, a new method that acts like a clever video editor and a creative director rolled into one.

The Problem with the Old Ways

Previous attempts to solve this were like trying to predict a whole movie by guessing every single frame from start to finish in one giant leap.

The "Importance Sampling" approach was like trying to fix a broken movie by multiplying the volume of every frame by a huge number. If the movie is long, the volume gets so loud it shatters the speakers (variance explodes).
The "Model-Based" approach was like trying to build a perfect simulator of the world. But if you make a tiny mistake in the first second, that mistake gets bigger and bigger every second after, until the simulation is completely wrong by the end.

The STITCH-OPE Solution: The "Lego" Strategy

STITCH-OPE changes the game by breaking the long movie into tiny, manageable chunks. Think of it like building a long Lego bridge. Instead of trying to snap the whole bridge together at once (which is hard and prone to snapping), you build it piece by piece, connecting the end of one small section to the start of the next.

Here is how it works, step-by-step:

The "Short-Clip" Editor:
Instead of training a model to generate a whole 10-minute video at once, STITCH-OPE trains a "diffusion model" (a fancy AI that learns to turn random static noise into clear pictures) to generate only short clips of the game. Let's say it learns to generate 8-second clips. It learns to take a messy, noisy version of a clip and clean it up until it looks like a real move Old Bot would make.
The "Stitching" Trick:
To make a long video, STITCH-OPE doesn't generate it all at once. It generates one 8-second clip, stops, and then asks the AI: "Okay, you ended at this specific spot. Now, generate the next 8-second clip starting from exactly where you left off." It stitches these short clips together end-to-end.
- Why this is cool: If the AI makes a small mistake in the first clip, it doesn't ruin the whole movie. It just fixes the next clip based on where the last one actually ended. This stops the "error explosion" that plagues other methods.
The "Director's Guidance" (The Secret Sauce):
Now, we need to make sure these clips look like what New Bot would do, not just Old Bot. The AI has a "score function" (a way of knowing what New Bot likes to do).
- The Trap: If you just tell the AI "Do what New Bot likes," it might get confused and create impossible moves because New Bot's style is totally different from Old Bot's.
- The Fix: STITCH-OPE uses a "negative guidance" trick. It tells the AI: "Do what New Bot likes, BUT subtract the things Old Bot loved to do."
- The Analogy: Imagine you are directing a scene. Old Bot loves to dance in a circle. New Bot wants to run in a straight line. If you just say "Run in a straight line," the actor might get stuck. But if you say "Don't dance in a circle, and then run," the actor knows exactly what to avoid. This prevents the AI from getting stuck in a "comfort zone" of Old Bot's moves and forces it to explore New Bot's style.

What the Numbers Say

The authors tested this on some famous robot control tasks (like making a robot hop, walk, or run) using datasets called D4RL and OpenAI Gym. They set the game length to 768 steps for the big robots and 256 or 196 steps for the smaller ones.

The results were promising:

Accuracy: STITCH-OPE beat almost every other method in 11 out of 15 specific test cases.
Ranking: It was really good at figuring out which robot was the "best" one, even if the robots were very different from the one in the recordings.
The "Window" Size: They found that generating clips of length 8 (the "window size" $w$ ) was the sweet spot. It balanced the need to stitch things together without making too many mistakes.

What They Don't Know (Yet)

The paper is careful to say this isn't a magic wand for everything.

The "Impossible" Moves: If New Bot tries to do something that Old Bot never did (like jumping off a cliff when Old Bot only walked on the ground), the AI might get confused and make up a fake, impossible move. The math assumes Old Bot has seen at least some of what New Bot does.
Real-World Complexity: While it works great on these computer simulations, the authors admit they aren't sure yet if it will work perfectly on messy, real-world problems where the data is incomplete or the robot can't see everything clearly.

The Bottom Line

STITCH-OPE suggests that by breaking long, scary problems into tiny, stitchable pieces and using a clever "don't do that" rule to guide the AI, we can much better predict how a new robot will perform without ever letting it touch the real world. It's a big step forward for safely testing robots in high-stakes situations, but the authors remind us that it's still a simulation-based success, and real-world testing is the next frontier.

Technical Summary: STITCH-OPE

Problem Statement

Off-Policy Evaluation (OPE) aims to estimate the performance of a target policy ( $\pi$ ) using offline data collected from a different behavior policy ( $\beta$ ). This is critical in domains like robotics and healthcare where online interaction is costly or unsafe. However, existing OPE methods struggle with high-dimensional, long-horizon problems due to two primary failure modes:

Importance Sampling (IS): Suffers from the "curse of horizon," where variance grows exponentially with the evaluation horizon.
Model-Based Methods: While sample-efficient, they suffer from compounding errors when learning dynamics models for long trajectories, leading to inaccurate value estimates.
Distribution Shift: Differences between $\beta$ and $\pi$ lead to inaccurate estimates, making it difficult to trust policy selection before deployment.

Recent attempts using model-free density ratio estimation have shown poor empirical performance when the behavior and target policies differ significantly.

Methodology: STITCH-OPE

The authors propose STITCH-OPE (Sub-Trajectory Importance-Weighted Trajectory Composition for Long-Horizon OPE), a model-based generative framework leveraging denoising diffusion models. The core innovation lies in generating synthetic trajectories for the target policy by guiding a diffusion model trained on behavior data, while avoiding the compounding errors of full-trajectory generation.

1. Guided Diffusion for Distribution Shift Correction

Instead of retraining a model for every new target policy, STITCH-OPE pre-trains a diffusion model on the behavior dataset ( $D_\beta$ ) to learn the behavior trajectory distribution $p_\beta(\tau)$ . During inference, it generates target policy trajectories ( $p_\pi$ ) by guiding the denoising process.

The guidance function $g(\tau)$ is derived from the likelihood ratio between the target and behavior policies. By interpreting the likelihood ratio as a classification density in diffusion guidance, the authors formulate the guidance as the difference between the score functions:
$g(\tau) = \nabla_\tau \sum_{t=0}^{T-1} \log \pi(a_t|s_t) - \nabla_\tau \sum_{t=0}^{T-1} \log \beta(a_t|s_t)$
This allows the model to shift the sampling distribution from $p_\beta$ toward $p_\pi$ without retraining.

2. Negative Behavior Guidance

A key technical contribution is the inclusion of the negative behavior score term ( $-\nabla_\tau \log \beta$ ).

Problem: Using only the target policy score ( $\nabla_\tau \log \pi$ ) can cause the guided sampling distribution to collapse into high-density regions of the behavior policy, even if those regions are unlikely under the target policy (over-regularization).
Solution: Subtracting the behavior score prevents this collapse, ensuring the model generalizes out-of-data and correctly addresses distribution shift. The authors generalize this with coefficients $\alpha$ and $\lambda$ (where typically $\lambda < \alpha$ ) to control the relative importance of the policies.

3. Trajectory Stitching with Conditional Diffusion

To address the compounding error of long-horizon generation, STITCH-OPE does not generate full trajectories in one pass. Instead:

Sub-Trajectory Generation: The diffusion model is trained on short sub-trajectories of length $w \ll T$ .
Conditional Sampling: Generation is performed in a semi-autoregressive manner. The model generates a sub-trajectory conditioned on the final state of the previously generated sub-trajectory.
Stitching: These partial rollouts are stitched end-to-end to form a full-length trajectory.

This approach treats the behavior trajectory distribution as a mixture over sub-trajectories. By conditioning only on the immediate previous state (rather than the full history), the model maintains higher entropy and broader coverage of the behavior data modes, significantly improving compositionality and generalization to new initial states.

Theoretical Guarantees

The paper provides a theoretical bound on the Mean Squared Error (MSE) of the STITCH-OPE estimator (Theorem 3.3).

Under mild assumptions (bounded likelihood ratio and bounded error in the learned behavior model), the MSE is bounded.
Crucially, the bound depends on the sub-trajectory length $w$ rather than the total horizon $T$ .
This results in an exponential reduction in variance compared to standard Importance Sampling (which scales as $O(\exp(T))$ ) and a more favorable polynomial dependence ( $O(T^2)$ ) compared to full-trajectory diffusion methods.

Empirical Results

The authors evaluated STITCH-OPE on the D4RL benchmark (Hopper, Walker2d, HalfCheetah) and OpenAI Gym (Pendulum, Acrobot).

Performance: STITCH-OPE outperformed state-of-the-art baselines (including Fitted Q-Evaluation, Doubly Robust, Importance Sampling, and Policy-Guided Diffusion) in 11 out of 15 instances across various metrics:
- Log RMSE: Significant reduction in error.
- Rank Correlation: Improved ability to correctly rank target policies.
- Regret@1: Lower regret in selecting the best policy.
Diffusion Policies: The method successfully evaluated complex target policies defined as diffusion models, a task where other estimators struggle due to the lack of explicit probability distributions.
Ablation: Experiments confirmed that optimal performance occurs with a small window size ( $w=8$ ) and a guidance configuration where the negative behavior term is present ( $0 < \lambda < \alpha$ ).

Significance and Claims

The paper claims STITCH-OPE is the first work to demonstrate robust off-policy evaluation in high-dimensional, long-horizon tasks. Its significance lies in:

Bridging the Gap: It successfully bridges the gap between model-based OPE (sample efficiency) and full-trajectory offline diffusion (generative capability) by mitigating compounding errors through trajectory stitching.
Scalability: It demonstrates that diffusion models can be applied to OPE in high-dimensional spaces without suffering from the exponential variance explosion typical of IS methods.
Generalization: The combination of conditional diffusion and negative behavior guidance allows for robust evaluation of target policies that deviate significantly from the behavior policy, including those with different initial states or action distributions.

The authors remain modest regarding limitations, noting that the method relies on the assumption that the behavior data covers the target policy's support (Assumption 3.1) and that evaluating trajectory fidelity in complex, partially observable states remains challenging. They also note that the applicability to domain-specific problems outside of robotics remains an open question.