AgentFly RL Framework for LM Agents

• AgentFly is a modular, extensible framework designed for scalable reinforcement learning with LM agents that execute multi-turn interactions.
• It employs a decorator-based API and token-level masking to adapt classic RL algorithms for continuous and asynchronous tool-invoked environments.
• The framework supports high-throughput rollouts across diverse environments and tools, driving research in advanced agent behaviors and efficiency.

AgentFly is a modular, extensible framework for scalable reinforcement learning (RL) with LLM (LM) agents, specifically targeting the Agent-RL paradigm where LM agents autonomously complete tasks through orchestrated multi-turn interactions with environments, tools, and APIs. The framework addresses the unique requirements of supporting RL over extended token-action sequences, asynchronous tool invocation, and high-throughput multi-environment execution, operationalized through a decorator-based API and token-level RL adaptation. AgentFly aims to unify the study of LM agent reinforcement learning, overcoming both engineering and methodological challenges prevalent in earlier approaches (Wang et al., 20 Jul 2025).

1. Motivation and System Objectives

Key challenges in Agent-RL include supporting multi-turn agent-environment interactions—characterized by recursive sequences of prompt, generation, tool/API invocation, observation, and further response—whereas standard RL libraries assume single-turn or fully episodic trajectories. Additionally, the high rollout costs inherent in LM-based agents are compounded by slow, sometimes blocking, external tool or environment calls, motivating significant advances in concurrent execution and resource coordination. There exists a strong demand for a system that is simultaneously modular—permitting researchers to easily plug in new tools, environments, and custom reward functions—and efficient at the scale of modern RL experiments.

AgentFly is designed to:

• Enable multi-turn RL training by adapting classic RL algorithms for token-level operations.
• Provide a decorator interface for seamless definition of tools (APIs, environments, or user-defined functions) and reward functions.
• Execute asynchronous, chain-based rollouts, coordinated via a centralized resource manager.
• Ship with prebuilt tools and benchmark environments, accelerating the development and evaluation of new Agent-RL approaches (Wang et al., 20 Jul 2025).

2. System Architecture

AgentFly's architecture is bifurcated into two core modules: the RL Trainer (based on Verl) and the Agent Rollout Module. The RL Trainer is responsible for management of policy and value networks (anchored in the LM backbone), gradient-based optimization, logging, and checkpointing. It utilizes a user-defined "RolloutSampler" to gather task trajectories and calculate policy advantages.

The Agent Rollout Module encapsulates the multi-turn agent logic, iterating between prompt-context LM generation, parsing and asynchronously dispatching tool calls, aggregating tool observations, and continuing token generation. Central to scalable execution is the Centralized Resource Manager, which maintains finite pools per stateful tool/environment class. Each tool call with a new session ID triggers the allocation of an idle environment; instances are reset and recycled upon session termination, and requests are queued if no idle resources are available.

A conceptual interaction diagram is as follows:

• [RL Trainer] ←→ [RolloutSampler] ←→ [Agent Core]
• The Agent Core uses the [Tool & Reward Interface] for extensibility.
• The [Central Resource Manager] coordinates pools for PythonSandboxEnv, WebShopEnv, etc.

This templated resource control and chain-based agent execution enable high-throughput, multi-agent rollouts across diverse tasks and environments (Wang et al., 20 Jul 2025).

3. Token-Level Masking for Multi-Turn Reinforcement Learning

Traditional RL with LMs, such as policy optimization via PPO, computes losses over contiguous, single-agent token streams:

$$L_{PPO}(\theta) = \sum_{t=1}^L f(\pi_\theta(a_t|s_t), \pi_{\theta_{old}}(a_t|s_t), \hat{A}_t)$$

In Agent-RL, token sequences interleave agent generations ($r_i$) with environment-observed fragments ($o_i$). To restrict reward attribution and policy updates to LM-emitted tokens, AgentFly introduces a binary token mask $M_t$ over the full token sequence:

$$M_t = \begin{cases} 1 &\text{if token } t \text{ belongs to any } r_i,\ 0 &\text{otherwise} \end{cases}$$

The masked PPO loss is:

$$L_{PPO}(\theta) = \sum_{t=1}^T M_t\, f\left(\pi_\theta(a_t|s_t), \pi_{\theta_{old}}(a_t|s_t), \hat{A}_t\right)$$

Only LM-generated tokens contribute to returns and advantages, ensuring proper credit assignment across complex, multi-turn interactions. This adaptation is central to faithfully training LM agents with RL algorithms in practical, tool-augmented environments (Wang et al., 20 Jul 2025).

4. Extensible Decorator-Based Interface

AgentFly exposes a Python decorator mechanism to facilitate rapid extension of the framework by registering tools and reward functions. Decorators abstract away environment allocation, session management, and concurrency, covering both stateless and stateful tools. Examples include:

• Stateless tool (calculator):

from agentfly.api import tool @tool(name="calculator", description="Evaluate a mathematical expression.") def calculate(expr: str) -> str: return str(eval(expr))

• Stateful tool (sandboxed Python):

from agentfly.api import tool from agentfly.envs import PythonSandboxEnv @tool(name="code_interpreter", env_cls=PythonSandboxEnv, pool_size=8, description="Run Python code snippets.") async def code_interpreter(code: str, env: PythonSandboxEnv) -> str: return await env.step(code)

• Stateless reward:

from agentfly.api import reward @reward(name="qa_f1") def qa_f1_reward(pred: str, gold: str) -> dict: f1 = f1_score(pred, gold) return {"reward": f1, "f1": f1}

• Environment-based reward:

@reward(name="webshop_success", env_cls=WebShopTextEnv, pool_size=8) async def webshop_reward(pred: str, env: WebShopTextEnv, task_id: int) -> float: out = await env.step("get_reward", task_id) return out["reward"]

This interface enables efficient, type-safe integration of new tasks and evaluation metrics without compromising throughput or resource control (Wang et al., 20 Jul 2025).

5. Asynchronous Execution Strategy and Scalability

AgentFly's agent rollout algorithm is fundamentally asynchronous, combining vLLM-based token generation, asyncio-dispatched tool calls, and parallel post hoc reward computation. The workflow follows:

$r_i$0

Scalability is empirically validated: on an H200 node with 8 GPUs, generating 16 chains per query, a full PPO run on code-interpreter tasks requires approximately 200 GPU-hours. Throughput increases linearly with the number of environment worker instances, up to the limits of compute and memory bandwidth (Wang et al., 20 Jul 2025).

6. Supported Environments, Tools, and Algorithms

AgentFly ships with a curated set of environments and tools:

Tool/Environment	Description
Code Interpreter	Sandboxed Python REPL
Calculator	Evaluation of mathematical expressions with eval
Search	Serper-based Google search (Redis-backed cache)
Retrieve	Wikipedia retriever (leveraging Search-R1)
WebShopTextEnv	Simulated e-commerce navigation, ~12k instructions
ALFWorldEnv	Embodied text-based household tasks
ScienceWorldEnv	Interactive science sandbox, 30 topics

Supported RL algorithms include PPO, REINFORCE++ (Hu et al.), GRPO (Shao et al.), and RLOO (Ahmadian et al.). The training pipeline sequentially initializes LM policy and baseline, samples prompts, executes chain-based rollouts, computes masked per-token returns, and optimizes parameters via Verl (Wang et al., 20 Jul 2025).

7. Empirical Results, Limitations, and Ongoing Directions

• Empirical findings on Qwen2.5-Instruct 3B (Code Interpreter, 16 chains/query) reveal rapid initial reward improvements across all algorithms (~50 steps to plateau), with PPO, GRPO, and RLOO outperforming REINFORCE++ by 5–10%. Scaling from 3B to 7B model sizes yields consistent 10–20% higher final rewards on ALFWorld, WebShop, and ScienceWorld environments, though long-horizon tasks increase learning instability. Ablation on turn limits demonstrates that reward convergence is similar irrespective of the max-turn setting, but shorter rollouts produce more stable gradients and smoother curves. Tool-use dynamics in ALFWorld indicate that agents shift from random actions to informed querying, decreasing hallucinated (invalid) tool calls from ~30% to <5% as training progresses.

Identified limitations include equal-weight credit assignment to all LM emissions—suggesting that finer-grained schemes may aid stability—and persistent optimization challenges for sparse-reward or long-horizon tasks. Offline RL and replay buffer strategies are not presently supported but could reduce rollout burden. The meta-learning of tool discovery remains an open direction, as does the extension to multi-agent and adversarial environments (Wang et al., 20 Jul 2025).