Efficient Reinforcement Learning for Long-Horizon Tool-Use Agentic Tasks

# Efficient Reinforcement Learning for Long-Horizon Tool-Use Agentic Tasks
Source: [https://arxiv.org/html/2608.10357](https://arxiv.org/html/2608.10357)
###### Abstract

Long\-horizon tool\-using agents must reason over user goals, domain policies, tool calls, simulator state, and delayed verifiable rewards\. Reinforcement learning \(RL\) is a natural fit for this setting, but multi\-turn on\-policy rollouts create long contexts, while model\-specific attention layers may require custom masks and learned sink normalization\. We presentSinkFlex\-RL, a modular training system for RL in dual\-control tool\-use environments\. The system combines a Gymnasium\-compatible environment wrapper, a VERL\-style rollout dataflow, group\-relative policy optimization without a separate value model, and a sink\-aware FlexAttention path designed to preserve model\-specific sink scaling under causal and sliding\-window masks\. In a preliminaryτ2\\tau^\{2\}\-Bench retail run, validation reward \(mean@1\) rises from0\.250\.25early in training to0\.440\.44later in the observed training window, while training\-score and trajectory\-reward proxies also trend upward\. In a fixed\-configuration memory benchmark, the optimized attention path reduces peak VRAM from28\.0628\.06GB to22\.5222\.52GB at40964096tokens, a19\.7%19\.7\\%reduction, and runs the measured81928192\-token configuration using25\.5325\.53GB where the eager baseline runs out of memory\. These results illustrate the value of integrating environment interfaces, RL dataflow, and attention\-kernel design for memory\-feasible long\-horizon agent training\.

## 1Introduction

Language\-model agents increasingly operate in interactive environments: they converse with users, inspect and update state through tools, follow domain policies, and receive sparse success signals only after long trajectories\. Benchmarks such asτ\\tau\-Bench andτ2\\tau^\{2\}\-Bench formalize this setting by evaluating agents in customer\-service\-style domains with tool APIs, policy constraints, and simulated users\(Yaoet al\.,[2024](https://arxiv.org/html/2608.10357#bib.bib10); Barreset al\.,[2025](https://arxiv.org/html/2608.10357#bib.bib1)\)\. These settings stress both reasoning quality and training\-system capacity\. A single trajectory may contain many dialogue turns, tool calls, and environment observations; on\-policy RL multiplies this cost by requiring fresh rollouts, reward checking, and repeated policy updates\.

This paper studies the systems side of memory\-feasible agentic RL for a large open\-weight mixture\-of\-experts \(MoE\) transformer\. MoE architectures can increase model capacity at lower per\-token feed\-forward cost by routing tokens to a subset of experts\(Feduset al\.,[2022](https://arxiv.org/html/2608.10357#bib.bib4)\), but they do not remove the long\-context attention bottleneck\. As multi\-turn trajectories grow, eager attention can exhaust high\-bandwidth memory before backpropagation\. Our current system evaluation covers sequence lengths through81928192tokens; longer target workloads motivate the design but are not evaluated here\. Existing fused attention kernels improve memory movement and throughput\(Daoet al\.,[2022](https://arxiv.org/html/2608.10357#bib.bib2); Dao,[2024](https://arxiv.org/html/2608.10357#bib.bib3)\), but production models may require attention variants with learned sink parameters, heterogeneous masks, or backward behavior that is not exposed by a fixed kernel interface\.

We therefore treat agentic RL as a systems\-integration problem and make three contributions\. First, we wrap dual\-control environments behind a Gymnasium\-style interface so that rollouts, tools, user simulators, and reward checkers can connect to a standard RL dataflow\. Second, we use group\-relative policy optimization \(GRPO\) to update the actor from groups of trajectories without training a separate value model, following the memory\-motivated rationale introduced in DeepSeekMath\(Shaoet al\.,[2024](https://arxiv.org/html/2608.10357#bib.bib7)\)and a VERL\-style post\-training dataflow\(Shenget al\.,[2024](https://arxiv.org/html/2608.10357#bib.bib15)\)\. Third, we implement a sink\-aware FlexAttention path that composes causal and sliding\-window masks with differentiable sink scaling while avoiding eager attention\-state materialization\. We evaluate the integrated pipeline with a preliminaryτ2\\tau^\{2\}\-Bench retail training run and a peak\-memory study through81928192tokens\. The core contribution of this work is the integration of these components and the resulting memory feasibility\.

## 2Related Work

Agentic benchmarks\.Evaluating language models as interactive agents requires environments in which models issue API calls, manage changing state, and process multi\-turn feedback\. AgentBench\(Liuet al\.,[2024](https://arxiv.org/html/2608.10357#bib.bib11)\)motivated broad interactive evaluation, while WebArena\(Zhouet al\.,[2024](https://arxiv.org/html/2608.10357#bib.bib12)\), GAIA\(Mialonet al\.,[2024](https://arxiv.org/html/2608.10357#bib.bib13)\), and SWE\-bench\(Jimenezet al\.,[2024](https://arxiv.org/html/2608.10357#bib.bib14)\)target web navigation, general assistant capabilities, and software engineering, respectively\. For stateful tool use,τ\\tau\-Bench evaluates executable state changes rather than text\-only judgments\(Yaoet al\.,[2024](https://arxiv.org/html/2608.10357#bib.bib10)\)\. We useτ2\\tau^\{2\}\-Bench\(Barreset al\.,[2025](https://arxiv.org/html/2608.10357#bib.bib1)\), whose dual\-control setting allows both the agent and a simulated user to influence task progress\. This provides a representative testbed for gathering missing information, recovering from invalid actions, and maintaining consistency over multi\-turn interactions\.

RL fine\-tuning for multi\-step tool use\.In complex tool\-use environments, annotating every intermediate action is costly, which motivates training with programmatically verifiable episode\-level rewards\(Zhenget al\.,[2025](https://arxiv.org/html/2608.10357#bib.bib16); Donget al\.,[2026](https://arxiv.org/html/2608.10357#bib.bib17)\)\. PPO\(Schulmanet al\.,[2017](https://arxiv.org/html/2608.10357#bib.bib6)\)commonly uses a learned value function, adding memory and compute overhead for large language models\. GRPO\(Shaoet al\.,[2024](https://arxiv.org/html/2608.10357#bib.bib7)\)removes the separate value model and normalizes rewards within groups of sampled outputs\. Frameworks such asverl\(Shenget al\.,[2024](https://arxiv.org/html/2608.10357#bib.bib15)\)organize rollout generation, reward computation, and policy updates as a distributed post\-training dataflow\. Our work adopts these algorithmic and systems components and focuses on integrating them with long\-context, model\-compatible attention for multi\-turn tool\-use rollouts\.

Attention mechanisms for long\-context training\.Transformers use attention as their core sequence\-mixing operation\(Vaswaniet al\.,[2017](https://arxiv.org/html/2608.10357#bib.bib8)\), but dense attention has quadratic memory and compute in sequence length\. FlashAttention and FlashAttention\-2 reduce memory traffic by tiling exact attention and avoiding materialized attention matrices\(Daoet al\.,[2022](https://arxiv.org/html/2608.10357#bib.bib2); Dao,[2024](https://arxiv.org/html/2608.10357#bib.bib3)\)\. StreamingLLM shows that retaining attention sinks can help stabilize windowed long\-context inference\(Xiaoet al\.,[2024](https://arxiv.org/html/2608.10357#bib.bib9)\)\. PyTorch FlexAttention exposes a programmable interface for composing masks and score modifications\(PyTorch Contributors,[2026](https://arxiv.org/html/2608.10357#bib.bib5)\)\.SinkFlex\-RLcombines these ideas by treating efficient attention as both a memory problem and a model\-compatibility problem when sink\-aware normalization and heterogeneous masks must remain in the differentiable training path\.

## 3Agentic RL Setting

##### Dual\-control environment\.

In single\-control tool\-use benchmarks, the agent is often the only actor that changes the environment through tools while the user supplies information passively\. In dual\-control environments, both sides can affect the world state: the agent chooses dialogue and API actions, while the simulated user can confirm choices, provide missing information, or take user\-side actions\. A task instance can be represented as

ℰ=\(g,πSOP,𝒜tool,s0,u\),\\mathcal\{E\}=\(g,\\pi\_\{\\mathrm\{SOP\}\},\\mathcal\{A\}\_\{\\mathrm\{tool\}\},s\_\{0\},u\),\(1\)whereggis the user goal,πSOP\\pi\_\{\\mathrm\{SOP\}\}denotes domain\-policy or standard\-operating\-procedure constraints,𝒜tool\\mathcal\{A\}\_\{\\mathrm\{tool\}\}is the agent tool set,s0s\_\{0\}is the initial shared state, anduuis the user simulator\. A trajectory is

τ=\{\(ot,at,rt,dt,it\)\}t=1T,\\tau=\\\{\(o\_\{t\},a\_\{t\},r\_\{t\},d\_\{t\},i\_\{t\}\)\\\}\_\{t=1\}^\{T\},\(2\)whereoto\_\{t\}is an observation,ata\_\{t\}is an agent action,rtr\_\{t\}is a reward or diagnostic signal,dtd\_\{t\}is a termination flag, anditi\_\{t\}contains metadata\. The trainer consumes a trajectory\-level rewardRiR\_\{i\}produced by the benchmark’s programmatic checker\. Other dashboard diagnostics are logged separately and are not outputs from a separately trained critic network\.

##### Why this setting stresses the training system\.

The environment creates three systems pressures\. First, rewards are delayed and often verifiable only after the final state is checked, increasing the number of sampled tokens per useful gradient\. Second, trajectories are multi\-turn and tool\-heavy, so the policy must retain user messages, tool outputs, and domain\-policy constraints in context\. Third, the system must coordinate environment execution, rollout generation, reward checking, and policy optimization while maintaining consistent episode state and policy versions\.

Figure[1](https://arxiv.org/html/2608.10357#S3.F1)summarizes the dual\-control task structure used in this work: a task specification initializes a shared world state, the user and agent take alternating actions, and the resulting trajectory is scored by a programmatic checker before the rollout group is used for a GRPO\-style update\.

![Refer to caption](https://arxiv.org/html/2608.10357v1/agentic_rl_framework.png)Figure 1:A dual\-control agentic RL episode\. A task specification initializes the environment, user–agent interaction produces a multi\-turn trajectory, and a programmatic checker scores the resulting rollout group for a GRPO\-style policy update\.

## 4Training Pipeline

Figure[2](https://arxiv.org/html/2608.10357#S4.F2)summarizes the training dataflow\. We describe the components abstractly to avoid dependence on deployment\-specific infrastructure\.

![Refer to caption](https://arxiv.org/html/2608.10357v1/verl_tau2.png)

Figure 2:Training\-system dataflow\. The actor samples multi\-turn trajectories through rollout workers and a Gymnasium\-compatible wrapper; trajectory rewards and metadata flow to a GRPO trainer, which updates and redistributes policy weights\.##### Environment wrapper\.

Each benchmark domain is exposed through a common reset/step interface\. On reset, the wrapper samples a task, initializes the shared state, and constructs the first policy observation\. On step, it routes model outputs to one of three handlers: natural\-language response, tool call, or termination\. The wrapper then advances the user simulator and state backend and returns the next observation together with reward metadata\. This separation keeps benchmark\-specific user simulation, tool behavior, and checking logic outside the trainer and is intended to support additional executable environments through the same interface\.

##### Rollout worker\.

The rollout worker owns the agent loop\. At each turn, it formats the observation, samples from the current policyπθ\\pi\_\{\\theta\}, parses the response into an action, and appends the transition to the trajectory buffer\. The worker repeats until the environment terminates or a maximum\-turn budget is reached\. The trainer receives token\-level log probabilities, action masks, trajectory rewards, and episode metadata without directly depending on benchmark\-specific parsing or tool implementations\.

##### Policy update\.

Given a prompt or task contextxxand a group ofGGsampled rolloutsy1,…,yGy\_\{1\},\\ldots,y\_\{G\}from the old policy, GRPO computes a group\-normalized trajectory advantage

A^i=Ri−μ​\(R1:G\)σ​\(R1:G\)\+ϵA,\\hat\{A\}\_\{i\}=\\frac\{R\_\{i\}\-\\mu\(R\_\{1:G\}\)\}\{\\sigma\(R\_\{1:G\}\)\+\\epsilon\_\{A\}\},\(3\)whereRiR\_\{i\}is the programmatically computed trajectory reward andϵA\\epsilon\_\{A\}is a small numerical constant\. Under outcome\-level supervision, the same normalized advantage is assigned to every optimized token in rolloutyiy\_\{i\}\. For tokenyi,ty\_\{i,t\}, the importance\-sampling ratio is

ρi,t​\(θ\)=πθ​\(yi,t∣x,yi,<t\)πθold​\(yi,t∣x,yi,<t\)\.\\rho\_\{i,t\}\(\\theta\)=\\frac\{\\pi\_\{\\theta\}\(y\_\{i,t\}\\mid x,y\_\{i,<t\}\)\}\{\\pi\_\{\\theta\_\{\\mathrm\{old\}\}\}\(y\_\{i,t\}\\mid x,y\_\{i,<t\}\)\}\.\(4\)Define the clipped ratio

ρ¯i,t​\(θ\)=clip⁡\(ρi,t​\(θ\),1−ϵc,1\+ϵc\)\.\\bar\{\\rho\}\_\{i,t\}\(\\theta\)=\\operatorname\{clip\}\\\!\\left\(\\rho\_\{i,t\}\(\\theta\),1\-\\epsilon\_\{c\},1\+\\epsilon\_\{c\}\\right\)\.\(5\)The loss minimized by the trainer is

𝒥i,t​\(θ\)\\displaystyle\\mathcal\{J\}\_\{i,t\}\(\\theta\)=min⁡\(ρi,t​\(θ\)​A^i,ρ¯i,t​\(θ\)​A^i\),\\displaystyle=\\min\\\!\\left\(\\rho\_\{i,t\}\(\\theta\)\\hat\{A\}\_\{i\},\\bar\{\\rho\}\_\{i,t\}\(\\theta\)\\hat\{A\}\_\{i\}\\right\),\(6\)ℒGRPO​\(θ\)\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{GRPO\}\}\(\\theta\)=−1∑iTi​∑i=1G∑t=1Ti𝒥i,t​\(θ\)\+β​DKL​\(πθ∥πref\),\\displaystyle=\-\\frac\{1\}\{\\sum\_\{i\}T\_\{i\}\}\\sum\_\{i=1\}^\{G\}\\sum\_\{t=1\}^\{T\_\{i\}\}\\mathcal\{J\}\_\{i,t\}\(\\theta\)\+\\beta D\_\{\\mathrm\{KL\}\}\(\\pi\_\{\\theta\}\\\|\\pi\_\{\\mathrm\{ref\}\}\),\(7\)whereTiT\_\{i\}is the optimized token length of rolloutii,ϵc\\epsilon\_\{c\}is the clipping radius, and the KL term regularizes the policy toward the reference model\. This is the minimization form of the usual PPO/GRPO clipped maximization objective\(Schulmanet al\.,[2017](https://arxiv.org/html/2608.10357#bib.bib6)\)\. The baseline design does not train a separate critic or value network\.

## 5Sink\-Aware FlexAttention

##### Memory wall\.

A trajectory of lengthnninduces ann×nn\\times nattention score structure for each relevant batch/head component when implemented eagerly\. The number of score positions therefore grows quadratically with context length; for example,n=20,000n=20\{,\}000corresponds to4×1084\\times 10^\{8\}positions before accounting for activations, optimizer state, MoE routing, or rollout batching\. Fused attention methods address this by computing attention in tiles and reducing high\-bandwidth\-memory traffic\(Daoet al\.,[2022](https://arxiv.org/html/2608.10357#bib.bib2)\)\. A model\-specific attention layer, however, may also require custom score modifications, mixed causal and sliding\-window masks, or learned sink logic\.

##### Attention sinks\.

Attention sinks are tokens or learned mechanisms that absorb attention mass and can stabilize long\-context behavior\. Prior work shows that retaining initial sink tokens can recover quality under windowed attention in streaming inference\(Xiaoet al\.,[2024](https://arxiv.org/html/2608.10357#bib.bib9)\)\. Our setting differs from inference\-only KV\-cache management: the sink behavior is part of the trainable attention computation and therefore must remain connected to the forward and backward paths\.

##### Zero\-value\-sink equivalence\.

To avoid materializing an explicit sink token in the key and value caches, we use the algebraic form of a zero\-valued sink\. Letsηs\_\{\\eta\}be a learned sink logit and let its value vector bevsink=𝟎v\_\{\\mathrm\{sink\}\}=\\mathbf\{0\}\. For a queryqq, the attention output with the explicit sink is

Osink\\displaystyle O\_\{\\mathrm\{sink\}\}=∑iexp⁡\(q⋅ki\)​vi\+exp⁡\(sη\)​𝟎∑iexp⁡\(q⋅ki\)\+exp⁡\(sη\)\\displaystyle=\\frac\{\\sum\_\{i\}\\exp\(q\\cdot k\_\{i\}\)v\_\{i\}\+\\exp\(s\_\{\\eta\}\)\\mathbf\{0\}\}\{\\sum\_\{i\}\\exp\(q\\cdot k\_\{i\}\)\+\\exp\(s\_\{\\eta\}\)\}\(8\)=\(∑iexp⁡\(q⋅ki\)∑iexp⁡\(q⋅ki\)\+exp⁡\(sη\)\)​Ostd\.\\displaystyle=\\left\(\\frac\{\\sum\_\{i\}\\exp\(q\\cdot k\_\{i\}\)\}\{\\sum\_\{i\}\\exp\(q\\cdot k\_\{i\}\)\+\\exp\(s\_\{\\eta\}\)\}\\right\)O\_\{\\mathrm\{std\}\}\.\(9\)Writingℓ=log​∑iexp⁡\(q⋅ki\)\\ell=\\log\\sum\_\{i\}\\exp\(q\\cdot k\_\{i\}\)gives

αsink=exp⁡\(ℓ\)exp⁡\(ℓ\)\+exp⁡\(sη\)=σ​\(ℓ−sη\),Osink=αsink​Ostd\.\\alpha\_\{\\mathrm\{sink\}\}=\\frac\{\\exp\(\\ell\)\}\{\\exp\(\\ell\)\+\\exp\(s\_\{\\eta\}\)\}=\\sigma\(\\ell\-s\_\{\\eta\}\),\\qquad O\_\{\\mathrm\{sink\}\}=\\alpha\_\{\\mathrm\{sink\}\}O\_\{\\mathrm\{std\}\}\.\(10\)Under the zero\-value assumption, explicit sink materialization is therefore algebraically equivalent to scaling the standard attention output by a factor computed from the log\-sum\-exp statistic\.

importtorch

importtorch\.nn\.functionalasF

scores\_with\_sink=torch\.cat\(\[sink\_score,scores\],dim=\-1\)

v\_with\_sink=torch\.cat\(\[torch\.zeros\_like\(v\[:,:,:1,:\]\),v\],dim=\-2\)

out\_explicit=F\.softmax\(scores\_with\_sink,dim=\-1\)@v\_with\_sink

out\_std=F\.softmax\(scores,dim=\-1\)@v

lse=torch\.logsumexp\(scores,dim=\-1,keepdim=True\)

out\_scaled=out\_std\*torch\.sigmoid\(lse\-sink\_score\)

asserttorch\.allclose\(out\_explicit,out\_scaled,atol=1e\-6\)

Figure 3:Illustration of the zero\-value\-sink identity\. The production implementation uses fused attention statistics rather than materializing the dense score tensor shown in this small reference example\.
##### Implementation pattern\.

The implementation uses FlexAttention as a programmable attention substrate\(PyTorch Contributors,[2026](https://arxiv.org/html/2608.10357#bib.bib5)\)\. A mask function composes causal constraints with an optional sliding\-window constraint and an always\-visible prefix\. The mask is compiled into a block\-sparse structure so that fully masked blocks can be skipped\. When requested, the attention call returns both the output and an auxiliary log\-sum\-exp statistic\. The sink path then applies the model\-specific scaling function to the output:

Mb,h,q,k\\displaystyle M\_\{b,h,q,k\}=𝟙​\[k≤q\]∧𝟙​\[q−k≤w∨k<p\],\\displaystyle=\\mathbb\{1\}\[k\\leq q\]\\wedge\\mathbb\{1\}\[q\-k\\leq w\\;\\vee\\;k<p\],\(11\)\(z,ℓ\)\\displaystyle\(z,\\ell\)=FlexAttention​\(Q,K,V;BlockMask​\(M\)\),\\displaystyle=\\mathrm\{FlexAttention\}\(Q,K,V;\\mathrm\{BlockMask\}\(M\)\),\(12\)αsink\\displaystyle\\alpha\_\{\\mathrm\{sink\}\}=fη​\(ℓ\),\\displaystyle=f\_\{\\eta\}\(\\ell\),\(13\)z′\\displaystyle z^\{\\prime\}=z⊙αsink,\\displaystyle=z\\odot\\alpha\_\{\\mathrm\{sink\}\},\(14\)wherewwis the local\-window size,ppis the number of always\-visible prefix positions, andfηf\_\{\\eta\}is the learned sink\-scaling function\. For the zero\-value\-sink construction above,fη​\(ℓ\)=σ​\(ℓ−sη\)f\_\{\\eta\}\(\\ell\)=\\sigma\(\\ell\-s\_\{\\eta\}\)\. The prefix budgetppand the learned sink logitsηs\_\{\\eta\}are distinct:ppcontrols token visibility in the mask, whereassηs\_\{\\eta\}reallocates softmax mass through output scaling\.

##### Gradient flow and autograd integration\.

The data\-dependent sink scale couples the gradients of the attention outputzzand the log\-sum\-exp statisticℓ\\ell\. Letdddenote the head dimension\. By the chain rule,

∇zℒ\\displaystyle\\nabla\_\{z\}\\mathcal\{L\}=∇z′ℒ⊙αsink,\\displaystyle=\\nabla\_\{z^\{\\prime\}\}\\mathcal\{L\}\\odot\\alpha\_\{\\mathrm\{sink\}\},\(15\)∇αsinkℒ\\displaystyle\\nabla\_\{\\alpha\_\{\\mathrm\{sink\}\}\}\\mathcal\{L\}=∑j=1d\(∇z′ℒ⊙z\)j,\\displaystyle=\\sum\_\{j=1\}^\{d\}\\left\(\\nabla\_\{z^\{\\prime\}\}\\mathcal\{L\}\\odot z\\right\)\_\{j\},\(16\)∇ℓℒ\\displaystyle\\nabla\_\{\\ell\}\\mathcal\{L\}=∇αsinkℒ⊙fη′​\(ℓ\),\\displaystyle=\\nabla\_\{\\alpha\_\{\\mathrm\{sink\}\}\}\\mathcal\{L\}\\odot f\_\{\\eta\}^\{\\prime\}\(\\ell\),\(17\)and the sink\-parameter gradient is

∇ηℒ=∇αsinkℒ⊙∂fη​\(ℓ\)∂η\.\\nabla\_\{\\eta\}\\mathcal\{L\}=\\nabla\_\{\\alpha\_\{\\mathrm\{sink\}\}\}\\mathcal\{L\}\\odot\\frac\{\\partial f\_\{\\eta\}\(\\ell\)\}\{\\partial\\eta\}\.\(18\)A fixed fused\-kernel interface may not expose the auxiliary gradient path throughℓ\\ell\. Our implementation composes FlexAttention with the sink\-scaling operation under AOTAutograd andtorch\.compile, allowing the compiler to generate forward and backward code forQQ,KK,VV, and the sink parameters without materializing anO​\(n2\)O\(n^\{2\}\)Jacobian\.

##### Why a programmable kernel interface is needed\.

Optimized attention kernels are essential, but a fixed interface can omit model\-specific masking, score modification, or auxiliary\-gradient behavior\. In on\-policy RL, a small attention mismatch can be repeatedly amplified across long trajectories and policy updates\. The FlexAttention path therefore prioritizes explicit control over the model’s attention semantics while retaining a fused, block\-sparse execution path\.

##### Memory\-oriented implementation optimizations\.

Profiling identified two avoidable sources of memory allocation: eager materialization in the sink\-rescaling computation and replication of mask metadata across batch and head dimensions\. We address them as follows\.

##### Compilation and fusion\.

We applytorch\.compileto the composed FlexAttention and sink\-scaling path\. PyTorch Inductor can fuse eligible pointwise operations with surrounding generated code, reducing the number and lifetime of materialized intermediate tensors\.

##### Mask broadcasting\.

When the attention pattern is shared across batch elements and heads, block\-mask construction omits explicit batch and head dimensions and relies on kernel\-side broadcasting\. This avoids storing repeated copies of the same sparse mask metadata\. The aggregate effect of the optimized path is evaluated through peak\-VRAM measurements in Section[6](https://arxiv.org/html/2608.10357#S6); we do not report a separate per\-optimization ablation\.

## 6Experiments

We report two complementary measurements: preliminary policy\-learning trends in theτ2\\tau^\{2\}\-Bench retail domain and peak\-memory scaling for the attention implementation\. The first documents behavior over an observed training window; it is not an algorithmic comparison\. The second measures peak VRAM only and does not establish throughput or wall\-clock improvements\.

### 6\.1Experimental Design

##### Experiment 1: Preliminary retail training\.

We evaluate the integrated GRPO pipeline on the retail domain ofτ2\\tau^\{2\}\-Bench\(Barreset al\.,[2025](https://arxiv.org/html/2608.10357#bib.bib1)\)\. The agent interacts with a dynamic user simulator, executes multi\-step API calls, and follows domain\-specific standard operating procedures\. The purpose of this experiment is to document whether validation reward and associated training diagnostics improve over the observed run, not to isolate the causal effect of GRPO relative to another optimizer\.

We track three metrics:

- •Validation Reward \(mean@1\):the benchmark’s programmatically verified task\-success score for a single sampled validation trajectory per task\.
- •Training Score Proxy:a dashboard diagnostic monitored during training\. The dashboard historically labels this trace as a “critic\-style score,” but the baseline does not train a separate critic or value network\.
- •Trajectory Reward Proxy:a rolling summary of raw episode rewards used to monitor the direction and variability of the observed training run\.

##### Experiment 2: Peak\-memory scaling\.

We compare the model\-native eager attention reference path with the optimized sink\-aware FlexAttention path under the same fixed model, batch, and training configuration\. The optimized path is configured to reproduce the same causal/sliding\-window policy and sink scaling as the reference path\. We measure peak high\-bandwidth\-memory usage \(Peak VRAM in GB\) at sequence lengths of10241024,20482048,40964096, and81928192tokens\. Because the current evaluation reports neither throughput nor a forward/backward numerical\-equivalence test, the conclusions are limited to observed peak memory and execution feasibility in this configuration\.

### 6\.2Experimental Results

#### 6\.2\.1Preliminary Retail Training Results

Table[1](https://arxiv.org/html/2608.10357#S6.T1)summarizes values estimated from the available dashboard screenshots\. The retail validation reward rises from0\.250\.25early in training to0\.440\.44later in the observed training window\. These values represent two portions of the same run rather than an untrained baseline and a final converged model\.

Table 1:Preliminary retail training trends\.Values are visually estimated from dashboard screenshots and should be interpreted as approximate\. Higher values are better\.![Refer to caption](https://arxiv.org/html/2608.10357v1/tau2_retail_reward.png)

![Refer to caption](https://arxiv.org/html/2608.10357v1/Tau2_train_reward.png)

Figure 4:Approximate learning curves\.The left panel shows retail validation reward; the right panel shows the training\-score diagnostic whose dashboard label is “critic\-style score\.” Dots denote noisy per\-evaluation readings estimated from screenshots, and thick lines show the displayed smoothed trends\.Individual evaluations are variable, as expected for multi\-turn, state\-dependent tasks\. Over the displayed window, the smoothed validation trace trends upward, while the training\-score and trajectory\-reward proxies rise from0\.180\.18to0\.400\.40and0\.390\.39, respectively\. These observations are consistent with the integrated pipeline receiving a useful learning signal under the implemented GRPO update\. Without multiple seeds, an optimizer baseline, or exported scalar logs, they do not isolate the effect of GRPO, establish variance reduction, or support a statistical significance claim\.

#### 6\.2\.2Peak\-Memory Results

Table[2](https://arxiv.org/html/2608.10357#S6.T2)reports peak VRAM for the two attention paths\. The optimized path uses less memory at every sequence length for which both paths complete, while the eager path runs out of memory at81928192tokens\.

Table 2:Peak\-memory comparison\.Peak VRAM during the measured training configuration\. The optimized sink\-aware FlexAttention path completes the81928192\-token configuration for which the eager reference path runs out of memory\.At40964096tokens, the optimized path reduces peak VRAM from28\.0628\.06GB to22\.5222\.52GB, a reduction of5\.545\.54GB or19\.7%19\.7\\%\. At81928192tokens, it completes the measured configuration with a peak allocation of25\.5325\.53GB, whereas the eager reference path encounters an out\-of\-memory error\. Thus, in this fixed configuration, the optimized path removes the81928192\-token memory failure observed for the eager baseline\. Because the experiment measures only peak VRAM, it does not establish improvements in throughput, latency, total training time, or accelerator utilization\.

## 7Discussion and Limitations

First, our system is particularly well suited to agentic tasks that operate in controlled environments with programmatically verifiable outcomes\. This setting is naturally instantiated byτ2\\tau^\{2\}\-Bench, where the policy interacts with a user simulator, invokes domain tools, and receives feedback based on the resulting environment state\. By grounding optimization in auditable actions and objective task outcomes, the framework enables scalable and reproducible policy learning\. The same approach has strong potential to extend to other domains that provide structured tool interfaces, executable environments, or clearly defined success criteria\. For more open\-ended applications, programmatic verification can be complemented with human feedback or learned reward models, providing a natural path toward broader task coverage\.

Second, the preliminary retail\-domain experiment provides encouraging evidence for the effectiveness of the proposed policy\-learning pipeline\. Over the observed training window, the displayed reward\-associated traces exhibit a clear upward trend from the early to later stages of training, indicating that the policy is able to benefit from the optimization signal\. These results serve as a promising proof of concept for applying GRPO to tool\-using agentic systems\. They also motivate broader evaluations across domains, random seeds, model scales, and training horizons to further characterize the robustness and generality of the observed improvements\.

Third, programmatic rewards provide an efficient and scalable source of supervision by directly connecting policy behavior to verifiable task outcomes\. This foundation can support increasingly rich reward functions that jointly assess final\-state correctness, policy compliance, tool\-use quality, and dialogue naturalness\. GRPO is especially well matched to settings in which rollout groups contain diverse behavioral outcomes, allowing relative advantages to identify and reinforce stronger trajectories\. Techniques such as diversity\-aware sampling, curriculum design, adaptive grouping, and reward shaping offer promising directions for maintaining informative reward variation and improving optimization efficiency\.

Forth, the memory evaluation demonstrates the feasibility of executing the proposed training approach at sequence lengths of up to81928192tokens under the tested configuration\. In particular, the sink\-aware execution path provides a practical foundation for scaling agentic policy learning to longer interaction trajectories while retaining the intended eager reference semantics\. The current peak\-VRAM measurements establish an important systems proof of concept, and the same evaluation framework can be extended to characterize throughput, latency, training time, and accelerator utilization\. Forward\- and backward\-equivalence tests, together with targeted validation of sink handling, masking, position indexing, and log\-sum\-exp statistics, can provide additional confirmation of semantic consistency\. Overall, these results establish a strong foundation for developing efficient, reliable, and scalable training systems for long\-horizon tool\-using agents\.

## 8Conclusion

We presentedSinkFlex\-RL, a modular RL training system that integrates a Gymnasium\-compatible dual\-control environment interface, VERL\-style rollout dataflow, GRPO updates without a separate value model, and a sink\-aware FlexAttention path\. In the current evidence, retail validation and training diagnostics trend upward over one preliminary run, while the optimized attention path reduces peak memory and completes the measured81928192\-token configuration where the eager reference path runs out of memory\. These results support the narrower claim that environment, RL\-dataflow, and attention\-kernel integration can improve the memory feasibility of long\-horizon agent training\. They do not yet establish algorithmic superiority, broad generalization, exact implementation equivalence, or end\-to\-end computational speedup\.

## Ethics Statement

The system trains agents that can call tools and update environment state\. Such agents should be evaluated for policy compliance, user deception, unsafe tool use, privacy leakage, and simulator overfitting before deployment\. The present work focuses on system feasibility and preliminary training behavior and does not claim deployment readiness\.

## References

- V\. Barres, H\. Dong, S\. Ray, X\. Si, and K\. Narasimhan \(2025\)τ2\\tau^\{2\}\-Bench: evaluating conversational agents in a dual\-control environment\.arXiv preprint arXiv:2506\.07982\.Cited by:[§1](https://arxiv.org/html/2608.10357#S1.p1.2),[§2](https://arxiv.org/html/2608.10357#S2.p1.2),[§6\.1](https://arxiv.org/html/2608.10357#S6.SS1.SSS0.Px1.p1.1)\.
- T\. Dao, D\. Fu, S\. Ermon, A\. Rudra, and C\. Ré \(2022\)Flashattention: fast and memory\-efficient exact attention with io\-awareness\.Advances in neural information processing systems35,pp\. 16344–16359\.Cited by:[§1](https://arxiv.org/html/2608.10357#S1.p2.1),[§2](https://arxiv.org/html/2608.10357#S2.p3.1),[§5](https://arxiv.org/html/2608.10357#S5.SS0.SSS0.Px1.p1.4)\.
- T\. Dao \(2024\)Flashattention\-2: faster attention with better parallelism and work partitioning\.InInternational Conference on Learning Representations,Vol\.2024,pp\. 35549–35562\.Cited by:[§1](https://arxiv.org/html/2608.10357#S1.p2.1),[§2](https://arxiv.org/html/2608.10357#S2.p3.1)\.
- Q\. Dong, G\. Sahin, P\. Wang, Z\. Cai, R\. Shrestha, H\. Yang, and D\. Modolo \(2026\)Visual reasoning through tool\-supervised reinforcement learning\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,pp\. 8993–9002\.Cited by:[§2](https://arxiv.org/html/2608.10357#S2.p2.1)\.
- W\. Fedus, B\. Zoph, and N\. Shazeer \(2022\)Switch transformers: scaling to trillion parameter models with simple and efficient sparsity\.Journal of Machine Learning Research23\(120\),pp\. 1–39\.Cited by:[§1](https://arxiv.org/html/2608.10357#S1.p2.1)\.
- C\. E\. Jimenez, J\. Yang, A\. Wettig, S\. Yao, K\. Pei, O\. Press, and K\. Narasimhan \(2024\)Swe\-bench: can language models resolve real\-world github issues?\.InInternational Conference on Learning Representations,Vol\.2024,pp\. 54107–54157\.Cited by:[§2](https://arxiv.org/html/2608.10357#S2.p1.2)\.
- X\. Liu, H\. Yu, H\. Zhang, Y\. Xu, X\. Lei, H\. Lai, Y\. Gu, H\. Ding, K\. Men, K\. Yang,et al\.\(2024\)Agentbench: evaluating llms as agents\.InInternational Conference on Learning Representations,Vol\.2024,pp\. 52989–53046\.Cited by:[§2](https://arxiv.org/html/2608.10357#S2.p1.2)\.
- G\. Mialon, C\. Fourrier, T\. Wolf, Y\. LeCun, and T\. Scialom \(2024\)Gaia: a benchmark for general ai assistants\.InInternational Conference on Learning Representations,Vol\.2024,pp\. 9025–9049\.Cited by:[§2](https://arxiv.org/html/2608.10357#S2.p1.2)\.
- PyTorch Contributors \(2026\)PyTorch flexattention documentation\.Note:[https://docs\.pytorch\.org/docs/stable/nn\.attention\.flex\_attention\.html](https://docs.pytorch.org/docs/stable/nn.attention.flex_attention.html)Accessed 2026\-06\-21Cited by:[§2](https://arxiv.org/html/2608.10357#S2.p3.1),[§5](https://arxiv.org/html/2608.10357#S5.SS0.SSS0.Px4.p1.9)\.
- J\. Schulman, F\. Wolski, P\. Dhariwal, A\. Radford, and O\. Klimov \(2017\)Proximal policy optimization algorithms\.arXiv preprint arXiv:1707\.06347\.Cited by:[§2](https://arxiv.org/html/2608.10357#S2.p2.1),[§4](https://arxiv.org/html/2608.10357#S4.SS0.SSS0.Px3.p1.10)\.
- Z\. Shao, P\. Wang, Q\. Zhu, R\. Xu, J\. Song, X\. Bi, H\. Zhang, M\. Zhang, Y\. Li, Y\. Wu,et al\.\(2024\)Deepseekmath: pushing the limits of mathematical reasoning in open language models\.arXiv preprint arXiv:2402\.03300\.Cited by:[§1](https://arxiv.org/html/2608.10357#S1.p3.2),[§2](https://arxiv.org/html/2608.10357#S2.p2.1)\.
- G\. Sheng, C\. Zhang, Z\. Ye, X\. Wu, W\. Zhang, R\. Zhang, Y\. Peng, H\. Lin, and C\. Wu \(2024\)HybridFlow: a flexible and efficient rlhf framework\.arXiv preprint arXiv: 2409\.19256\.Cited by:[§1](https://arxiv.org/html/2608.10357#S1.p3.2),[§2](https://arxiv.org/html/2608.10357#S2.p2.1)\.
- A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. N\. Gomez, Ł\. Kaiser, and I\. Polosukhin \(2017\)Attention is all you need\.Advances in neural information processing systems30\.Cited by:[§2](https://arxiv.org/html/2608.10357#S2.p3.1)\.
- G\. Xiao, Y\. Tian, B\. Chen, S\. Han, and M\. Lewis \(2024\)Efficient streaming language models with attention sinks\.InInternational Conference on Learning Representations,Vol\.2024,pp\. 21875–21895\.Cited by:[§2](https://arxiv.org/html/2608.10357#S2.p3.1),[§5](https://arxiv.org/html/2608.10357#S5.SS0.SSS0.Px2.p1.1)\.
- S\. Yao, N\. Shinn, P\. Razavi, and K\. Narasimhan \(2024\)τ\{\\tau\}\-Bench: a benchmark for tool\-agent\-user interaction in real\-world domains\.arXiv preprint arXiv:2406\.12045\.Cited by:[§1](https://arxiv.org/html/2608.10357#S1.p1.2),[§2](https://arxiv.org/html/2608.10357#S2.p1.2)\.
- Y\. Zheng, D\. Fu, X\. Hu, X\. Cai, L\. Ye, P\. Lu, and P\. Liu \(2025\)Deepresearcher: scaling deep research via reinforcement learning in real\-world environments\.InProceedings of the 2025 Conference on Empirical Methods in Natural Language Processing,pp\. 414–431\.Cited by:[§2](https://arxiv.org/html/2608.10357#S2.p2.1)\.
- S\. Zhou, F\. F\. Xu, H\. Zhu, X\. Zhou, R\. Lo, A\. Sridhar, X\. Cheng, T\. Ou, Y\. Bisk, D\. Fried,et al\.\(2024\)Webarena: a realistic web environment for building autonomous agents\.InInternational Conference on Learning Representations,Vol\.2024,pp\. 15585–15606\.Cited by:[§2](https://arxiv.org/html/2608.10357#S2.p1.2)\.