Test gpu runners

.github/workflows/precommit.yml

@@ -18,6 +18,7 @@ jobs:
       - uses: prefix-dev/setup-pixi@v0.9.3
         with:
           pixi-version: latest
+          environments: dev-mjwarp
           cache: true
 
       # ── Stubs cache ────────────────────────────────────────────────
@@ -30,14 +31,14 @@ jobs:
 
       - name: Generate stubs if cache miss
         if: steps.typings-cache.outputs.cache-hit != 'true'
-        run: pixi run -e dev pybind11-stubgen mujoco -o typings/
+        run: pixi run -e dev-mjwarp pybind11-stubgen mujoco -o typings/
 
       # ── Linting & type checking ────────────────────────────────────
       - name: Run Ruff Linter
-        run: pixi run -e dev ruff check judo/ tests/ --output-format=github
+        run: pixi run -e dev-mjwarp ruff check judo/ tests/ --output-format=github
 
       - name: Run Ruff Formatter
-        run: pixi run -e dev ruff format judo/ tests/ --diff
+        run: pixi run -e dev-mjwarp ruff format judo/ tests/ --diff
 
       - name: Run Pyright
-        run: pixi run -e dev pyright
+        run: pixi run -e dev-mjwarp pyright

.github/workflows/test_gpu.yml

@@ -0,0 +1,52 @@
+name: Test (GPU / MJWarp)
+
+# on:
+#   push:
+#     branches: [main]
+#   pull_request:
+#     branches: [main]
+
+permissions:
+  contents: read
+
+jobs:
+  test-mjwarp:
+    runs-on: github-gpu-runner
+    timeout-minutes: 30
+
+    steps:
+      - uses: actions/checkout@v6
+
+      - uses: prefix-dev/setup-pixi@v0.9.3
+        with:
+          pixi-version: latest
+          environments: dev-test-mjwarp
+          cache: true
+
+      # ── mujoco_extensions build cache ───────────────────────────────
+      - name: Restore mujoco_extensions build cache
+        id: ext-cache
+        uses: actions/cache@v5
+        with:
+          path: |
+            mujoco_extensions/build
+            mujoco_extensions/policy_rollout/*.so
+          key: mujoco-ext-gpu-${{ hashFiles('pixi.lock', 'mujoco_extensions/**/*.cpp', 'mujoco_extensions/**/*.h', 'mujoco_extensions/**/CMakeLists.txt') }}
+
+      - name: Build mujoco_extensions (cache miss)
+        if: steps.ext-cache.outputs.cache-hit != 'true'
+        run: pixi run -e dev-mjwarp build
+
+      # ── Download mesh assets (avoid xdist race) ────────────────────
+      - name: Download meshes from release
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: pixi run -e dev-mjwarp python -c "from judo.utils.assets import download_and_extract_meshes; from judo import MODEL_PATH; download_and_extract_meshes(str(MODEL_PATH))"
+
+      # ── Pre-warm robot_descriptions cache (avoid xdist race) ───────
+      - name: Cache mujoco_menagerie meshes
+        run: pixi run -e dev-mjwarp python -c "from robot_descriptions import spot_mj_description"
+
+      # ── Tests ───────────────────────────────────────────────────────
+      - name: Run MPC GPU tests
+        run: pixi run -e dev-mjwarp pytest tests/test_run_mpc.py -rsx -n=1 --timeout=300

.gitignore

@@ -199,3 +199,6 @@ wheelhouse/
 
 # byproduct
 meshes.zip
+
+# generated dataset
+*.h5

CHANGELOG.md

@@ -1,3 +1,16 @@
+# v0.1.0
+
+## Added
+* Added support for parallelization using MJWarp
+  * Parallel simulation backend using MJWarp
+  * GPU accelerated data collection in run_mpc
+
+## Changed
+  * Added an optional GPU test since GPU runners aren't currently available on repo
+
+## Dev
+  * Updated mujoco to `>=3.5.0` to support MJWarp
+
 # v0.0.7
 
 ## Added

README.md

@@ -122,6 +122,7 @@ if __name__ == "__main__":
 ```
 
 To see more information about the available tasks, please refer to [the task README](judo/tasks/README.md).
+For batched MPC trajectory generation, see [run_mpc/README.md](run_mpc/README.md).
 
 ## 3. Running `judo` as a Dependency
 You can easily install `judo` as a dependency in your own project. A few comments:

judo/controller/__init__.py

@@ -1,6 +1,11 @@
 # Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
 
-from judo.controller.controller import Controller, ControllerConfig, make_controller
+from judo.controller.controller import (
+    BatchedControllers,
+    Controller,
+    ControllerConfig,
+    make_controller,
+)
 from judo.controller.overrides import (
     set_default_caltech_leap_cube_overrides,
     set_default_cartpole_overrides,
@@ -20,6 +25,7 @@
 set_default_spot_overrides()
 
 __all__ = [
+    "BatchedControllers",
     "Controller",
     "ControllerConfig",
     "make_controller",