AstraFlow: Aprendizado por Reforço Orientado a Fluxo de Dados para LLMs Agênticos

Resumo

O aprendizado por reforço (RL) é cada vez mais utilizado para aprimorar as capacidades de raciocínio, codificação e uso de ferramentas de grandes modelos de linguagem (LLMs), mas o RL agentivo continua proibitivamente caro. Escalar o RL para LLMs agentivos requer o suporte a cargas de trabalho complexas, incluindo treinamento colaborativo com múltiplas políticas, ao mesmo tempo em que se utilizam eficientemente recursos computacionais elásticos, heterogêneos e de múltiplas regiões. Sistemas existentes de RL para LLMs oferecem suporte a algumas dessas capacidades, mas cada nova extensão frequentemente exige engenharia de sistema dedicada. Esse ônus decorre de arquiteturas de controle centradas no treinador e da falta de abstrações fundamentadas para os componentes dos sistemas de RL. Para superar essas limitações, propomos o AstraFlow, um sistema de RL orientado a fluxo de dados que substitui o controle convencional centrado no treinador por abstrações de componentes fundamentadas. No AstraFlow, serviços de coleta de experiência (rollout), gerenciamento de fluxo de dados e treinamento são desacoplados em componentes autônomos, permitindo que o sistema suporte nativamente cargas de trabalho complexas de RL agentivo com múltiplas políticas e explore eficientemente recursos computacionais diversificados. Avaliamos o AstraFlow em cargas de trabalho de matemática, código, busca e AgentBench, demonstrando que o mesmo sistema suporta treinamento com múltiplas políticas, escalabilidade elástica, execução heterogênea em múltiplas regiões e algoritmos de dados combináveis, sem necessidade de alterações no código do sistema. No treinamento colaborativo com múltiplas políticas, o AstraFlow alcança precisão comparável ou superior à de sistemas RL existentes, ao mesmo tempo que acelera o tempo de treinamento em 2,7 vezes.

English

Reinforcement learning (RL) is increasingly used to improve the reasoning, coding, and tool-use capabilities of large language models, but agentic RL remains prohibitively expensive. Scaling RL to agentic LLMs requires supporting complex workloads, including multi-policy collaborative training, while efficiently using elastic, heterogeneous, and cross-region compute resources. Existing LLM RL systems support some of these capabilities, but each new extension often requires dedicated system engineering. This burden arises from trainer-centered control architectures and the lack of principled abstractions for RL system components. To address these limitations, we propose AstraFlow, a dataflow-oriented RL system that replaces conventional trainer-centered control with principled component abstractions. In AstraFlow, rollout services, dataflow management, and training are decoupled into autonomous components, enabling the system to natively support complex multi-policy agentic RL workloads and efficiently exploit diverse compute resources. We evaluate AstraFlow across math, code, search, and AgentBench workloads, showing that the same system supports multi-policy training, elastic scaling, heterogeneous cross-region execution, and composable data algorithms without system-level code changes. In multi-policy collaborative training, AstraFlow achieves comparable or better accuracy than existing RL systems while speeding up training time by 2.7x.