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

2026-08-10

AI Summary

This work addresses the challenge of long-horizon tool-using agents, which must perform complex reasoning across user goals, domain-specific strategies, and delayed rewards—a setting where conventional reinforcement learning suffers from excessive memory consumption due to long contexts and attention mechanisms. The authors propose SINKFLEX-RL, the first framework to co-design the environment interface, reinforcement learning data pipeline, and sink-aware attention kernel. It integrates a Gymnasium-wrapped environment, a VERL-style data flow, grouped relative policy optimization, and the FlexAttention mechanism, enabling efficient training under both causal and sliding-window masking. Evaluated on the Tau2Bench retail task, SINKFLEX-RL improves reward from 0.25 to 0.44, reduces GPU memory usage by 19.7% (from 28.06 GB to 22.52 GB) at 4096 tokens, and achieves the first successful training at 8192 tokens—where baseline methods fail due to out-of-memory errors.

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 present SINKFLEX-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 Tau2Bench retail run, validation reward (mean@1) rises from 0.25 early in training to $0.44$ later 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 from 28.06GB to 22.52GB at 4096 tokens, a $19.7\%$ reduction, and runs the measured 8192-token configuration using $25.53$~GB 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.

Problem

• long-horizon
• tool-use
• reinforcement learning
• memory efficiency
• agent training

Innovation

• SINKFLEX-RL
• long-horizon reinforcement learning
• tool-use agentic tasks
• sink-aware attention
• memory-efficient training