Resumo Técnico: KAGE-Bench

Declaração do Problema

Agentes de aprendizado por reforço (RL) treinados em observações de pixels de alta dimensão são notoriamente frágeis quando confrontados com mudanças de distribuição puramente visuais, mesmo quando a dinâmica latente, os kernels de transição e as funções de recompensa permanecem inalterados. Embora os agentes possam convergir e performar bem em distribuição (in-distribution), eles frequentemente falham catastroficamente sob mudanças de aparência, iluminação ou ruído visual. Um obstáculo primário para diagnosticar essas falhas é a falta de benchmarks de avaliação que possam isolar fontes específicas de mudança visual. Benchmarks existentes (ex: Procgen, Distracting Control Suite) frequentemente entrelaçam múltiplas fontes de variação — como aparência de fundo, geometria, dinâmica e distratores — tornando impossível atribuir as lacunas de desempenho a fatores visuais específicos. Além disso, muitos ambientes de RL baseados em pixels são computacionalmente caros, dificultando estudos de ablação em larga escala e testes de hipóteses sistemáticos.

Metodologia

Os autores introduzem o KAGE-Env e o KAGE-Bench para abordar essas limitações através de uma filosofia de design de "eixos conhecidos" (known-axis).

KAGE-Env: Um Plataforma de Plataforma Fatorizada

O KAGE-Env é um ambiente de plataforma 2D nativo em JAX projetado para fatorizar o processo de observação em eixos visuais independentemente controláveis, mantendo o problema de controle latente fixo.

• Arquitetura: O ambiente é implementado inteiramente em JAX com compilação JIT de ponta a ponta e execução vetorizada (usando vmap e lax.scan). Isso permite um paralelismo massivo, alcançando até 33 milhões de passos de ambiente por segundo em uma única GPU com $2^{16}$ ambientes paralelos.
• Fatorização: O kernel de observação $O_\xi(\cdot|s)$ é parametrizado por uma configuração visual $\xi$. Crucialmente, o kernel de transição $P$ e a função de recompensa $r$ são invariantes a $\xi$. Isso garante que qualquer diferença de desempenho entre o treinamento e a avaliação surja apenas de como uma política baseada em pixels $\pi(a|o)$ responde a diferentes renderizações do mesmo estado latente $s$.
• Configuração: O ambiente é controlado via um único arquivo YAML, expondo 93 parâmetros explicitamente controláveis em categorias como imagens/cores de fundo, aparência do agente/sprites, NPCs, distratores, filtros fotométricos (brilho, matiz, ruído), efeitos de iluminação e layout.

KAGE-Bench: O Protocolo de Benchmark

Construído sobre o KAGE-Env, o KAGE-Bench define um protocolo de avaliação padronizado composto por seis suítes de "eixos conhecidos" e 34 pares de configuração de treino-avaliação.

• Isolamento: Em cada par, a configuração de treinamento ($\xi_{train}$) e a configuração de avaliação ($\xi_{eval}$) diferem apenas em um subconjunto designado de parâmetros de renderização (um eixo específico), enquanto todos os outros parâmetros e a dinâmica subjacente permanecem idênticos.
• Fundamentação Teórica: O artigo formaliza que a generalização visual sob estas condições é equivalente a uma mudança na política de estado induzida $\pi_\xi$. Ao marginalizar o kernel de observação, a política de pixels $\pi(a|o)$ induz uma distribuição de ação condicionada ao estado $\pi_\xi(a|s)$. O benchmark demonstra que as lacunas de desempenho são exatamente as diferenças entre estas políticas induzidas agindo no mesmo MDP latente.
• Métricas: Além do retorno episódico padrão, o benchmark reporta métricas de nível de trajetória: distância percorrida, progresso normalizado em direção ao objetivo e sucesso binário da tarefa. Esta distinção é vital, pois o retorno sozinho pode ocultar falhas onde o agente faz progresso frontal, mas falha em completar a tarefa.

Principais Contribuições

1. KAGE-Env: Um ambiente de RL de alto rendimento, nativo em JAX, com 93 parâmetros visuais controláveis, capaz de simulação vetorizada a até 33M de passos/segundo em uma única GPU.
2. KAGE-Bench: Um protocolo de benchmark que isola mudanças de distribuição visual por construção, apresentando seis suítes (Agente, Fundo, Distratores, Efeitos, Filtros, Layout) e 34 pares específicos de treino-avaliação.
3. Diagnóstico Empírico: Uma avaliação sistemática da generalização visual usando uma baseline padrão PPO-CNN, quantificando como diferentes eixos visuais impactam o desempenho e identificando classes específicas de mudanças que induzem degradação severa.

Resultados

Usando uma baseline PPO-CNN padrão treinada a partir de frames RGB únicos, os autores observaram falhas dependentes de eixo fortes:

• Degradação Severa: Filtros (ex: mudança de matiz, saturação, brilho) e Efeitos (ex: intensidade de iluminação, contagem) causaram as colisões de desempenho mais significativas. As taxas de sucesso caíram 86,8% e 80,5% respectivamente, muitas vezes enquanto o agente mantinha algum movimento frontal.
• Mudanças de Fundo: Mudanças na imagem ou cor do fundo levaram a uma queda de 53,3% na taxa de sucesso e uma queda de 30,5% na distância percorrida, indicando que mudanças de fundo prejudicam tanto a locomoção quanto a conclusão da tarefa.
• Mudanças Benignas: Mudanças na aparência do agente (ex: mudar o sprite de um esqueleto para um palhaço) foram comparativamente benignas, causando apenas uma queda de 21,1% na taxa de sucesso.
• A Lacuna "Movimento vs. Conclusão": Várias mudanças (notadamente Distratores e Layout) preservaram o movimento frontal (pequenas lacunas de distância de ~3-4%), mas causaram quedas massivas na conclusão da tarefa (lacunas de Taxa de Sucesso de 30-60%). Isso destaca que métricas baseadas em retorno sozinhas podem mascarar falhas de generalização onde o agente se move, mas falha em resolver a tarefa.
• Baselines de Robustez: Os autores testaram vários métodos de robustez (RAD-PPO com cutout-color, crop-resize, CURL, DrQ, SVEA). Embora o RAD-PPO com cutout-color tenha mostrado alguma melhoria em configurações difíceis, métodos que dependem de aumentos geométricos (crop-resize) falharam em convergir, provavelmente porque interromperam as posições espaciais absolutas necessárias para o tempo preciso de pulo e desvio de obstáculos neste cenário de plataforma.

Significância e Alegações

O artigo afirma que o KAGE-Bench fornece um framework rápido, reprodutível e diagnóstico para avaliar RL baseado em pixels sob variação visual controlada.

• Atribuição: Ao fatorizar o processo de observação, o benchmark permite a atribuição precisa de falhas de generalização a fatores visuais específicos, removendo os efeitos de confusão de dinâmicas entrelaçadas ou mudanças na estrutura da tarefa encontrados em outros benchmarks.
• Escalabilidade: A implementação nativa em JAX aborda o gargalo computacional da avaliação de generalização visual, tornando viáveis varreduras exaustivas sobre parâmetros visuais.
• Utilidade Diagnóstica: Os resultados demonstram que a dificuldade de generalização visual não é uniforme; ela varia substancialmente entre os eixos. O benchmark serve como uma ferramenta para identificar quais fatores visuais específicos (ex: mudanças fotométricas vs. aparência do agente) são mais críticos para uma política aprender invariâncias.

Os autores posicionam o KAGE-Bench não como um proxy abrangente para todos os desafios do mundo real do RL, mas como uma ferramenta de diagnóstico controlada para apoiar uma análise mais sistemática da robustez visual e para guiar o desenvolvimento de algoritmos capazes de lidar com mudanças visuais de eixos conhecidos específicos.