PocketDP3: Efficient Pocket-Scale 3D Visuomotor Policy

Resumo Técnico: PocketDP3: Política Visuomotora 3D Eficiente em Escala de Bolso

Declaração do Problema
Avanços recentes em manipulação robótica têm aproveitado políticas de difusão baseadas em visão 3D para aprender habilidades complexas a partir de dados. No entanto, os autores identificam um desajuste arquitetônico crítico nos modelos atuais de estado da arte (SOTA), como o DP3. Esses modelos tipicamente combinam um codificador de nuvem de pontos leve e eficiente com um decodificador de backbone excessivamente grande, frequentemente um U-Net condicional. Embora os U-Nets sejam adequados para síntese de imagem de alta resolução onde a reconstrução de pixels detalhada é necessária, os autores argumentam que esse design é ineficiente para o controle visuomotor 3D. Neste contexto, o codificador já produz uma representação de cena compacta e rica em informações, e o alvo de predição é uma trajetória de ação de baixa frequência. Consequentemente, o decodificador massivo introduz redundância substancial de parâmetros, aumenta a pegada de memória e dificulta a implantação em tempo real em robôs com recursos limitados, sem proporcionar ganhos proporcionais no desempenho do controle. Além disso, a amostragem de difusão padrão requer muitos passos de denoising iterativos, criando gargalos de latência, o que muitas vezes exige técnicas complexas de destilação de consistência para acelerar a inferência.

Metodologia
O artigo propõe o PocketDP3, uma política de difusão 3D em escala de bolso projetada para resolver o desequilíbrio arquitetônico e acelerar a inferência.

• Arquitetura (Diffusion Mixer - DiM): Em vez de um pesado U-Net condicional, o PocketDP3 emprega um decodificador Diffusion Mixer (DiM) leve construído sobre blocos MLP-Mixer. A arquitetura mantém o codificador de nuvem de pontos eficiente do DP3 (usando Farthest Point Sampling e MLPs para gerar uma representação 3D compacta), mas substitui o decodificador por $K$ blocos DiM empilhados.

  • Mecanismo de Fusão: Os blocos DiM realizam uma fusão de informações eficiente alternando entre mistura temporal (agregando informações ao longo da dimensão temporal da trajetória) e mistura de canais (agregando através das dimensões de características).
  • Condicionamento: O modelo incorpora representações de cena e estados do robô via uma projeção linear e injeta essa informação de condicionamento nos blocos DiM usando camadas FiLM (Feature-wise Linear Modulation) combinadas com conexões residuais com portão (gated).
  • Eficiência de Parâmetros: Os autores propõem duas configurações, "Tiny" (0,53M de parâmetros) e "Base" (1,73M de parâmetros), que representam menos de 1% dos parâmetros usados por métodos anteriores como o DP3 (~250M de parâmetros).

• Estratégia de Inferência: O método utiliza amostragem do tipo DDIM. Notavelmente, os autores demonstram que o PocketDP3 suporta inferência de dois passos (2 avaliações de função) sem exigir destilação de consistência adicional ou objetivos de treinamento especializados (ex: One-Step Diffusion ou Modelos de Consistência). Isso reduz significativamente a latência de inferência enquanto mantém o desempenho.

• Treinamento: O modelo é treinado usando um objetivo de predição $x_0$, minimizando o erro quadrático médio entre a trajetória limpa prevista e a trajetória especialista real, condicionado à observação e ao passo de tempo.

Principais Contribuições

1. Insight Arquitetônico: O artigo identifica que parear codificadores 3D compactos com decodificadores U-Net massivos é uma ineficiência generalizada em políticas de difusão 3D, levando ao desperdício de parâmetros.
2. Arquitetura PocketDP3: A introdução de uma política em escala de bolso substituindo o U-Net pelo Diffusion Mixer (DiM), permitindo fusão temporal e de canais eficaz com um orçamento de parâmetros mínimo.
3. Amostragem Rápida Sem Destilação: Demonstração de que a inferência de alto desempenho em dois passos é alcançável sem destilação de consistência, simplificando o pipeline de treinamento e permitindo a implantação de baixa latência.
4. Validação Empírica: Avaliação abrangente mostrando que o PocketDP3 iguala ou excede o desempenho SOTA enquanto reduz drasticamente o tamanho do modelo e o tempo de inferência.

Resultados Experimentais
O método foi avaliado em três benchmarks de simulação (RoboTwin2.0, Adroit e MetaWorld) e experimentos no mundo real.

• Desempenho de Simulação:

  • RoboTwin2.0: O PocketDP3-Base alcançou uma taxa de sucesso média de 71,6%, superando o SOTA anterior (DP3) em 20,8%. O PocketDP3-Tiny também superou o DP3 em 15,2%.
  • Adroit & MetaWorld: O PocketDP3-Base alcançou uma taxa de sucesso média de 77,4%, superando o DP3 (73,0%) e o FlowPolicy (72,6%).
  • Estudos de Ablação: Comparações com um "Vanilla-MLP" (que falhou em convergir) e um "Vanilla-U-Net" reduzido (que teve desempenho inferior) confirmaram que o mecanismo de fusão específico do bloco DiM é crítico para o desempenho sob orçamentos estritos de parâmetros.
  • Passos de Inferência: A ablação no Número de Avaliações de Função (NFE) mostrou que o desempenho estabiliza após 2 passos; aumentar para 5 ou 10 passos não gerou ganhos significativos.

• Eficiência:

  • Tamanho do Modelo: O PocketDP3-Base usa 1,73M de parâmetros, comparado a 255,1M do DP3 (uma redução de >99%).
  • Latência: Com 2 passos de inferência, o PocketDP3-Base alcança uma latência de 4,80ms em uma RTX 5880 Ada, significativamente mais rápido que o DP3 (51,4ms com 10 passos) e comparável ao FlowPolicy (7,04ms com 1 passo), porém com dois ordens de magnitude a menos de parâmetros.

• Implantação no Mundo Real:

  • Testado em um robô AgileX Piper com uma câmera RealSense D455.
  • O PocketDP3-Base demonstrou forte transferência sim-to-real, superando o baseline DP3 em 15,3% em média em tarefas como posicionamento de objetos e empilhamento de blocos.

Significância e Alegações
O artigo afirma que o PocketDP3 oferece um caminho prático para o controle visuomotor 3D eficiente. Ao desafiar a suposição de que decodificadores grandes são necessários para políticas baseadas em difusão quando a representação de entrada já está comprimida, os autores mostram que "menor é melhor" neste contexto específico. O método alcança desempenho de estado da arte com <1% dos parâmetros de métodos anteriores e permite inferência em tempo real sem treinamento de destilação complexo. Os autores posicionam este trabalho como um passo em direção à implantação de políticas robóticas sofisticadas em hardware com recursos limitados.

Limitações e Trabalhos Futuros
Os autores reconhecem que a arquitetura DiM pode enfrentar um "gargalo espacial" para tarefas que exigem raciocínio geométrico local refinado (ex: inserção submilimétrica) devido à falta de vieses espaciais alinhados ao pixel inerentes aos U-Nets. Eles também observam que o mecanismo subjacente que permite a inferência de alta qualidade em dois passos permanece obscuro e requer investigação adicional. Sugere-se que trabalhos futuros foquem no design de codificadores de condicionamento mais eficientes e ricos em informações, além de estudar sistematicamente os compromissos entre capacidade de condicionamento, complexidade da tarefa e escala do decodificador.