TL;DR: I ported all 130 LIBERO robot manipulation tasks to MuJoCo Warp and ran them on one RX 9070 XT ($700, 16 GB). Physics hits 17,137 env-steps/s at 2,048 worlds, against 38 env-steps/s on CPU. A BC transformer scores 42.5% on Warp against 50% on CPU. Code: poad42/libero_mjx.

LIBERO is the benchmark I keep coming back to. It’s 130 robot manipulation tasks across five suites (spatial, object, goal, scene10, and scene90), each with human demonstrations and a language goal, and it’s the standard testbed for language-conditioned imitation learning. The reference implementation runs on robosuite with CPU MuJoCo, one environment per process. Ten 600-step episodes take 92.6 seconds that way, one environment at a time, which is 38.4 environment-steps per second.

I wanted that on a GPU. Not a datacenter GPU; a $700 one.

The reason is rollouts. Behavioral cloning needs them. You run the policy, watch where it fails, and generate labels from the failures. The counterfactual-regret work I wrote about in July needed 18,000 labels, and the CPU vectorized environment would have chewed on that for hours. The AMD Instinct blog used the same trick for XArm6 sim2real, on MI300X and MI355X: run hundreds of worlds at once on the GPU with MuJoCo Warp. I co-wrote that post; it was my project over the summer at AMD. I wanted to know if the consumer card could do it.

There is a second reason, and it is the one that surprised me. Generating the demonstrations for one LIBERO suite, 10 tasks, took me 8 to 9 hours on a 16-core/32-thread consumer CPU. All five suites is a multi-day job. The same work on the GPU is minutes, which turns dataset generation from something you do once into something you can iterate on. That matters if you are producing rollouts, counterfactual labels, or synthetic demos rather than just evaluating a policy.

So I ported all 130 LIBERO tasks to MuJoCo Warp, and as far as I can tell it’s the first MJX+Warp port of LIBERO. I trained and evaluated a BC transformer on it too, on an RX 9070 XT, which I believe is the first time that’s been done on RDNA 4 outside AMD. It works. Some of it doesn’t. I’ll get to both.

Everything is at github.com/poad42/libero_mjx.

The GPU

The RX 9070 XT is a $700 consumer card: RDNA 4, gfx1201, 16 GB of VRAM, 304 W. The FLOPS don’t matter much here; the 16 GB and the ROCm stack do. Warp compiles HIP kernels for gfx1201 at runtime, JAX sees the card as rocm:0, and PyTorch sees it as cuda. Same card renders the scene and runs the policy.

16 GB is the annoying part. With two 128x128 cameras per environment, the Warp renderer allocates ray-tracing buffers per world, and my eval runs out of memory at 100 environments. Fifty fit, and give 66.4 environment-steps per second. It’s a memory wall, not a compute wall; the physics kernels are half asleep at that batch size.

What the port actually is

Three pieces.

The 130 tasks come from 131 MJCF XMLs pulled out of robosuite. The XMLs are scene structure only; the meshes and textures stay in the LIBERO and robosuite packages and get resolved by path. I wrote a small resolver so the same XMLs load in a container, on my workstation, or on a cluster without rewriting 131 files by hand.

LiberoEnv wraps any of the 130 tasks behind one reset/step, batched with jax.vmap, with a flax.struct.PyTreeNode holding the mjx.Data. The controller is a JAX port of robosuite’s OSC_POSE: Jacobian at the gripper site, densify the mass matrix, solve for the torques that push the end-effector to a target pose. That part was mostly transcription.

WarpRenderer does the images, on the GPU, batched. That part was not transcription.

The Warp port

Warp is NVIDIA’s. The AMD HIP/ROCm port for Strix Halo was open-sourced by Tomas Thoresen, validated on gfx1151. It works great on RDNA 4 as well, which is what this whole post runs on. If you want Warp on an AMD GPU, that is the branch to clone: AMD-Ecosystem/warp, or Tomas’s own fork.

The Instinct blog from the intro is this same stack on Instinct MI300X and MI355X.

How fast it goes

The first working Warp path was slower than the CPU. I kept the OSC controller in JAX and ran 25 physics substeps per control step, and got 18.9 environment-steps per second at 10 environments. The JAX physics step launches dozens of small kernels, and at 10 worlds the launch overhead eats everything.

Cutting the substeps from 25 to 5 (a 0.01 s simulation timestep instead of 0.002 s) cut the step time by 2.9x and got me to 40.9 env-steps/s. Then I ported the OSC controller from JAX to Torch and called mujoco_warp.step() directly, which removed the JAX dispatch entirely: 51.7 env-steps/s, 1.35x the CPU, at 60% average success across three seeds.

environment-steps / s 38.4 CPU 1 env 18.9 Warp JAX 25 sub 40.9 Warp JAX 5 sub 51.7 Warp Torch 5 sub
LIBERO spatial task 0, 10 environments, 600 max steps. My first Warp path was slower than the CPU. Moving the controller to Torch and cutting the substeps got it to 1.35x. Physics-only, at scale, it hits 17,137 env-steps/s at 2,048 worlds.

The 10-environment numbers are launch-bound. When I profile physics alone and crank the batch, the card wakes up: 1,568 env-steps/s at 768 environments, and 17,137 at 2,048, where the solver finally becomes compute-bound. The July post quoted 9,838 steps/s across 1,024 environments, a 49x speedup over the CPU SubprocVectorEnv, which is what made those 18,000 labels take 7 minutes instead of hours.

For a lighter scene, the state-based PandaPickCube from MuJoCo Playground, the same card does 212,445 env-steps/s at 2,048 worlds, 2.53x the JAX backend. Brax PPO on that task runs 2.16x faster on Warp (64.7 s vs 140.0 s for 2 million timesteps).

MuJoCo Playground PandaPickCube, 15 warmup and 50 timed steps. Hover a point for the exact number, and toggle log or linear. Warp is ahead at 64, 1024, 2048 and 4096, JAX wins at 256, and 512 is a tie. Warp peaks at 2048 (212,445 env-steps/s) and degrades at 4096; JAX keeps scaling. The 256 point is reproducibly bad for Warp (21.78 ms/step against 7.00 at 64) and I never explained it.

The BC transformer

The policy is LIBERO’s BC transformer: a ResNet image encoder, a 4-layer temporal transformer (6 heads, embedding 64, sequence length 10), and a 5-mode Gaussian-mixture head. It takes two 128x128 cameras (agentview and robot0_eye_in_hand) plus joint and gripper state, and spits out a 7-DoF action.

Fifty epochs on spatial task 0 gets 50% success on CPU episodes. On Warp, with the renderer fixes below, the same checkpoint gets 42.5% across four seeds (50, 40, 30, 50). Without the fixes, it gets 0%.

0% 20% 40% 60% success rate 50 CPU 50 Warp 42 40 Warp 123 30 Warp 999 50 Warp 7 0 Warp, no fix
LIBERO spatial task 0, one 50-epoch BC checkpoint, 10 episodes per seed. The Warp numbers need the renderer fixes; without them the policy scores 0% because it is looking at images it never trained on.

Across the spatial suite the 50-epoch checkpoints average 51.5% (task 8 is 95%, task 2 is 0%). The object suite averages 56.5%.

The interesting number isn’t 42.5%. It’s that a policy trained on CPU pixels works at all on GPU pixels, which means the renderer is close enough that the policy recognizes what it’s looking at.

The renderer was the whole fight

The physics port was mechanical. The renderer was not. A BC policy is a fixed function of the pixels it was trained on, and the Warp ray tracer is not MuJoCo’s CPU renderer. Every difference in shading, orientation, or texture moves the policy off its training distribution. My first Warp eval scored 0%, and it took a while to figure out why.

The first one was dumb. The Warp renderer and MuJoCo’s OpenGL renderer both use a bottom-left origin; robosuite’s training images are top-left. The policy was looking at an upside-down scene. Flip the image and it works.

The second was a constant. The Warp megakernel kept 30% of the diffuse and specular light on shadowed geometry (visible = 0.3). MuJoCo zeroes it and gives shadowed geometry ambient only. Changing it to 0.0 dropped the Warp-vs-EGL RMSE from 34.4 to 22.2.

The third was brightness. The ray tracer is about 85% as bright as the CPU renderer, and not uniformly (0.856 at center, 0.929 at the edges), so it’s not a clean exposure scale. A 1.15x multiplier on the output moved average success from 30% to 42.5%. That one fix did more than the kernel changes combined. 1.10, 1.15, and 1.20 all gave the same result on the seed I checked, so the exact value isn’t critical.

The fourth was cube maps. Warp sampled every texture as 2D. MuJoCo supports cube textures stored as six faces, and a cube material sampled as 2D reads a horizontal line across the face strip. LIBERO’s table uses a cube texture for its wood grain, so the table rendered as a flat tan smear. Reproducing MuJoCo’s cube texgen, for both the 6:1 strip layout and the square texture MuJoCo repeats on all six faces, cut the eye-in-hand RMSE from 29.9 to 20.4 and brought the grain back.

CPU (EGL) reference render Warp ray-traced render
CPU (EGL) Warp
Drag to compare. Eye-in-hand camera, spatial task 0. Before the fix the table is flat tan and the stove has no burner rings; after, the wood grain and rings are back. Eye-in-hand RMSE went from 29.9 to 20.4.
CPU (EGL) reference render Warp ray-traced render
CPU (EGL) Warp
The same comparison for the agentview camera.

The fifth one cost me a day. The renderer fixes rewrite mujoco_warp’s source on disk, so they have to run before mujoco_warp is imported. I had them running after. The file changed under the compiler, and warp resolved the nested render megakernel against stale line numbers, which fails as “kernels cannot return values”. Patch first, import second.

Fix Effect
Vertical flip 0% → working
Shadow constant 0.3 → 0.0 RMSE 34.4 → 22.2
Brightness 1.15x Success 30% → 42.5%
Cube-map sampling Eye-in-hand RMSE 29.9 → 20.4

The gap that’s left between Warp (42.5%) and CPU (50%) is the residual brightness mismatch, plus the ray tracer having no alpha blending for the transparent end-effector target geoms. Both are in the mujoco_warp render kernel, not in my code.

What I didn’t finish

Two things, and I want to be straight about them.

The loop isn’t closed. The Warp path collects data; it doesn’t train on that data in the loop. The 50-epoch policy trains offline on the LIBERO HDF5 demos in PyTorch, no simulator in the process. The Warp rollouts generate the counterfactual labels the repair head learns from and evaluate the trained policy, but the gradient step never sees the renderer. Rendering in the training loop is a different problem from rendering for evaluation, and I haven’t done it for LIBERO.

CPU rendering is broken on my reproducible base. I rebuilt the container on the TheRock AlmaLinux image, and its Mesa (23.1) is too old for gfx1201 EGL. The GPU Warp renderer doesn’t care, but the CPU (robosuite) renderer does, so eval_bc.py and the CPU half of render_comparison.py need a newer Mesa. I run those on a host that has one. The comparison image above came from there.

Where the in-loop version already exists

Rendering inside the training loop is not new, and I didn’t invent it. MuJoCo Playground does it with DeferredVisionWrapper: the env builds a Warp render context once (mjx.create_render_context) and, every step, calls mjx.refit_bvh, mjx.render, and mjx.get_rgb under jit and vmap. Physics, rendering, and the policy end up in one compiled program. PandaPickCubeCartesian is the reference env.

That path runs on one R9700, the 32 GB big brother of the 9070 XT, same RDNA 4 silicon. PandaPickCubeCartesian trained to 15,052,800 steps with the renderer in the loop, and the XArm6 pick-place vision runs reached 150,200,320 steps. The XArm6 run was a single card at 1,024 environments, and it did roughly 5,900 env-steps/s over 6.6 hours with rendering on the critical path. Those are PPO runs on the XArm6, not BC on LIBERO, but they show the pattern works on this hardware.

What this port does: collect offline Warp physics + Warp render dataset (HDF5 / labels) offline train PyTorch, no sim policy What Playground and the cluster do: train in the loop Warp physics + Warp render policy forward (CNN / transformer) PPO update on device one compiled step, no host round trip
The two pipelines. Top is what this port does today. Bottom is what the Playground and cluster vision runs do, and what I'd need to build for LIBERO.

How this actually went

It wasn’t a straight line, and there were long gaps where I was working on other things. Roughly:

  • May to June: the counterfactual-regret and flow-matching work. That’s where the first 10 LIBERO_SPATIAL tasks got ported, because I needed fast rollouts for gate labels.
  • July 7: the CFR blog post.
  • July 8 to 14: the real port. All 130 tasks, the parallel init, the Warp eval pipeline.
  • July 29 to 31: a long push on performance and rendering. The Torch OSC port, the profiling, the shadow and brightness fixes, the CPU-vs-Warp comparison images.
  • August to September: other projects. The port sat.
  • October: I came back, rebuilt the container on public sources, and fixed the cube-map and patch-ordering bugs.

So the repo has bursts of work with gaps in between. That’s just how it went.

Making it reproducible

The port is public now, and the container builds from public sources only. TheRock manylinux base, ROCm from the pip SDK, PyTorch and JAX from the TheRock index, warp from the public cu-basil/warp fork, MuJoCo pinned on PyPI, and LIBERO from the public cu-basil/LIBERO fork. No private wheels, no local docker commit.

Two warp bugs I hit on the way, the LLVM 23 build failure and the nested-kernel source extraction, are upstream in cu-basil/warp now, so the container build applies no patches at all. A third is a runtime workaround in libero_mjx/warp_kernel_patch.py, and I will explain why it is not upstream yet. docker/verify_stack.py checks the GPU backend, the MJX-Warp wiring, and a reset-and-step on the real task XMLs.

The CCD kernel name drift

The example below surfaced a third warp bug. mujoco_warp’s CCD path calls wp.get_suggested_block_size, which loads the module at the default block_dim of 256, and then launches the same kernel at block_dim=64. For a nested module="unique" kernel, warp hashes the resolved static expressions, and adj.build() repopulates those when the module is compiled for the second block_dim. So the live kernel’s hash drifts away from the name the module was cached under, and ModuleExec.get_kernel_hooks looks up a mangled name the metadata does not have:

KeyError: 'ccd_kernel_builder__locals__ccd_kernel_<hash>_cuda_kernel_forward_smem_bytes'

It is intermittent, depending on hashing order, which is why it took a while to pin down. warp_kernel_patch.py wraps get_kernel_hooks so that when the current mangled name is missing it recovers the name the module was actually built with from the metadata. The match is exact (a single <key>_<8 hex> entry), so a module that legitimately holds several kernels under one key still fails loudly. It is a no-op where the hash does not drift, and it works on both the warp 1.13 and 1.17 APIs used here. It is a workaround, not the fix; the real fix belongs in warp, where a kernel’s hash should be stable across block_dim variants.

Try it

There’s a minimal example in the repo, examples/libero_quickstart.py. It loads one task, runs a batched rollout on Warp, renders both cameras on the GPU, and prints the throughput:

python examples/libero_quickstart.py --suite spatial --task-id 0 --n-envs 8 --steps 20

[quickstart] 100 init states from the LIBERO init file
[quickstart] spatial task 0, 8 envs, 128x128 cameras, impl=warp
  step   10      19.8 env-steps/s  success=0.00
  step   20      20.0 env-steps/s  success=0.00
[quickstart] 8 envs x 20 steps in 8.0s = 19.9 env-steps/s
[quickstart] success: 0/8 = 0%
[quickstart] wrote /tmp/frame.png (agentview | eye-in-hand)

Zero actions, so 0% success is the expected result; the point is that the env steps and the renderer work. At 8 environments the number is launch-bound, the same reason the 10-environment eval above is slow. --batch-bench sweeps the batch size and shows the curve:

  n_envs   env-steps/s   ms/step
      16          38.5    415.61
      64         151.9    421.27
     256         302.6    846.02

The --checkpoint flag runs a trained BC policy instead of zero actions. The 50-epoch spatial task 0 checkpoint (21.6 MB) ships in the repo, so you don’t have to train before trying it; it still needs that task’s demo HDF5 for the shape metadata. --save-image writes one agentview | eye-in-hand frame. Inside the container it’s:

./scripts/docker_run.sh python examples/libero_quickstart.py --suite spatial --task-id 0 --save-image /tmp/frame.png

Takeaways

  • A $700 consumer GPU runs all 130 LIBERO tasks in parallel. Physics alone hits 17,137 env-steps/s on LIBERO at 2,048 worlds, and 212,445 on a lighter Playground scene. The full loop, with rendering and the BC policy, runs 51.7 env-steps/s at 10 environments.
  • It is also a dataset generator. One suite of demonstrations took 8 to 9 hours on a 16-core/32-thread CPU; the full 130 tasks is a multi-day job. On the GPU it is minutes, so you can iterate instead of running it once.
  • The renderer was the hard part, not the physics. Four fixes, and the brightness one alone was worth 12.5 points of success rate.
  • The loop isn’t closed here. The Warp path collects data; training is offline. In-loop vision training works on this hardware (MuJoCo Playground, and the cluster’s 150M-step XArm6 runs), and porting it to LIBERO is the next thing I’d do.
  • Everything builds from public sources. If you want to try it, the repo is poad42/libero_mjx.

Work with me

I’m looking for a role starting May 2027, in ML systems, GPU/edge deployment, or robotics infrastructure. If you’re hiring, or just want to talk about any of this, my email is in the footer. The code and the rest of my projects are public.

References