⚡ electrical engineering

Amortising Trajectory Optimisation for Residual MPC via Implicit Contact Differentiation

Este artigo introduz um método de simulação diferenciável eficiente baseado no Teorema da Função Implícita para otimização de trajetórias ricas em contatos que reduz drasticamente o uso de memória em comparação com a diferenciação automática desenrolada, e o combina com destilação de otimizador para melhorar significativamente a taxa de sucesso do MPC residual em tarefas robóticas complexas.

Autores originais: Daniel Layeghi, Thomas Corbères, Calum Arnott, Aditya Kamireddypalli, Hashim Al-Obaidi, Steve Tonneau, Michael Mistry

Publicado 2026-07-29

📖 7 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Daniel Layeghi, Thomas Corbères, Calum Arnott, Aditya Kamireddypalli, Hashim Al-Obaidi, Steve Tonneau, Michael Mistry

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você está tentando ensinar um robô a andar, fazer malabarismo ou jogar futebol. Para fazer isso, o robô precisa de um "cérebro" que possa simular o futuro em sua cabeça, testando milhões de movimentos diferentes para ver qual funciona melhor. Isso é chamado de otimização de trajetória. A parte difícil é o contato. Quando o pé de um robô atinge o chão, uma bola bate em uma parede ou uma mão agarra uma xícara, a física torna-se caótica e imprevisível. É como tentar prever exatamente como uma pilha de blocos Jenga cairá quando você puxa um deles; pequenas mudanças na força levam a mudanças enormes no resultado.

Para fazer essas previsões, os cientistas usam simulação diferenciável. Pense nisso como um motor de jogo de vídeo superpoderoso que não apenas mostra o próximo quadro, mas também diz exatamente como o jogo mudaria se você desse um pequeno toque nos controles. Essa "sensibilidade ao toque" permite que o roba aprenda com seus erros instantaneamente. No entanto, há uma pegadinha: calcular esses toques para tarefas com muito contato é incrivelmente caro. É como tentar filmar um filme em câmera lenta, mas toda vez que a câmera dá zoom em uma colisão, o rolo do filme fica mais longo e longo, eventualmente enchendo seu disco rígido antes que você consiga terminar a cena. Este artigo aborda esse problema de memória e mostra como fazer os robôs aprenderem habilidades de contato complexas de forma muito mais rápida e confiável.

O Problema: O "Monstro da Memória" nos Cérebros dos Robôs

Imagine que você está tentando resolver um labirinto. A maneira padrão de ensinar um robô a resolver um labirinto é deixá-lo caminhar pelo labirinto, bater em uma parede e depois rebobinar a fita para ver exatamente onde ele errou. No mundo da física de robôs, esse "rebobinar a fita" é chamado de diferenciação automática desenrolada (unrolled automatic differentiation).

O problema surge quando o robô atinge uma parede (ou um chão, ou outro objeto). Para entender a física desse impacto, o computador tem que executar um cálculo complexo muitas vezes, como um detetive verificando e rechecando pistas até que a resposta seja perfeita. Se o computador precisar verificar as pistas 10 vezes para acertar, a "fita de rebobinar" precisa armazenar a memória de todas as 10 verificações. Se você quiser que a resposta seja ainda mais perfeita, pode precisar de 100 verificações. De repente, a fita de memória torna-se 100 vezes mais longa.

Isso cria um equilíbrio terrível. Se você quer que o robô seja preciso (verificar as pistas 100 vezes), você só pode executar alguns robôs por vez porque seu computador fica sem memória. Se você quiser executar milhares de robôs ao mesmo tempo para aprender mais rápido, terá que encurtar as verificações (talvez apenas 5 vezes), o que significa que o robô aprenderá com uma resposta imprecisa e desleixada. É como tentar aprender uma dança assistindo apenas aos primeiros cinco segundos do vídeo; você pode pegar os passos, mas perderá o giro crucial no final.

A Solução: O "Instantâneo Mágico"

Os autores deste artigo, trabalhando com o simulador de física MuJoCo (uma ferramenta popular na pesquisa de robótica), encontraram uma maneira inteligente de contornar o monstro da memória. Em vez de rebobinar toda a fita das 100 verificações do detetive, eles usaram um truque matemático chamado Teorema da Função Implícita (IFT).

Pense da seguinte forma: Imagine um detetive que resolveu um mistério. Em vez de mostrar a você as 100 páginas de notas que ele tomou para chegar lá, ele apenas lhe entrega o arquivo do caso resolvido e um "instantâneo mágico" da solução. Esse instantâneo diz exatamente como a solução mudaria se você alterasse um detalhe minúsculo, sem a necessidade de ver as notas bagunçadas.

Em termos técnicos, o artigo introduz um método que diferencia o resíduo de estacionariedade (uma forma elegante de dizer "o ponto onde a matemática diz que terminamos") em vez dos passos dados para chegar lá.

O Jeito Antigo (AD Desenrolada): Armazena cada passo do solver. Se você passar de 1 passo para 10 passos, o uso de memória salta 10,6 vezes.
O Jeito Novo (IFT): Armazena uma quantidade de memória quase constante. Mesmo que você aumente o esforço do solver de 1 para 10 passos, o uso de memória muda em menos de 4%.

Isso é uma mudança de jogo. Significa que o computador pode exigir uma resposta superprecisa (verificando as pistas 100 vezes) sem ficar sem memória. Na verdade, quando o artigo testou isso com 256 contatos ativos (como um robô com muitos dedos tocando uma mesa), o novo método usou 20 vezes menos memória do que o método antigo. Com 16 contatos e um modelo de robô complexo, usou 6 vezes menos memória.

O Resultado: Ensinando Robôs a "Destilar" Sabedoria

Com esta ferramenta nova e eficiente em termos de memória, os autores não pararam apenas em tornar a matemática mais rápida; eles a usaram para ensinar robôs melhor. Eles criaram um sistema que chamam de Destilação de Otimizador (Optimiser Distillation).

Imagine um chef mestre (o "professor") que passa horas aperfeiçoando uma receita complexa. Este chef é lento, mas incrivelmente preciso. Depois, você tem um subchefe (o "aluno" ou política) que é rápido, mas precisa de orientação.

O Professor: O computador executa uma otimização de horizonte longo (como o chef mestre planejando toda a refeição) para encontrar a sequência perfeita de movimentos. Isso é feito em lotes, graças ao novo truque de economia de memória.
O Aluno: O robô aprende com essas sequências perfeitas, criando uma "política" (um conjunto de instintos) que conhece o plano geral.
O Híbrido: Quando o robô está realmente realizando a tarefa, ele não segue apenas cegamente a política. Ele usa a política para o panorama geral (o plano de longo prazo), mas adiciona uma correção "residual" rápida e local (um otimizador de horizonte curto) para lidar com solavancos ou escorregões repentinos.

O artigo testou isso em três robôs diferentes:

Finger: Um pequeno braço girando um pião.
Franka: Um braço grande empurrando uma caixa.
Unitree: Um robô quadrúpede semelhante a um cachorro correndo.

Os resultados foram impressionantes. Quando o horizonte de planejamento (o quão longe o robô olha à frente) era curto (apenas 6 passos), o método padrão (iLQR) frequentemente falhava. Mas com a política "destilada" guiando o robô, a taxa de sucesso saltou dramaticamente:

Em todas as três tarefas (Finger, Franka e Unitree): O sucesso aumentou entre 28 a 98 pontos percentuais em comparação ao iLQR padrão.

Para o robô Franka empurrando uma caixa, o robô padrão de visão curta mal conseguia ter sucesso, enquanto o novo robô híbrido teve sucesso com muito menos "olhar à frente", provando que a política forneceu a estratégia de longo prazo enquanto o otimizador local lidou com os momentos complicados de contato.

Por Que Isso Importa

Este artigo não apenas sugere uma ideia teórica; ele fornece uma ferramenta funcional e de código aberto que muda as regras do jogo. Ao provar que é possível obter derivadas de contato de alta precisão sem o enorme custo de memória, os autores removeram um grande gargalo no aprendizado de robôs. Eles mostraram que você não precisa escolher entre "rápido mas desleixado" e "lento mas preciso". Você pode ter ambos.

Os autores estão confiantes em suas descobertas, tendo validado-as contra métodos numéricos padrão (diferenças finitas) e mostrado que seu novo método iguala a precisão dos métodos antigos usando apenas uma fração dos recursos. Eles também disponibilizaram seu código ao público, convidando outros a construir robôs mais rápidos, inteligentes e destros que possam lidar com a física bagunçada do mundo real do contato sem ficarem presos em um loop de memória.

Resumo Técnico: Otimização de Trajetória Amortizada para MPC Residual via Diferenciação de Contato Implícita

Declaração do Problema

A simulação diferenciável é essencial para a otimização e o aprendizado de trajetórias ricas em contatos, pois expõe as sensibilidades locais dos resultados das tarefas em relação aos controles. No entanto, a escala deste pipeline é dificultada pelo custo computacional de diferenciar solvers de contato iterativos. As abordagens existentes enfrentam um compromisso fundamental entre convergência e paralelismo:

Diferenças Finitas (FD): Caro e sensível ao tamanho do passo.
Diferenciação Automática (AD) Desdobrada (Unrolled): Diferencia todo o traço do solver. Embora precisa, armazena um rastro de computação crescente para cada iteração do solver. Apertar a solução de contato (aumentando as iterações) consome significativamente mais memória, limitando o número de trajetórias paralelas que podem ser processadas.
Derivações de Sensibilidade KKT: Requerem derivações manuais intrincadas e específicas do solver.

O desafio central é que a AD de modo reverso padrão diferencia o traço executado do solver em vez da solução convergida. Isso cria um gargalo de memória onde o esforço do solver (para melhorar a precisão) reduz diretamente a capacidade de lote (batch), impedindo a otimização eficiente de horizonte total.

Metodologia

Os autores propõem um framework que desacopla o solver de propagação (forward) de seu derivativo de retropropagação (backward) usando o Teorema da Função Implícita (IFT).

1. Diferenciação Implícita via IFT

Em vez de desdobrar as iterações do solver, o método diferencia a residual de estacionalidade do problema de contato suavizado e regularizado na solução de convergência de tolerância.

Formulação: O problema de contato no MuJoCo (MJX) é formulado como uma otimização convexa regularizada. A solução $a^\star$ é definida como a raiz do residual de estacionalidade $F(a, \theta) = 0$ , onde $\theta$ representa os parâmetros do sistema (estado, controles, fricção, etc.).
Derivação: Ao aplicar o IFT ao residual $F$ , a derivada da solução em relação aos parâmetros é derivada como:
$\frac{da^\star}{d\theta} = -F_a^{-1} F_\theta$
onde $F_a$ é o Jacobiano do residual em relação à aceleração.
Implementação: O método utiliza jax.lax.custom_root para implementar este passo de retropropagação. Ele requer apenas uma resolução linear transposta e um produto vetor-Jacobiano. Crucialmente, ele não armazena os estados intermediários das $K$ iterações do solver usadas no passo de propagação.
Escopo: A derivada é local a regiões de contato suave fixas e pode mudar através de transições de conjunto de contato ou modo de fricção, mas dentro de um regime fixo, fornece gradientes exatos da solução convergida.

2. Destilação de Otimizador para MPC Residual

Para aproveitar a escalabilidade da derivada implícita, os autores introduzem um pipeline de destilação de otimizador para Controle Preditivo de Modelo (MPC):

Professor (Teacher): Um iLQR (Controlador Linear Quadrático Iterativo) de horizonte total e em lote utiliza a dinâmica MJX diferenciável implícita para gerar trajetórias de alta qualidade e ricas em contatos.
Aluno (Student): Uma política de rede neural é treinada para mimetizar as ações do professor iLQR de horizonte total.
Implantação (Controlador Híbrido): Em tempo de execução, o sistema utiliza uma abordagem de MPC Residual:
- A política destilada fornece a ação nominal de longo horizonte ( $\bar{u}_t$ ).
- Um iLQR residual de curto horizonte calcula uma correção local ( $r_t$ ) e ganhos de feedback baseados no modelo de contato diferenciável.
- O controle final é $u_t = \text{sg}[\pi_\theta(o_t)] + r_t$ , onde $\text{sg}$ denota uma operação de parada de gradiente (stop-gradient), impedindo que a política seja atualizada pelo otimizador de curto horizonte durante a implantação.

Principais Contribuições

Passo de Retropropagação Implícito Assistido por AD: Os autores introduzem um passo de retropropagação para o residual de estacionalidade do solver de contato regularizado do MJX. Esta abordagem deixa o solver de propagação inalterado, evita sistemas KKT montados manualmente e remove o armazenamento de memória dependente da iteração do solver exigido pela AD desdobrada.
Validação e Escalabilidade: O método é validado contra diferenças finas centrais e AD desdobrada. Os resultados mostram que o IFT iguala a precisão da AD desdobrada em altas contagens de iterações, enquanto supera substancialmente esta em termos de memória temporária compilada.
- Escalabilidade de Memória: A memória do IFT permanece quase constante conforme as iterações do solver aumentam (mudando <4% de 1 para 10 iterações), enquanto a AD desdobrada cresce 10,6×.
- Escalabilidade de Contato/Dimensão: O IFT usa 20× menos memória com 256 contatos ativos e 6× menos com 16 contatos e 96 Graus de Liberdade (DoF) em comparação com a AD desdobrada.
MPC Residual Guiado por Política: O artigo demonstra um pipeline de destilação onde ações otimizadas de horizonte total são destiladas em uma política. Esta política guia um iLQR residual de curto horizonte, permitindo a execução bem-sucedida em tarefas complexas onde o iLQR de curto horizonte padrão falha.

Resultados Experimentais

O framework foi avaliado em três tarefas ricas em contatos: Finger (girar um spinner), Franka (empurrar uma caixa) e Unitree (aceleração de quadrupede).

Otimização de Trajetória: Em otimização de horizonte total em lote, o iLQR usando a derivada implícita alcançou custos médios de execução menores e intervalos de percentil mais estreitos em comparação com métodos de primeira ordem (Adam) em todas as tarefas.
Desempenho de Destilação:
- O controlador híbrido (Política + iLQR Residual de Curto Horizonte) superou significativamente o iLQR de curto horizonte padrão.
- Com um horizonte de planejamento de 6 passos, a abordagem híbrida melhorou as taxas de sucesso em 28 a 98 pontos percentuais sobre o iLQR padrão nas três tarefas.
- Franka: O iLQR de curto horizonte padrão alcançou sucesso próximo de zero, enquanto o controlador híbrido obteve sucesso com consideravelmente menos lookahead. A política forneceu a estratégia de empurrar de longo horizonte, enquanto o iLQR lidou com as correções de contato local.
- Finger & Unitree: A política permitiu que o sistema tivesse sucesso com lookaheads mais curtos do que o exigido pelo iLQR puro. Em certos regimes (ex: Finger), um controlador apenas de política superou horizontes residuais muito curtos, pois o otimizador residual poderia perturbar uma política já correta.

Significância e Alegações

O artigo afirma que a diferenciação implícita baseada em resíduos torna prática a diferenciação de contato do MJX com tolerância convergida em escala de aceleradores. Ao remover o gargalo de memória associado às iterações do solver, o método permite:

Convergência Mais Estrita: Os solvers podem ser executados para maior precisão sem sacrificar o paralelismo de lote.
Supervisão de Horizonte Total: A capacidade de executar iLQR de horizonte total em lote como um professor online, o que era anteriormente inviável devido às restrições de memória.
Amortização Eficaz: O controlador híbrido resultante alcança taxas de sucesso substancialmente maiores do que o iLQR padrão em horizontes de planejamento curtos, efetivamente amortizando o cálculo caro específico da tarefa entre condições iniciais e objetivos.

Os autores disponibilizam a implementação como código aberto, fornecendo um padrão modular para simulação de contato diferenciável que separa o solver de propagação de seu derivativo sem exigir derivações analíticas específicas do solver.