Neural Operators for Multi-Task Control and Adaptation

Título: Operadores Neurais para Controle Multi-Tarefa e Adaptação

Autores: David SeWell, Xingjian Li, Stepan Tretiakov, Krishna Kumar, David Fridovich-Keil (Universidade do Texas em Austin).

1. Problema e Motivação

O controle ótimo em aplicações reais frequentemente envolve resolver não um único problema de otimização, mas uma família de problemas relacionados (ex: navegar para diferentes locais, rastrear trajetórias com parâmetros de veículo variáveis, ou evitar obstáculos em configurações diferentes).

• Desafio Central: Pequenas mudanças nas especificações da tarefa (dinâmica, função de custo ou restrições) podem induzir políticas ótimas drasticamente diferentes.
• Limitações das Abordagens Atuais:
  • Aprendizado de Política Separada: Treinar uma política distinta para cada tarefa não escala bem e ignora estruturas compartilhadas.
  • Embeddings de Tarefa: Representar tarefas como vetores de dimensão finita cria um gargalo representacional, sendo sensível à resolução e ordenação dos dados.
  • Meta-Aprendizado (ex: MAML): Embora promissor, o MAML não modela explicitamente o mapeamento entre espaços de funções e exige computação adicional na fase de implantação (inferência) para adaptação.
• Perspectiva Proposta: O controle multi-tarefa deve ser formulado como um mapeamento entre espaços de funções infinitas: da função que define a tarefa (ex: função de custo $c(x,u)$) para a lei de controle ótima (política $\pi(x)$).

2. Metodologia

Os autores propõem o uso de Operadores Neurais (Neural Operators - NOs), especificamente a arquitetura SetONet (baseada em DeepONet), para aprender esse mapeamento.

2.1. Arquitetura SetONet

O modelo utiliza uma decomposição Branch-Trunk (Ramo-Tronco):

• Branch Network (Ramo): Recebe avaliações pontuais da função que define a tarefa (ex: pares estado-ação-custo). Utiliza um codificador de conjuntos permutation-invariant (invariante à permutação), permitindo lidar com entradas de tamanho variável e desordenado. Mapeia a tarefa para coeficientes $\{c_k\}$.
• Trunk Network (Tronco): Recebe as localizações de consulta (estados e tempos) e gera uma base de funções $\{b_k(y)\}$ aprendidas.
• Saída: A política é aproximada como uma combinação linear:
  $$\pi(y) \approx \sum_{k=1}^{p} c_k(\text{tarefa}) \cdot b_k(y)$$
  Essa estrutura permite que o operador generalize para diferentes resoluções de dados e tamanhos de conjuntos de contexto sem re-treinamento.

2.2. Estratégias de Aprendizado e Adaptação

O trabalho explora três abordagens principais:

1. Clonagem Comportamental Multi-Tarefa (Zero-Shot): O operador é treinado para imitar políticas expert geradas por solvers ótimos (LQR, iLQR, NLP) ou RL (SAC) em uma distribuição de tarefas. O modelo aprende a mapear diretamente a função de custo/dinâmica para a política.
2. Adaptação Estruturada (Fine-tuning): Para tarefas fora da distribuição de treinamento (OOD) ou com poucos dados, o modelo pré-treinado é adaptado. Devido à estrutura Branch-Trunk, são propostas estratégias de ajuste parcial:
   • Full-network: Ajuste de todos os parâmetros.
   • Branch-only: Ajuste apenas dos coeficientes da tarefa (mantendo a base fixa).
   • Last-layer: Ajuste apenas das camadas finais.
   • Adaptação via Feedback de Custo: Ajuste direto minimizando o custo da trajetória sem necessidade de demonstrações expert, apenas usando o modelo de dinâmica e a função de custo.
3. Meta-Treinamento (Few-Shot): Propõem duas variantes inspiradas no MAML para otimizar a inicialização do operador para adaptação rápida:
   • SetONet-Meta: Otimiza apenas a inicialização do Branch (coeficientes), mantendo o Trunk (base) fixo durante o loop interno. Eficiente em dados limitados.
   • SetONet-Meta-Full: Otimiza a inicialização de toda a rede (Branch + Trunk), permitindo adaptação de funções de base para tarefas OOD complexas.

3. Contribuições Principais

1. Operadores Neurais para Controle Multi-Tarefa: Estabelecem os NOs como uma classe de modelo eficaz para controle, demonstrando que podem aproximar com precisão o operador solução sobre uma distribuição de tarefas, superando a necessidade de embeddings fixos.
2. Adaptação Estruturada Eficiente: Demonstram que a decomposição Branch-Trunk permite estratégias de adaptação leves (ex: apenas a camada final ou apenas o branch) que alcançam precisão comparável ao fine-tuning completo com custo computacional muito menor.
3. Meta-Treinamento para Adaptação Rápida: Introduzem variantes meta-treinadas que superam consistentemente o baseline MAML em todos os ambientes, especialmente em cenários de few-shot (poucos dados) e adaptação rápida online.
4. Invariância à Resolução de Tarefa: O modelo generaliza para quantidades de dados de contexto não vistas durante o treinamento (ex: prever políticas com 16 pontos de contexto quando treinado com 32).

4. Resultados Experimentais

Os métodos foram avaliados em quatro ambientes de controle ótimo paramétrico (P2P-Cost, P2P-Dynamics, Quadrotor Planar, Evitação de Obstáculos) e um ambiente de locomoção de alta dimensão (HalfCheetah-v3 do iMuJoCo).

• Aproximação Zero-Shot: O operador pré-treinado alcança baixo erro em tarefas próximas à distribuição de treinamento. Em tarefas mais complexas (dinâmica variável, obstáculos), a adaptação é necessária.
• Comparação com MAML:
  • O MAML sofreu consistentemente, apresentando erros altos e falhando em melhorar sua inicialização em muitos ambientes.
  • As variantes SetONet (especialmente Last-Branch e SetONet-Meta) superaram o MAML em todos os cenários, muitas vezes com apenas 1 passo de gradiente.
• Adaptação com Poucos Dados (Few-Shot):
  • No ambiente HalfCheetah-v3, com poucos exemplos (1-5 demonstrações), os modelos meta-treinados (SetONet-Meta e SetONet-Meta-Full) superaram o fine-tuning padrão (SetONet-FT).
  • À medida que a quantidade de dados aumenta, o fine-tuning padrão converge para o melhor desempenho, mas os métodos meta-treinados oferecem uma vantagem crítica em cenários de recursos limitados.
• Adaptação OOD (Fora da Distribuição): O SetONet-Meta-Full demonstrou capacidade de adaptar tanto os coeficientes quanto as funções de base para tarefas OOD (ex: Quadrotor com alvo fora da região de treinamento), onde o modelo pré-treinado puro falhava.
• Adaptação via Custo: Foi possível refinar o operador para evitar colisões e atingir objetivos sem nenhuma demonstração expert, apenas diferenciando através do modelo de dinâmica e função de custo.

5. Significado e Conclusão

O trabalho demonstra que os Operadores Neurais fornecem um framework unificado e eficiente para controle multi-tarefa.

• Vantagens sobre Métodos Existentes: Eliminam o gargalo de representações de tarefas fixas e reduzem a sobrecarga computacional na implantação em comparação com métodos baseados em gradiente como o MAML.
• Flexibilidade: A capacidade de lidar com entradas de tamanho variável e a invariância à permutação tornam o modelo robusto a variações na quantidade de dados disponíveis em tempo de execução.
• Aplicabilidade Prática: As estratégias de adaptação estruturada permitem que operadores neurais sejam usados em cenários de implantação real onde o tempo de computação e a quantidade de dados de adaptação são limitados, oferecendo uma alternativa superior aos métodos tradicionais de controle e aprendizado por reforço multi-tarefa.

Em suma, o artigo valida que aprender o mapeamento direto entre espaços de funções (tarefa $\to$ política) é uma abordagem superior para sistemas de controle que exigem generalização e adaptação rápida.