Understanding the Behaviors of Environment-aware Information Retrieval

LLMs Learn to 'Code-Switch' Query Styles to Match Specific Retriever Characteristics

Different search engines (retrievers) work better with different types of questions. This research shows that by using reinforcement learning (RL), AI can learn to rewrite a single question into different styles. It might use long descriptions for one engine or short keywords for another. This helps get the best possible search results.

In the current landscape of Retrieval-Augmented Generation (RAG), we typically treat the retriever as a static tool. Whether you are hitting a vector database with dense embeddings or a legacy Elasticsearch cluster with BM25 (a keyword-based retrieval algorithm), the assumption is the same. We assume a well-phrased natural language question is "good enough." We optimize the generator to handle whatever the retriever spits out. However, we rarely optimize the interface between the two.

This paper argues that this "one-size-fits-all" approach is a missed opportunity. The authors find that different retrieval backends require fundamentally different "dialects" to achieve peak performance. A query that works for a dense retriever like Contriever might be poor for a sparse keyword matcher like BM25.

The mismatch between intent and formulation

The core problem identified by the authors is "Structural Drift." In most RAG pipelines, the semantic intent remains constant. If you ask "Who wrote God's Plan?", the goal does not change regardless of the backend. However, the optimal formulation of that intent varies wildly.

Current agentic RAG approaches often treat retrieval as a uniform "tool call." They are often optimized for a single commercial search engine. This breaks down in production environments. You might swap backends or deal with local, specialized repositories. As the authors demonstrate, a strategy optimized for one retriever cannot be effectively transferred to another. They quantify this using a metric called Retrieval Environment MMD (RE-MMD). This metric distinguishes between "Semantic Drift" (changes in what the user wants) and "Structural Drift" (changes in how the query is styled).

As shown in, semantic drift remains consistently low across different retriever pairs.

In contrast, structural drift is an order of magnitude higher. This proves that transfer failures are stylistic rather than intentional. The agent fails because it cannot "code-switch" into the required style, such as moving from declarative phrasing to dense keyword clusters.

Teaching agents to speak retriever

To solve this, the researchers frame query adaptation as a reinforcement learning (RL) problem. They use Group Relative Policy Optimization (GRPO). This is an RL method that allows for stable learning without a separate value-function critic (a secondary model used to estimate rewards). The goal is to train an LLM-based rewriter policy $\pi_\theta$ to generate an optimal query $q$. This query aims to maximize the $nDCG@10$ (Normalized Discounted Cumulative Gain, a metric that rewards both relevance and high ranking) of the retrieved documents.

The implementation moves through several logical stages:

1. Environment Interaction: The rewriter receives a user question and generates a query. This query is sent to a "black-box" retriever (the environment), which returns a set of documents.
2. Reward Assignment: The system calculates the reward based on how well the retrieved documents match the ground truth.
3. Branching Rollout for Multi-Turn Stability: For iterative, multi-step retrieval, the authors introduce a "Branching Rollout" strategy. This prevents training collapse. Instead of running independent linear trajectories, they use a $4 \times 4$ structure. They sample $M$ first-turn actions and branch each into $N$ second-turn completions.

This branching mechanism is critical for credit assignment (the process of determining which action led to a reward). As explained in, the branching structure isolates the learning signal for the second turn.

By grouping $N$ branches that share the same history, the baseline reward becomes specific to that subgroup. This ensures the model is rewarded for the quality of the second rewrite independently of the first. Mathematically, this treats the first-turn reward as a constant within the subgroup. This reduces the variance that typically destabilizes multi-turn RL training.

Scaling beyond human heuristics

The empirical results are significant. The authors report that RL effectively teaches LLMs to tailor queries to specific retriever characteristics. For instance, Contriever performs best with long-form, descriptive "hypothetical documents." Meanwhile, BM25 favors concise, keyword-heavy strings.

The most striking takeaway involves model scaling. Small models (4B or 8B) often hit performance plateaus. They may also suffer from "policy collapse," where they settle on a mediocre, safe strategy. However, the authors found that scaling to a 14B model allowed the agent to break through these limits. It can discover non-intuitive strategies. As illustrated in, a 14B agent tasked with querying Contriever moved past traditional question-style rewrites.

It discovered "statement-style" queries. For the question "Is it possible to be white and Latino?", the agent produced the concise statement "Being white and Latino." This achieved an $nDCG@10$ of 0.965. This is a very high score for retrieval effectiveness.

The impact on end-to-end RAG is also measurable. In experiments using the HotpotQA dataset, the authors show that retriever-specific rewriting improved full RAG accuracy to 46.9%. This compares to 41.4% for vanilla RAG and 42.3% for general rewrites using Gemini-2.5-flash [Table 4].

Limits of the current approach

While the results are compelling, there are clear boundaries. First, the study is strictly limited to text-only retrievers and text-based corpora. Modern production stacks often involve multimodal RAG (searching through images, video, or audio). It is unclear if these stylistic "dialects" follow the same patterns in those spaces.

Second, the multi-step exploration is capped at a two-turn horizon. The authors note that the agent generalizes to longer sequences in a zero-shot manner (as seen in the case study in Appendix E). However, the training itself does not account for deep, 5- or 10-step reasoning chains.

Finally, the reliance on GRPO and the branching rollout implies a non-trivial computational cost during training. The inference-time cost is simply the cost of an extra LLM call for the rewrite. However, the "offline" cost to discover these optimal policies involves managing multiple parallel rollout branches. This will require careful orchestration of GPU resources.

The verdict: Build for the backend

If you are building a RAG system that uses a single, static vector database, this paper might not change your workflow. But if you operate in a heterogeneous environment, "query awareness" is no longer optional. You might use a hybrid of BM25 for keyword precision and a dense model for semantic breadth.

The research provides a clear signal: stop treating the retriever as a generic tool. Automating the discovery of retriever-specific dialects via RL is a viable path. It allows you to squeeze more performance out of your existing infrastructure without upgrading your underlying embedding models.

For those ready to move beyond basic prompting, the code and resources are available at LCO-Embedding/Envs-aware-Information-Retrieval. If you have the compute to spare, prototyping a branching-based rollout for your multi-turn agents is a high-leverage move.

Figures from the paper