ARM-FM: Automated Reward Machines via Foundation Models for Compositional Reinforcement Learning

O artigo apresenta o ARM-FM, um framework que utiliza modelos de fundação para gerar automaticamente máquinas de recompensa a partir de especificações em linguagem natural, permitindo o design composicional de recompensas e a generalização zero-shot em aprendizado por reforço.

O essencial

Imagine que você está tentando ensinar um robô a cozinhar um jantar complexo, como um bolo de chocolate.

O Problema:
Se você apenas disser ao robô: "Faça um bolo" e der um ponto de recompensa apenas quando o bolo estiver pronto e perfeito, o robô vai ficar perdido. Ele vai tentar coisas aleatórias, queimar a farinha, derrubar ovos e nunca entender por que está falhando. Isso é o que acontece na Inteligência Artificial quando usamos recompensas muito raras (chamadas de "recompensas esparsas"). O robô não tem feedback suficiente para aprender.

A Solução Antiga (e chata):
Antes, os cientistas tinham que escrever manualmente um guia passo a passo para o robô: "Primeiro, pegue a farinha. Se pegar, ganhe 1 ponto. Agora, pegue os ovos. Se pegar, ganhe 2 pontos..." Isso é como escrever um manual de instruções gigante e tedioso para cada tarefa nova. Se o robô precisa fazer algo diferente, você tem que reescrever tudo.

A Nova Ideia (ARM-FM):
Os autores deste paper criaram uma ferramenta chamada ARM-FM. Pense nela como um Arquiteto de Missões Automático que usa a inteligência de modelos de linguagem gigantes (como o GPT-4, mas muito focado em raciocínio lógico).

Aqui está como funciona, usando analogias do dia a dia:

1. O Tradutor Mágico (O Modelo de Fundação)

Em vez de você escrever o manual, você apenas conversa com o robô em linguagem natural: "Preciso que você pegue a chave amarela, abra a porta azul e depois pegue a caixa vermelha."

O ARM-FM pega essa frase simples e a transforma automaticamente em um Mapa de Tesouro (chamado de "Máquina de Recompensas").

• Antes: O robô só sabia que o "Tesouro" (a recompensa final) existia, mas não sabia onde estava.
• Agora: O robô recebe um mapa com marcos claros: "Ponto de Check-in 1: Pegou a chave amarela (+1 ponto). Ponto de Check-in 2: Porta aberta (+1 ponto). Ponto de Check-in 3: Caixa pega (+1 ponto)."

Isso transforma uma tarefa impossível em uma série de pequenas vitórias fáceis de alcançar.

2. O GPS de Voz (Embeddings de Linguagem)

A parte mais genial é como o robô "lê" esse mapa.
Imagine que cada etapa do mapa tem um etiqueta de voz.

• Quando o robô está na etapa "Pegar a chave amarela", ele ouve mentalmente: "Ok, agora preciso pegar a chave amarela".
• Se a tarefa mudar para "Pegar a chave vermelha", o robô ouve: "Ok, agora preciso pegar a chave vermelha".

Como o robô entende que "chave amarela" e "chave vermelha" são coisas muito parecidas (ambas são chaves), ele usa o que aprendeu com a amarela para aprender a pegar a vermelha muito mais rápido. É como se ele tivesse um GPS de voz que entende o contexto e adapta a rota sem precisar de um novo manual inteiro.

3. O Treinamento (A Prática)

O robô começa a treinar.

• Ele tenta pegar a chave. Se conseguir, o "Mapa" avança para o próximo passo e ele ganha pontos.
• Se ele soltar a chave, o "Mapa" avisa: "Ops, você perdeu o progresso" e ele perde pontos.
• Isso acontece em frações de segundo, milhares de vezes, até que o robô se torna um mestre em completar a tarefa.

Por que isso é incrível?

O paper mostrou que esse sistema funciona em lugares muito difíceis:

• Mundos 2D (MiniGrid): Como labirintos simples.
• Mundos 3D (Minecraft): Onde o robô precisa minerar madeira, pedra e ferro antes de conseguir um diamante. Sem o mapa, o robô ficaria preso minerando pedras para sempre. Com o mapa, ele sabe a ordem exata.
• Robótica Real: Braços robóticos que precisam pegar objetos e colocá-los em lugares específicos.

O Resultado Final

O ARM-FM é como ter um instrutor de pilotagem que:

1. Escuta o que você quer fazer em português.
2. Cria instantaneamente um plano de voo detalhado com checkpoints.
3. Ensina o piloto (o robô) a usar o que já sabe para voar em novas rotas, sem precisar reiniciar o treinamento do zero.

Isso permite que robôs aprendam tarefas complexas e longas que antes eram impossíveis, transformando a "fala humana" em "ação inteligente" de forma automática e eficiente.

Resumo técnico

1. O Problema

O Aprendizado por Reforço (RL) enfrenta um desafio central: a especificação de funções de recompensa eficazes para tarefas complexas.

• Recompensas Esparsas: Em muitos ambientes, a recompensa só é dada ao final da tarefa (ex: chegar ao objetivo), o que fornece um sinal de aprendizado insuficiente e torna a exploração ineficiente.
• Recompensas Densas Manuais: Recompensas densas criadas manualmente são propensas a "hacking de recompensa" (onde o agente explora falhas na especificação sem cumprir o objetivo real) e exigem esforço intensivo de especialistas.
• A Lacuna dos Modelos de Base (FMs): Embora os Modelos de Base (Foundation Models - FMs) sejam excelentes em interpretar instruções em linguagem natural e decompor tarefas, eles falham em traduzir esse entendimento abstrato em sinais de recompensa estruturados e executáveis necessários para o treinamento de agentes de RL. Planos de alto nível gerados por FMs muitas vezes não se "ancoram" (grounding) adequadamente no controle de baixo nível.

2. Metodologia: ARM-FM

O trabalho propõe o ARM-FM (Automated Reward Machines via Foundation Models), um framework que automatiza o design de recompensas compostas utilizando a capacidade de raciocínio dos FMs. A abordagem baseia-se em Máquinas de Recompensa (RMs), uma formalização baseada em autômatos que divide tarefas complexas em sub-objetivos sequenciais.

O framework possui dois componentes principais:

A. Máquinas de Recompensa Alinhadas à Linguagem (LARMs)

O ARM-FM utiliza FMs para gerar automaticamente RMs completas a partir de descrições de tarefas em linguagem natural e observações visuais.

• Geração Iterativa: O sistema emprega um loop de auto-aperfeiçoamento com pares de FMs (gerador e crítico) para refinar a especificação da RM. Humanos podem intervir opcionalmente para correções.
• Estrutura da LARM: Cada estado da máquina de recompensa ($u \in U$) é associado a:
  1. Uma instrução em linguagem natural ($l_u$) descrevendo o sub-objetivo.
  2. Um vetor de incorporação (embedding) ($z_u = \phi(l_u)$) gerado a partir dessa instrução.
  3. Funções de rotulagem (labeling functions) em Python que mapeiam estados do ambiente para eventos de transição na RM.
• Espaço de Habilidades Semântico: Ao injetar os embeddings das instruções de linguagem no espaço de estados do agente, cria-se um espaço de habilidades onde sub-tarefas semanticamente similares (ex: "pegar chave azul" vs. "pegar chave vermelha") estão próximas no espaço vetorial, permitindo transferência de conhecimento.

B. Treinamento de RL com LARMs

O agente de RL é treinado em um MDP aumentado ($M'$), que é o produto cartesiano dos estados do ambiente ($S$) e os estados da RM ($U$).

• Política Condicionada: A política do agente $\pi(s_t, z_{u_t})$ recebe não apenas a observação do ambiente, mas também o embedding do estado atual da RM. Isso torna o agente ciente do seu sub-objetivo ativo.
• Sinal de Recompensa: A recompensa total é a soma da recompensa do ambiente ($R_t$) e a recompensa estruturada da RM ($R^{RM}_t$). A RM fornece recompensas densas e incrementais para cada sub-objetivo concluído, guiando o agente através de tarefas de longo horizonte.

3. Contribuições Principais

1. Framework de Geração Automática de Especificações: Desenvolvimento de um método para gerar especificações completas de tarefas (estrutura do autômato, funções de rotulagem executáveis e instruções em linguagem natural) diretamente a partir de prompts de linguagem natural, sem necessidade de demonstrações humanas.
2. Espaço de Habilidades Compartilhado: Introdução de uma técnica que utiliza embeddings de linguagem dos estados da RM para permitir a reutilização de experiências e transferência de políticas entre tarefas relacionadas, criando um espaço de habilidades semanticamente fundamentado.
3. Validação Empírica Abrangente: Demonstração de que o método resolve tarefas complexas de longo horizonte em diversos domínios (Grid Worlds, ambientes 3D procedurais e robótica de controle contínuo), superando métodos de RL padrão e abordagens baseadas em intrinsic motivation.

4. Resultados Experimentais

Os autores avaliaram o ARM-FM em quatro ambientes distintos:

• MiniGrid & BabyAI (Tarefas Esparsas): O agente com LARM superou consistentemente baselines como DQN puro, DQN com motivação intrínseca (ICM) e agentes ReAct (LLM puro). Em tarefas complexas como UnlockToUnlock e KeyCorridor, onde baselines falharam completamente, o ARM-FM alcançou recompensas quase perfeitas.
• Craftium (Ambiente 3D Procedural): Em um ambiente estilo Minecraft para coleta de recursos (madeira, pedra, ferro, diamante), um agente PPO padrão falhou em fazer progresso. O agente guiado por LARM gerada automaticamente completou a sequência inteira de tarefas com sucesso, demonstrando escalabilidade em ambientes visualmente complexos.
• Meta-World (Manipulação Robótica): Em tarefas de controle contínuo, o ARM-FM automatizou a engenharia de recompensas, superando agentes que aprendiam apenas com recompensas esparsas, alcançando altas taxas de sucesso em tarefas como montagem e colocação em prateleiras.
• Generalização Zero-Shot (XLand-MiniGrid):
  • Ablação: O estudo mostrou que tanto as recompensas estruturadas quanto os embeddings de estado são essenciais para o aprendizado multi-tarefa.
  • Generalização: Um agente treinado em um conjunto de tarefas (A e B) conseguiu resolver uma nova tarefa composta (C) sem treinamento adicional, desde que os sub-objetivos de C fossem semanticamente familiares (compartilhando embeddings próximos no espaço de habilidades).

5. Significado e Conclusão

O ARM-FM preenche uma lacuna crítica entre o raciocínio semântico de alto nível dos Modelos de Base e o controle de baixo nível do Aprendizado por Reforço.

• Interpretabilidade e Controle: Ao gerar RMs baseadas em linguagem, o sistema torna as políticas de RL interpretáveis e permite que humanos inspecionem e refinem as especificações da tarefa.
• Eficiência de Amostra: A decomposição automática de tarefas em sub-objetivos com recompensas densas resolve o problema da esparsidade de recompensa, permitindo que agentes aprendam tarefas que seriam intratáveis para métodos convencionais.
• Composicionalidade: A abordagem estabelece uma nova classe de agentes de RL capazes de traduzir intenções humanas abstratas em comportamentos competentes, generalizáveis e interpretáveis, através de uma biblioteca de habilidades reutilizáveis.

Em suma, o trabalho demonstra que a automação do design de recompensas via FMs não apenas acelera o desenvolvimento de agentes de RL, mas também habilita a generalização zero-shot e a transferência de habilidades em domínios complexos e variados.