Resumo Técnico: RoboSSM – Aprendizado de Imitação In-context Escalável via Modelos de Espaço de Estados

1. Definição do Problema

O Aprendizado de Imitação In-context (ICIL) permite que robôs se adaptem a novas tarefas usando demonstrações de poucos disparos (few-shot) sem atualizações de parâmetros durante a implantação. Embora métodos recentes de ICIL tenham adotado com sucesso arquiteturas baseadas em Transformers para formular o aprendizado de imitação como um problema de modelagem de sequência, eles enfrentam limitações significativas de escalabilidade. Os Transformers exibem complexidade de tempo quadrática ($O(L^2)$) em relação ao comprimento da sequência e lutam para extrapolar além dos comprimentos de prompt observados durante o treinamento. Consequentemente, os métodos de ICIL existentes degradam-se drasticamente ao lidar com prompts de teste mais longos (por exemplo, mais demonstrações ou tarefas de longo horizonte) do que aqueles usados no treinamento, limitando sua aplicabilidade em cenários do mundo real complexos que exigem contexto extenso.

2. Metodologia: RoboSSM

Os autores propõem o RoboSSM, um framework de ICIL escalável que substitui os Transformers por Modelos de Espaço de Estados (SSMs), utilizando especificamente o Longhorn como arquitetura de base (backbone).

Arquitetura e Codificação

• Processamento de Entrada: O RoboSSM codifica observações multimodais (dados visuais via CNNs e dados proprioceptivos via MLPs) em embeddings. Notavelmente, a representação de entrada exclui tokens de ação e instruções de linguagem de tarefa para forçar o modelo a inferir a intenção da tarefa apenas a partir das demonstrações in-context.
• Backbone Longhorn: A sequência de embeddings de observação é processada pelo Longhorn, que mantém uma matriz de estado de memória $s_t$ atualizada via recorrência linear. Diferente do mecanismo de autoatenção dos Transformers, o Longhorn utiliza uma regra de atualização recorrente derivada de programação convexa online:
  $$s_t = \text{argmin}_{s} \left( \|s - s_{t-1}\|_F^2 + \|s k_t - x_t\|_{\text{diag}(\beta_t)}^2 \right)$$
  Esta formulação equilibra a retenção de informações de estado prévio com a incorporação de novos inputs.
• Escalonamento $\beta$: Para incentivar especificamente a atenção aos prompts de demonstração no cenário de ICIL, os autores introduzem um mecanismo de escalonamento em tempo de teste onde o vetor de ponderação $\beta_t$ é escalonado por um fator $\gamma \in (0, 1]$. Isso modula o equilíbrio entre a retenção de estado e a integração de nova informação.

Treinamento e Inferência

• Treinamento: O modelo é treinado em um conjunto de dados multitarefa onde a entrada consiste em um prompt (um conjunto de $N_{train}$ trajetórias de demonstração) e uma trajetória de consulta (query). O modelo prevê ações para a trajetória de consulta condicionada ao prompt.
• Inferência: No tempo de teste, a política se adapta a tarefas não vistas usando um prompt de $N_{test}$ demonstrações. O modelo prevê iterativamente as ações com base na observação inicial e no contexto fornecido pelo prompt, sem quaisquer atualizações de parâmetros.

3. Principais Contribuições

1. Framework de ICIL Escalável: A introdução do RoboSSM, a primeira abordagem que utiliza SSMs para aprendizado de imitação in-context, permitindo inferência de tempo linear ($O(L)$) e fortes capacidades de extrapolação.
2. Extrapolação de Longo Contexto: Demonstração de que o RoboSSM pode processar efetivamente prompts de teste até 16 vezes mais longos do que os vistos durante o treinamento, uma capacidade na qual os métodos baseados em Transformer falham.
3. Inferência Eficiente: Alcançando tempos de execução de inferência significativamente menores em comparação com as bases de comparação (baselines) de Transformer à medida que o comprimento do prompt aumenta, devido à complexidade linear dos SSMs versus a complexidade quadrática dos Transformers.
4. Generalização para Tarefas de Longo Horizonte: Mostrando que o modelo pode generalizar de distribuições de treinamento de curto horizonte para tarefas de longo horizonte e múltiplos estágios, e cenários de dilatação temporal.

4. Resultados Experimentais

Experimentos foram conduzidos no benchmark LIBERO, cobrindo diversas tarefas de manipulação (suites Object, Goal, Spatial, Long e 90).

• Generalização de Comprimento de Prompt:
  • Quando treinado com apenas 2 demonstrações ($N_{train}=2$), o RoboSSM manteve ou melhorou as taxas de sucesso conforme o número de demonstrações de teste ($N_{test}$) aumentava até 32.
  • Em contraste, o baseline baseado em Transformer (ICRT) experimentou uma queda acentuada de desempenho e colapsou para 0% de sucesso quando $N_{test} > N_{train}$.
• Longo Horizonte e Dilatação Temporal:
  • O RoboSSM resolveu com sucesso tarefas de longo horizonte (ex: "Coloque tanto o Objeto A quanto o Objeto B no Objeto C") e tarefas com demonstrações de tempo dilatado (até 16x de estiramento temporal), enquanto o ICRT falhou completamente.
• Desempenho In-Distribution:
  • Mesmo quando $N_{train} = N_{test}$, o RoboSSM superou consistentemente o ICRT na maioria dos cenários, particularmente nas cenas LIBERO-90 Study e Living Room.
• Comparação com Aprendizado Multitarefa (MTL):
  • RoboSSM e ICRT superaram significativamente os baselines padrão de MTL (que exigem instruções de linguagem e não utilizam prompting in-context), confirmando a eficácia do paradigma ICIL.
• Análise de Espaço Latente:
  • Visualizações das representações latentes mostraram que o RoboSSM gera trajetórias que permanecem próximas aos clusters de demonstração do prompt, indicando uma extrapolação estável. Por outro lado, as trajetórias do ICRT derivaram significativamente para fora do manifold do prompt conforme o comprimento do contexto aumentava.
• Tempo de Execução (Runtime):
  • O RoboSSM demonstrou um escalonamento quase linear com o comprimento do prompt, enquanto o tempo de execução do ICRT aumentou quadraticamente, criando uma lacuna de desempenho crescente a favor do RoboSSM para contextos longos.

5. Significância e Alegações

O artigo afirma que o RoboSSM representa uma mudança de paradigma para o aprendizado de imitação in-context ao estabelecer os Modelos de Espaço de Estados como um backbone eficiente e escalável. A principal significância reside em superar os gargalos computacionais e de extrapolação dos Transformers, permitindo que robôs:

• Lidem com janelas de contexto substancialmente mais longas no tempo de teste sem necessidade de retreinamento.
• Adaptem-se a tarefas não vistas, de longo horizonte e temporalmente variáveis usando prompts de poucos disparos (few-shot).
• Suportem a adaptação contínua para aprendizado vitalício (lifelong learning) simplesmente ingerindo novos prompts de demonstração sem atualizações de parâmetros específicas da tarefa.

Os autores concluem que, embora o RoboSSM demonstre o potencial dos SSMs para o ICIL de longo contexto, trabalhos futuros podem exigir corpora de treinamento mais amplos e diversos para aumentar ainda mais a generalização. Eles também observam que, embora a validação em robôs reais não seja a contribuição central deste trabalho, ela permanece como uma via potencial para futuras demonstrações de generalização.