Task Tokens: A Flexible Approach to Adapting Behavior Foundation Models

Resumo Técnico: Task Tokens

1. O Problema

Os avanços recentes em aprendizado por imitação levaram ao surgimento de Modelos Fundamentais de Comportamento (BFMs) baseados em transformadores para controle de humanoides. Especificamente, os Modelos Fundamentais de Comportamento Condicionados a Objetivos (GC-BFMs), como o MaskedMimic, demonstraram uma notável capacidade de generalização zero-shot, gerando comportamentos robustos e naturais a partir de objetivos de alto nível (prompts).

No entanto, existem lacunas fundamentais para a aplicação prática desses modelos:

• Engenharia de Prompts Limitada: Definir prompts de alto nível precisos para tarefas complexas é difícil e muitas vezes resulta em comportamentos subótimos (ex: um personagem andando de costas para atingir um objetivo).
• Dificuldade com Recompensas: O design de funções de recompensa para tarefas de longo horizonte é propenso a erros e pode levar a comportamentos não naturais.
• Compromisso entre Especialização e Generalização: Métodos tradicionais de ajuste fino (fine-tuning) de todo o modelo para uma tarefa específica tendem a destruir o conhecimento pré-treinado (causando catastrophic forgetting), resultando em movimentos menos robustos e naturais. Por outro lado, métodos puramente baseados em recompensa (RL puro) não aproveitam o vasto conhecimento de movimento humano pré-existente.

O objetivo é criar um paradigma unificado que permita adaptar BFMs a tarefas específicas com alta precisão, mantendo a robustez, a naturalidade e a capacidade de generalização do modelo original.

2. Metodologia: Task Tokens

Os autores propõem o Task Tokens, uma abordagem que integra o controle baseado em objetivos com otimização orientada a recompensas dentro da arquitetura de transformadores dos GC-BFMs.

• Conceito Central: Em vez de ajustar todo o modelo, o método treina um codificador de tarefa (Task Encoder) leve e específico para cada tarefa. Este codificador gera um "Token de Tarefa" que é injetado na entrada do modelo BFM congelado.
• Arquitetura Híbrida:
  1. Tokens de Prioridade (Priors): O usuário fornece objetivos de alto nível (texto ou condições de juntas) que são codificados pelos encoders pré-treinados do BFM.
  2. Token de Tarefa (Task Token): O Task Encoder (uma rede neural feed-forward) processa as observações do objetivo atual ($g_t$) e gera um token especializado ($\tau_t$).
  3. Estado: O estado atual do ambiente é também tokenizado.
  4. Fusão: Todos esses tokens são concatenados e alimentados no BFM congelado, que gera a ação final.
• Treinamento:
  • O BFM permanece congelado (sem atualização de parâmetros), preservando seu conhecimento de movimento humano.
  • Apenas o Task Encoder é treinado usando Otimização de Política Próxima (PPO).
  • O gradiente da função de recompensa da tarefa específica flui através do BFM congelado para atualizar apenas os parâmetros do Task Encoder.
• Eficiência: O método requer apenas ~200 mil parâmetros treináveis por tarefa, comparado a milhões em métodos de ajuste fino ou RL puro.

3. Contribuições Principais

• Adaptação Específica sem Ajuste Fino: Propõe uma nova abordagem para adaptar o MaskedMimic a tarefas específicas via tokens, preservando as capacidades zero-shot e a robustez do modelo base.
• Eficiência de Parâmetros e Convergência: O método é altamente eficiente, exigindo até 125x menos parâmetros treináveis e convergindo até 6x mais rápido do que métodos de base (como fine-tuning completo ou PULSE).
• Paradigma de Controle Híbrido: Permite a combinação perfeita de priors definidos pelo usuário (ex: "andar de frente") com otimização aprendida por recompensa (ex: "atingir o alvo com impacto máximo").
• Generalização Robusta: Ao manter o BFM congelado, o método herda a robustez do modelo pré-treinado contra perturbações fora da distribuição (OOD), como variações de gravidade e atrito, superando métodos que ajustam o modelo completo.

4. Resultados Experimentais

Os autores avaliaram o método em cinco tarefas de controle de humanoides simulados (Reach, Direction, Steering, Strike, Long Jump):

• Desempenho em Tarefas Específicas: O Task Tokens atingiu taxas de sucesso comparáveis ou superiores ao fine-tuning completo e ao RL puro na maioria das tarefas (ex: 94,88% em Reach, 99,26% em Direction).
• Eficiência de Amostragem: O método convergiu em cerca de 50 milhões de passos, enquanto o método PULSE precisou de ~300 milhões de passos para atingir desempenho similar.
• Generalização (OOD): Em testes com perturbações de atrito e gravidade, o Task Tokens manteve um desempenho superior, enquanto o fine-tuning completo sofreu degradação significativa, perdendo a robustez inerente ao modelo base.
• Estudo Humano: Em uma avaliação com ~100 participantes, os movimentos gerados pelo Task Tokens foram considerados mais "humanos" do que os gerados pelo MaskedMimic padrão (apenas com condicionamento de juntas) e comparáveis ou superiores a outros métodos de ponta, evitando comportamentos estranhos (como andar de costas).
• Prompting Multimodal: O método demonstrou capacidade de combinar prompts textuais e condicionamento de juntas para corrigir comportamentos indesejados (ex: forçar o personagem a olhar para o alvo enquanto ataca).

5. Significado e Impacto

O trabalho apresenta uma solução escalável para a adaptação de modelos fundamentais de comportamento. Ao introduzir uma camada leve de adaptação (Task Tokens) que não altera o núcleo do modelo, os autores resolvem o dilema entre generalização e especialização.

• Para Animação e Robótica: Permite criar personagens e robôs que podem executar tarefas complexas e específicas sem perder a naturalidade dos movimentos humanos aprendidos durante o pré-treinamento.
• Viabilidade Prática: A redução drástica no número de parâmetros e no tempo de treinamento torna viável a adaptação de grandes modelos para múltiplas tarefas em cenários do mundo real, onde o custo computacional é um fator limitante.
• Futuro: O trabalho abre caminho para sistemas de controle que combinam a intuição humana (via prompts) com a otimização automática (via recompensas), superando as limitações de abordagens puramente baseadas em recompensa ou puramente baseadas em prompts.

Em resumo, Task Tokens oferece um mecanismo elegante e eficiente para "afinar" modelos fundamentais de comportamento para tarefas específicas, mantendo a essência e a robustez do modelo original.