Resumo Técnico: Latent Memory Palace (LMP)

Definição do Problema

A tomada de decisão humana exibe um espectro que vai de ações reflexas a raciocínios deliberados de múltiplas etapas, onde o tempo e o esforço investidos escalam com a complexidade da tarefa. Embora os Grandes Modelos de Linguagem (LLMs) tenham demonstrado com sucesso capacidades semelhantes de "raciocínio" adaptativo por meio de geração iterativa de tokens (ex: Cadeia de Pensamento/Chain-of-Thought), transferir essa capacidade para políticas de controle robótico contínuo permanece desafiador. O raciocínio direto no espaço da linguagem muitas vezes carece da granularidade necessária para compreensão espacial e controle de movimento preciso. Além disso, as abordagens padrão de inferência variacional para controle tipicamente dependem de variáveis latentes de dimensão fixa, que não conseguem adaptar sua profundidade computacional às exigências específicas de cada entrada. O problema central é como permitir que políticas robóticas realizem computação iterativa e adaptativa em um espaço latente aprendido de ponta a ponta para predição de ações, visando equilibrar a tomada de decisão eficiente com uma melhor generalização.

Metodologia

Os autores propõem o Latent Memory Palace (LMP), um framework que formula o raciocínio para controle como inferência variacional autorregressiva.

1. Formulação Central

O LMP trata uma ação $a$ não como uma predição direta de uma observação $o$, mas como o resultado de um cálculo latente intermediário $z$. A política é definida por um limite inferior variacional (ELBO):
$$\log p(a | o) \geq \mathbb{E}_{z \sim q_\theta(z|o,a)} [\log p_\phi(a | o, z)] - D_{KL}(q_\theta(z | o, a) \parallel p_\theta(z | o))$$
Ao contrário das abordagens padrão, a variável latente $z$ não é um único vetor, mas uma sequência autorregressiva de comprimento variável $z = (z_1, z_2, \dots, z_{T(z)})$ extraída de um vocabulário que inclui um token de Fim de Sequência (EOS). A sequência termina quando o EOS é gerado, permitindo que o modelo aloque computação de forma adaptativa com base na complexidade da decisão.

2. Raciocínio Adaptativo via Compressão

Para elicitar um comportamento adaptativo onde o modelo utiliza mais etapas conforme necessário, o LMP introduz um mecanismo de compressão. O decodificador de ação $p_\phi$ é parametrizado como uma Gaussiana isotrópica onde a variância $\sigma(T(z))$ decai exponencialmente com o comprimento da sequência latente $T(z)$.
$$\sigma(T(z)) = \gamma^{T(z)} \sigma_0, \quad \gamma \in (0, 1)$$
Isso impõe uma penalidade de comprimento: cada etapa latente adicional torna a distribuição Gaussiana mais nítida. Consequentemente, o modelo só se "permite" gerar outra etapa se a precisão da ação prevista justificar a redução da variância. Isso força a política a gastar mais computação em decisões complexas e menos em decisões simples ou irremediavelmente incertas (ex: o momento exato do fechamento de uma garra).

3. Otimização Tratável via RL no Espaço Latente

A otimização do ELBO através de amostragem autorregressiva não diferenciável é intratável via backpropagation padrão. Os autores derivam um objetivo de Aprendizado por Reforço (RL) sobre trajetórias latentes:

• Episódio Latente: Uma trajetória é amostrada da posterior $q_\theta$ até o EOS.
• Recompensa: Composta por uma recompensa de reconstrução ($\log p_\phi(a|o,z)$) e uma penalidade KL contra a prior.
• Redução de Variância: O termo KL é decomposto em recompensas por etapa usando a estrutura autorregressiva, permitindo a atribuição de crédito por etapa.
• Otimização: Os autores utilizam um objetivo substituto truncado (clipped surrogate objective) do tipo PPO para estabilizar o treinamento e evitar o colapso latente, utilizando rollouts on-policy coletados em um buffer.

4. Tokenização de Ação (LMP-tok)

Ao remover o condicionamento da observação do framework variacional, a mesma arquitetura instancia o LMP-tok, um tokenizador de ação de comprimento variável. Isso comprime blocos de ações contínuas em tokens discretos latentes, que podem então ser usados para treinar políticas autorregressivas subsequentes. Diferente de tokenizadores que exigem a reconstrução exata de transformações discretas (ex: FAST), o LMP-tok permite que sequências de comprimentos arbitrários sejam decodificadas em ações válidas.

Principais Contribuições

1. Latent Memory Palace (LMP): Uma formulação inovadora do raciocínio como inferência variacional com uma distribuição latente autorregressiva, permitindo computação iterativa e adaptativa para controle robótico.
2. LMP-$\pi$ (Política): Uma política de controle que gera traços de "raciocínio" de comprimento variável durante o teste, alocando computação de forma adaptativa com base na dificuldade e incerteza da tarefa.
3. LMP-tok (Tokenizador): Um tokenizador de ação de comprimento variável que melhora significativamente o desempenho de políticas autorregressivas subsequentes ao comprimir ações em espaços latentes discretos.
4. Framework de Otimização: Um método baseado em RL tratável para otimizar o limite inferior variacional para amostragem autorregressiva de comprimento variável, incluindo técnicas para estabilizar o treinamento contra o colapso latente.

Resultados Experimentais

Os autores avaliaram o LMP tanto em benchmarks do mundo real (DROID) quanto simulados (LIBERO, D3IL, RoboMimic).

• Desempenho da Política (LMP-$\pi$):

  • Zero-Shot & Fine-tuning: Na plataforma DROID, o LMP-$\pi$ superou o Diffusion Policy (DP) em configurações zero-shot e alcançou taxas de sucesso mais altas (ex: 95% em limpeza de múltiplos objetos) em tarefas de alta precisão com fine-tuning.
  • Generalização: No LIBERO-90, o LMP-$\pi$ mostrou um "piso" de desempenho significativamente mais alto nas 10 tarefas mais difíceis em comparação ao DP, indicando redução de interferência entre tarefas e melhor capacidade de recuperação.
  • Iterativo vs. Não-Iterativo: O LMP-$\pi$ superou consistentemente baselines não-iterativos (VAE Policy, VQ-BeT), demonstrando que representações latentes de comprimento fixo carecem da granularidade necessária para manipulação complexa.

• Desempenho do Tokenizador (LMP-tok):

  • Ao ser usado para treinar políticas autorregressivas subsequentes, o LMP-tok superou significativamente os baselines (VQ-VAE, FAST, OAT) em tarefas do RoboMimic e LIBERO.
  • Notavelmente, o LMP-tok foi o único método que alcançou desempenho não trivial na tarefa de alta precisão "tool-hang" do RoboMimic, onde outros tokenizadores falharam devido a restrições de precisão ou reconstrução.

• Análise de Computação Adaptativa:

  • Interpretabilidade: Visualizações mostram que o número de etapas de raciocínio latente correlaciona-se com a complexidade da tarefa. O modelo usa menos etapas durante movimentos de garra (onde o tempo é irremediavelmente variável) e mais etapas durante tarefas de alinhamento preciso.
  • Compressão: Estudos de ablação confirmaram que uma compressão mais forte (decaimento de variância mais rápido) leva a uma melhor generalização, desde que seja equilibrada para evitar erro de reconstrução excessivo.
  • Busca de Modo (Mode-Seeking): Truncar as etapas latentes durante o teste não degradou significativamente o desempenho; em vez disso, o modelo exibiu comportamento de "busca de modo", comprometendo-se com modos de ação específicos válidos em vez de fazer a média entre eles.

Significância e Alegações

O artigo afirma estabelecer a inferência variacional autorregressiva como um framework prático e principiado para computação iterativa e adaptativa no aprendizado robótico. Ao traçar uma analogia ao mnemônico "palácio da memória", o trabalho demonstra que organizar informações em um espaço latente de comprimento variável permite que as políticas:

1. Imitem a Variabilidade de Decisão Humana: Escalar o esforço computacional com a complexidade da decisão.
2. Melhorem a Generalização: Alcançar desempenho robusto através de diversas tarefas e ambientes ao aprender a alocar computação eficientemente.
3. Melhorem a Tokenização: Fornecer um método superior para discretizar ações contínuas para modelagem autorregressiva.

Os autores posicionam o LMP como uma ponte entre as capacidades de raciocínio dos modelos de fundação e os requisitos de precisão do controle robótico, oferecendo uma nova perspectiva onde o "raciocínio" é um processo latente implícito otimizado através de inferência variacional e aprendizado por reforço.