Multi-Stage RL Pipelines with Tool-Conditioned Control

• Multi-stage RL pipelines with tool-conditioned control are frameworks that decompose complex tasks into sequential stages to optimize tool invocations.
• They integrate supervised fine-tuning, RL-based policy optimization, and hierarchical refinement to improve tool usage strategies across diverse applications.
• The approach enhances performance in program synthesis, robotics, and automated systems management while addressing challenges like reward sparsity and compliance constraints.

Multi-stage reinforcement learning (RL) pipelines with tool-conditioned control constitute a foundational paradigm for training agentic systems—particularly LLMs, vision-LLMs (VLMs), and domain-specialized agents—to interact with external APIs, code interpreters, robotics modules, and complex software stacks. These pipelines are designed to hierarchically decompose workflows into discrete stages, each with its own action space over callable “tools” or functions, and to use advanced RL algorithms to optimize both the sequence of tool invocations and the parameters for each call. This approach supports generalization to unseen toolkits, workflows, and real-world environments in mathematical reasoning, embodied AI, program synthesis, and automated systems management.

1. Formal Structure and Problem Formulation

Multi-stage RL pipelines with tool-conditioned control are typically formulated as Markov Decision Processes (MDPs) or Partially Observable MDPs (POMDPs), where agent actions correspond to tool selection and invocation. The agent’s state integrates both task context (problem inputs, prior outputs) and a structured history of tool calls, observations, or API responses. The action space is defined by the catalog of available tools and their input/output schemas; each action is a function call (often serialized in JSON or specialized markup), with parameters specified by the agent (Du et al., 22 Sep 2025, Dong et al., 22 May 2025, Chen et al., 3 Dec 2025, Lu et al., 28 Oct 2025, Feng et al., 15 Apr 2025, Chang et al., 17 Sep 2025, Liu, 2 Mar 2026, Soni et al., 15 Jan 2026).

The environments used in these pipelines may be synthetic (e.g., coding gyms with verifiable tasks), simulated (e.g., CI/CD pipeline emulators or robot manipulation simulators), or real-world (financial APIs, knowledge-graph backends). Each step’s execution returns partial or full observations and, depending on the environment, may yield sparse or decomposed rewards shaped by end-task success, tool usage quality, or compliance metrics.

2. Pipeline Decomposition and Staging

A characteristic feature is the explicit decomposition of the training pipeline into sequential stages, each designed to address a distinct aspect of the agent’s tool-use fluency:

1. Data Generation and Supervised Fine-Tuning: Cold-start with curated or synthesized traces demonstrating correct tool use, including thoughts, tool calls, and corresponding outputs. Supervised fine-tuning (SFT) is performed via standard cross-entropy objectives to teach trajectories and tool-invocation syntax, ensuring general well-formedness (Feng et al., 15 Apr 2025, Du et al., 22 Sep 2025, Chen et al., 3 Dec 2025, Dong et al., 22 May 2025, Liu, 2 Mar 2026).
2. RL-Based Policy Optimization: RL is applied to enable agents to discover, compose, and optimize tool-usage strategies beyond what can be imitated in SFT data. Techniques include Proximal Policy Optimization (PPO), Group Relative Policy Optimization (GRPO), Q-learning, or policy gradients adapted to the tool-conditioned action space (Du et al., 22 Sep 2025, Soni et al., 15 Jan 2026, Lu et al., 28 Oct 2025).
3. Hierarchical and Stepwise Fine-Tuning: Some pipelines use hierarchical RL to separately target trajectory-level (end-to-end correct solutions) and step-level (tool invocation accuracy or code execution success) objectives, combining their gradients for joint optimization (Chang et al., 17 Sep 2025).
4. Direct Preference Optimization (DPO) and Self-Critique: An additional stage overlays human-labeled preference pairs or self-critic feedback to refine the policy’s alignment with nuanced compliance regimes, style, or edge-case handling (Dong et al., 22 May 2025, Liu, 2 Mar 2026).
5. Self-Correction and On-the-Fly Backtracking: In deployment, agents may iteratively revise failed tool invocations by leveraging runtime feedback (e.g., code execution errors) to resample actions or backward-edit outputs (Chang et al., 17 Sep 2025).

3. Tool-Conditioned Action and Observation Modeling

The core of tool-conditioned control lies in the agent’s ability to selectively invoke from a palette of tools, leveraging structured prompts or system messages that enumerate each tool’s signature, documentation, and allowed schema. Models are trained to emit action tokens for tool names, parameters, and wrapper markup, as well as to parse and condition subsequent reasoning on downstream tool responses (Du et al., 22 Sep 2025, Chen et al., 3 Dec 2025, Lu et al., 28 Oct 2025, Liu, 2 Mar 2026).

Internally, tools are often represented as learned embeddings. In vision or spatial reasoning settings, the policy may attend to both visual and textual context, factorizing the joint action as $$\pi_\theta(a_t, z_t \mid s_t) = \pi_\theta(z_t \mid s_t) \,\pi_\theta(a_t \mid s_t, z_t)$$ (Chen et al., 3 Dec 2025). For LLMs, system-level coordination can involve autoregressive generation interleaved with external function calls, structured semantic actions, and parsing of partial environment states.

The following table summarizes representative tool-conditioned action paradigms from the literature:

Pipeline	Action Format	Tool Context Integration
CodeGym (Du et al., 22 Sep 2025)	JSON-wrapped function call	Tool list in prompt; action parsed/executed
ToolRLA (Liu, 2 Mar 2026)	JSON object (tool, params)	Tool schema in prompt/internal embedding
SpaceTools (Chen et al., 3 Dec 2025)	<tool_call> block + reasoning	Learned tool and image embeddings
OrchDAG (Lu et al., 28 Oct 2025)	Token-level JSON for DAG node	Prompt-embedded, no GNN module

4. Reward Design and Optimization Objective

Reward design in multi-stage RL pipelines is notably diverse, tailored to task requirements:

• Sparse terminal rewards (binary success/failure at episode termination) are common in environments with easily verifiable outcomes, such as code generation or mathematical reasoning. This presents significant exploration challenges; dense or decomposed rewards are often added to counteract training brittleness (Feng et al., 15 Apr 2025, Du et al., 22 Sep 2025).
• Fine-grained composite rewards decompose performance into orthogonal components: format validity, correct tool selection, invocation efficiency, and compliance, as in the multiplicative reward function of ToolRLA (Liu, 2 Mar 2026):

$$R(\tau) = R_\mathrm{fmt}(\tau) + R_\mathrm{cor}(\tau) + R_\mathrm{eff}(\tau) + R_\mathrm{cpl}(\tau)$$

with $R_\mathrm{cor}$ multiplicative over name correctness, tool coverage, and parameter accuracy.

• Hierarchical or multi-level rewards reward intermediate tool-execution correctness as well as full-trajectory outcomes (answer correctness), as in THOR (Chang et al., 17 Sep 2025):

$$J(\theta) = E_{τ∼π_θ}\Big[\sum_t r^{\text{traj}}(τ)+\sum_k r^{\text{step}}(s_k,a_k)\Big].$$

• Weighted graph-edit rewards for DAG-structured tool workflows, where edit distance between predicted and ground-truth DAGs supplies a graded reward signal (Lu et al., 28 Oct 2025).
• Efficiency and safety: Terms penalizing excess tool invocations, latency, or regulatory violations (e.g., large negative penalty $\lambda=10$ for compliance failure in financial APIs) (Liu, 2 Mar 2026).

5. Environments, Applications, and Generalization Properties

Applications of multi-stage RL pipelines with tool-conditioned control span a broad spectrum, enabled by the abstraction of agent actions as externally parameterized tool calls. Key domains and empirical accomplishments include:

• Mathematical reasoning and program synthesis: Integration of code interpreters (e.g., Python sandboxes) for arithmetic, symbolic computation, and algorithmic reasoning. State-of-the-art results are achieved on mathematical Olympiad, MATH500, and AIME benchmarks, with large models (32B, 72B) showing superior out-of-distribution generalization (Du et al., 22 Sep 2025, Feng et al., 15 Apr 2025, Chang et al., 17 Sep 2025).
• Multi-modal embodied reasoning: In spatial environments (e.g., “SpaceTools”), VLMs orchestrate perception and manipulation tools for robotic control, achieving improved spatial understanding and real-world manipulation (e.g., Kinova Jaco arm) (Chen et al., 3 Dec 2025).
• Workflow optimization in software engineering pipelines: RL-optimized CI/CD pipelines dynamically modulate tool invocations (build, test, deploy), enhancing throughput by up to 30% while maintaining defect rates (Soni et al., 15 Jan 2026).
• Orchestration of compositional toolchains: DAG-based infrastructure for testing agentic tool-use over flexible, multi-turn, multi-dependency workflows, as with OrchDAG and related benchmarks (Lu et al., 28 Oct 2025).
• Safety- and compliance-critical deployment: ToolRLA demonstrates >90% task completion and ~14% invocation error rate in a real-world financial advisor setting, with explicit enforcement of regulatory constraints via reward design and DPO (Liu, 2 Mar 2026).

6. Limitations, Lessons, and Future Directions

Several challenges and open questions persist:

• Pipeline robustness depends on environment synthesis quality; LLM-generated environments may introduce errors or shortcuts. Reward sparsity remains a bottleneck, especially for small models, accelerating premature convergence but impeding exploration (Du et al., 22 Sep 2025).
• Graph-edit or dense metric shaping improves sample efficiency but may not fully capture semantic tool-calling correctness (Lu et al., 28 Oct 2025).
• Current tool APIs and backend schemas are often static and curated. Generalization to dynamically discovered or open-world tools (with noisy, unstructured I/O) remains limited.
• Real-world deployment demands finely tuned compliance, latency, and efficiency tradeoffs, often solved by explicit reward component weighting and post-hoc alignment (Liu, 2 Mar 2026).

Research directions include curriculum learning over tool complexity, integration of simulated adversaries or user simulators, contrastive RL for shortcut-avoidance, and theoretical analysis of environment hardness or reward horizon properties (Du et al., 22 Sep 2025, Chen et al., 3 Dec 2025, Lu et al., 28 Oct 2025).

7. General Framework and Abstraction

Across agentic domains, a canonical multi-stage RL pipeline for tool-conditioned control is now emerging:

1. Supervised initialization on tool-augmented examples to facilitate correct syntax and primitive tool selection.
2. RL-based discovery and optimization with non-trivial reward shaping, allowing agents to learn compositional, context-sensitive tool usage.
3. Compliance, alignment, and preference optimization ensuring real-world acceptability.
4. Iterative refinement via self-correction and deployment feedback for robust long-horizon workflows.

This pattern is effective across modalities, task regimes, and deployment settings, providing a scalable pathway to robust agentic tool use in both synthetic and real-world environments (Du et al., 22 Sep 2025, Liu, 2 Mar 2026, Chen et al., 3 Dec 2025).