🤖 AI

Edit-R2: Context-Aware Reinforcement Learning for Multi-Turn Image Editing

This paper introduces Edit-R2, a reinforcement learning framework that enhances multi-turn image editing by reconstructing session intent to mitigate context dilution and state contamination, alongside the MICE-Bench benchmark for systematic evaluation.

Original authors: Yuxiao Ye, Haoran He, Fangyuan Kong, Xintao Wang, Pengfei Wan, Kun Gai, Ling Pan

Published 2026-06-05

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yuxiao Ye, Haoran He, Fangyuan Kong, Xintao Wang, Pengfei Wan, Kun Gai, Ling Pan

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 working with a very talented but slightly forgetful digital artist. You want to edit a photo, so you give them a simple instruction: "Change the dog's collar to blue." They do it perfectly.

But then, you want to keep going. You say, "Now, make the dog's leash purple." They do that too. Then you say, "Actually, change the dog's eyes to green."

Here is the problem: Most current AI artists get confused as the conversation gets longer. They might forget that you asked for a purple leash earlier, or they might think "it" in your latest sentence refers to the whole picture instead of just the dog. They lose the thread of the conversation, and the final result becomes a messy mix of your instructions and their own mistakes.

This paper introduces a new system called Edit-R2 that solves this problem. Think of it as giving that digital artist a super-powered notepad and a strict coach.

The Two Big Problems

The authors identified two main reasons why AI fails at long editing sessions:

The "Long-Context Dilution" (The Forgetful Artist): As you add more instructions and images, the AI gets overwhelmed. It's like trying to remember a grocery list while someone is shouting new items at you every second. The early instructions (like "make everything purple") get drowned out by the noise of the new instructions.
The "State Contamination" (The Snowball Effect): If the AI makes a small mistake in step one (e.g., it changes the wrong object), that mistake carries over to step two. By step three, the image is so messed up that the AI can't recover, and the whole session fails.

The Solution: Edit-R2

The authors created Edit-R2, which uses a technique called Reinforcement Learning (think of it as a video game where the AI learns by playing, failing, and trying again until it wins).

Here is how Edit-R2 works, using a simple analogy:

1. The "Notepad" (IC-CoT)

Before the AI touches the image, it has to write a note. This is called In-Context Chain-of-Thought.

What it does: Instead of just looking at the latest instruction, the AI pauses and summarizes the entire conversation so far.
The Analogy: Imagine you are editing a photo, and before you make a change, you read a summary note that says: "Remember, the user wants everything added to be purple. The previous object was a ribbon. The current instruction is to add a book."
The Result: This note "distills" all the scattered history into one clear, compact instruction. It prevents the AI from forgetting the "purple rule" or getting confused about what "it" refers to.

2. The "Coach" (Trajectory Filtering)

During the training process, the AI tries many different ways to edit the photo. Sometimes, it makes a huge mistake early on.

What it does: The system has a safety mechanism. If the AI makes a mistake that breaks the "rules" of the session (like ignoring the purple color rule), the system immediately stops that attempt. It throws away the bad attempt and doesn't let the AI learn from that specific failure.
The Analogy: It's like a coach watching a basketball player. If the player trips on the first step, the coach blows the whistle and says, "Stop. That run is over. Let's try again from the start." This prevents the player from practicing bad habits.

3. The "Unified Goal"

Usually, AI models are trained to be good at thinking (writing the note) and doing (drawing the picture) separately. Edit-R2 trains them together.

The Analogy: It's like training a writer and an illustrator to work as a single team. The writer knows that if they write a confusing note, the illustrator will fail. The illustrator knows that if they draw the wrong thing, the writer's note was useless. They learn to optimize for the whole session, not just the current turn.

The New Test: MICE-Bench

To prove their system works, the authors built a new test called MICE-Bench.

The Analogy: Before this, tests were like asking the AI, "Can you draw a cat?" (Single turn). MICE-Bench is like a multi-round interview: "Draw a cat. Now make it blue. Now add a hat to the blue cat. Now remove the hat but keep the blue color."
The test measures three things:
- Instruction Following: Did they do what you asked?
- Content Consistency: Did they keep the parts of the image you didn't touch?
- Global Awareness: Did they remember the "rules" you set at the very beginning (like "everything must be purple")?

The Results

When they tested Edit-R2 against other top AI models:

It remembered better: It didn't forget the "purple rule" or get confused by pronouns like "it" or "them."
It stayed consistent: Even after three or four rounds of editing, the image still looked like the original photo, just with the requested changes.
It outperformed the competition: It beat other strong models, including some from big tech companies, especially in those long, multi-step editing sessions.

In Summary

The paper argues that for AI to truly be a helpful editing tool, it needs to stop treating every instruction as a fresh start. Edit-R2 teaches the AI to remember the conversation, summarize its own thoughts before acting, and learn from its mistakes without letting those mistakes ruin the whole project. It turns a forgetful, one-trick pony into a reliable, long-term creative partner.

Technical Summary: Edit-R2 – Context-Aware Reinforcement Learning for Multi-Turn Image Editing

Problem Statement

Current text-guided image editing systems, driven by diffusion and unified multimodal foundation models, predominantly operate in single-turn settings. In these scenarios, a model receives an input image and a single instruction to produce one edited output. This formulation fails to capture realistic multi-turn in-context editing workflows, where users iteratively refine an image through a sequence of instructions.

In such multi-turn sessions, a model faces two coupled failure modes:

Long-Context Dilution: As the session progresses, the interleaved history of images and text grows. Sparse textual constraints from early turns become difficult to recover due to attention decay and the heavy visual-token budget required to encode generated images.
State Contamination: Errors made in early turns (e.g., incorrect object placement or attribute modification) are carried forward, degrading subsequent generations and making the training rollout unreliable.

Existing benchmarks and training paradigms are largely history-agnostic, treating each turn as an isolated task dependent only on the immediately preceding image and current instruction, ignoring accumulated session-level constraints.

Methodology: Edit-R2

The authors propose Edit-R2, a novel reinforcement learning (RL) post-training framework designed for unified multimodal models (specifically built upon BAGEL). Edit-R2 formulates multi-turn in-context editing as a session-level Markov Decision Process (MDP) and introduces three core mechanisms to address context loss and state contamination:

1. Session-State Reconstruction via In-Context Chain-of-Thought (IC-CoT)

To mitigate long-context dilution, Edit-R2 does not pass the raw interleaved history directly to the generator. Instead, it employs an IC-CoT module (an "understanding expert") that performs reasoning over the full session history before each generation step.

Function: The module reconstructs the operative session intent, consolidating scattered historical constraints (e.g., global rules like "all added objects must be purple") and resolving pronouns (e.g., "change its color") into an explicit, enhanced instruction ( $\tilde{T}_t$ ).
Benefit: This creates a compact reasoning trace that shortens the distance between the generation module and the critical information it needs, providing a stable state representation.

2. Unified Multi-Turn RL Objective

Edit-R2 jointly optimizes the IC-CoT reasoning trace and the visual generation process.

Coupled Spaces: The framework couples the discrete text space of intent reconstruction with the continuous latent space of the flow-matching image generator.
Reward Structure: A compound reward function ( $r_t$ $r_{t}$ ) is defined at each turn, combining:
- Instruction Following (IF): Did the model execute the current instruction?
- Content Consistency (CC): Is non-target content preserved?
- Global Awareness (GA): Does the output adhere to session-level constraints established in previous turns?
Optimization: The policy is updated using Unified-GRPO, which treats the CoT token sequence and the flow trajectory as a joint decision, ensuring the reasoning module is optimized specifically for the downstream visual generation capabilities.

3. Trajectory Filtering (Prefix-Valid Advantage Refinement)

To address state contamination, Edit-R2 implements a trajectory filtering mechanism.

Mechanism: If the reward for a specific turn falls below a threshold ( $\delta$ ), the rollout is terminated immediately. Subsequent turns in that trajectory are discarded.
Advantage Refinement: The advantage estimates for policy updates are normalized exclusively over active rollouts (those that have not been terminated). This prevents corrupted, low-quality trajectories from distorting the normalization baseline and contaminating the training signal.

Benchmark: MICE-Bench

To enable systematic evaluation, the authors introduce MICE-Bench (Multi-turn In-Context Editing Benchmark), a large-scale dataset comprising 720 multi-turn instances (3 turns each).

Tasks:
- Content Memory: Tests adherence to session-level constraints (e.g., "all added objects must be red").
- Content Understanding: Tests the resolution of pronominal coreferences across turns (e.g., "remove it").
Metrics: Introduces Global Awareness (GA) as a new metric to quantify compliance with accumulated session-level constraints, alongside standard Instruction Following (IF) and Content Consistency (CC).

Key Results

Experiments on MICE-Bench demonstrate that Edit-R2 substantially outperforms existing open-source and closed-source baselines:

Performance Gains: Compared to the base model BAGEL, Edit-R2 achieves significant improvements, particularly in later turns where context loss is most severe:
- Turn 2: +18% IF, +18% GA.
- Turn 3: +8% IF, +15% GA.
Ablation Studies:
- Removing IC-CoT leads to performance degradation in Turn 2 and 3, confirming its role in mitigating context dilution.
- Removing Text-GRPO (joint optimization) reduces gains, highlighting the necessity of optimizing the reasoning and generation modules together.
- Removing Prefix-Valid Advantage Refinement results in lower overall performance, validating the importance of filtering corrupted rollouts.
Qualitative Analysis: Attention visualizations show that Edit-R2 maintains sharply focused attention on key object tokens from earlier turns, whereas baseline models (like BAGEL) exhibit diffuse, misdirected attention due to RoPE distance decay.

Significance and Claims

The paper claims that Edit-R2 represents a significant step forward in making image editing systems capable of session-level consistency.

Paradigm Shift: It moves beyond isolated local edits to a session-level RL problem, acknowledging that the value of an edit depends on its ability to preserve context for future turns.
First of its Kind: Edit-R2 is presented as one of the first RL post-training frameworks specifically designed for in-context editing, addressing the unique challenges of long-horizon dependencies and state contamination.
Evaluation Standard: By introducing MICE-Bench and the GA metric, the work provides the first systematic, automated protocol for evaluating multi-turn in-context editing, a capability previously underexplored in both models and datasets.

The authors conclude that while MICE-Bench covers two representative scenarios, the framework lays the groundwork for future exploration of broader user intents in real-world interactive editing. They acknowledge the risks of controllable image generation (e.g., disinformation) and commit to releasing usage guidelines alongside the model weights.