---
title: Trajectory Purification in Reinforcement Learning
url: https://www.emergentmind.com/topics/trajectory-purification
type: topic
---

# Trajectory Purification in Reinforcement Learning

Trajectory purification refers to methods for systematically eliminating or repairing failure segments and associated errors from agentic reinforcement learning (RL) trajectories, producing “purified” paths that yield more informative gradients and efficient policy optimization. In the context of large language models (LLMs) augmented with tool use—such as Python code execution—trajectory purification is fundamental for mitigating spurious credit assignment and variance induced by noisy feedback. The Similarity-Aware Adaptive Rollback (SAAR) mechanism, introduced in the CLEANER framework, offers an influential and rigorously evaluated approach to this problem, combining self-correction and semantic similarity heuristics to generate training data comprised only of successful, self-consistent steps [2601.15141].

## 1. Formalization and Notation

Trajectory purification operates on RL rollouts (trajectories) comprising alternating sequences of reasoning, code, and observations. Let $x$ be the user query, $\mathcal{E}$ the code execution environment, and $\pi_\theta$ the agent’s policy. Each trajectory $\tau$ is a sequence:
\[
\tau = \bigl[x,\, (r_0, c_0, o_0),\, (r_1, c_1, o_1),\,\ldots,\, (r_T, c_T, o_T)\bigr]
\]
where $r_t$ is a text-based chain-of-thought, $c_t$ the generated code, and $o_t$ the outcome, with $o_t^+$ for success and $o_t^-$ for failure.

A failure segment at step $t$ is given by $f_t = (r_t,\,c_t,\,o_t^-)$. The model may attempt self-correction, producing $s_t = (r'_{t,\mathrm{aux}},\,c'_t,\,o'_t)$ with $o'_t = o^+$. The result of trajectory purification via SAAR is a new sequence $\tau_\mathrm{purified}$ where each $f_t$ successfully corrected is replaced by $(r_{t, \mathrm{final}}, c'_t, o'_t)$.

## 2. Mechanism: Similarity-Aware Adaptive Rollback (SAAR)

The SAAR mechanism is applied during data collection whenever an execution failure $o_t^-$ occurs, in two phases:

- **Phase I (Lookahead Correction):** Augmenting the context with $(r_t, c_t, o_t^-)$, the agent generates up to $K$ corrective attempts under policy $\pi_\theta$ until a successful execution $o'_t$ is produced.

- **Phase II (Similarity-Aware Replacement):** Compute the semantic similarity $\mathrm{Sim}(c_t, c'_t)$ between the failed and corrected code segments, using difflib.SequenceMatcher. With a threshold $\gamma \in (0,1)$ (default $\gamma=0.5$), the “rollback rule” specifies the replacement:
  \[
  r_{t, \mathrm{final}} =
    \begin{cases}
      r_t, & \mathrm{if}\; \mathrm{Sim}(c_t, c'_t) \geq \gamma \ \text{(shallow repair)} \\
      r'_{t, \mathrm{aux}}, & \mathrm{if}\; \mathrm{Sim}(c_t, c'_t) < \gamma \ \text{(deep repair)}
    \end{cases}
  \]
  The tuple $(r_{t, \mathrm{final}}, c'_t, o'_t)$ then enters the purified trajectory.

Additionally, after substitution, $\log\pi_\theta(c'_t \mid h_t \oplus r_{t, \mathrm{final}})$ is recomputed (with RadixAttention) to ensure correct importance weighting.

## 3. Adaptive Granularity and Purification Strategies

SAAR adaptively selects replacement granularity according to semantic similarity:

- **Shallow Repair ($\mathrm{Sim} \geq \gamma$):** Minor corrections (e.g., syntactic changes), retaining the original reasoning $r_t$.
- **Deep Repair ($\mathrm{Sim} < \gamma$):** Substantive corrections (e.g., alternative algorithmic logic), replacing reasoning with $r'_{t, \mathrm{aux}}$.

This granular replacement ensures that error-recovery patterns are excised and only valid, direct reasoning chains are encoded in purified trajectories.

## 4. Algorithmic Description

The SAAR algorithm can be summarized as follows:

1. Detect failure $(o_t^-)$ at step $t$.
2. For up to $K$ retries, generate corrections $(r'_{t,\mathrm{aux}}, c'_t)$ and evaluate $o'_t = \mathcal{E}(c'_t)$.
3. If $o'_t = o^+$, compute $\mathrm{Sim}(c_t, c'_t)$.
4. If $\mathrm{Sim} \geq \gamma$, construct tuple with original $r_t$ (shallow repair); else, substitute $r'_{t,\mathrm{aux}}$ (deep repair).
5. Substitute $(r_{t, \mathrm{final}}, c'_t, o'_t)$ into $\tau_{\mathrm{purified}}$.
6. Update log-probabilities as needed.

This procedure results in trajectories for RL training that are free from superfluous error-correcting steps and systematically align credit assignment with actually successful reasoning and action choices.

## 5. Theoretical Implications

Trajectory purification via SAAR addresses critical credit assignment and variance reduction problems in RL with outcome-only reward $R(\tau) \in \{-1,1\}$. Standard RL assigns positive reward to all steps in any successful $\tau$, including failed actions if subsequently self-corrected. Purification ensures that only directly successful actions are credited, yielding unbiased policy gradients. The formal claim is:
\[
\text{For $\hat{g}_{\mathrm{base}}$ (noisy) and $\hat{g}_{\mathrm{SAAR}}$ (purified),}\quad \mathrm{Var}\left[\hat{g}_{\mathrm{SAAR}}\right] \leq \mathrm{Var}\left[\hat{g}_{\mathrm{base}}\right]
\]
Purification thereby improves sample efficiency, especially in parameter-constrained model regimes.

## 6. Empirical Results and Experimental Context

CLEANER, using SAAR, was evaluated on AIME24/25, GPQA, and LiveCodeBench with 4B- and 7B-parameter LLMs. Key experimental settings are:
- Retry limit $K = 3$
- Similarity threshold $\gamma = 0.5$
- Curriculum mixing: SAAR on 70% of rollouts, 30% raw

Comparison to DAPO and SOTA baselines is summarized:

| Method               | AIME24 | AIME25 | GPQA | LiveCode | RL Steps |
|----------------------|--------|--------|------|----------|----------|
| DAPO-baseline (4B)   | 66.7   | 59.4   | 56.9 | 49.5     | 250      |
| CLEANER-4B (SAAR)    | 72.7   | 67.1   | 60.2 | 54.9     | 250      |
| DemyAgent-4B (SOTA)  | 72.6   | 70.0   | 58.5 | 51.7     | 750      |

Ablation with Qwen3-4B-Instruct demonstrates absolute accuracy gains of +6–7% on math (AIME24/25), +3% on GPQA, and +5% on LiveCodeBench versus the baseline. Notably, CLEANER with SAAR achieves or exceeds SOTA with only one-third the RL training steps. Figure 1 in the original text displays suppressed error rates and accelerated accuracy improvements over training.

## 7. Impact and Prospective Significance

The retrospective, self-purified trajectory construction enabled by SAAR ensures policies internalize correct reasoning patterns while minimizing error-recovery loops. The main effects are: (i) elimination of noisy and misleading credit assignments, (ii) reduction of policy gradient variance, and (iii) marked improvements in both accuracy and sample efficiency, particularly for compute-constrained LLMs. These outcomes have been corroborated on competitive mathematical, scientific, and code-based benchmarks [2601.15141].

A plausible implication is that trajectory purification, distinct from external filtering or reward shaping, constitutes a scalable paradigm for high-accuracy, efficient agentic RL with realistic tool-augmented LLMs.

Source: https://www.emergentmind.com/topics/trajectory-purification