Hướng dẫn NVIDIA: Tăng tốc mô phỏng robot MuJoCo lên 2048 môi trường song song với Warp và MJWarp

MuJoCo cổ điển cung cấp mô phỏng robot dựa trên CPU tốc độ cao để phát triển, thử nghiệm và điều khiển robot, đồng thời có khả năng song song hóa việc lấy mẫu trên các nhân CPU. Tuy nhiên, khi khối lượng công việc học máy tăng lên, câu hỏi chuyển từ việc một thế giới có thể chạy nhanh đến mức nào sang việc có bao nhiêu thế giới có thể chạy cùng một lúc. Tăng tốc GPU giúp thúc đẩy các thế giới đó theo từng lô lớn trong khi vẫn giữ dữ liệu mô phỏng và học máy gần với thiết bị.

MuJoCo Warp (MJWarp), được xây dựng trên NVIDIA Warp, đưa các mô hình MuJoCo tương thích vào quy mô GPU đó. Trong bài viết này, chúng ta sẽ chuyển cánh tay robot SO-101 từ quy trình làm việc MuJoCo quen thuộc sang tối đa 2.048 môi trường MJWarp song song và xem xét các công nghệ cũng như các bước xác thực giúp quá trình chuyển đổi này trở nên khả thi.

Hình 1. Cách MJWarp kết nối Python với mô phỏng GPU. MuJoCo tải và biên dịch mô hình MJCF; MJWarp triển khai vật lý trong NVIDIA Warp, vốn biên dịch các nhân CUDA để thúc đẩy các trạng thái mô phỏng trên GPU NVIDIA.

Đây là bài viết thứ hai trong loạt bài về Trạng thái Mô phỏng cho AI Vật lý (State of Simulation for Physical AI). Bài viết đầu tiên đã vạch ra bối cảnh mô phỏng robot. Ở đây, chúng ta chuẩn bị và mở rộng quy mô môi trường mô phỏng; chúng ta không huấn luyện chính sách (policy). Các phần tiếp theo về Newton và Isaac Lab sẽ đề cập đến các lớp tích hợp kế tiếp.

Tổng hợp lại

Lối tắt quyết định:

Bắt đầu với một Warp Kernel hữu ích

NVIDIA Warp là một framework Python để viết các nhân (kernel) hiệu năng cao, được tăng tốc bởi GPU. Warp cho phép các nhà phát triển tạo ra các nhân định kiểu tĩnh bằng Python và biên dịch chúng để thực thi trên CPU hoặc CUDA. Lần khởi chạy đầu tiên sẽ xây dựng và lưu vào bộ nhớ đệm một mô-đun gốc; các lần khởi chạy sau sẽ tái sử dụng nó. Ngôn ngữ nhân là một tập con của Python tập trung vào hiệu năng, trong khi Python thông thường vẫn chịu trách nhiệm về cấu hình, cấp phát và điều phối khởi chạy.

Nhân nhỏ hướng tới robot này thúc đẩy các vị trí điểm dưới tác động của trọng lực. Một luồng logic xử lý một điểm, vì vậy cùng một mã nguồn có thể mở rộng từ hai điểm lên hàng triệu điểm mà không cần đưa thuật ngữ GPU vào luồng điều khiển.

Ba giá trị cốt lõi của Warp là:

import numpy as np
import warp as wp

@wp.kernel
def integrate(
   positions: wp.array[wp.vec3],
   velocities: wp.array[wp.vec3],
   dt: float,
):
   i = wp.tid()
   velocities[i] += wp.vec3(0.0, 0.0, -9.81) * dt
   positions[i] += velocities[i] * dt

wp.init()
device = "cuda:0" if wp.is_cuda_available() else "cpu"
start = np.array([[0.0, 0.0, 0.5], [0.2, 0.0, 0.5]], dtype=np.float32)
positions = wp.array(start, dtype=wp.vec3, device=device)
velocities = wp.zeros_like(positions)

wp.launch(
   integrate,
   dim=len(start),
   inputs=[positions, velocities, 0.01],
   device=device,
)
wp.synchronize_device(device)
print(positions.numpy())

Ba đặc tính làm cho điều này trở nên hữu ích trong robot học:

• Công việc song song tường minh. wp.tid xác định điểm, tiếp điểm, vật thể hoặc thế giới thuộc sở hữu của luồng logic hiện tại.
• Mảng thiết bị tường minh. Một mảng nằm trên thiết bị đã chọn. Gọi.numpy trên một mảng CUDA sẽ đồng bộ hóa và sao chép nó vào bộ nhớ CPU; đây không phải là đường dẫn không sao chép (zero-copy). Đối với đường ống PyTorch hoặc JAX nằm trên thiết bị, hãy sử dụng các bộ điều hợp framework của Warp hoặc chia sẻ tương thích DLPack thay thế.
• Khởi chạy nhân có thể kết hợp. Một chương trình có thể khởi chạy một chuỗi các nhân tập trung và nắm bắt công việc CUDA được hỗ trợ vào một đồ thị để giảm chi phí điều phối lặp lại. Việc nắm bắt đồ thị sẽ phát lại các lần khởi chạy dựa trên các bộ đệm hiện có; nó không hợp nhất các nhân tùy ý.

Tính vi phân và tính tất định.

Hai khả năng khác của Warp đáng để biết, mặc dù không được sử dụng trong quy trình SO-101 trong bài viết này. Các nhân Warp có thể vi phân: một wp.Tape ghi lại các lần khởi chạy nhân tiến (forward) được thực hiện trong ngữ cảnh của nó và phát lại các phần phụ (adjoint) của chúng theo thứ tự ngược lại khi gọi backward, đó là lý do tại sao các nhóm xây dựng hình học vi phân, CFD và vật lý tùy chỉnh trong Warp, bao gồm cả quy trình CAE cho mô phỏng và tối ưu hóa thiết kế. Warp cũng hỗ trợ thực thi tất định, được giới thiệu trong Warp 1.15: Các nguyên tử GPU (GPU atomics) mặc định phụ thuộc vào bộ lập lịch, vì vậy các lần khởi chạy lặp lại của cùng một nhân có thể khác biệt đôi chút, và các chế độ tất định tùy chọn sẽ đánh đổi một phần hiệu năng để có thứ tự tái lập trong mô phỏng, xác thực và kiểm thử hồi quy. Đây là các khả năng của Warp, không phải là sự đảm bảo về tính vi phân hoặc tính tất định cho toàn bộ quá trình triển khai MJWarp. Xem tài liệu Warp về tính vi phân và thực thi tất định để biết chi tiết.

Thử Warp: pip install warp-lang (≥ 1.15 cho tính tất định GPU), sau đó python -m warp.examples.browse, hoặc các notebook hướng dẫn.

MuJoCo Warp (MJWarp) là gì?

Một trình mô phỏng robot liên tục tính toán điều gì sẽ xảy ra tiếp theo: dựa trên các vị trí khớp, vận tốc, điều khiển và tiếp điểm hiện tại, nó thúc đẩy cảnh mô phỏng thêm một bước thời gian nhỏ. Trong bài viết này, một "thế giới" có nghĩa là một bản sao độc lập của cảnh đó và trạng thái của nó. Một thế giới có thể chứa cánh tay SO-101 đang với lấy một khối lập phương; thế giới khác có thể chứa cùng cánh tay đó bắt đầu từ một tư thế hơi khác.

MuJoCo và MJWarp có thể chạy cùng một robot và tác vụ tương thích, nhưng chúng tổ chức công việc khác nhau. MuJoCo phù hợp tự nhiên với việc phát triển và kiểm tra một hoặc vài thế giới trên CPU. MJWarp là một triển khai NVIDIA Warp của đường ống vật lý MuJoCo, đặt mô hình và một lô các trạng thái độc lập trên GPU NVIDIA; một lệnh gọi mjw.step sẽ thúc đẩy toàn bộ lô đó.

Giá trị của MJWarp không nhất thiết là bước mô phỏng nhanh hơn cho một thế giới. Đó là khả năng thúc đẩy hàng trăm hoặc hàng nghìn thế giới cùng lúc, cung cấp cho GPU đủ công việc song song để cải thiện thông lượng tổng hợp, tức là tổng số bước thế giới hoàn thành mỗi giây. Điều đó ưu tiên cho học tăng cường và lấy mẫu quy mô lớn, nơi việc thu thập kinh nghiệm quan trọng hơn việc giảm độ trễ của một môi trường đơn lẻ.

Blog này bao gồm các nội dung sau:

1. xác thực một thế giới MuJoCo,
2. chuyển nó sang MJWarp, tạo thành một lô,
3. xác minh nó và đo lường nó một cách chính xác.

Việc tinh chỉnh bộ giải, biểu diễn Jacobian và các chủ đề chuyên biệt về đa GPU hoặc tính tất định không bắt buộc cho quá trình di chuyển này và có thể được đề cập riêng.

Khi đó, sự khác biệt là chính xác:

• Độ trễ là thời gian thực tế (wall-clock time) cho một bước mô phỏng.
• Thông lượng tổng hợp là tổng số bước thế giới hoàn thành trên mỗi giây thời gian thực tế được đo lường.

Cách sử dụng cơ bản: cấu trúc (structs), kích thước lô và một bước tối thiểu

• Quá trình chuyển đổi API cốt lõi rất nhỏ:

Sử dụng mjw.make_data khi muốn tạo trạng thái mặc định/mới. Sử dụng mjw.put_data khi trạng thái MuJoCo đã khởi tạo chính xác cần phải vượt qua ranh giới di chuyển.

Việc cấp phát tài nguyên theo lô yêu cầu xác định các tham số sau (tham khảo Kích thước lô):

Tinh chỉnh hiệu suất

1. Ghi lại CUDA graph: mjw.step thực hiện nhiều lần khởi chạy kernel; hãy ghi lại một lần và phát lại thường xuyên:

with wp.ScopedCapture() as capture:
     mjw.step(m, d)
wp.capture_launch(capture.graph)

2. Định kích thước nconmax / naconmax / njmax vừa khít: bộ nhớ và khối lượng công việc sẽ thay đổi theo các thông số này. Tinh chỉnh bằng mjwarp-testspeed: --measure_alloc và theo dõi tình trạng tràn trong mjwarp-viewer.

Bài gốc còn tiếp — xem tiếp tại bài gốc ↗