Resumo Técnico: MuJoCoUni

Declaração do Problema

Sistemas de aprendizado robótico integram cada vez mais simuladores de física diretamente no loop de treinamento, exigindo execução em alta frequência, em lote, de etapas de física, leituras de sensores e reinicializações de ambiente. Embora o MuJoCo upstream forneça ativos e solucionadores maduros, sua interface Python padrão incorre em sobrecarga relacionada ao gerenciamento de tempo de vida de objetos, chamadas de interface e formatação de saída em altas frequências.

Soluções existentes apresentam trade-offs:

• mujoco.rollout upstream: Eficiente para geração de trajetórias em malha aberta (planejamento, identificação de sistema), mas carece de persistência de ambiente com estado, capacidades de reinicialização esparsa e ciclos de vida de randomização por ambiente necessários para aprendizado por reforço (RL) online.
• Simuladores Residuais em GPU (ex.: Brax, MJX, Isaac Gym): Oferecem paralelismo massivo, mas frequentemente exigem que modelos e layouts de dados se ajustem a restrições específicas de aceleradores, potencialmente sacrificando cobertura total de recursos, tratamento complexo de contatos ou a capacidade de reutilizar ativos MuJoCo otimizados para CPU existentes.
• Runtimes Vetorizados em CPU Padrão: Frequentemente carecem de integração profunda com as semânticas específicas do solucionador MuJoCo ou exigem engenharia significativa para manter a persistência de modelos e a randomização de domínio de forma eficiente.

Há uma necessidade de um runtime baseado em CPU, em lote, que retenha a fidelidade semântica completa do MuJoCo upstream (modelos, sensores, contatos, restrições) enquanto fornece execução paralela de alto rendimento e gerenciamento de ambiente com estado para RL online e avaliação de física em lote.

Metodologia

MuJoCoUni é uma distribuição downstream leve do MuJoCo que concentra suas adições de engenharia estritamente na camada de ligação Python, deixando o núcleo de física, solucionador, integrador e modelo de contatos inalterados.

Arquitetura Central: BatchEnvPool

O componente central é o BatchEnvPool, um executor C++/pybind11 que gerencia:

1. Recursos Persistentes: Possui $N$ cópias de mjModel (uma por ambiente) e $W$ trabalhadores mjData (um por thread).
2. Pool de Threads: Um pool de threads interno distribui índices de ambiente para threads de trabalhadores para execução paralela.
3. Execução com Estado: Diferentemente do rollout, que é sem estado por chamada, o BatchEnvPool mantém o estado dos ambientes entre chamadas, permitindo reinicializações esparsas e variantes de modelo persistentes.

Primitivas Chave

A API expõe várias primitivas projetadas para treinamento e avaliação de robôs online:

• step: Realiza $n$ etapas de integração (mj_step) em todos os ambientes. Retorna apenas o estado final (e opcionalmente dados finais de sensores), evitando a sobrecarga de retornar trajetórias completas.
• forward: Executa mj_forward para atualizar sensores com base no estado atual sem avançar a dinâmica, útil para construção de observações.
• reset: Implementa reinicialização esparsa, atuando apenas em um subconjunto de ambientes terminados (env_ids). Suporta randomização de domínio com ciclo de vida de reinicialização, aplicando patches de campo (ex.: massa, inércia, atrito) e chamando mj_setConst apenas para os ambientes afetados.
• compute_site_jacobians: Cálculo em lote de jacobianos translacionais e rotacionais para locais específicos, executando o prefixo cinemático mínimo (mj_kinematics, mj_comPos) seguido de mj_jacSite.
• sample_hfield_height: Amostra bilinearmente geometrias de campos de altura do MuJoCo em lote, suportando vários alinhamentos de quadros (mundo, corpo, guinada).

Limites de Design

A implementação evita modificar a árvore de origem do MuJoCo. Baseia-se em mj_copyModel para instanciação de ambiente e mj_setConst para randomização. O design suporta dois modos:

1. Modelo Único Compartilhado: Todos os ambientes compartilham uma cópia de modelo (eficiente em memória, menos flexibilidade de randomização).
2. Variantes de Modelo: Cada ambiente possui uma cópia distinta de modelo, permitindo randomização ao nível de geometria (ex.: comprimentos de elos, escalas de malha) pré-compilada no momento da construção.

Contribuições Principais

1. Runtime em Lote com Estado: Introduz o BatchEnvPool para preencher a lacuna entre a geração de trajetórias em malha aberta e os requisitos de ambientes de RL online e com estado.
2. Reinicialização Esparsa e Randomização Eficientes: Fornece um mecanismo para reinicializar apenas ambientes terminados e aplicar randomização de domínio (tanto patches de campo quanto variantes de modelo) com sobrecarga mínima, escalando linearmente com o número de ambientes reinicializados.
3. Consultas de Física em Lote: Expõe implementações de alto desempenho, em lote, de passagens diretas de sensores, jacobianos de locais e amostragem de campos de altura sem avançar a dinâmica.
4. Preservação de Semânticas: Garante que o comportamento físico permaneça idêntico ao MuJoCo em CPU upstream, mantendo compatibilidade com ativos XML/MJB existentes, ferramentas de depuração e definições de sensores.

Resultados

Os autores validaram o MuJoCoUni em quatro modelos de robôs (Unitree Go1, Wonik Allegro, Franka Panda, CMU Humanoid) usando uma CPU Intel i9-14900HX com 16 threads.

• Vazão: A vazão satura em torno de 256–512 ambientes. Na saturação, o sistema alcança:
  • Wonik Allegro: ~1,8M etapas/segundo.
  • Unitree Go1: ~1,2M etapas/segundo.
  • Franka Panda: ~410k etapas/segundo.
  • CMU Humanoid: ~290k etapas/segundo.
• Desempenho de Reinicialização: O caminho C++ BatchEnvPool é significativamente mais rápido que loops Python. Para uma reinicialização completa de 4096 ambientes, leva 3,5 ms comparado a 53 ms para um loop for em Python (~15× de aceleração).
• Desempenho de Jacobianos: Calcular jacobianos de locais para 4096 ambientes leva 0,53 ms em C++ versus 11,9 ms em Python (~22× de aceleração).
• Amostragem de Campos de Altura: Amosstrar campos de altura para 4096 ambientes leva 0,52 ms em C++ versus 290 ms em Python (~555× de aceleração).
• Variantes de Modelo: A sobrecarga de usar cópias de modelo distintas por ambiente (modo variante de modelo) versus um único modelo compartilhado é negligenciável na saturação, com a vazão tornando-se essencialmente idêntica.

Significado e Escopo

O artigo posiciona o MuJoCoUni como uma solução de engenharia complementar, e não como um substituto para simuladores residuais em GPU ou o rollout upstream.

• Complemento à GPU: Alinha-se a cargas de trabalho onde a cobertura de recursos (semânticas completas do MuJoCo, contatos complexos, depuração) e a portabilidade entre plataformas são priorizadas sobre a residência bruta em aceleradores. Permite que sistemas aproveitem CPUs multicore modernas para geração de dados de alto rendimento, mantendo a estabilidade do solucionador MuJoCo em CPU.
• Complemento ao rollout: Atende às necessidades específicas de RL robótico online e loops de avaliação com estado, que exigem ambientes persistentes, reinicializações esparsas e consultas de estado atual, tarefas para as quais o rollout (projetado para geração de trajetória completa) não é otimizado.
• Integração de Sistema: Serve como um backend de simulação que desacopla a execução de física de lógica de treinamento de nível superior, registro e agendamento. É particularmente adequado para treinamento sim-to-real que exige randomização de domínio, locomoção consciente do terreno e otimização em lote (ex.: estratégias evolutivas).

Os autores enfatizam que a contribuição é orientada para engenharia, fornecendo um caminho concreto e reproduzível para execução em lote em CPU que preserva o ecossistema MuJoCo upstream.