Praxis-1

praxis é o adaptador do LeRobot para o Praxis-1, o World Action Model de pesos abertos da Runway para robôs. O Praxis-1 é um modelo de vídeo-ação: um backbone de difusão de vídeo com causalidade por blocos (um transformer mixture-of-experts com 27 bilhões de parâmetros, cerca de 3 bilhões ativos por token) remove o ruído dos próximos 1,6 s de vídeo das duas câmeras do rig, enquanto um especialista de ações, lendo os estados ocultos do backbone, remove o ruído dos 24 comandos cartesianos correspondentes a 15 Hz. O autoencoder de vídeo Wan2.2 e o codificador de texto Qwen3.5-2B são congelados e vêm de seus pesos públicos. O checkpoint liberado, praxis-1-droid, controla um rig DROID de braço único: um braço Franka com uma câmera externa e uma câmera no pulso. Um checkpoint DROID anterior, não liberado, tem a mesma arquitetura com outros pesos e outro canvas de câmera; o config.json do checkpoint os distingue (veja “The checkpoints”). praxis-1-yam, o checkpoint bimanual para dois braços YAM, é distribuído no mesmo layout (veja “The bimanual YAM checkpoint”). praxis-1-base, o backbone de vídeo e o tronco do especialista de ações sem as cabeças de um rig, é o ponto de partida para um fine-tune de um novo rig de braço único (veja “Fine-tune the base checkpoint”).

Recursos:

O código do modelo está na árvore do repositório, em lerobot/policies/praxis/praxis_core (vendado do repositório praxis no commit 8dad1a4, mais a regra canvas de praxis-1-droid, com os caminhos de importação reescritos; os cabeçalhos dos arquivos indicam os commits originais). A simulação MuJoCo é o pacote praxis-sim no PyPI (desenvolvido em packages/praxis-sim do repositório praxis, sem torch); o LeRobot não possui um ambiente para ele.

Instalação

pip install -e ".[praxis,dataset]"   # the policy; `lerobot-eval` and `lerobot-train` also need `dataset`

Os checkpoints convertidos estão no Hub: --policy.path=runwayml/praxis-1-droid-lerobot carrega o checkpoint de braço único lançado (54 GB, baixado uma vez). runwayml/praxis-1-yam-lerobot é o checkpoint bimanual. runwayml/praxis-1-base-lerobot é a base, que precisa de init_heads.py primeiro.

A política é executada nas versões que o LeRobot 0.7 fixa (0.7.1 no momento da escrita: torch 2.15.0 com CUDA 13.2, transformers 5.22.2, triton 3.9). Os resultados abaixo foram medidos nessa pilha (com transformers 5.22.1 e 5.22.2, que produzem os mesmos comandos) e na do LeRobot 0.6 (torch 2.11.0 com CUDA 12.8, transformers 5.5.4); os comandos são iguais bit a bit aos do repositório praxis em cada uma.

Kernels opcionais da configuração de validação do checkpoint. flash-attn-4 fornece o backend de atenção flash da configuração do checkpoint (atenção flexível com o kernel CuTe do FlashAttention-4); sem ele, a política usa o kernel Triton da atenção flexível. flash-linear-attention e causal-conv1d são os kernels rápidos das camadas de atenção linear do Qwen3.5; sem eles, o transformers executa sua implementação em PyTorch. Os três alteram os últimos bits das saídas, não o modelo:

pip install "flash-attn-4[cu13]==4.0.0b36"        # the CUDA 13 build for torch 2.15; "flash-attn-4==4.0.0b36" on a CUDA 12 torch
pip install "flash-linear-attention==0.6.0"       # 0.5.0 does not run on triton 3.9
CAUSAL_CONV1D_FORCE_BUILD=TRUE MAX_JOBS=16 pip install --no-build-isolation "causal-conv1d==1.6.1"   # compiles against torch; the praxis repository's setup

A inferência requer uma GPU com 80 GB de memória (os pesos ocupam 54 GB em bfloat16).

Em GPUs Blackwell de datacenter (B200 e GB200 são SM100, B300 e GB300 são SM103; Jetson Thor, SM110, tem o mesmo FlashAttention-4 forward):

• B300 e GB300 exigem um build do torch para CUDA 12.9 ou posterior, como o torch 2.15.0+cu132 que o LeRobot 0.7 fixa. Os especialistas do MoE chamam torch._grouped_mm, e o build do torch 2.11.0+cu128 (LeRobot 0.6) não possui kernel SM103 para ele: a primeira chamada de política falha com RuntimeError: cutlass cannot run, error 7, como ocorre no pipeline do repositório praxis. O B200 funciona em ambas as pilhas.
• FlashAttention-4 aceita máscaras block-sparse apenas em blocos de 256 tokens (ele levanta ValueError: Block sparse tensorsforward have block size 128, which must be a multiple of 256). Com o backend flash, a política constrói as máscaras do backbone em blocos de 256 tokens nestes GPUs, em vez do backbone.attention.flex_block_size de 128 do checkpoint. O tamanho do bloco apenas agrupa os pares de tokens: a função de máscara ainda decide cada par. O repositório praxis executa o mesmo com "flex_block_size": 256 no config.json do checkpoint. O backend triton mantém o 128 do checkpoint.

Em GPUs RTX Blackwell (capacidade de computação 12.x; a RTX PRO 6000 Blackwell possui os 96 GB que os pesos necessitam), o FlashAttention-4 não possui forward block-sparse em nenhum tamanho de bloco (com os blocos de 128 tokens do checkpoint, ele levanta Block sparsity requires sparse_block_size[1]=128 to match tile_n=64), portanto o backend flash não pode ser executado ali. attention_backend="auto" seleciona triton nessas GPUs, mesmo com flash-attn-4 instalado, e "flash" levanta um ValueError quando a política é construída. O repositório praxis falha na sua primeira chamada com o flash do checkpoint e executa o mesmo com "backend": "triton" no config.json do checkpoint. Os dois backends diferem nos últimos bits, portanto compare execuções nessas GPUs com execuções de triton em outros locais. Nada mais muda em uma RTX PRO 6000: as compilações torch de ambos os stacks abaixo incluem kernels SM120 (incluindo torch._grouped_mm do MoE), e a inferência com as configurações de praxis-1-droid atinge um pico de cerca de 62 GB.

Essas regras seguem as próprias verificações do FlashAttention-4. Os resultados de paridade abaixo foram medidos em um GB300, um B200 e um RTX PRO 6000 Blackwell.

transformers 5.4.0 a 5.5.4 (LeRobot 0.6; corrigido em 5.6.0, e LeRobot 0.7 fixa 5.22) carregam o checkpoint do Hub Qwen/Qwen3.5-2B com todos os 320 tensores do modelo de linguagem reportados como ausentes (ele renomeia os nomes armazenados model.language_model.* antes de fazer o pareamento) e os inicializa silenciosamente de forma aleatória. O codificador de texto da política detecta isso e carrega os tensores sob seus nomes armazenados, com um aviso no log. Um codificador de texto deixado aleatório resulta em uma política que ainda se move, mas raramente tem sucesso (7 de 56 episódios de guarda na nossa primeira execução).

Entradas e saídas

A política mantém seu próprio histórico: um buffer circular com as últimas 25 observações por ambiente do lote (o chunk executado mais o frame atual) alimenta a próxima chamada como contexto, de modo que n_obs_steps permanece 1 e select_action deve ser chamado a cada tick de 15 Hz com a observação daquele tick. A primeira chamada de uma janela de contexto retorna 21 comandos, e cada chamada subsequente retorna 24. Com n_action_steps abaixo de 24, cada chamada é um início a frio sem contexto, como no praxis. drop_queued_actions (por exemplo, uma nova instrução no meio do rollout) também reinicia o contexto temporal, pois o modelo condicionaliza no fato de todo o chunk anterior ter sido executado. Um chamador que executa chunks inteiros (predict_action_chunk) passa os frames executados desde sua chamada anterior como uma pilha, (B, T, 3, H, W) câmeras com (B, T, 7) estados, do mais antigo para o mais recente; o próximo chunk é então o de um rollout com chamada por tick. Uma chamada de chunk com um frame reinicia o contexto novamente.

Processadores

O pré-processador prepara as entradas da mesma forma que o código de dados e de serviço da praxis o faz, como etapas salvas e inspecionáveis:

Entre eles, a política vê tiles e linhas pose8 e emite poses de wire, exatamente o que a sessão Session da praxis consome e produz. A normalização de ação q99, a normalização proprio de média/desvio padrão e a decodificação de ação relativa permanecem dentro da sessão: eles são ancorados à pose no início de cada chunk executado, que a sessão rastreia ao longo do rollout, e suas estatísticas são as do checkpoint (PraxisConfig.model_config["normalization"]); as etapas do normalizador padrão são identidade. A política também aceita a pose bruta de 7 valores, os frames brutos da câmera e a string task, para que possa ser chamada sem o pipeline: frames no tamanho do tile são tratados como tiles, outros tamanhos passam pela regra de canvas do checkpoint (a regra de fit de praxis-1-droid não é idempotente, então um tile nunca passa por ela duas vezes), e uma task sem tokens é tokenizada pela política. Os tiles permanecem no dispositivo em que o pipeline os coloca até o autoencoder lê-los, e os comandos voltam na CPU: nada cruza a fronteira do dispositivo duas vezes.

Os checkpoints

O canvas vem do config.json do checkpoint (video.tile_height, video.antialias) por meio do conversor; as configurações do sampler são campos de PraxisConfig, portanto uma execução do checkpoint DROID anterior passa os seus próprios (--policy.t_shift=8.0 --policy.cfg_video_scale=1.0 --policy.max_context_chunks=3).

Converter o checkpoint

Um checkpoint praxis é um diretório com config.json, 17 shards safetensors, model.safetensors.index.json e SHA256SUMS. Converta-o em um diretório de política LeRobot:

python examples/praxis/convert_checkpoint.py \
    --checkpoint /path/to/praxis-1-droid --out outputs/praxis-1-droid --link

O rig vem do embodiment do checkpoint: o checkpoint bimanual recebe suas três chaves de câmera (--camera-keys para renomeá-las) e as configurações de serve do seu checkpoint. O checkpoint base recebe um diretório sem processadores que o init_heads.py completa. O runwayml/praxis-1-droid-lerobot e os outros checkpoints convertidos são a saída deste script.

--link cria links simbólicos para os shards em vez de copiar 54 GB. O diretório contém o config.json do LeRobot, os shards, o índice e seu próprio SHA256SUMS (sha256sum -c passa nele) (as políticas do LeRobot leem e escrevem safetensors em shards no layout do transformers; PraxisPolicy.save_pretrained escreve shards de 5 GB), e os dois arquivos de processador. --vae-id e --text-encoder-id definem de onde vêm as partes congeladas (ids do Hub por padrão).

Resultados

Os números de laço fechado abaixo vêm do rig MuJoCo do repositório praxis (o pacote praxis-sim: Franka + Robotiq 2F-85, câmeras 512 x 288, 15 Hz) e de seus conjuntos de avaliação fixos; a semente LeRobot de um episódio é a semente da cena dele. O ambiente que executa os conjuntos por meio de lerobot-eval não faz parte do LeRobot.

Durante a amostragem, a política fixa as configurações numéricas com as quais o praxis é executado (sem o modo de benchmark do cuDNN, multiplicações de matrizes em float32 para o roteador do MoE, TF32 permitido no cuDNN) e restaura as configurações do host em seguida, de modo que --tf32, torch.backends.cudnn.benchmark e o tamanho do lote não alteram suas saídas.

A suíte droid_guard (7 tarefas DROID x 8 cenas, 56 episódios, semente do sampler 0) com o checkpoint DROID anterior, não lançado (tiles 256 x 416), nas configurações do sampler do repositório praxis na época (2 passos Euler, flow shift 1.0, action guidance 3.0, sem video guidance, 1 chunk de contexto) e seu backend de atenção flash, em uma H100 80 GB (cerca de 10 minutos). Os mesmos episódios foram executados de cinco maneiras: o próprio runner do repositório praxis em seu próprio ambiente (torch 2.13, transformers 5.2.0, causal-conv1d), o mesmo runner e lerobot-eval na pilha do LeRobot 0.6 (torch 2.11, transformers 5.5.4), e ambos novamente na pilha que o LeRobot 0.7 fixa (torch 2.15 com CUDA 13.2, triton 3.9, flash-linear-attention 0.6.0; transformers 5.16.1 na época). Repetir as execuções com torch 2.15 usando transformers 5.22.1 produz comandos idênticos bit a bit e os mesmos resultados de 46/56 episódios tanto para o runner do praxis quanto para lerobot-eval, e a versão 5.22.2 (a versão que o LeRobot 0.7.1 fixa) produz os mesmos comandos para praxis-1-droid e o checkpoint YAM: a versão do transformers não entra nos números.

Apareados ao longo dos 56 episódios (praxis_sim.compare), lerobot-eval menos a referência:

Todos os intervalos incluem zero: a política e o ambiente do LeRobot reproduzem os resultados do repositório praxis, e nem a pilha de software (kernels, versão do transformers) nem as alterações no harness os modificam além do ruído episódio a episódio. O próprio runner do praxis se desloca na mesma magnitude entre as pilhas (44 a 46 de 56). press_button é difícil para este checkpoint em todas as configurações.

Com entradas idênticas, a política é igual bit a bit ao pipeline praxis em ambos os stacks (veja abaixo); entre os dois stacks, os comandos de entradas idênticas diferem em até 0,04 (flash) e 0,02 (triton), o que é o que move os episódios individuais.

praxis-1-droid , o checkpoint publicado

Os mesmos 56 episódios com praxis-1-droid (tiles antialiased de 224 x 416) em suas configurações recomendadas (2 passos Euler, flow shift 12.5, video guidance 4.0, sem action guidance, 1 context chunk) e o backend flash, em uma H100 80 GB: o runner do repositório praxis (seu main com o sampler atual, no ambiente praxis: torch 2.13, transformers 5.2.0), lerobot-eval na stack do LeRobot 0.6 (torch 2.11, transformers 5.5.4) e tanto o runner quanto o lerobot-eval na stack que o LeRobot 0.7 fixa (torch 2.15 com CUDA 13.2, transformers 5.22.1).

Apareados nos 56 episódios, lerobot-eval menos o runner praxis é +0,0 pontos em ambos os stacks (intervalo bootstrap de 95 % [-7,1 %, +7,1 %], venceu 2 / perdeu 2, McNemar p 1,00 cada). Os episódios que mudam não são os mesmos nos dois stacks: os kernels do torch 2.15 movem os comandos em até 0,02, e alguns episódios de fronteira seguem. praxis-1-droid com suas configurações resolve press_button, o que o checkpoint anterior do DROID com suas configurações falhou na maioria das vezes.

Paridade com o repositório praxis

• Para observações idênticas em um processo, a política LeRobot e a Session do praxis produzem embeddings de texto, latentes do VAE, dados proprioceptivos, ações amostradas e comandos iguais bit a bit (mesmos pesos carregados). tests/policies/praxis/test_praxis.py repete essa verificação em um modelo aleatório pequeno na CPU.
• Entre processos, a execução do pipeline praxis por meio do pacote original praxis e a política do LeRobot geram comandos iguais bit a bit no checkpoint real, tanto no backend flash quanto no backend triton (os dois backends diferem entre si em até 0.055 nos comandos), na pilha do LeRobot 0.6 (torch 2.11.0+cu128, transformers 5.5.4, flash-attn-4 4.0.0b36) e na pilha que o LeRobot 0.7 fixa (torch 2.15.0+cu132, transformers 5.22.1 e 5.22.2, triton 3.9, flash-attn-4 4.0.0b36). A própria referência varia entre as duas pilhas (em até 0.019 no flash, 0.015 no triton); a política a acompanha bit a bit em cada uma.
• praxis-1-droid: a verificação entre processos contra o pipeline praxis do repositório praxis, na branch main (com seu sampler atual), é igual bit a bit nos dois backends e nas duas stacks em oito maneiras de chamar a política: através de seus pipelines de processador salvos (como o lerobot-eval faz), em frames brutos e a pose de 7 valores sem os pipelines, com dois ambientes em um único lote (instruções e sementes de cena diferentes), através de predict_action_chunk (um cold start, e o próximo chunk dado os frames executados como uma stack), com n_action_steps=12 (um chunk a cada 12 comandos, condicionado aos frames executados), com orientação de ação ativada (cfg_scale=3), e sobre 70 comandos com três chunks de contexto (profundidade 0 a 3). A verificação se mantém sob as flags padrão de TF32 e benchmark do cuDNN do lerobot-eval: a política fixa a numérica do praxis enquanto amostra e devolve as flags ao host. O passo de tile computa o canvas na CPU, independentemente do dispositivo dos frames: o redimensionamento antialiased da CUDA arredonda um nível diferente em cerca de um terço dos pixels, o que moveu os comandos em até 0.02 antes que o passo o fizesse.
• Treinamento: nos pesos liberados, PraxisPolicy.forward e praxis.train.objective do repositório praxis produzem a mesma perda (e os termos de vídeo e ação) bit a bit para a mesma janela e extrações aleatórias, e os mesmos gradientes em todos os 200 tensores do especialista de ação. A métrica de equilíbrio MoE reportada varia em um ou dois ulps de float32 entre execuções em qualquer um dos lados (uma redução não determinística na estatística de equilíbrio; ela entra na perda com peso 1e-4).
• GB300 (SM103, um host aarch64): na pilha, o LeRobot 0.7 fixa (torch 2.15.0+cu132, cuDNN 9.26, transformers 5.22.1, triton 3.9, flash-attn-4 4.0.0b36, flash-linear-attention 0.6.0), a verificação entre processos contra o repositório praxis, na main, é igual bit a bit nos dois backends (flash com blocos de máscara de 256 tokens em ambos os lados): para praxis-1-droid em sete maneiras de chamar a política (os pipelines salvos sob as flags de lerobot-eval, quadros brutos de 720 x 1280, dois ambientes em um lote, predict_action_chunk com pilhas de quadros, n_action_steps=12, cfg_scale=3 e 94 comandos com três blocos de contexto), para praxis-1-yam através dos pipelines, em quadros brutos e com dois ambientes, e para o rig DROID reconstruído a partir de praxis-1-base. O forward fornece a perda do objetivo do praxis e os gradientes de todos os 200 tensores de especialistas de ação bit a bit nos pesos liberados.
• RTX PRO 6000 Blackwell (SM120, 96 GB, host x86_64), triton nos dois lados: na pilha do LeRobot 0.6 (torch 2.11.0+cu128, cuDNN 9.19, transformers 5.5.4, triton 3.6, flash-linear-attention 0.5.0) e na pilha que o LeRobot 0.7 fixa (torch 2.15.0+cu132, cuDNN 9.26, transformers 5.22.1, triton 3.9, flash-linear-attention 0.6.0 e causal-conv1d 1.6.1, para que o Qwen3.5 execute seus kernels), ambos com flash-attn-4 4.0.0b36 instalado, a verificação entre processos contra o main do repositório praxis é igual bit a bit nos mesmos termos do GB300: praxis-1-droid nos sete (463 comandos, profundidades de contexto de 0 a 3), praxis-1-yam nos três, e o rig DROID reconstruído a partir de praxis-1-base (seus seis tensores de cabeçalho e sua configuração são iguais aos do init_heads do repositório praxis). O forward fornece a perda do objetivo praxis e os gradientes de todos os 200 tensores de especialistas de ação bit a bit em ambas as pilhas, com o backbone congelado nos dois lados (seus 54 GB de gradientes não cabem ao lado dos pesos). Na pilha do LeRobot 0.6, o lado praxis também precisa da correção do Qwen3.5 para o aviso acima: o praxis-1 fixa o transformers 5.2.0, e sob o 5.5.4 seu codificador de texto permanece aleatório. Loop fechado no droid_guard com praxis-1-droid na pilha do LeRobot 0.6: lerobot-eval 53/56 (94,6 %), o runner praxis 54/56 (96,4 %, como no H100 e no GB300); em pares, -1,8 pontos, intervalo de 95 % [-8,9, +5,4], venceu 2 / perdeu 3, p de McNemar 1,00. Um fine-tune LoRA através do lerobot-train (r 16, alpha 32, alvos padrão, uma janela por passo) roda na única GPU: 5,5M de parâmetros treináveis, 1,1 s por passo, 62 GB.
• B200 (SM100, host x86_64): na pilha do LeRobot 0.6 (torch 2.11.0+cu128, cuDNN 9.19, transformers 5.5.4, triton 3.6) e na pilha fixada pelo LeRobot 0.7 (torch 2.15.0+cu132, cuDNN 9.26, transformers 5.22.1, triton 3.9, flash-linear-attention 0.6.0), ambos com flash-attn-4 4.0.0b36, a verificação entre processos contra o main do repositório praxis é igual bit a bit nos dois backends (flash com blocos de máscara de 256 tokens em ambos os lados) nas mesmas condições do GB300: praxis-1-droid nos sete (463 comandos por backend e pilha, profundidades de contexto de 0 a 3), praxis-1-yam nos três, e o rig DROID reconstruído a partir de praxis-1-base (seus seis tensores de cabeçalho e sua configuração são iguais aos do init_heads do repositório praxis). Os dois backends diferem entre si em até 0.19 nos comandos. O forward fornece a perda do objetivo praxis, seus termos de vídeo e ação e os gradientes de todos os 200 tensores de especialistas de ação bit a bit nos dois backends e nas duas pilhas, com o backbone treinado. Um fine-tune LoRA através do lerobot-train roda na única GPU (1,2 s por passo, 62 GB), e o adaptador e seu checkpoint mesclado carregam no lerobot-eval.
• Loop fechado no B200 (droid_guard, praxis-1-droid, flash): o runner praxis pontua 54/56 no stack torch 2.15 e 55/56 no torch 2.11, lerobot-eval 53/56 e 51/56 (pareado: -1.8 pontos, intervalo de 95 % [-8,9, +5,4]; e -7.1 pontos, [-14,3, -1,8], venceu 0 / perdeu 4, McNemar p 0.125). Cada harness se repete episódio por episódio, então a diferença não é ruído entre execuções, e também não é a política: lerobot-eval carrega o estado e os comandos como os 7 valores do DROID em float32 (PraxisSimGymEnv, as etapas do processador de pose e comando), enquanto o runner praxis passa as poses e os comandos da sessão como estão. A ida e volta desloca uma pose em cerca de 1e-7, e ao longo de um episódio a diferença cresce para outra trajetória. O runner praxis com seu estado e comandos enviados através da mesma conversão float32 produz o resultado do lerobot-eval em todos os 56 episódios, nos dois stacks.
• Em relação ao próprio ambiente do repositório praxis (torch 2.13.0+cu130, transformers 5.2.0, causal-conv1d), as saídas não são iguais bit a bit: os kernels diferem, o embedding do Qwen3.5 difere em até 0.125 em bfloat16, e os últimos bits das ações seguem essa tendência. Portanto, a verificação entre repositórios é a avaliação fechada em laço pareada acima.
• A simulação é renderizada com 4x multisampling, e alguns pixels por quadro diferem entre renderizações da mesma cena, de modo que dois episódios em laço fechado da mesma política divergem após o primeiro chunk, mesmo com pesos idênticos. Episódios em pares são comparados estatisticamente, não quadro a quadro.

O checkpoint bimanual do YAM

praxis-1-yam ("embodiment": "bimanual", runwayml/praxis-1-yam-lerobot) controla dois braços YAM com uma câmera superior e duas câmeras de pulso (visões ZED 16:10, 1920 x 1200). O mesmo modelo o executa; o checkpoint difere nas larguras da cabeça de especialista (20 valores de ação: [xyz, rot6d, gripper] por braço), no canvas de três tiles e na normalização. O adaptador LeRobot segue o praxis/bimanual.py do repositório praxis (vendored como praxis_core/bimanual.py):

As configurações de serve do checkpoint (o arquivo configs/serve_yam.yaml do repositório praxis: 2 passos de Euler, flow shift 1.0, sem orientação de texto, 1 chunk de contexto, o rig bimanual não possui modo de controle) são o que o conversor escreve na configuração, portanto --policy.path=runwayml/praxis-1-yam-lerobot as executa. Os prompts são enviados ao codificador de texto verbatim; o checkpoint foi treinado com legendas como “Pick up the pen from the center of the table, then place the pen in the bag”, sem ponto final.

Paridade: no checkpoint real, a política LeRobot e o pipeline praxis BimanualServeSession geram comandos iguais bit a bit para as mesmas três visões, poses e instrução (22 ticks: o chunk de inicialização a frio e o primeiro comando do chunk de profundidade 1) através dos pipelines salvos, em frames brutos sem eles e com dois ambientes de instruções diferentes em um único lote. O teste de modelo pequeno repete a verificação na CPU contra a Session do praxis.

Limitações: não há simulação do rig YAM, portanto não há suíte de laço fechado; forward (fine-tuning) é apenas de braço único, pois o repositório praxis ainda não possui receita de treinamento bimanual.

Treinamento

lerobot-train faz o fine-tuning da política com o objetivo do repositório praxis (praxis/train): a perda de flow-matching conjunta do vídeo e das ações, o dropout de texto e proprio, o termo de balanceamento de sequência do MoE e a atualização do routing-bias após cada passo do otimizador. Dado a mesma janela e os mesmos sorteios aleatórios, PraxisPolicy.forward é igual bit a bit ao objetivo do treinador praxis no checkpoint liberado. Uma janela é de 93 frames consecutivos (chunks de 21, 24, 24, 24) de ambas as câmeras, com a pose observada e comandada da flange de cada frame e a instrução; a política declara isso por meio de observation_delta_indices / action_delta_indices (93 frames) e mantém as janelas dentro dos episódios com drop_n_last_frames (93, para que o sampler sorteie os inícios de janela da receita). Os tiles da câmera seguem a regra de canvas do checkpoint (video.antialias), como no módulo de dados praxis. forward roda sob a numérica praxis, independentemente do que o launcher configurar: sem autocast (as pontuações do roteador MoE permanecem em float32), sem TF32 em multiplicações de matriz em float32. O número de janelas de treinamento vistas é um buffer da política (train_counter): ele semeia os sorteios, viaja com o checkpoint e o adaptador, e uma execução retomada continua a sequência; uma avaliação com dados retidos não o avança.

Dados

A receita treina no DROID. lerobot/droid_1.0.1 mantém a pose cartesiana em seis mais uma colunas; as colunas de 7-D observation.state.cartesian e action.cartesian de PR #11 pré-montam elas (use --dataset.revision=refs/pr/11 até que seja mesclado). O próprio observation.state e action do conjunto de dados são no espaço de juntas; o mapa de renomeação coloca as colunas cartesianas em seu lugar (uma chave renomeada substitui uma existente) e as duas câmeras sob as chaves que a política lê:

RENAME='{"observation.state.cartesian":"observation.state","action.cartesian":"action",
"observation.images.exterior_1_left":"observation.images.exterior","observation.images.wrist_left":"observation.images.wrist"}'

lerobot-train aplica o mapa ao lote bruto e novamente dentro do pipeline do processador, portanto, ele não deve trocar as chaves (mapear observation.state para outro nome seria desfeito).

Metade dos episódios do DROID não possui instrução e alguns são mais curtos que uma janela; a receita os exclui. examples/praxis/select_droid_episodes.py escreve a lista para --dataset.episodes:

python examples/praxis/select_droid_episodes.py --repo-id lerobot/droid_1.0.1 \
    --revision refs/pr/11 --out droid_episodes.json

A cópia no Hub armazena vídeos de 180x320; o checkpoint foi treinado nos gravações DROID de 1280x720. Os tiles são ampliações com menos detalhes na cópia do Hub, o que é aceitável para um teste de fumaça. Para um fine-tune fiel, crie uma cópia local com os vídeos brutos, conforme descrito no praxis/train/README.md do repositório praxis, e aponte --dataset.root para ela.

Ajuste fino completo

Os pesos são bfloat16; o preset do otimizador praxis_adamw é AdamW com stochastic rounding (um AdamW comum arredonda os updates de 3e-6 para zero), o scheduler é um linear warmup e um cosine decay para 3e-7, e cada bloco do transformer recalcula seu forward no backward pass (activation_checkpointing, ativado por padrão). Pesos, gradientes e momentos do otimizador ocupam 8 bytes por parâmetro, 218 GB para o 27.2B, distribuídos entre as GPUs por FSDP: um fine-tune completo requer no mínimo 4 x 80 GB. A receita praxis é 96 janelas por passo para 1.200 passos; em 8 GPUs, isso é --batch_size=1 --accelerator.gradient_accumulation.steps=12. lerobot-train conta micro-batches (um passo do loop é um batch em cada GPU) e avança o scheduler por micro-batch, então os 1.200 updates são --steps=14400, com o warmup e o decay em micro-passos.

torchrun --nproc-per-node=8 $(which lerobot-train) \
    --dataset.repo_id=lerobot/droid_1.0.1 --dataset.revision=refs/pr/11 \
    --dataset.episodes="$(cat droid_episodes.json)" --rename_map="$RENAME" \
    --policy.path=runwayml/praxis-1-droid-lerobot --policy.device=cuda --policy.push_to_hub=false \
    --parallelism.dp_shard=8 --accelerator.fsdp.mixed_precision_policy.reduce_dtype=float32 \
    --batch_size=1 --accelerator.gradient_accumulation.steps=12 --steps=14400 \
    --policy.scheduler_warmup_steps=600 --policy.scheduler_decay_steps=14400 \
    --num_workers=8 --save_freq=100 --output_dir=outputs/train/praxis_droid

Um único save contém o modelo em shards de 54 GB (pretrained_model/, carregável por PraxisPolicy.from_pretrained e lerobot-eval) e cerca de 108 GB de estado do otimizador. PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True evita fragmentação do alocador, como no praxis trainer.

Medido em 8 x H100 80 GB com praxis-1-droid, uma janela por GPU e --dataset.episodes de 256 episódios DROID (os vídeos do Hub), 1.000 passos do otimizador: 3,7 s por passo após o primeiro (que compila os kernels), 48,5 GB de memória de GPU por rank, o data loader acompanhando com 8 workers por rank. A perda vai de 0,216 nos primeiros cem passos para 0,190 nos últimos cem (vídeo 0,191 para 0,171, ação 0,025 para 0,019). Um checkpoint salva o modelo shardado e o estado do otimizador em cerca de dois minutos através do FSDP. O checkpoint fine-tuned pontua 50 de 56 no droid_guard contra 54 do base: os vídeos do Hub são upscales de 180x320, então um fine-tune neles afasta a política das gravações 1280x720 em que ela foi treinada.

LoRA

Adaptadores por meio do PEFT (pip install -e ".[peft]"). Os alvos padrão são o grupo de expert do praxis, as projeções de atenção e MLP do expert de ação (PraxisPolicy.PEFT_TARGETS); as projeções de saída da atenção do backbone formam o segundo grupo, passado com --peft.target_modules. Os pesos dos experts do MoE não são camadas lineares e não recebem adaptadores. Com os pesos base congelados, uma execução com LoRA cabe em menos GPUs:

torchrun --nproc-per-node=4 $(which lerobot-train) ... --parallelism.dp_shard=4 \
    --peft.r=16 --peft.lora_alpha=32 --output_dir=outputs/train/praxis_lora

Deixe --peft.full_training_modules sem definir: o padrão de modules_to_save da política é o contador de treinamento, que o adaptador precisa carregar (a configuração rejeita um adaptador sem ele); os únicos outros módulos que ela aceita são os três módulos de cabeça do rig, para um fine-tune a partir do checkpoint base (abaixo). Os vieses de roteamento do MoE pertencem à base e permanecem fixos sob PEFT. Medido com r 16, alpha 32 e os alvos padrão em 4 x H100 (--parallelism.dp_shard=4): 5,5M de parâmetros treináveis, 2,3 s por passo com uma janela por GPU, 30 GB de memória de GPU por rank, um adaptador de 22 MB; --parallelism.dp_replicate com dp_shard (HSDP) também funciona. A perda total permanece próxima de 0,22 ao longo de 500 passos porque o backbone congelado detém o termo de vídeo; o termo de ação vai de 0,026 para 0,022. O adaptador dessa execução pontua 52 de 56 em droid_guard, carregado sobre a base por lerobot-eval.

O checkpoint contém o adaptador (adapter_config.json, adapter_model.safetensors); o lerobot-eval o carrega sobre o modelo base por meio de --policy.path. Para obter um checkpoint simples para a simulação, o servidor e from_pretrained, incorpore o adaptador nos pesos:

python examples/praxis/merge_praxis_adapter.py --base runwayml/praxis-1-droid-lerobot \
    --adapter outputs/train/praxis_lora/checkpoints/last/pretrained_model --out outputs/praxis-1-droid-lora

Diferenças em relação ao praxis trainer

O objetivo, a receita da janela, o otimizador e o agendamento são os da praxis. Os dados vêm de LeRobotDataset em vez do leitor de arquivos da praxis, as extrações aleatórias são semeadas a partir da semente do LeRobot, da contagem de chamadas, do rank e da posição no lote (de modo que as perdas sejam reproduzíveis, mas não sejam idênticas, passo a passo, às de uma execução da praxis), e os checkpoints são checkpoints do LeRobot. Não portado: o modo de overfit com extrações fixas e no_sync.

Ajuste fino do checkpoint base

praxis-1-base ("embodiment": "base", runwayml/praxis-1-base-lerobot) é o backbone de vídeo e o tronco do especialista de ação no final do pré-treino, sem os seis tensores de cabeça de um rig (action_encoder.w1, embed_proprio, head, pesos e vieses) e sem o estado do otimizador. Ele não é carregado como uma política; dois scripts o transformam no ponto de partida de um fine-tune para o seu rig de braço único, a forma LeRobot dos scripts compute_norm_stats.py e init_heads.py do repositório praxis:

# the normalization statistics of your data: 1 % / 99 % quantiles of the relative actions and the
# mean and std of the proprio over random training windows (no video is read)
python examples/praxis/compute_norm_stats.py --repo-id lerobot/droid_1.0.1 --revision refs/pr/11 \
    --state-key observation.state.cartesian --action-key action.cartesian \
    --episodes droid_episodes.json --n-windows 2000 --out stats.json
# the base plus six new head tensors (`--seed`) and the rig fields: a policy directory that loads
python examples/praxis/init_heads.py --base runwayml/praxis-1-base-lerobot --out outputs/praxis-1-myrig \
    --stats stats.json --tile-height 224 --tile-width 416 --antialias --link
# then the full fine-tune command above with --policy.path=outputs/praxis-1-myrig, or LoRA with the
# six new head tensors trained in full next to the adapters:
torchrun --nproc-per-node=4 $(which lerobot-train) ... --policy.path=outputs/praxis-1-myrig \
    --parallelism.dp_shard=4 --peft.r=16 --peft.lora_alpha=32 \
    --peft.full_training_modules='["train_counter","action_expert.action_encoder.w1","action_expert.embed_proprio","action_expert.head"]'

As colunas de estado e ação contêm a pose do DROID e o [x, y, z, roll, pitch, yaw, gripper] do garra por quadro (as próprias colunas do conjunto de dados, ou um mapa de renomeação). --stats-from <single-arm checkpoint> copia um bloco de normalização existente (o rig do DROID). O rig do DROID reconstruído a partir da base e das estatísticas do checkpoint liberado tem exatamente a configuração de praxis-1-droid, e os tensores da cabeça são os do init_heads.py do repositório praxis para a mesma semente, bit a bit. No subconjunto DROID das execuções acima, compute_norm_stats.py reproduz as estatísticas do checkpoint com precisão de 0,01 m nos quantis xyz, com a garra fixada em 0 e 1. O próprio fine-tune de base do repositório praxis (configs/finetune_base.yaml, um dia em 8 GPUs) atinge 62,5 % em droid_guard a partir dessa base, contra 91 % para praxis-1-droid treinado com cerca de 40 vezes mais janelas: a base é um ponto de partida didático para um novo rig, não um substituto do checkpoint liberado. O fine-tune completo treina todos os pesos em optimizer_lr; sob LoRA, os adaptadores e os seis tensores da cabeça são treinados (PraxisPolicy.PEFT_HEAD_MODULES são os únicos módulos, além do contador, que full_training_modules aceita) e o restante do tronco permanece congelado; os passos de aquecimento apenas da cabeça e a escala de taxa de aprendizado da cabeça da receita praxis não foram portados.

Limitações conhecidas

• Dois rigs: o checkpoint de braço único DROID (embodiment: single_arm) e o YAM bimanual (bimanual). O fine-tuning e a simulação são de braço único.
• O Praxis gera alvos de pose cartesiana para um braço com um controlador de impedância; as classes de robô do LeRobot não incluem um Franka, portanto, o uso em robôs reais é feito por meio de um cliente que executa alvos cartesianos.
• Um lote de select_action executa seus ambientes sequencialmente (uma chamada de modelo por ambiente).
• O serviço remoto não é suportado. O executor de blocos padrão (lerobot/inference/policy_runner.py) recusa uma política que condicione em observações passadas, e o Praxis condiciona nos quadros dos blocos executados (current_observation_only=False em chunk_inference_spec); a etapa de tile também redimensiona os quadros da câmera para o canvas do checkpoint, de modo que eles não corresponderiam às formas verificadas por validate_chunk_input_features. Execute a política em processo (lerobot-eval, um cliente de robô).