Warp და MJWarp: 2048 პარალელურ სამყაროში რობოტების სიმულაცია ერთ GPU-ზე

NVIDIA-მ Physical AI-ს სიმულაციის სერიის მეორე სტატია გამოაქვეყნა. მასში აღწერილია MuJoCo Warp (MJWarp), MuJoCo-ს ფიზიკის GPU რეალიზაცია, რომელიც NVIDIA Warp-ის ბირთვების ჩარჩოზეა აგებული. მაგალითში SO-101 რობოტის მკლავი CPU პროცესიდან ერთ GPU-ზე 2048 პარალელურ სამყარომდე გადადის. ავტორების თქმით, მთავარი აღარ არის, რამდენად სწრაფად მუშაობს ერთი სამყარო, არამედ ის, რამდენი სამყარო მუშაობს ერთდროულად.

Warp არის Python-ის ჩარჩო GPU-ზე აჩქარებული ბირთვების დასაწერად. დეველოპერები სტატიკურად ტიპიზებულ ბირთვებს Python-ით წერენ, Warp კი მათ CPU-სთვის ან CUDA-სთვის კომპილირებს. ჩარჩო გვთავაზობს CUDA-ის დონის სიჩქარეს JIT კომპილაციითა და CUDA Graphs-ით, ჩაშენებულ ვექტორებს, მატრიცებსა და კვატერნიონებს, აგრეთვე დიფერენცირებად ბირთვებს DLPack-ის თავსებადობით, რაც სიმულაციის ML წვრთნის ციკლში ჩასმას იძლევა.

როგორ ალაგებს MJWarp ფიზიკას პარტიებად

MJWarp მოდელის ფორმატს ინარჩუნებს: MJCF ფაილი MuJoCo-ში იტვირთება, mjw.put_model კი მას მოწყობილობაზე ატვირთავს, სადაც ერთი mjw.step გამოძახება დამოუკიდებელი სამყაროების მთელ პარტიას ამოძრავებს. ტევადობას ოთხი პარამეტრი განსაზღვრავს: nworld (პარალელური გარემო), nconmax (კონტაქტები თითო სამყაროზე), naconmax (კონტაქტების გლობალური ზღვარი, უპირატესი განსაზღვრისას) და njmax (შეზღუდვების ზღვარი თითო სამყაროზე). მოგება ერთი სამყაროს დაყოვნება არ არის, არამედ ჯამური გამტარუნარიანობა, რაც განმტკიცებითი სწავლისთვის მნიშვნელოვანია.

ჯერ შემოწმება, შემდეგ გაზომვა

ავტორები ეტაპობრივ შემოწმებებს გადიან: CPU-ზე საბაზისო შედეგს აყალიბებენ, შემდეგ ერთ სამყაროს GPU-ზე გადააქვთ და თანხვედრას ამოწმებენ, ბუფერებს ამოცანის ყველაზე დატვირთული მომენტით ზომავენ, ბოლოს 2048 სამყაროზე გადადიან. მაგალითში SO-101 მკლავს წითელი 44 მმ კუბი ლურჯზე უნდა დააწყოს. წარმატება ორი პირობით იზომება: ჰორიზონტალური ცდომილება 15 მმ-ზე ნაკლები ან ტოლი და ვერტიკალური დაშორება 35-55 მმ შუალედში. გადავსება შეცდომად არ აღიქმება: MJWarp ლიმიტის გაზრდას გვთავაზობს და დაზარალებულ სამყაროებს აღნიშნავს, ამიტომ ასეთი გაშვება ბენჩმარკისთვის სანდო არ არის.

გაზომვა ბოლო ეტაპია. GPU-ის გაშვებები ასინქრონულია, ამიტომ უბრალო ტაიმერი ზომავს, რამდენად სწრაფად ალაგებს Python სამუშაოს და არა, რამდენად სწრაფად დაასრულა GPU-მ. სტატია ჯერ გახურებას ატარებს და გაზომვამდე და მის შემდეგ სინქრონიზაციას აკეთებს; შედეგი წამში შესრულებული სამყარო-ნაბიჯებითა და პარტიის ზომით არის მოცემული. Warp 1.15-ში დაემატა დეტერმინირებული რეჟიმი, სერიის შემდეგი ნაწილი იმავე გარემოს Newton-ზე გადატანას დაეთმობა.