Search papers, labs, and topics across Lattice.
This paper introduces SINKFLEX-RL, a modular training system designed for reinforcement learning in long-horizon tool-use tasks, addressing the challenges of multi-turn on-policy rollouts and memory constraints. By integrating a Gymnasium-compatible environment wrapper with a novel sink-aware FlexAttention mechanism, the approach optimizes policy training while maintaining memory efficiency. Experimental results show significant improvements in validation rewards and a notable reduction in peak VRAM usage, highlighting the effectiveness of the proposed system in enhancing agent performance and resource management.
Validation rewards increased by 76% as SINKFLEX-RL tackles the memory limitations of long-horizon reinforcement learning tasks.
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.
