MJWarp της NVIDIA επιταχύνει ταυτόχρονες προσομοιώσεις ρομποτικής σε GPU

Νέα τεχνολογία επιτρέπει να εκτελούνται έως 2.048 παράλληλες προσομοιώσεις ρομποτικής σε GPUs για γρήγορη εκπαίδευση και δοκιμές AI.

Η NVIDIA ανακοίνωσε το MJWarp, μια νέα υλοποίηση του δημοφιλούς physics engine MuJoCo προσαρμοσμένη για GPUs, με στόχο να επιταχύνει τον αριθμό των ταυτόχρονων προσομοιώσεων ρομποτικών περιβαλλόντων. Το MJWarp βασίζεται στην πλατφόρμα προγραμματισμού NVIDIA Warp, που φέρνει εγγενή υποστήριξη για CUDA και JIT compilation, επιτρέποντας υψηλές ταχύτητες εκτέλεσης.

Σε αντίθεση με την παραδοσιακή CPU-βασισμένη εκτέλεση MuJoCo, το MJWarp έχει σχεδιαστεί για να διαχειρίζεται δεκάδες έως και χιλιάδες παράλληλες προσομοιώσεις, εκμεταλλευόμενο τη GPU για να αυξήσει σημαντικά τον συνολικό ρυθμό επεξεργασίας (aggregate throughput). Σύμφωνα με την NVIDIA, το MJWarp μπορεί να τρέξει έως 2.048 παράλληλα περιβάλλοντα σε συμβατά μοντέλα MuJoCo.

Η αρχιτεκτονική του MJWarp διαχειρίζεται με ειδικές παραμέτρους τον αριθμό των αναμενόμενων επαφών (nconmax) και των περιορισμών (njmax) μέσα στις προσομοιώσεις. Για παράδειγμα, στη σκηνή SO-101 που χρησιμοποιείται για ρομποτικό pick-and-place, οι παράμετροι ορίζονται σε nconmax=128 και njmax=300 ώστε να αποφεύγονται προβλήματα υπερχείλισης στη μνήμη.

Επιπλέον, υποστηρίζεται η τεχνική CUDA graph capture, που βελτιστοποιεί την απόδοση καταγράφοντας πολλαπλές εκτελέσεις kernels σε ένα γράφο, ώστε να αναπαράγονται πιο γρήγορα, μειώνοντας το κόστος του προγραμματισμού σε GPU.

Το MJWarp είναι διαθέσιμο ως πακέτο Python και μπορεί να εγκατασταθεί απλά με την εντολή pip install mujoco-warp. Αυτό καθιστά προσβάσιμη τη χρήση του σε περιβάλλοντα που εστιάζουν σε προχωρημένες εφαρμογές τεχνητής νοημοσύνης και εκπαίδευσης ρομποτικών μοντέλων μέσω reinforcement learning.

Σημαντικό είναι να αναφερθεί πως παρόλο που το MJWarp υποστηρίζει την παράλληλη εκτέλεση χιλιάδων προσομοιώσεων, ο χρόνος εκτέλεσης σε κάθε βήμα μπορεί να είναι μεγαλύτερος σε σχέση με την κλασική CPU έκδοση, γεγονός που το καθιστά λιγότερο κατάλληλο για εφαρμογές σε πραγματικό χρόνο που απαιτούν χαμηλή καθυστέρηση.

Η NVIDIA Warp, η βάση του MJWarp, προσφέρει δυνατότητες όπως differentiability και deterministic execution, που ευνοούν προχωρημένες μεθόδους βελτιστοποίησης και ανάλυσης, αλλά δεν ισχύουν απαραίτητα για όλο το MJWarp.