NVIDIA tăng tốc mô phỏng robot lên hơn 2.000 môi trường song song bằng Warp và MJWarp

Giải quyết nút thắt thông lượng trong mô phỏng vật lý cho robot

Trong phát triển robot và trí tuệ nhân tạo thể nghiệm (Physical AI), công cụ mô phỏng MuJoCo truyền thống trên CPU vốn là tiêu chuẩn quen thuộc nhờ độ chính xác cao khi kiểm thử các thuật toán điều khiển như MPC hay teleoperation. Tuy nhiên, khi chuyển sang bài toán học tăng cường (reinforcement learning), nhu cầu thực tế đã thay đổi: câu hỏi không còn là "một thế giới mô phỏng chạy nhanh đến mức nào", mà là "có thể chạy bao nhiêu thế giới cùng một lúc". Việc chạy hàng loạt môi trường trên CPU dần bộc lộ giới hạn về băng thông và thông lượng tính toán.

Để giải quyết bài toán này, nhóm kỹ sư của NVIDIA gồm Johnny Nuñez Cano, Asier Arranz, Rishabh Chadha và Ben Oliveri đã công bố quy trình ứng dụng NVIDIA Warp cùng MuJoCo Warp (MJWarp). Bằng cách tái hiện toàn bộ pipeline vật lý của MuJoCo dựa trên ngôn ngữ kernel Warp, MJWarp biên dịch các phép tính sang CUDA để xử lý hàng loạt môi trường song song trực tiếp trên GPU. Dữ liệu trạng thái và mô hình học máy được giữ trọn vẹn trong bộ nhớ VRAM, loại bỏ hoàn toàn độ trễ trao đổi dữ liệu qua cổng kết nối PCIe.

Quy trình mở rộng cánh tay robot SO-101 lên 2.048 môi trường độc lập

Để chứng minh năng lực của giải pháp, nhóm nghiên cứu đã thử nghiệm trên bài toán gắp - đặt (pick-and-place) bằng cánh tay robot SO-101. Nhiệm vụ yêu cầu cánh tay gắp một khối lập phương màu đỏ kích thước 44 mm và xếp chồng lên khối màu xanh đặt trên bàn, tuân thủ các ngưỡng sai số khắt khe: khoảng cách tâm mặt phẳng ngang $xy\_err \le 0{,}015\text{ m}$ và độ chênh lệch cao độ thẳng đứng $0{,}035\text{ m} \le dz \le 0{,}055\text{ m}$.

Quy trình chuyển đổi từ CPU sang GPU được chuẩn hóa qua các bước kỹ thuật cụ thể:

• Đồng bộ hóa tần số điều khiển và vật lý: Hệ thống thiết lập tần số bộ điều khiển ở mức 50 Hz và thực hiện 10 bước đệm vật lý (substeps) cho mỗi khung hình điều khiển. Bước thời gian vật lý được cố định chính xác ở mức 0.002 giây (frame_dt / sim_substeps) trên cả CPU và GPU trước khi nạp mô hình qua hàm mjw.put_model, tránh sai lệch thời gian thực trong quá trình triển khai mô hình học được ra thực tế.
• Đối chiếu tính tương đương trên một thế giới: Ban đầu, mô hình chỉ chạy thử trên 1 môi trường GPU (nworld = 1) nhưng đồng bộ dữ liệu vị trí (qpos) và vận tốc (qvel) ngược về CPU sau mỗi bước để quan sát trên viewer và so khớp trực tiếp kết quả với mô phỏng MuJoCo gốc.
• Tối ưu hóa bộ nhớ va chạm và ràng buộc: MJWarp cấp phát trước bộ đệm trên GPU cho va chạm và ràng buộc. Với SO-101, thông số khởi điểm được cấu hình ở mức nconmax = 128 (số tiếp xúc tối đa mỗi môi trường) và njmax = 300 (giới hạn trên về ràng buộc). Các giá trị này phải được đo lường tại thời điểm va chạm phức tạp nhất trong tác vụ—khi cả hai hàm kẹp và mặt bàn cùng chạm vào khối hộp. Nếu vượt ngưỡng, hệ thống sẽ cảnh báo tràn bộ đệm (narrowphase overflow) và làm sai lệch quỹ đạo, buộc kỹ sư phải điều chỉnh lại trước khi đo điểm chuẩn. Nhóm nghiên cứu cũng cung cấp thêm hồ sơ cho robot reBot với cấu hình lớn hơn (nconmax = 256, njmax = 500).
• Mở rộng quy mô lên 2.048 môi trường bằng CUDA Graph: Khi tính tương đương được xác nhận, hệ thống loại bỏ hoàn toàn thao tác đọc dữ liệu về CPU và nhân bản trạng thái ban đầu lên nworld = 2.048. Để triệt tiêu độ trễ điều phối kernel từ Python, toàn bộ chuỗi bước tính mjw.step được đóng gói bằng wp.ScopedCapture() thành một đồ thị CUDA Graph, cho phép tái thực thi liên tục trên GPU với hiệu năng phần cứng tối đa.
• Đo lường thông lượng chính xác: Do các lệnh chạy GPU diễn ra bất đồng bộ, quá trình đo chuẩn đòi hỏi chạy làm nóng (warm-up) 10 bước để hoàn tất biên dịch kernel JIT và cấp phát bộ nhớ, sau đó sử dụng wp.synchronize() ngay trước và sau 200 bước đo thực tế nhằm ghi nhận đúng thời gian phần cứng hoàn thành tác vụ tính theo số bước thế giới mỗi giây (world-steps/second).

Vị trí của MJWarp trong bức tranh Physical AI

NVIDIA Warp đóng vai trò là tầng ngôn ngữ kernel viết bằng Python nhưng biên dịch thành mã máy native CUDA hoặc CPU. Framework này hỗ trợ tính toán khả vi (differentiable) thông qua cơ chế wp.Tape ghi nhận lượt chạy xuôi và phát lại tiếp hợp (adjoint) theo chiều ngược lại, đồng thời hỗ trợ thực thi tất định (deterministic execution) từ phiên bản 1.15 để phục vụ kiểm thử hồi quy.

Nhờ khả năng tương thích thông qua DLPack, dữ liệu mô phỏng từ MJWarp có thể kết nối trực tiếp với các framework huấn luyện học máy phổ biến như PyTorch hay JAX. Các nhà phát triển có thể tận dụng kiến trúc này qua các nền tảng như mjlab, MuJoCo Playground (thông qua backend MJX), hoặc đưa vào Newton và Isaac Lab trong các giai đoạn tích hợp sâu hơn của quy trình phát triển robot thông minh.