BATON, a new training framework for AI agents, tackles two overlooked variables in how those agents learn from trial and error.
Researchers behind the paper point out that reinforcement learning for AI agents involves two separate decisions: how to credit individual actions within a single attempt at a task, and how to weigh entire attempts against each other when updating the model. BATON handles the first with a Bayesian method that estimates which actions likely caused a good or bad outcome. It handles the second with trajectory mass normalization, which gives every complete attempt equal weight during training regardless of how many steps it took. The team tested BATON on top of two existing training algorithms, GRPO and GiGPO, across three benchmark environments: the text-based household simulator ALFWorld, the online shopping simulator WebShop, and the search task SearchQA.
Both changes improved performance on their own, and combining them produced the best results across every model size tested. That is a useful data point for anyone building agentic AI systems, where reinforcement learning is a common way to teach a model to use tools and take multi-step actions. Most existing training tweaks address only one of these two variables at a time, so the paper's split-and-recombine framing is a reasonable engineering fix rather than a wholesale rewrite of how agent training works.
Still, this is a preprint tested only on toy text environments, not production-scale agents pulling real levers in the real world.