NVIDIA Warp ve MJWarp: tek GPU'da 2.048 paralel robot simülasyonu

NVIDIA, State of Simulation for Physical AI serisinin ikinci yazısını yayımladı. Yazıda, NVIDIA Warp çekirdek çatısı üzerine kurulu MuJoCo fizik uygulaması MuJoCo Warp (MJWarp) ele alınıyor. Örnekte SO-101 robot kolu, alışılmış CPU akışından tek GPU üzerinde 2.048 paralel simülasyon dünyasına taşınıyor. Yazarlara göre asıl soru artık tek bir dünyanın ne kadar hızlı ilerlediği değil, kaç dünyanın aynı anda ilerleyebildiği.

Warp, yüksek performanslı ve GPU hızlandırmalı çekirdekler yazmak için bir Python çatısıdır. Geliştiriciler statik tipli çekirdekleri Python ile yazar; Warp bunları CPU veya CUDA için derleyip önbelleğe alır. Çatı; JIT derleme ve CUDA Graphs ile CUDA düzeyinde hız, hazır vektör, matris, kuaterniyon ve hash grid yapıları, ve DLPack uyumlu türevlenebilir çekirdekler sunar; böylece simülasyon eğitim döngüsüne yerleşebilir.

MJWarp fiziği nasıl partiler halinde çalıştırıyor

MJWarp model formatını korur: MJCF dosyası MuJoCo'da yüklenir, mjw.put_model onu aygıta aktarır ve tek bir mjw.step çağrısı bağımsız dünyalardan oluşan partinin tamamını ilerletir. Kapasiteyi dört parametre belirler: nworld, paralel ortam sayısı; nconmax, dünya başına beklenen temas sayısı; naconmax, tanımlandığında öncelik kazanan genel temas sınırı; ve njmax, dünya başına kısıt tavanı. Kazanç tek dünyanın gecikmesi değil toplam verimdir; bu da pekiştirmeli öğrenme ve geniş ölçekli örneklemeye uygundur.

Önce doğrulama, sonra ölçüm

Yazarlar sıralı kontrollerden geçiyor: CPU'da bir temel sonuç oluşturuyor, tek dünyayı GPU'ya taşıyıp eşleşmeyi doğruluyor, temas ve kısıt tamponlarını görevin en yoğun temas anına göre boyutlandırıyor ve ardından 2.048 dünyaya geçiyor. Örnekte SO-101 kolu 44 mm'lik kırmızı küpü mavi küpün üzerine koymalı. Başarı iki koşulla ölçülüyor: yatay hata en fazla 15 mm ve dikey aralık 35 ile 55 mm arasında. Tampon taşması hata olarak yükseltilmez; bu yüzden taşma sonrası süren bir koşu performans ölçümü için güvenilir sayılmaz.

Ölçüm son aşamadır. GPU çalıştırmaları eşzamansız olduğu için saf bir zamanlayıcı Python'un işi ne kadar hızlı kuyruğa aldığını ölçer, GPU'nun ne zaman bitirdiğini değil. Yazı önce ısınma turu yapar ve ölçülen aralığın öncesi ile sonrasında eşitler. Sonuçlar, parti boyutuyla birlikte saniyedeki dünya adımı ve parti başına milisaniye olarak raporlanır. Warp 1.15 belirlenimci yürütme kipi ekledi; sonraki bölüm aynı ortamı Newton'a taşıyacak.