Batched Differentiable Rigid Body Dynamics in PyTorch for GPU-Accelerated Robot Learning

ملخص تقني: ديناميكا الأجسام الجاسئة المجمعة والقابلة للتفاضل في PyTorch (bard)

بيان المشكلة
مع انتقال التحكم في الروبوتات نحو التعلم المعزز (RL) واسع النطاق الذي يتضمن حساب الديناميكا داخل الحلقة (in-loop)، يواجه المجال اختناقاً حرجاً في معدل الإنتاجية. تعتمد خطوط أنابيب التدريب الحالية، مثل تلك الموجودة في Isaac Lab وMuJoCo Playground، على PyTorch للتنفيذ المجمع على وحدة معالجة الرسومات (GPU). ومع ذلك، فإن مكتبة الديناميكا القياسية، Pinocchio، مقيدة بوحدة المعالجة المركزية (CPU) وتعمل بشكل تسلسلي. تفرض هذه البنية عبء نقل بيانات خطي بين الـ CPU والـ GPU لكل عملية تقييم، مما يترك نوى الـ GPU خاملة أثناء عمليات التتبع الشجري التسلسلية. وفي حين توفر محاكيات الـ GPU الحالية مثل Brax وMuJoCo MJX إنتاجية عالية، إلا أنها مرتبطة بأنظمة غير تابعة لـ JAX، أو تعمل كمحاكيات متكاملة لا توفر مؤثرات مستقلة وقابلة للتفاضل (مثل المصفوفات اليعقوبية أو مصفوفات الكتلة) المطلوبة لدوال الخسارة المخصصة أو المكافآت المستندة إلى الفيزياء. وعلى العكس من ذلك، فإن المكتبات المتوافقة مع PyTorch والقابلة للتفاضل مثل ADAM تضحي بإنتاجية الـ GPU من أجل عمومية الخلفية، وتفشل في استغلال دمج النواة (kernel fusion) أو تجنب التحكم في التدفق المعتمد على البيانات الذي يكسر تتبع الرسم البياني لمترجم PyTorch.

المنهجية
يقدم المؤلفون bard (الديناميكا المجمعة للأجسام الجاسئة المفصلية)، وهي مكتبة pure-PyTorch تنفذ خوارزميات Featherstone لديناميكا الأجسام الجاسئة (FK، RNEA، CRBA، ABA، واليعقوبيات). تم تصميم المكتبة للتقييم المجمع على الـ GPU والتوافق الكامل مع torch.compile. يتم تحقيق مكاسب الأداء من خلال أربع استراتيجيات تصميم محددة:

1. التقييم الكسول المتدرج (Tiered Lazy Evaluation): بدلاً من حساب جميع الكميات الكينماتيكية في كل استدعاء، تقوم bard بتقسيم الحالة الكينماتيكية إلى ثلاث طبقات محمية بأعلام صلاحية: (1) تحويلات الأب-الابن المحلية، (2) الملحقات المكانية (spatial adjoints) (المطلوبة فقط لخوارزميات الديناميكا)، و(3) وضعيات إطار العالم (المطلوبة فقط لاستعلامات FK/Jacobian). يضمن ذلك أن حلقة RL التي تستدعي RNEA فقط تتجنب التكلفة الحسابية لنشر إطار العالم.
2. تحويلات المفاصل الخالية من ضرب المصفوفات: للقضاء على عبء الإطلاق لآلاف عمليات ضرب المصفوفات $4 \times 4$ الصغيرة على الـ GPU، تقوم المكتبة بحساب ثوابت دوران رودريغز (Rodrigues rotation constants) مسبقاً عند بناء النموذج. في وقت التشغيل، يتم استرداد تحويلات الأب-الابن باستخدام عمليات ضرب قياسية-مصفوفية فقط مع $\sin(q_i)$ و$\cos(q_i)$، مما يسمح لمترجم PyTorch بدمج العمليات عبر جميع المفاصل وعناصر الدفعة في نواة واحدة.
3. الانتشار الشجري المتوازي على المستويات (Level-Parallel Tree Propagation): بينما تعتبر خوارزميات Featherstone تسلسلية بطبيعتها، تقوم bard بإعادة تنظيم عمليات التتبع الشجري إلى عمليات مجمعة متوازية على مستوى المستوى. من خلال تجميع العقد عند نفس عمق الشجرة، تقلل المكتبة عدد الخطوات التسلسلية $n_b$ (عدد الأجسام) إلى $d$ (عمق الشجرة) من الخطوات المجمعة. وبينما تظل عمليات التمرير الخلفي (تراكم القوى) تسلسلية بسبب التبعيات في البيانات، يتم معالجة جميع البيئات $B$ في وقت واحد عند كل عقدة.
4. البناء المباشر لأعمدة اليعقوبي (Direct Jacobian Column Construction): بالنسبة للمفاصل ذات درجة الحرية الواحدة (single-DOF)، تتجاوز المكتبة بناء مصفوفات ملحقة (adjoint matrices) كاملة $6 \times 6$. بدلاً من ذلك، تقوم بحساب أعمدة اليعقوبي مباشرة باستخدام الإزاحات المحسوبة مسبقاً والوضعيات المتاحة في إطار العالم، مستبدلةً بناء الملحقات بضربات تقاطع (cross products) وعمليات ضرب مصفوفة-متجه فعالة.

تستخدم المكتبة تخطيط ذاكرة يعتمد على ترتيب الدفعات أولاً (batch-first) وتخزن التحويلات المكانية المنقولة بشكل متصل لتجنب هدر خط الكاش (cache-line). وهي تدعم الوضع الفوري (eager mode) (تخطي تتبع التدرج لتحقيق أقصى قدر من الإنتاجية) والوضع الوظيفي (functional mode) (التفاضل التلقائي الكامل) وهي متوافقة تماماً مع torch.compile باستخدام التحكم في التدفق الثابت وتنسيرات الفهرس الثابتة.

المساهمات الرئيسية

1. bard: مكتبة ديناميكا pure-PyTorch مفتوحة المصدر تتيح التقييم المجمع لخوارزميات Featherstone على الـ GPU.
2. تقنيات على مستوى النظام: طرق مبتكرة لتشغيل الحسابات الشجرية المتكررة على الـ GPU، بما في ذلك تخزين الكينماتيكا المؤقت، ومساحات العمل المحجوزة مسبقاً، والانتشار المتوازي على المستويات.
3. التحقق: اختبار شامل على تحديد الأنظمة المستند إلى التدرج والتعلم المعزز واسع النطاق مع 4,096 بيئة متوازية.

النتائج التجريبية
قيم المؤلفون bard على خمسة نماذج روبوتية (تتراوح من 7 إلى 23 درجة حرية) مقابل Pinocchio وADAM على وحدة معالجة رسومات NVIDIA H200:

• الدقة العددية: تطابق bard مكتبة Pinocchio ضمن أخطاء التقريب العائم. كانت أقصى الفروقات قريبة من قيمة الآلة ($\approx 10^{-16}$) لـ FK واليعقوبيات، وأقل من $10^{-13}$ لـ RNEA وCRBA. أظهرت ABA تباينات أكبر قليلاً (تصل إلى $10^{-10}$) بسبب تراكم الخطأ في هيكلها ثلاثي المراحل، لكنها ظلت ضمن الحدود العملية.
• الإنتاجية: عند حجم دفعة قدره 4,096، حققت bard إنتاجية أعلى بـ 64 ضعفاً من Pinocchio لـ Forward Kinematics وبـ 63 ضعفاً لـ Jacobians. بالنسبة لخوارزميات الدينمايكا، تراوحت نسب التسريع بين 3 إلى 7 أضعاف.
• تحديد الأنظمة: في مهمة تحديد مستندة إلى التدرج على ذراع آلية بـ 7 درجات حرية مع ضوضاء عزم دوران بنسبة 5%، نجحت bard في استعادة كتل الوصلات بمتوسط خطأ قدره 1.24%.
• التكامل مع RL: عند دمجها في خط تدريب Isaac Lab AMP لـ quadruped بـ 11 درجة حرية و4,096 بيئة متوازية، حققت bard إنتاجية أعلى بـ 8.5 ضعفاً من Pinocchio و2.0 ضعفاً أعلى من ADAM.

الأهمية والادعاءات
يدعي البحث أن bard تسد فجوة حرجة في منظومة الروبوتات من خلال توفير مكتبة ديناميكا نمطية عالية الأداء وأصلية لـ PyTorch. فهي تتيح تدريب RL واسع النطاق مع ديناميكا داخل الحلقة على وحدة GPU واحدة، وهي مهمة كانت تعيقها سابقاً اختناقات الـ CPU أو عدم توافق الأنظمة. يشير المؤلفون إلى أن ميزة الإنتاجية تكون أكثر وضوحاً في مهام العمل التي تهيمن عليها الكينماتيكا عند أحجام الدفعات الكبيرة. كما يحذرون من أن torch.compile قد يقلل الأداء في مهام العمل الكثيفة بالديناميكا على وحدات الـ GPU الأقل قوة (مثل NVIDIA L4) بسبب تسريب السجلات (register spilling)، موصين باستخدام الوضع الفوري (eager mode) في هذه الحالات. يوضح العمل أن إعادة تنظيم الخوارزميات التسلسلية للتنفيذ المتوازي على الـ GPU يمكن أن يحقق مكاسب كبيرة في الأداء دون التضحية بالدقة العددية أو القابلية للتفاضل.