🤖 machine learning

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

This paper introduces SINKFLEX-RL, a modular training system that integrates environment interfaces, RL dataflow, and a sink-aware FlexAttention path to enable memory-efficient, long-horizon reinforcement learning for tool-using agents, demonstrating improved validation rewards and significant VRAM reductions in preliminary benchmarks.

Original authors: Zelei Cheng, Amritansh Mishra, Sambit Sahu, William Campbell

Published 2026-08-12

📖 8 min read🧠 Deep dive

CC BY 4.0

Original authors: Zelei Cheng, Amritansh Mishra, Sambit Sahu, William Campbell

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a super-smart robot to play a very complicated, multi-day board game. This isn't a game where you just roll dice and move a piece; it's a game where the robot has to talk to other players, use a toolbox of special gadgets, follow a strict rulebook, and wait days to find out if it actually won. This is the world of "Agentic AI," where computer programs act like little agents that can take actions in a digital world. The big problem is that these games get so long and complex that the robot's brain (its memory) starts to overflow. It's like trying to remember every single word of a 10-hour conversation while also doing math; eventually, you just run out of space to think. Scientists are trying to figure out how to train these robots to get better at these long games without their brains melting from the sheer amount of information they have to hold onto.

This paper introduces a new training system called SINKFLEX-RL, which is like a clever trick to help the robot remember the most important parts of the game without needing a bigger brain. The researchers built a system that combines a standard game interface, a smart way to learn from mistakes (called "Group-Relative Policy Optimization" or GRPO), and a special memory-saving technique for the robot's attention. Instead of trying to remember every single detail of the past 8,000 steps in a conversation, the system uses a "sink" mechanism. Think of this like a magical sponge in the robot's mind: it soaks up the overwhelming noise of old information but keeps the essential "sinks" (the starting points) that help the robot stay stable. By using this sponge, the robot can focus on the current move without getting crushed by the weight of the past.

The team tested this system in a simulated retail store environment where the robot has to help a customer, use tools to check inventory, and follow store policies. In these early tests, the robot's performance score went up, rising from a 0.25 to a 0.44 as it learned. But the real magic happened when they tested how much computer memory (VRAM) the system needed. When the conversation got very long (8,192 tokens), the old, standard way of thinking ran out of memory and crashed. However, the new SINKFLEX-RL system kept going, using only 25.53 GB of memory. Even at a slightly shorter length of 4,096 tokens, the new system saved a massive 19.7% of memory compared to the old way. The authors suggest that this proves it is possible to train these long-horizon agents without needing super-expensive hardware, though they note this is just a preliminary look and not a final, perfect solution.

The Problem: The Robot's Memory Leak

Imagine you are the robot in this story. You are playing a game where you have to help a customer buy a shirt. You ask them their size, they tell you, you check the inventory, they ask for a different color, you check again, and so on. After 50 turns, you have to remember the very first thing they said to make sure you didn't forget their size.

In the world of AI, this is called a "long-horizon" task. The problem is that as the conversation gets longer, the amount of memory the computer needs to hold all those words grows incredibly fast. It's like trying to carry a backpack that gets heavier every time you take a step. If the conversation gets too long (like 8,000 words), the backpack gets so heavy that the robot collapses. This is called running out of "VRAM" (Video Random Access Memory), which is the high-speed memory the computer uses to think.

The researchers noticed that standard ways of training these robots were hitting a "memory wall." The robots were getting stuck because they were trying to remember everything explicitly, which is too expensive. They needed a way to be efficient without losing the ability to reason.

The Solution: The Magic Sponge and the Team Huddle

The paper proposes a three-part solution to fix this memory leak, which they call SINKFLEX-RL.

1. The Gymnasium Wrapper (The Universal Remote)
First, they built a standard interface, like a universal remote control, that lets the robot talk to different game environments. Whether the robot is in a store, a bank, or a travel agency, this wrapper makes sure the robot knows how to ask for help, use tools, and get feedback. It's like giving the robot a standard set of rules so it doesn't have to learn a new language for every single game it plays.

2. The Group Huddle (GRPO)
Next, they changed how the robot learns. Usually, to learn, a robot might need a separate "coach" (a value model) to tell it how good a move was. But that coach takes up extra memory. Instead, this system uses a method called Group-Relative Policy Optimization (GRPO).

Imagine the robot plays the same game 10 times in a row, but with slightly different choices each time. Instead of asking a coach, the robot looks at all 10 versions of itself. It says, "Hey, version 3 got a better score than version 1, so I'll copy version 3's moves." It compares itself to its own group to figure out what worked best. This is like a study group where students compare their answers to see who got it right, without needing a teacher to grade every single paper. This saves a huge amount of memory because they don't need to train a separate coach.

3. The Magic Sponge (Sink-Aware FlexAttention)
This is the most creative part. The robot's "attention" is how it decides which words in the conversation are important. In a long conversation, the robot usually tries to look at every single word, which is slow and memory-heavy.

The researchers realized that in long conversations, the beginning of the chat acts like a "sink"—a place where the robot's attention naturally pools to stay stable. They created a special math trick (using something called FlexAttention) that lets the robot treat these "sink" words differently.

Think of it like this: If you are reading a 100-page book, you don't need to hold the entire book in your hands at once. You can hold the first page (the sink) and the current page you are reading, and let the middle pages fade away until you need them. The "sink" is like a bookmark that keeps the story's context alive without you having to carry the whole book. The system uses a "zero-value sink," which is a mathematical way of saying, "We don't need to store the actual data for the old words, we just need to know they were there to keep our balance." This allows the robot to handle long conversations (up to 8,192 tokens) without running out of memory.

The Results: Does It Work?

The team tested this new system in a simulated retail store environment. They watched the robot learn over a period of time.

Learning Progress: At the start of the training, the robot's score for helping customers was 0.25. By the end of the observed training window, the score rose to 0.44. The team also saw that the "training score" and "trajectory reward" (which are like internal checkmarks for how well the robot is doing) went up from 0.18 to 0.40 and 0.39 respectively. This suggests the robot is actually learning and getting better, though the authors are careful to say this is just a preliminary look and not a final proof that the method is perfect.
Memory Savings: The most exciting result was about memory. They tested the system with different conversation lengths:
- At 4,096 tokens, the old way needed 28.06 GB of memory. The new SINKFLEX-RL system only needed 22.52 GB. That is a savings of 5.54 GB, or 19.7%.
- At 8,192 tokens, the old way crashed completely (it ran out of memory, or "OOM"). The new system, however, kept running and only used 25.53 GB of memory.

What This Means (and What It Doesn't)

The paper shows that by combining a standard game interface, a smart group-learning method, and a memory-saving attention trick, it is possible to train robots for very long, complex tasks without needing a supercomputer. The "sink" trick acts like a sponge, soaking up the memory pressure so the robot can keep going.

However, the authors are very honest about what they haven't proven yet. They didn't test if this is the best way to learn, or if it works for every single type of robot. They didn't measure how much faster the robot learns, just that it can learn without crashing. They also noted that this is a "proof of concept"—a successful test run that shows the idea works, but it's not a finished product ready for the real world.

In short, SINKFLEX-RL is a promising new tool that helps AI agents remember long conversations without their brains exploding. It suggests that with the right tricks, we can build robots that are smart enough to handle complex, multi-day tasks, as long as we give them a way to manage their memory.

Technical Summary: SINKFLEX-RL for Efficient Long-Horizon Agentic Training

Problem Statement
Long-horizon tool-using agents operate in interactive environments where they must reason over user goals, domain policies, tool calls, and simulator states to achieve sparse, verifiable rewards only after extended trajectories. While Reinforcement Learning (RL) is a natural fit for this setting, training such agents faces significant systems challenges. Multi-turn on-policy rollouts generate long contexts that exacerbate memory bottlenecks. Specifically, standard "eager" attention mechanisms require quadratic memory growth relative to sequence length, often exhausting high-bandwidth memory (VRAM) before backpropagation can occur. Furthermore, production models often require model-specific attention variants, such as learned sink normalization or heterogeneous masks (e.g., combining causal and sliding-window constraints), which are not always exposed by fixed fused-kernel interfaces. This creates a tension between the need for memory-efficient attention kernels and the requirement for model-specific differentiable behaviors in the training path.

Methodology
The paper introduces SINKFLEX-RL, a modular training system designed to integrate environment interfaces, RL dataflows, and attention-kernel design for memory-feasible long-horizon agent training. The system is built around three core components:

Dual-Control Environment Wrapper: The system wraps dual-control environments (where both the agent and a simulated user influence the state) behind a Gymnasium-compatible interface. This allows rollouts, tools, user simulators, and reward checkers to connect to a standard RL dataflow without the trainer depending on benchmark-specific parsing logic.
Group-Relative Policy Optimization (GRPO): To update the policy without the memory overhead of a separate value model (critic), the system employs GRPO. This approach normalizes rewards within groups of sampled trajectories to compute advantages, following the rationale introduced in DeepSeekMath and the VERL-style post-training dataflow. The loss function minimizes a clipped objective regularized by a KL term against a reference model.
Sink-Aware FlexAttention Path: To address the memory wall while preserving model-specific semantics, the authors implement a custom attention path using PyTorch FlexAttention.
- Zero-Value Sink Equivalence: Instead of materializing explicit sink tokens in key/value caches (which increases memory), the system utilizes an algebraic equivalence where a zero-valued sink is represented by scaling the standard attention output. The scaling factor is derived from the log-sum-exp statistic and a learned sink logit ( $s_\eta$ ).
- Programmable Masking: The attention path composes causal constraints, sliding-window constraints, and always-visible prefixes via a block-sparse mask function.
- Gradient Flow: The implementation couples the attention output with the log-sum-exp statistic using AOTAutograd and torch.compile. This allows the compiler to generate forward and backward code for the sink parameters without materializing an $O(n^2)$ Jacobian, ensuring the sink behavior remains differentiable and connected to the training path.

Key Contributions
The paper makes three primary contributions:

System Integration: It wraps dual-control environments in a standard RL interface, enabling the connection of complex agentic tasks to a unified dataflow.
Algorithmic Adaptation: It adopts GRPO to eliminate the need for a separate value model, reducing memory and compute overhead during multi-turn rollouts.
Attention Kernel Design: It implements a sink-aware FlexAttention path that preserves learned sink scaling under causal and sliding-window masks. This design avoids eager attention-state materialization while maintaining the specific attention semantics required by the model.

Experimental Results
The authors evaluate the system through two complementary measurements:

Preliminary Retail Training (τ2-Bench): In a retail domain experiment, the validation reward (mean@1) increased from 0.25 in early training to 0.44 later in the observed window. Training score and trajectory-reward proxies also trended upward (from 0.18 to 0.40 and 0.39, respectively). The authors note these results document behavior over a single run and do not isolate the causal effect of GRPO against other optimizers or establish statistical significance.
Peak-Memory Scaling: The optimized attention path was tested against an eager baseline across sequence lengths up to 8192 tokens.
- At 4096 tokens, the optimized path reduced peak VRAM from 28.06 GB to 22.52 GB (a 19.7% reduction).
- At 8192 tokens, the optimized path completed the configuration using 25.53 GB, whereas the eager baseline ran out of memory (OOM).
- The authors explicitly state these results measure peak VRAM only and do not establish improvements in throughput, latency, or total training time.

Significance and Claims
The paper claims that the integration of environment interfaces, RL dataflow, and attention-kernel design improves the memory feasibility of training long-horizon tool-using agents. The results serve as a "systems proof of concept," demonstrating that it is possible to train agents on sequences up to 8192 tokens using a sink-aware path that retains the intended eager reference semantics.

The authors maintain a modest tone regarding their claims:

They do not claim algorithmic superiority over other RL methods.
They do not claim broad generalization across domains or model scales.
They do not claim end-to-end computational speedups or exact implementation equivalence beyond the specific memory and feasibility metrics reported.
They emphasize that the work focuses on system feasibility and preliminary training behavior, not deployment readiness.

The primary significance lies in demonstrating that a modular system combining GRPO and a programmable, sink-aware attention kernel can overcome the memory barriers that typically prevent the training of long-context agentic policies.

The Problem: The Robot's Memory Leak

The Solution: The Magic Sponge and the Team Huddle

The Results: Does It Work?

What This Means (and What It Doesn't)

Technical Summary: SINKFLEX-RL for Efficient Long-Horizon Agentic Training

More like this