mjbatch

Roda milhares de simulações do MuJoCo em paralelo na CPU

Licença

Apache-2.0

Visão geral

O mjbatch é uma biblioteca Python que roda milhares de simulações do MuJoCo em paralelo na CPU. Um pool de threads em C++ faz o avanço com o GIL liberado, e o lote é exposto ao Python como visões NumPy vivas: bind devolve um array sobre um campo do MjData em todas as simulações, de modo que escrever em ctrl e ler qpos dispensa cópia e laço por simulação. expand faz o mesmo com campos do MjModel, o que dá a cada simulação sua própria física — atritos diferentes, comprimentos de elo diferentes — com set_const recalculando as constantes derivadas depois. O repositório traz seis exemplos autocontidos cobrindo aprendizado por reforço, controle preditivo, identificação de sistemas e coprojeto de hardware.

O núcleo é deliberadamente pequeno. Você constrói um Batch a partir de um único MjModel e de uma contagem de simulações, o pool de threads assume por padrão todas as CPUs lógicas, e batch.step() avança todas elas. O que o torna mais que um laço é que os arrays são vivos em vez de coletados: qpos, ctrl = batch.bind("qpos"), batch.bind("ctrl") entrega visões que um controlador pode escrever sobre o lote inteiro numa única atribuição, e que se atualizam no lugar conforme o lote avança. Os parâmetros de modelo por simulação são a outra metade — batch.expand("geom_friction") devolve um array que se pode aleatorizar por simulação, e isso já é aleatorização de domínio e identificação de sistemas sem maquinaria adicional.

Os exemplos são o argumento, porque cada um é um método completo e não uma demonstração: um pêndulo sobre carrinho erguido com iLQR, um MPC de pêndulo sobre carrinho por amostragem preditiva, um humanoide G1 seguindo um mortal de costas de referência com iLQR de horizonte deslizante, um controlador de joystick para o quadrúpede Go1 treinado com PPO, um CEM otimizando em conjunto proporções, engrenagens e controle de um braço arremessador, e um ajuste de Gauss-Newton amortecido dos parâmetros inerciais de um braço Rizon. O autor relata que o controlador do Go1 aprende a andar em menos de um minuto num notebook M1 de cinco anos. Instala pelo PyPI, fixa o mujoco 3.13.0, compila via scikit-build-core com nanobind, e traz CI, uma suíte pytest e um benchmark de escalabilidade.

Por que importa

Simulação massivamente paralela virou uma história de GPU — MJX, Isaac, Warp — e isso elevou em silêncio o custo de entrada para quem não tem uma placa de estação de trabalho. O mjbatch devolve esse tipo de trabalho ao alcance de um notebook usando a CPU direito em vez de um núcleo por vez, e faz isso através do MuJoCo puro em vez de uma reimplementação, então os modelos e a física continuam sendo aqueles em que você já confia. Para um catálogo cujo público em geral tem um notebook e um braço barato, e não um cluster, esse é o lado útil do compromisso.

Problema que resolve

Conseguir milhares de rollouts do MuJoCo costumava obrigar a migrar para uma pilha de GPU e aceitar sua conta de hardware e sua reimplementação da física. O mjbatch mantém o MuJoCo original e extrai o paralelismo da CPU que já está ali, o que basta para trabalho de controle, identificação e coprojeto em escala de notebook.

Pontos principais

• Um pool de threads em C++ avança o lote inteiro com o GIL liberado, então o paralelismo é real e não Python intercalado
• bind expõe campos do MjData como arrays vivos sobre todas as simulações — o controlador escreve uma vez em ctrl e lê qpos no lugar, sem etapa de coleta
• expand dá a cada simulação seus próprios parâmetros de MjModel, com set_const recalculando as constantes derivadas — aleatorização de domínio e identificação sem maquinaria extra
• Seis exemplos trabalhados, cada um um método completo: iLQR, MPC por amostragem preditiva, PPO, coprojeto de hardware com CEM e ajuste inercial de Gauss-Newton amortecido
• O autor relata um controlador de joystick do Go1 que aprende a andar em menos de um minuto num notebook M1 de cinco anos
• Instala pelo PyPI, fixa o mujoco 3.13.0 e traz CI, suíte pytest e benchmark de escalabilidade

Para quem é

• Pesquisadores de MPC, otimização de trajetória ou CEM que precisam de muitos rollouts por passo
• Quem treina controladores pequenos por reforço sem acesso a um cluster de GPU
• Profissionais que fazem identificação de sistemas ou aleatorização de domínio sobre parâmetros do modelo
• Usuários do MuJoCo que querem simulação em lote sem migrar para MJX ou Warp

Limitações

• Somente CPU por projeto — para lotes em escala de GPU o caminho continua sendo MJX ou MuJoCo Warp, e isto não os substitui
• Seis dias de vida na hora em que isto foi escrito, em v0.1.1 e classificado como Beta; as 521 estrelas refletem tanto o histórico do autor quanto a biblioteca
• Fixa mujoco==3.13.0 exatamente, então só se move quando o autor sobe a versão
• Precisa de uma toolchain C++ para compilar do código, e os exemplos que abrem janela precisam de display a menos que rodem com --headless
• A superfície de API é deliberadamente estreita — Batch, bind, expand, set_const — então sair dessa forma significa descer direto ao MuJoCo
• Nos classificadores publicados constam apenas macOS e Linux; não há wheels para Windows