NVIDIA Warp և MJWarp. 2048 զուգահեռ ռոբոտային սիմուլյացիա մեկ GPU-ով

NVIDIA-ն հրապարակել է State of Simulation for 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-ով, ներկառուցված վեկտորներ, մատրիցներ ու hash grid կառույցներ, ինչպես նաև տարբերակելի միջուկներ DLPack ինտերֆեյսով, ինչը թույլ է տալիս սիմուլյացիան ներառել ուսուցման մեջ։

Ինչպես է MJWarp-ը ֆիզիկան խմբավորում խմբաքանակներով

MJWarp-ը պահպանում է մոդելի ձևաչափը՝ MJCF ֆայլը բեռնվում է MuJoCo-ում, իսկ mjw.put_model-ը այն տեղափոխում է սարք, որտեղ mjw.step-ի մեկ կանչը առաջ է տանում անկախ աշխարհների ամբողջ խմբաքանակը։ Տարողությունը սահմանում են չորս պարամետր՝ nworld-ը՝ զուգահեռ միջավայրերի թիվը, nconmax-ը՝ մեկ աշխարհի սպասվող կոնտակտները, naconmax-ը՝ կոնտակտների գլոբալ սահմանը, որը սահմանվելու դեպքում գերակա է, և njmax-ը՝ մեկ աշխարհի սահմանափակումների վերին շեմը։ Շահույթը մեկ աշխարհի ուշացումը չէ, այլ ընդհանուր թողունակությունը, ինչը կարևոր է ամրապնդող ուսուցման և ընտրանքահանման համար։

Նախ ստուգում, հետո չափում

Հեղինակները հաջորդական ստուգումներ են անցնում՝ CPU-ի վրա հիմնական արդյունք, մեկ աշխարհի տեղափոխում GPU և համընկնման ստուգում, բուֆերների չափավորում առաջադրանքի ամենածանր պահով, ապա անցում 2048 աշխարհի։ Օրինակում SO-101 թևը 44 մմ կարմիր խորանարդը պետք է դնի կապույտի վրա։ Հաջողությունը չափվում է երկու պայմանով՝ հորիզոնական սխալը՝ 15 մմ-ից ոչ ավելի, և ուղղահայաց հեռավորությունը՝ 35-55 մմ սահմանում։ Բուֆերի գերազանցումը սխալ չի առաջացնում, ուստի դրանից հետո շարունակված գործարկումը հուսալի չէ։

Չափումը վերջին փուլն է։ GPU-ի գործարկումները ասինխրոն են, ուստի պարզ ժամաչափը չափում է Python-ի հերթագրումը, ոչ թե GPU-ի ավարտը։ Հոդվածը սկզբում տաքացնում է և համաժամացնում չափվող հատվածից առաջ ու հետո։ Արդյունքները ներկայացվում են վայրկյանում կատարված աշխարհ-քայլերով և խմբաքանակի քայլի միլիվայրկյաններով։ Warp 1.15-ում ավելացել է դետերմինացված ռեժիմ, իսկ հաջորդ մասում միջավայրը կտեղափոխվի Newton։