Resumo Técnico: Modelos Causais Centrados em Objetos para Planejamento com Busca em Árvore Monte Carlo

Declaração do Problema

O aprendizado por reforço baseado em modelo (MBRL) alcançou um sucesso significativo ao aprender a dinâmica do ambiente para gerar experiências imaginadas, reduzindo assim a necessidade de interações no mundo real. No entanto, aprender modelos de mundo precisos continua sendo um desafio em ambientes de alta dimensão e não estacionários compostos por múltiplos objetos interagentes.

Abordagens de RL visual existentes geralmente dependem de codificadores de Redes Neurais Convolucionais (CNN) que produzem uma representação holística única de uma imagem de entrada. Essas representações monolíticas frequentemente falham em capturar a estrutura e as interações ao nível de objeto, que são críticas para a tomada de decisão. Em cenas complexas, a presença de fundos dinâmicos, entidades irrelevantes ou objetos pequenos, mas relevantes para a tarefa, pode degradar significativamente o desempenho do agente. Enquanto os humanos percebem o mundo como composto de entidades discretas, permitindo um raciocínio eficiente, os agentes de RL padrão lutam para isolar essas entidades sem supervisão explícita ou máscaras de segmentação de instância.

Metodologia: COMET

Os autores introduzem o COMET (Causal Object-centric Model for Efficient Tree search), um algoritmo de aprendizado por reforço baseado em modelo que realiza Busca em Árvore Monte Carlo (MCTS) dentro de um espaço latente estruturado por slots. O COMET integra três componentes primários:

1. Codificador Centrado em Objetos Congelado

O COMET utiliza um codificador centrado em objetos não supervisionado e congelado (ex: SLATE, DINOSAUR ou Slot Contrast) para mapear observações visuais em um conjunto de representações latentes centradas em objetos, ou "slots" ($\bar{s}_t = \{s^1_t, \dots, s^n_t\}$).

• Inicialização de Slots: Para abordar a invariância de permutação e a natureza estocástica da atenção de slots em modelos de imagens estáticas, os autores inicializam as representações de slots no tempo $t+1$ usando os slots do tempo $t$. Isso garante a consistência temporal das atribuições de objetos dentro de um episódio.
• Congelamento: O codificador é pré-treinado em dados coletados via uma política aleatória e permanece congelado durante a fase de RL para fornecer características estáveis ao nível de objeto.

2. Modelo de Mundo Baseado em Transformer com Fusão Ação-Slot

O núcleo do COMET é um modelo de mundo baseado em transformer que opera nesses slots. Diferente das abordagens padrão do tipo MuZero que processam um vetor de estado monolítico, o COMET processa uma sequência de slots de objetos.

• Vinculação Ação-Objeto: Um mecanismo inovador funde ações com slots de objetos. Em vez de comprimir todos os latentes de objetos em um único vetor para prever o próximo estado (o que cria um gargalo), o embedding da ação é concatenado com cada slot de forma independente. Esses embeddings de slot condicionados à ação são passados através de um projetor MLP compartilhado e, então, por um backbone de transformer para prever os próximos slots de estado ($\hat{s}_{t+1}$) e recompensas.
• Arquitetura: O modelo segue um pipeline de treinamento estilo UniZero usando um backbone de transformer baseado em nanoGPT. Ele prevê futuros slots e recompensas com base no conjunto de slots atual e na ação tomada.

3. Atenção Causal de Objeto para Política e Valor

As cabeças de política e valor utilizam uma arquitetura de transformer modificada equipada com atenção causal de objeto.

• Scores de Causalidade: Para cada slot latente $h^k_t$, a rede estima um score de causalidade $\alpha^k_t \in [0, 1]$, representando a probabilidade de que o objeto correspondente seja relevante para a decisão atual.
• Modulação de Atenção: Esses scores são usados para modular a atenção de produto escalar escalonada. Uma matriz de pesos aprendida $W_t$ eleva uma estrutura de grafo causal predefinida ao nível de token, permitindo que o mecanismo de atenção se concentre em entidades relevantes para a tarefa (objetos causais) enquanto suprime as irrelevantes (objetos não causais). Isso permite que o agente foque seu planejamento no subconjunto específico de objetos que atualmente interagem com o ambiente.

Principais Contribuições

1. Algoritmo COMET: A introdução de um algoritmo de MBRL baseado em MCTS que combina representações de nível de objeto congeladas com um modelo de mundo baseado em transformer para planejamento em um espaço latente estruturado.
2. Mecanismo de Fusão Ação-Slot: Um método inovador para vincular ações a slots centrados em objetos dentro de um backbone de transformer unificado. Isso permite a modelagem de mundo e a predição de política/valor centradas em objetos sem colapsar a informação do objeto em um único vetor.
3. Atenção Causal de Objeto: A integração de scores de causalidade aprendidos para modular interações de tokens, permitindo que a política foque dinamicamente em entidades relevantes para a tarefa durante a tomada de decisão.
4. Avaliação Empírica: Avaliação abrangente em oito tarefas visual e dinamicamente diversas, incluindo o benchmark de RL visual centrado em objetos, ManiSkill, Robosuite e VizDoom.

Resultados Experimentais

O COMET foi avaliado contra várias linhas de base (baselines):

• Baselines Centrados em Objetos: OCRL (PPO com codificador de objeto), OCDreamer (RSSM centrado em objeto) e SOLD (dinâmica baseada em slots não supervisionada).
• Baselines Monolíticos: UniZero (planejamento latente padrão estilo MuZero).

Descobertas de Desempenho:

• Eficiência de Amostragem: Através do conjunto diversificado de tarefas, o COMET alcança uma Pontuação Normalizada Média superior durante os estágios iniciais de treinamento em comparação tanto com métodos fortes de MCTS monolíticos quanto com outras linhas de base de RL centradas em objetos.
• Especificidades da Tarefa:
  • Em tarefas visualmente simples (ex: Objetivo de Objeto, Comparação de Objetos), o COMET aproveita representações de objetos de alta qualidade e atenção causal para identificar precisamente objetos relevantes, levando a uma convergência mais rápida.
  • Em tarefas de controle contínuo (ex: Levantamento de Bloco, Empurrar Cubo), o COMET demonstra desempenho consistente, embora não supere universalmente todas as linhas de base em cada métrica específica (ex: não superou competidores em Levantamento de Bloco).
• Generalização: O método mostra robustez em ambientes com variadas complexidades visuais, desde formas 2D em fundos monocromáticos até manipulação robótica 3D complexa e cenários de tiro em primeira pessoa.

Significância e Alegações

O artigo postula que modelar explicitamente a importância de objetos individuais para a tomada de decisão melhora o aprendizado da política em ambientes orientados a objetos. Ao introduzir um viés indutivo explícito de nível de objeto no planejamento latente estilo MuZero, o COMET permite que o modelo raciocine sobre entidades individuais e suas interações, em vez de tratar a cena como um bloco holístico.

Os autores afirmam que o COMET aborda efetivamente o desafio de representar observações em ambientes complexos de múltiplos objetos. Ao desacoplar representações de objetos e usar atenção causal para filtrar a relevância, o algoritmo alcança maior eficiência de amostragem e uma tomada de decisão mais focada. O trabalho sugere que vieses indutivos centrados em objetos são um caminho viável e eficaz para melhorar o planejamento baseado em modelo em RL visual, particularmente quando a segmentação de instância não está disponível ou quando o aprendizado de representação não supervisionada é necessário.