Resumo Técnico: TaskSense

Definição do Problema

Os modelos de mundo para controle visual tipicamente dependem da reconstrução de observações completas para aprender estados latentes compactos. Embora eficaz nessa abordagem em benchmarks padronizados com fundos estáticos, ela incentiva implicitamente as representações latentes a preservar informações em toda a entrada visual, independentemente de sua relevância para a tarefa de controle. Em ambientes onde o conteúdo relevante para a tarefa ocupa apenas uma pequena fração da observação ou onde há ruído visual e distratores, esse objetivo de reconstrução consome uma capacidade representacional valiosa em conteúdo visual irrelevante para a tarefa. Esse descompasso entre o objetivo de reconstrução visual e os objetivos de controle enviesa a dinâmica latente, diluindo os sinais de aprendizado para características relevantes ao controle e degradando severamente o desempenho subsequente sob distrações visuais. Por outro lado, confiar exclusivamente na supervisão da tarefa (por exemplo, predição de recompensa) geralmente fornece um sinal fraco demais para aprender representações visuais ricas. As soluções existentes que eliminam totalmente a reconstrução sacrificam as capacidades generativas necessárias para o planejamento baseado em imaginação, enquanto métodos que aplicam atenção ou mascaramento após a extração de características falham em controlar qual informação entra no modelo de dinâmica latente em primeiro lugar.

Metodologia

Os autores propõem o TaskSense, um framework de modelagem de mundo centrado na tarefa que impõe a relevância da tarefa antes da codificação latente. O TaskSense aumenta a arquitetura DreamerV3 com um mecanismo de atenção espacial estocástica diferenciável condicionado ao estado de crença latente anterior.

Componentes Principais

1. Atenção Espacial Estocástica Condicionada ao Latente:

   • Em vez de prever a atenção para cada pixel, o modelo define a atenção sobre uma grade tênue $\Lambda$ que cobre a observação.
   • Um MLP leve prevê os logits de atenção baseando-se no estado latente anterior $s_{t-1}$ (especificamente o estado oculto $h_{t-1}$ e o latente estocástico $z_{t-1}$), permitindo que o modelo decida onde focar com base em sua compreensão evolutiva da tarefa, em vez de apenas na aparência visual instantânea.
   • O campo de atenção é modelado como uma coleção de variáveis estocásticas independentes usando uma distribuição Binary Concrete (Relaxed Bernoulli). Isso permite que múltiplas regiões sejam atendidas ou ignoradas simultaneamente, evitando as restrições competitivas da atenção softmax.
   • Durante o treinamento, máscaras diferenciáveis são amostradas; durante a inferência, utiliza-se a média determinística.

2. Codificação Consciente de Atenção (Attention-Aware Encoding):

   • A máscara de atenção amostrada é redimensionada (upsampled) e usada para modular a observação via multiplicação elemento a elemento, atenuando as regiões de baixa atenção.
   • Para resolver a ambiguidade entre pixels de baixa intensidade e pixels de alta intensidade suprimidos, o mapa de atenção redimensionado é concatenado como um canal de entrada adicional ao codificador.

3. Objetivo Auxiliar de Dinâmica Inversa:

   • Para evitar que o regularizador de esparsidade colapse a atenção para uma solução trivial (ignorar todas as observações), um objetivo auxiliar é introduzido.
   • O modelo prevê a ação $a_{t-1}$ executada entre dois passos de tempo consecutivos usando a diferença temporal das características codificadas ($f_t - f_{t-1}$).
   • Este objetivo fundamenta o mecanismo de atenção em regiões relevantes ao controle, incentivando a retenção de características visuais que são preditivas do comportamento do agente.

4. Reconstrução Condicionada à Atenção:

   • Em vez de reconstruir a observação completa, o decodificador reconstrói apenas as regiões atendidas ($o^{att}_t$).
   • O decodificador é condicionado ao mapa de atenção amostrado (via stop-gradient) para garantir que ele saiba quais regiões foram retidas, evitando que o estado latente seja penalizado por informações perdidas que foram intencionalmente filtradas.

Objetivo de Treinamento

A função de perda total combina as perdas padrão de dinâmica latente do DreamerV3 ($L_{dyn}$) com os novos componentes:
$$L_{WM} = L_{dyn} + L_{rec} + \lambda_{sparse}L_{sparse} + \lambda_{inv}L_{inv}$$
Onde $L_{rec}$ é a perda de reconstrução sobre as observações filtradas, $L_{sparse}$ incentiva mapas de atenção compactos, e $L_{inv}$ é a perda de dinâmica inversa.

Contribuições Principais

1. Framework TaskSense: Um framework de modelagem de mundo baseado em reconstrução que impõe a relevância da tarefa antes da codificação através de atenção espacial estocástica condicionada ao latente.
2. Mecanismo de Aprendizado Autossupervisionado: Uma abordagem inovadora que fundamenta a atenção estocástica em regiões relevantes ao controle, combinando supervisão de dinâmica inversa com reconstrução sobre observações filtradas.
3. Validação Empírica: Demonstração de que o TaskSense mantém um desempenho competitivo em benchmarks padrão, enquanto supera significativamente os baselines sob distrações visuais.

Resultados Experimentais

Os autores avaliaram o TaskSense no DeepMind Control Suite (DMC) e no Distracting Control Suite (DCS), que apresenta fundos de imagens naturais dinâmicos.

• Desempenho Padrão: No DMC padrão (fundos estáticos), o TaskSense iguala o desempenho do baseline DreamerV3 na maioria das tarefas, demonstrando que reconstruir apenas um subconjunto da observação é suficiente para um planejamento latente eficaz.
• Robustez a Distrações: No Distracting Control Suite, o TaskSense supera consistentemente o DreamerV3. Por exemplo, na tarefa Walker Run, o TaskSense alcançou um retorno médio de 531.2 comparado aos 286.2 do DreamerV3. Isso confirma que filtrar o conteúdo visual irrelevante antes da codificação melhora a robustez.
• Análise Qualitativa: Visualizações dos mapas de atenção aprendidos na tarefa Cheetah Run mostram que o TaskSense foca consistentemente em regiões relevantes ao controle (ex: as pernas do cheetah) enquanto suprime o torso e os fundos distrativos.
• Estudos de Ablação:
  • Remover o objetivo de dinâmica inversa leva a uma atenção difusa que falha em localizar regiões relevantes ao controle, causando uma degradação significativa no desempenho.
  • Substituir a atenção estocástica por atenção determinística resulta em uma queda significativa de desempenho, destacando a importância de expor o modelo a diversas observações parciais durante o treinamento.
  • Remover a reconstrução condicionada à atenção causa uma degradação menor, porém perceptível, indicando a necessidade de alinhar o alvo de reconstrução com a entrada filtrada.

Significância e Alegações

O artigo afirma que o TaskSense aborda um descompasso fundamental nos modelos de mundo atuais: a tendência de preservar toda a informação visual em vez de focar no que importa para o controle. Ao deslocar o foco para o filtragem ao nível da observação em vez de apenas o refinamento de representação pós-hoc, o TaskSense permite que modelos de mundo baseados em reconstrução aloquem a capacidade representacional de forma eficiente para informações relevantes ao controle.

Os autores sustentam que esta abordagem permite que os modelos de mundo mantenam os benefícios generativos da reconstrução (essenciais para o planejamento baseado em imaginação) ao mesmo tempo em que alcançam robustez contra distrações visuais que tipicamente degradam modelos padrão. O trabalho sugere que impor a relevância da tarefa na interface de observação é uma direção promissora para melhorar modelos de mundo visuais e o aprendizado de representação em aprendizado por reforço baseado em modelo.

O artigo reconhece uma limitação: assume-se que características visuais preditivas de ações imediatas (capturadas pela dinâmica inversa) são sinônimas das que são relevantes para o sucesso da tarefa a longo prazo. Em cenários onde elas divergem, o sinal de supervisão atual pode focar a atenção em características subótimas, uma direção identificada para trabalhos futuros.