Efficient reinforcement learning for long-horizon tool-use agentic tasks
A modular training system for RL in dual-control tool-use environments.
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 τ2-Bench 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.06 GB to 22.52 GB 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.
Latest publications
Know when to stop: Segment-level credit assignment for reducing overthinking
A benchmark which assigns segmentlevel credit based on whether each reasoning segment leads toward or away from correctness.
COLMDecomposing the delta: what do models actually learn from preference pairs?
An investigation of two distinct notions of quality delta in preference data.
COLMRefusal tokens: a simple way to calibrate refusals
Refusal tokens enable controlling a single model's refusal rates without the need of any further fine-tuning.
COLM