StepPO: Step-Aligned Policy Optimization for Agentic Reinforcement Learning

Resumo Técnico: StepPO

1. O Problema: A Inadequação do Paradigma Token-Cêntrico

O artigo identifica uma lacuna fundamental na aplicação de Aprendizado por Reforço (RL) para Agentes de Inteligência Artificial baseados em Grandes Modelos de Linguagem (LLMs).

• Contexto: Agentes gerais (como OpenClaw e Claude Code) realizam tarefas complexas de múltiplas voltas (multi-turn), envolvendo tomada de decisão, uso de ferramentas e interação com ambientes.
• A Falha Atual: Os métodos tradicionais de RL para LLMs (como RLHF e RLVR) utilizam um Processo de Decisão de Markov (MDP) em nível de token. Nesses modelos, a unidade de ação é a geração de um único token.
• Desafios Específicos de Agentes:
  • Recompensas Atrasadas e Esparsas: O sucesso de uma ação (ex: chamar uma ferramenta correta) só é recompensado após várias interações futuras.
  • Inconsistência de Retokenização: Em pipelines atuais, o histórico de interação é frequentemente armazenado como texto e depois re-tokenizado para treinamento. Como a mapeamento token-texto-token não é reversível em subpalavras, isso causa "deriva de retokenização", quebrando a consistência entre a geração (rollout) e o treinamento.
  • Granularidade Incorreta: Atribuir crédito (recompensa) a cada token individual é muito ruidoso e local, enquanto atribuir crédito a toda a trajetória é muito grosseiro para identificar decisões intermediárias cruciais.

2. Metodologia: A Abordagem StepPO

O StepPO propõe uma mudança de paradigma de um MDP em nível de token para um MDP em nível de etapa (step-level), alinhando a otimização com a estrutura natural da interação do agente.

• Reformulação do MDP (Nível de Etapa):

  • Em vez de $s_t$ e $a_t$ representarem um token, eles representam uma etapa completa de interação.
  • Uma etapa consiste em: Observação ($s_t$) $\rightarrow$ Ação Completa ($a_t$, que pode ser uma sequência de tokens, chamada de ferramenta ou raciocínio) $\rightarrow$ Recompensa ($r_t$) $\rightarrow$ Próxima Observação ($s_{t+1}$).
  • Isso separa claramente o que pertence à política (escolher a ação) e o que pertence ao ambiente (executar a ação e retornar feedback).

• Representação de Trajetória Estruturada:

  • Abandona-se a representação de texto plano ou fluxo de tokens achatado.
  • Adota-se uma Representação Estruturada em Nível de Etapa, onde cada passo é uma unidade autônoma contendo IDs de prompt, IDs de resposta, recompensa e metadados. Isso elimina a deriva de retokenização e preserva a consistência do espaço de tokens dentro de cada ação.

• Atribuição de Crédito em Nível de Etapa (Step-Level Credit Assignment):

  • O StepPO redefine a estimativa de vantagem (GAE - Generalized Advantage Estimation) para operar sobre etapas, não tokens.
  • A política pode ainda ser fatorada internamente por tokens (para geração), mas a propagação da recompensa e a atualização da política ocorrem no nível da etapa.
  • A função de perda do ator (actor) é calculada com base na razão de importância da etapa inteira ($w_t$) multiplicada pela vantagem da etapa ($A_t$).

• Arquitetura de Sistema:

  • O sistema suporta treinamento assíncrono, onde engines de rollout e treinamento operam de forma desacoplada.
  • Utiliza um gateway e um datapool para gerenciar dados de agentes heterogêneos (caixa-branca e caixa-preta).
  • Implementa otimização de eficiência através da reutilização de prefixos compartilhados entre etapas, reduzindo computação redundante em contextos longos.

3. Contribuições Principais

1. Proposta Teórica (StepPO): Defende que o MDP e a atribuição de crédito em RL para agentes devem ser alinhados à granularidade da "etapa de interação", não ao token.
2. Solução de Consistência de Dados: Introduz uma representação de dados nativa de etapa que resolve o problema da deriva de retokenização, garantindo que o treinamento reflita exatamente o comportamento gerado.
3. Sistema de Treinamento Escalável: Apresenta uma arquitetura de sistema (exemplificada nos projetos Agent-R1 e Claw-R1) que suporta agentes heterogêneos, gerenciamento de dados centralizado e treinamento assíncrono eficiente.
4. Validação Empírica: Demonstra que a otimização em nível de etapa supera a otimização em nível de token em tarefas complexas.

4. Resultados Experimentais

• Configuração: O modelo foi testado no conjunto de dados HotpotQA (respostas a perguntas que exigem raciocínio multi-estágio e busca de informações), comparando o StepPO com o PPO tradicional em nível de token.
• Desempenho:
  • O StepPO superou consistentemente o PPO token-level durante todo o treinamento.
  • A curva de aprendizado do StepPO atingiu picos de pontuação mais altos e manteve um platô mais forte nas fases média e tardia.
  • Conclusão: Alinhar a otimização com a estrutura de interação do agente fornece um sinal de aprendizado mais eficaz para tarefas de longo horizonte do que distribuir crédito apenas em nível de token.

5. Significado e Impacto

O artigo estabelece que o RL para Agentes (Agentic RL) não é apenas uma extensão do RL para LLMs, mas um problema distinto que exige uma reestruturação fundamental.

• Mudança de Paradigma: A transição de "token" para "etapa" é essencial para capturar a causalidade real do comportamento do agente (decisão $\rightarrow$ ação $\rightarrow$ feedback do ambiente).
• Viabilidade Prática: O StepPO oferece uma lente prática para construir agentes mais robustos, resolvendo problemas de estabilidade de treinamento e escalabilidade em sistemas reais.
• Futuro: O trabalho sugere que o futuro do treinamento de agentes dependerá de sistemas que integrem representações de dados estruturadas, otimização alinhada à etapa e infraestrutura assíncrona, como demonstrado pela evolução dos projetos Agent-R1 para Claw-R1.

Em resumo, o StepPO é um framework que alinha a teoria (MDP), a representação de dados e a otimização (RL) ao nível de interação semântica do agente, superando as limitações dos métodos baseados em tokens para tarefas complexas e multi-etapa.