Learning to Reason in Structured In-context Environments with Reinforcement Learning

1. Problem Statement

While Reinforcement Learning (RL) has significantly improved the reasoning capabilities of Large Language Models (LLMs) through environmental exploration, the training environment itself is often overlooked. Existing RL environments for LLMs suffer from three main limitations:

• Scalability: Mathematical and coding environments often rely on expensive expert annotations, making it difficult to construct large-scale training data.
• Generalizability: Skills learned in game-based environments are often too specialized and fail to transfer to general reasoning domains.
• Verifiability: Many environments lack clear, objective rules to verify the correctness of an answer, leading to reward hacking or ambiguous learning signals.

The authors argue that an ideal reasoning environment must possess scalability, generalizable reasoning patterns, and verifiability. They propose that structured data (e.g., Knowledge Graphs) is the missing link to satisfy these criteria but note that current methods do not effectively leverage it for RL fine-tuning.

2. Methodology: The Structured In-Context Environment (SIE)

The paper introduces the Structured In-Context Environment (SIE) framework, which automatically constructs reasoning environments from massive structured data (specifically Knowledge Graphs) and uses them as "soft constraints" within the LLM's prompt.

A. Automated Construction Pipeline

The framework transforms Knowledge Graph Question Answering (KGQA) tasks into a Markov Decision Process (MDP) where the LLM explores a subgraph to find an answer. The pipeline involves four steps:

1. Seed Subgraph Retrieval: Uses a bidirectional multi-hop search (from question entities and answer entities) to retrieve a seed subgraph containing potential reasoning paths, optimizing for efficiency.
2. Supporting Subgraph Extraction: Extracts the minimal set of triples (the "supporting subgraph") that form valid reasoning paths connecting the question to the answer using Dijkstra's algorithm.
3. Distractor Subgraph Filtering: Filters the remaining triples to create a "distractor subgraph." A two-stage semantic filtering process (using a cross-encoder) retains only the most relevant and challenging distractors to balance context length and difficulty.
4. Partial SIE Construction: Creates a series of environments with varying difficulty by controlling the retention ratio of the supporting subgraph (e.g., 100%, 75%, 50%, 25%, 0%). This allows the model to train under information-constrained conditions, forcing it to infer missing information.

B. RL Fine-Tuning Process

• Input: The LLM receives a prompt containing the question and the structured in-context (the shuffled mix of supporting and distractor triples).
• Action: The model performs implicit exploration within the provided context to generate a reasoning chain and a final answer.
• Algorithm: The authors employ GRPO (Group Relative Policy Optimization), which eliminates the need for a separate critic model by using group-relative scoring.
• Reward Mechanism:
  • Answer Reward: Binary reward (1.0 or 0.0) based on exact match with the ground-truth answer.
  • Format Reward: Positive reward for adhering to a specific reasoning format (e.g., <thought>...</thought><answer>...</answer>).
  • This rule-based verification ensures the model learns correct reasoning logic rather than just mimicking formats.

3. Key Contributions

1. SIE Framework: Formalized a new framework that automates the construction of scalable, verifiable, and generalizable reasoning environments from structured data (Knowledge Graphs).
2. Partial Environment Analysis: Introduced the concept of Partial SIEs (varying information retention). This allows for systematic study of how LLMs learn to reason when supporting information is incomplete, shifting the paradigm from shallow retrieval to deep compositional reasoning.
3. Empirical Validation: Demonstrated that RL fine-tuning within SIEs significantly outperforms both standard Chain-of-Thought (CoT) prompting and Supervised Fine-Tuning (SFT) on structured reasoning data.
4. Generalization Discovery: Showed that reasoning skills learned in structured KG environments transfer effectively to out-of-domain tasks, including mathematical reasoning (GSM8K, MATH500) and logical puzzles (Knights and Knaves).

4. Experimental Results

Experiments were conducted on Qwen (2.5-7B, 3-8B) and Llama (3.1-8B) models using the Freebase Knowledge Graph.

• In-Domain Performance (Structured Reasoning):

  • RL w/ SIE significantly outperformed RL w/o Context (no structured data) and CoT.
  • On the WebQSP dataset, RL w/ SIE improved performance by +33.7% over RL w/o Context for Qwen2.5-7B-Instruct.
  • Compared to SFT w/ SRD (Supervised Fine-Tuning on distilled reasoning data), RL w/ SIE showed superior efficiency, achieving an average improvement of ~53.7% over CoT, whereas SFT only improved by ~11.4%.

• Out-of-Domain Generalization:

  • Models trained in SIEs showed strong transfer to math and logic tasks.
  • On GSM8K, RL w/ SIE improved Qwen2.5-7B-Instruct from 29.2% (CoT) to 87.4%.
  • On MATH500 and KK-hard (logic puzzles), significant gains were observed, proving that the compositional reasoning strategies are domain-agnostic.

• Robustness in Partial Environments:

  • Even in SIE-0% (where all supporting facts are removed and only distractors remain), RL fine-tuning yielded substantial improvements (e.g., +55.0% on WebQSP).
  • Case Study: In SIE-0%, the model learned to recognize missing information in the context and successfully combined it with its intrinsic parametric knowledge to solve the problem, whereas the pre-trained model hallucinated incorrect answers.

• Ablation Studies:

  • Reward Verification: Removing the answer correctness reward (Format Only) caused performance to collapse, confirming that verifiable feedback is critical.
  • Starting Checkpoints: Starting RL from an SFT checkpoint improved generalization slightly but limited the maximum potential gain in structured reasoning, suggesting that pure RL exploration is more effective for learning the reasoning paradigm.
  • Scalability: The framework scaled effectively to larger models (Qwen2.5-14B), showing consistent gains.

5. Significance and Impact

This paper addresses a critical bottleneck in LLM alignment: the lack of scalable, high-quality reasoning environments. By leveraging structured data (Knowledge Graphs), the authors provide a solution that is:

• Automated and Scalable: Removes the dependency on human annotation for training data generation.
• Generalizable: Proves that reasoning patterns learned in one structured domain (KGs) can be transferred to entirely different domains (Math, Logic).
• Robust: Demonstrates that RL can teach models to "fill in the gaps" when information is missing, a crucial capability for real-world applications where data is often incomplete.

The SIE framework represents a shift from passive data consumption (SFT) to active environmental exploration (RL), offering a promising path toward developing LLMs with more robust, generalizable, and verifiable reasoning capabilities.