SPO++: Stream-Aligned Policy Optimization for Asynchronous Agentic RL

No campo em rápida evolução da inteligência artificial, pesquisadores estão ensinando programas de computador a resolver problemas complexos permitindo que eles tentem, falhem e aprendam com os resultados. Esse processo, conhecido como aprendizagem por reforço, é particularmente útil para treinar grandes modelos de linguagem para atuarem como agentes que podem usar ferramentas, navegar em ambientes virtuais ou resolver problemas matemáticos. Um desafio central neste trabalho é como medir o sucesso quando o caminho para uma resposta é longo e imprevisível. Métodos tradicionais frequentemente exigem que o sistema gere muitas tentativas diferentes para o mesmo problema simultaneamente, esperando que todas as tentativas terminem antes que ele possa aprender com os resultados. Isso cria um gargalo: se uma tentativa demora muito para ser concluída porque fica presa ou tenta muitas ferramentas, todo o processo de aprendizagem pausa, esperando que essa tentativa mais lenta alcance as demais.

Para resolver essa ineficiência, um método mais novo chamado Otimização de Política de Fluxo Único (Single-stream Policy Optimization) foi desenvolvido. Em vez de esperar por um grupo de tentativas, este método permite que o sistema aprenda com uma tentativa de cada vez, usando uma memória persistente de sucessos e falhas passados para guiar decisões futuras. No entanto, uma equipe de pesquisadores liderada por Kai Ruan e Jinghao Lin descobriu que, embora este método tenha removido o tempo de espera, ele introduziu um descompasso sutil, mas significativo, na forma como o computador calculava seu progresso. Eles descobriram que a maneira como o sistema media a média de suas recompensas não se alinhava com a forma como ele realmente processava as etapas de uma solução. Ao corrigir esse alinhamento, eles criaram uma versão aprimorada do método, que nomearam SPO++, que permite que a inteligência artificial aprenda de forma significativamente mais rápida e eficiente.

O cerne do problema residia em como o computador lidava com a extensão de suas respostas. No método original, o sistema calcularia uma pontuação única para toda uma tentativa, como se um robô tivesse sucesso ao limpar um quarto ou um resolvedor de matemática encontrasse a resposta correta. Ele então distribuiria essa pontuação única por cada palavra ou "token" que o modelo gerasse durante essa tentativa. Isso parecia lógico, mas criava uma distorção. Se uma tentativa fosse muito longa e outra curta, a tentativa longa diluiria sua pontuação por muitas palavras, enquanto a tentativa curta concentraria sua pontuação em menos palavras. Quando o sistema tentava aprender com essas pontuações, a extensão da resposta alterava silenciosamente o ponto central da aprendizagem, fazendo com que o modelo otimizasse as coisas erradas. Era como se o sistema estivesse tentando equilibrar uma balança, mas os pesos na balança mudassem dependendo de quantos itens eram colocados nela, em vez de dependerem do valor dos itens em si.

Os pesquisadores identificaram duas áreas específicas onde esse desalinhamento ocorria. Primeiro, o sistema estava rastreando quando recebia os resultados de uma tentativa, em vez de quando a tentativa era realmente gerada. Em um sistema assíncrono, onde as tarefas são enviadas e concluídas em velocidades diferentes, a ordem em que os resultados chegam é frequentemente aleatória e dependente da velocidade da rede ou da carga do computador. O método original usava essa ordem de chegada para atualizar sua memória, o que significava que o sinal de aprendizagem era influenciado pelo tempo do sistema de computação, e não pela lógica da tarefa em si. Segundo, e mais criticamente, o método de calcular a média das pontuações não levava em conta o fato de que o modelo aprende com cada palavra que gera, não apenas com o resultado final. Os pesquisadores perceberam que, para corrigir o processo de aprendizagem, precisavam padronizar as pontuações com base no número de palavras de ação geradas, garantindo que o sinal de recompensa correspondesse ao volume real de trabalho que o modelo estava realizando.

Para abordar esses problemas, a equipe introduziu o SPO++, que implementa duas mudanças principais no processo de aprendizagem. Primeiro, eles reorganizaram o sistema de memória para rastrear o "evento de política" (policy event), que é o momento específico em que uma solicitação foi enviada, em vez de quando o resultado retornou. Isso garante que a memória do sistema sobre sucessos passados esteja vinculada ao estado do modelo no momento em que a tarefa foi criada, tornando o sinal de aprendizagem consistente, independentemente de quanto tempo a tarefa levou para terminar. Segundo, eles mudaram a forma como as pontuações são calculadas em média. Em vez de tratar cada tentativa como uma unidade única, o novo método calcula a média das pontuações com base no número total de palavras de ação geradas em todas as tentativas. Isso garante que o sinal de aprendizagem esteja perfeitamente alinhado com a forma como o modelo realmente atualiza seu conhecimento, palavra por palavra.

Os resultados dessas mudanças foram testados em dois tipos diferentes de tarefas: um conjunto de 128 tarefas domésticas simuladas onde um agente deve mover objetos para locais específicos, e um conjunto de dados de 1.500 problemas matemáticos que requerem o uso de uma calculadora Python. Os pesquisadores realizaram experimentos usando dois tamanhos diferentes de modelos de linguagem, um com 0,8 bilhão de parâmetros e outro com 2 bilhões de parâmetros. Em todos os testes, o novo método, SPO++, aprendeu mais rápido que a abordagem original. Na tarefa de tarefas domésticas, a melhoria foi substancial, com o novo método alcançando uma recompensa total significativamente maior ao longo do treinamento. Nos problemas matemáticos, os ganhos foram menores, mas ainda consistentes, mostrando que o modelo atingiu um nível de desempenho mais alto mais rapidamente. Os pesquisadores descobriram que a parte mais poderosa do novo método foi a mudança na forma como as pontuações eram calculadas em média, que, sozinha, foi responsável pela maior parte da melhoria.

Essas descobertas sugerem que, no mundo complexo de treinamento de agentes de inteligência artificial, os detalhes de como os dados são processados podem ser tão importantes quanto os próprios dados. Ao garantir que a maneira como um sistema mede seu progresso corresponda à maneira como ele realmente aprende, pesquisadores podem desbloquear ganhos significativos de eficiência sem a necessidade de mais poder computacional ou modelos maiores. O trabalho demonstra que mesmo pequenos desalinhamentos no processo de aprendizagem podem atrasar o progresso, e que corrigir esses descompassos permite que o sistema foque sua energia em resolver o problema, em vez de compensar uma medição falha. À medida que a inteligência artificial continua a enfrentar tarefas mais difíceis e variadas, métodos como o SPO++ oferecem um caminho mais claro, garantindo que cada passo dado pelo modelo seja medido com precisão e contribua efetivamente para o seu crescimento.