libero-mjx

All 130 LIBERO manipulation tasks ported to MuJoCo Warp. 131 task XMLs, 5 suites, a BC transformer training pipeline, and two eval paths: CPU (robosuite) and GPU (Warp physics + Warp rendering).

LIBERO runs 130 manipulation tasks on robosuite with CPU MuJoCo. This repo ports all 130 to MuJoCo Warp, which runs MuJoCo on GPU via JAX & Warp. You get thousands of parallel environment instances on one GPU instead of one at a time on CPU.

A BC transformer trained on CPU demo data hits 50% success on CPU eval and 42.5% average on Warp eval (4 seeds: 50%, 40%, 30%, 50%). Without the rendering fixes in this repo, Warp eval gives 0%.

Suites

131 task XMLs ship in libero_mjx/assets/xml/.

Install

pip install -e .
pip install -e ".[torch,libero]"   # BC training & eval with robosuite
pip install -e ".[render]"        # GPU rendering

You need Python 3.10+, MuJoCo 3.1.0+, JAX 0.4.28+, and Warp 1.0+. A CUDA or ROCm GPU is required for Warp physics.

NVIDIA: pip install jax[cuda12]

AMD ROCm: use the pip ROCm SDK and the matching JAX wheels. For ROCm 10.x that is jax-rocm10-plugin + jax-rocm10-pjrt; pin them to the same version as jax. See docs/reproducibility.md or just use the container, which does this for you.

Docker (recommended)

The container is reproducible and vendor-parameterised. Everything comes from a public source: the TheRock manylinux base image, the pip ROCm SDK, the TheRock torch/JAX wheels, the public cu-basil/warp branch, and PyPI MuJoCo. See docs/reproducibility.md for the full pin table.

./scripts/docker_build.sh                          # AMD ROCm, gfx1201 (default)
GFX_TARGET=gfx942 ./scripts/docker_build.sh        # CDNA
ACCELERATOR=cuda ./scripts/docker_build.sh         # NVIDIA

./scripts/docker_run.sh python /opt/verify_stack.py
./scripts/docker_run.sh python tests/test_all_suites.py

scripts/docker_run.sh passes all arguments through to the container. It sets JAX_PLATFORMS=rocm, XLA_PYTHON_CLIENT_MEM_FRACTION=0.15, mounts the repo at /workspace/libero-mjx, and optionally mounts a host LIBERO checkout (LIBERO_BASIL_PATH), datasets (LIBERO_DATASETS) and robosuite assets (ROBOSUITE_ASSETS). The image already contains LIBERO and robosuite assets, so these mounts are only needed to override them.

Warp branch by target: amd-integration-halo (RDNA, gfx10xx/gfx11xx/gfx12xx) or amd-integration (CDNA, gfx9xx). scripts/docker_build.sh picks the right one from GFX_TARGET.

Rendering: the GPU Warp renderer is the supported path in the container. The CPU (robosuite) renderer needs Mesa 24.1+ for gfx1201 EGL, so scripts/eval_bc.py and the CPU side of scripts/render_comparison.py require a host with a newer Mesa. See docs/rendering.md.

If you add an entry point, call patch_render_kernel() (or import libero_mjx) before importing mujoco_warp; see docs/rendering.md.

Assets

The task XMLs reference meshes and textures by path; none are redistributed. LIBERO (MIT code) provides libero/libero/assets, robosuite (MIT) provides robosuite/models/assets. On a host:

pip install robosuite==1.5.1
./scripts/setup_assets.sh

A non-canonical layout is supported with LIBERO_ASSETS_ROOT and ROBOSUITE_ASSETS_ROOT; the loader rewrites the XML roots in memory. See docs/assets.md.

Datasets

BC training uses LIBERO human demonstrations, published under CC BY 4.0 at https://huggingface.co/datasets/yifengzhu-hf/LIBERO-datasets.

python scripts/download_datasets.py --suite spatial     # ~6.2 GB
python scripts/download_datasets.py --suite all         # ~100 GB

See docs/datasets.md for sizes and verification.

Quick start

Runnable example

examples/libero_quickstart.py loads one task, runs a batched Warp rollout, renders both cameras on the GPU, and prints the throughput and success:

python examples/libero_quickstart.py --suite spatial --task-id 0 --n-envs 8 \
    --steps 20 --save-image /tmp/frame.png
python examples/libero_quickstart.py --batch-bench

# The 50-epoch spatial task 0 BC checkpoint ships in the repo (21.6 MB).
# It needs the task demo HDF5 for the shape metadata; see docs/datasets.md.
python examples/libero_quickstart.py --suite spatial --task-id 0 \
    --checkpoint checkpoints/task0_model_50ep.pth --n-envs 8

--save-image writes one agentview | eye-in-hand frame. Inside the container:

./scripts/docker_run.sh python examples/libero_quickstart.py --suite spatial --task-id 0

Warp environment

import jax
import jax.numpy as jp
from libero_mjx.envs.libero import LiberoEnv

env = LiberoEnv(suite="spatial", task_id=0, impl="warp", n_envs=256)
state = env.reset(jax.random.PRNGKey(0))

action = jp.zeros((256, 7))
state = env.step(state, action)
print(state.metrics["success"])  # (256,) array

Train a BC policy

python scripts/train_bc.py --suite spatial --task-id 0 --epochs 50 \
    --batch-size 32 --save checkpoints/spatial_task0.pth

Evaluate

Two paths. CPU eval uses robosuite's OffScreenRenderEnv, same renderer that produced the training data. Warp eval runs Warp physics & Warp rendering on GPU with 10 parallel envs.

# CPU eval
python scripts/eval_bc.py --suite spatial --task-id 0 \
    --ckpt checkpoints/spatial_task0.pth --n-eval 20 --max-steps 600

# Warp eval
python scripts/eval_warp_only.py --suite spatial --task-id 0 \
    --ckpt checkpoints/spatial_task0.pth --n-eval 10 --max-steps 600

The Warp eval accepts --brightness 1.15 (default) to correct the Warp ray tracer's output brightness. Pass --brightness 1.0 to disable.

Architecture

libero_mjx/
  __init__.py              Package exports, auto GPU patch
  warp_gpu_patch.py        ROCm FFI, device detection, lbvh, GraphMode
  warp_kernel_patch.py     Warp kernel-name drift fix for nested unique modules
  robosuite_patch.py       robot_base_factory fallback for non-spatial suites
  texture_patch.py         Warp type code for Texture2D arrays
  render_kernel_patch.py   Patches mujoco_warp render kernel (shadow, haze)
  envs/
    base.py                LiberoMjxEnv: batched reset/step, state save/restore
    libero.py              LiberoEnv: unified env for all 5 suites
    spatial.py             Legacy spatial-only env
  controllers/
    osc.py                 OSC Cartesian impedance controller (JAX port)
  predicates/
    spatial.py             Success predicates: distance, on, in_region, is_open
  obs/
    __init__.py            State-only observation builder
  render.py                WarpRenderer: batched GPU rendering with DLPack
  assets/
    xml/                   131 task XMLs extracted from robosuite

examples/
  libero_quickstart.py     Batched Warp rollout + GPU render + throughput

scripts/
  train_bc.py              Train BC transformer on LIBERO demo data
  eval_bc.py               Evaluate BC via robosuite (CPU)
  eval_warp_only.py        Evaluate BC via Warp physics + rendering (GPU)
  render_comparison.py     Generate CPU vs Warp render comparison images for docs
  benchmark_eval.py        Benchmark CPU vs Warp eval throughput
  prof_eval.py             Per-phase profiling of the Warp eval loop
  eval_warp_native.py      Eval via pure Warp + Torch (no JAX) physics
  prof_warp_native.py      Per-phase profiling of the Warp native eval
  extract_all_xmls.py      Extract task XMLs from robosuite
  setup_assets.sh          Fetch LIBERO + link robosuite assets
  download_datasets.py     Download LIBERO demo datasets (CC BY 4.0)
  docker_build.sh          Build the reproducible image (ROCm or CUDA)
  docker_run.sh            Docker wrapper for GPU scripts

docker/
  verify_stack.py          Build- and run-time stack verification

tests/
  test_all_suites.py       Smoke test: all 5 suites load & step
  test_env_smoke.py        Basic env reset/step
  test_osc.py              OSC controller tests
  test_predicates.py       Predicate tests
  test_spatial.py          Spatial suite integration test
  test_parallel.py         Parallelism / vmap tests
  test_vmap_step.py        vmap step tests
  test_physics_compare.py  Warp vs CPU physics comparison
  validate_warp_render.py  Warp renderer validation

Rendering fixes

The Warp ray tracer (mujoco_warp) differs from MuJoCo's CPU / EGL renderer in ways that break a BC policy trained on CPU data. WarpRenderer in libero_mjx/render.py applies the fixes below. The kernel and texture fixes patch the installed mujoco_warp package on disk before import; the flip and brightness fixes run at render time.

Cube-map material textures (render_cube_patch.py)

Warp sampled every texture as 2D, so cube materials (the table wood grain, the stove top) rendered as a flat smear. The patch reproduces MuJoCo's classic cube texgen for both layouts: a 6:1 strip of six faces and a square texture that MuJoCo repeats on all six sides. On spatial task 0 this cuts eye-in-hand RMSE from 29.9 to 20.4. See docs/rendering.md.

Left: CPU (EGL). Right: Warp (ray trace) with all fixes applied, brightness 1.15x. Spatial task 0, init state 0, after 5 zero-action steps. Regenerate with python scripts/render_comparison.py.

Shadow fallback (render_kernel_patch.py)

The Warp render megakernel set visible = 0.3 for shadowed pixels. That constant, NO_LIGHT_AMBIENT_FALLBACK, keeps 30% of diffuse & specular light on geometry in shadow. MuJoCo's CPU renderer sets visible = 0.0. Shadowed geometry gets ambient light only, applied in a separate pass.

The patch changes the constant to 0.0 in _render_megakernel in the installed render.py. This dropped the Warp-vs-EGL RMSE from 34.4 to 22.2 on spatial task 0.

Haze blending (render_kernel_patch.py)

MuJoCo applies atmospheric haze: distant geometry blends toward the background color based on vis.map.haze, fogstart, and fogend. The Warp renderer had no haze. The patch adds haze blending after shading, before the pixel write.

For LIBERO scenes this has no visible effect. The fog starts at fogstart * extent = 3.0 * 10.61 = 31.83 units from the camera, but objects sit at 1 to 2 units. The haze factor is 0 for all visible geometry. The improvement comes from kernel recompilation: a different code path produces different floating-point intermediate values.

RenderContext fields (render_kernel_patch.py)

The RenderContext dataclass in types.py had no fields for haze parameters. The patch adds haze_amount, fogstart, fogend, and background_color_float. The create_render_context() function in io.py populates them from mjm.vis.map.haze, mjm.vis.map.fogstart * mjm.stat.extent, mjm.vis.map.fogend * mjm.stat.extent, and the background color.

Vertical flip (WarpRenderer)

OpenGL renders with a bottom-left origin. Robosuite & MuJoCo CPU output top-left. img.flip(dims=[1]) corrects the vertical axis.

Brightness boost (WarpRenderer)

The Warp ray tracer produces images at about 85% of the CPU renderer's brightness. The ratio is 0.856 at image center and 0.929 at edges, so it is not a uniform scale. The difference is likely a missing tone mapping or exposure step in the ray tracer. A 1.15x multiplier on the output RGB closes the gap. This single fix moved average success from 30% to 42.5% across 4 seeds.

The --brightness flag controls the multiplier. Values of 1.10, 1.15, and 1.20 all produced 50% success on seed 42. The default is 1.15.

What does not work

Replacing cube map textures (type=1) with flat average colors destroys rendering. Warp eval drops to 0%. The Warp renderer samples cube maps as 2D vertical strips, which is wrong, but the wrong result still carries enough texture information for the policy. Flat colors carry none.

Removing transparent geoms (the EEF target spheres & boxes at alpha 0.5 and 0.8) hurts success rate. Warp renders them as opaque because the ray tracer has no alpha blending. Keeping them visible, even as opaque shapes, matches the training data better than removing them.

Results

Spatial task 0, 50-epoch BC checkpoint

10 envs per eval. The policy samples actions stochastically, so success rates vary 30 to 50 percentage points across seeds with this sample size.

CPU vs Warp throughput

Spatial task 0, 10 episodes, 600 max steps, AMD RX 9070 XT:

The JAX-based Warp eval (25 substeps) is 0.5x CPU throughput. The JAX physics step (75% of time) runs 25 substeps of OSC controller + mjx.step, each launching dozens of small kernels on 10 envs. The Warp native env (libero_mjx/warp_env.py) ports the OSC controller to Torch and runs physics via direct mujoco_warp.step() calls, eliminating JAX overhead. With 5 substeps, it achieves 51.7 env-steps/s (1.35x CPU) with 60% average success. See docs/architecture.md for the full profiling breakdown.

Spatial suite, CPU eval, 20 episodes per task

Object suite, CPU eval

Goal, scene10, and scene90 CPU eval results are pending training.

BC training & evaluation

Training uses LIBERO's BC transformer: a ResNet image encoder, a temporal transformer, and a GMM policy head. It reads demo data from HDF5 files. Training runs on PyTorch & GPU. No robosuite or Warp needed.

python scripts/train_bc.py --suite spatial --task-id 0 --epochs 50 --save ckpt.pth

# With domain randomization (brightness, contrast, noise, blur)
python scripts/train_bc.py --suite spatial --task-id 0 --epochs 50 --save ckpt.pth --augment

Download demo data

python scripts/download_datasets.py --suite spatial

The datasets are CC BY 4.0 and come from LIBERO's HuggingFace hub, https://huggingface.co/datasets/yifengzhu-hf/LIBERO-datasets. Sizes, the pinned revision, the original Box fallback links, and a content check are in docs/datasets.md. LIBERO's own downloader still works:

from libero.libero.utils.download_utils import libero_dataset_download
libero_dataset_download(datasets="libero_spatial", use_huggingface=True)

Patches

The repo patches three external packages at import time. Each patch exists because the upstream package has a bug or missing feature that blocks LIBERO tasks on GPU.

warp_gpu_patch.py

JAX on ROCm reports its platform as rocm (lowercase). Warp registers FFI targets for ROCM (uppercase). The XLA compiler looks up rocm and finds nothing. The patch calls jax.ffi.register_ffi_target for both casings.

MJX's io module checks for CUDA GPUs with has_cuda_gpu_device. On ROCm, JAX reports devices under the gpu backend, but MJX's check looks for cuda by name. The patch rewrites the check to use jax.devices('gpu').

HIP does not support the cubql BVH constructor. The patch forces lbvh for mesh & heightfield BVH builds.

Some mujoco_warp versions ship GraphMode as an int instead of an enum. MJX expects GraphMode.WARP. The patch injects a compatible enum.

robosuite_patch.py

Robosuite 1.5.1's robot_base_factory returns a string for unknown base names like NullBase, which LIBERO uses for floor & kitchen scenes. Downstream code calls the return value as a class constructor and crashes with TypeError: 'str' object is not callable. The patch falls back to NullMount for unknown names.

This is required for object, goal, scene10, and scene90. Spatial works without it because it uses RethinkMount.

texture_patch.py

Warp 1.13.0's get_type_code does not recognize wp.array[wp.Texture2D] when hashing kernel arguments. It raises TypeError: Unrecognized type. The patch adds type codes "tex2d" and "atex2d" for Texture2D and arrays of Texture2D.

render_kernel_patch.py

Patches three files in the installed mujoco_warp package on disk: render.py, types.py, and io.py. See Rendering fixes above. Must run before any import mujoco_warp statement, because importing libero_mjx triggers mujoco.mjx.warp, which caches mujoco_warp._src.types in sys.modules. Once cached, the patch cannot take effect.

Documentation

• docs/architecture.md: environment, controller, and renderer design
• docs/rendering.md: CPU vs Warp rendering differences and fixes
• docs/physics.md: Warp physics verification & known differences
• docs/api.md: public API reference for LiberoEnv, WarpRenderer, and OscController
• docs/reproducibility.md: container pin table, build args, verification
• docs/assets.md: where the 3D assets come from and how to link them
• docs/datasets.md: where the demo datasets come from and how to fetch them

Attribution & Third-Party Licenses

This project builds on several open-source projects. See LICENSE for the full third-party license texts.

No binary assets are redistributed. The 131 task XML files in libero_mjx/assets/xml/ are derived from LIBERO's robosuite task definitions but contain only scene structure (object placement, robot config, contact parameters). All 3D meshes, textures, and robot models are referenced by path and are not copied into this repository. They come from:

• LIBERO assets — the cu-basil/LIBERO fork of https://github.com/Lifelong-Robot-Learning/LIBERO (libero/libero/assets), MIT
• robosuite assets — https://github.com/ARISE-Initiative/robosuite (robosuite/models/assets), MIT

scripts/setup_assets.sh fetches and links them; docs/assets.md documents the roots and the LIBERO_ASSETS_ROOT / ROBOSUITE_ASSETS_ROOT overrides.

Demonstration datasets used for BC training are downloaded from LIBERO's HuggingFace hub, https://huggingface.co/datasets/yifengzhu-hf/LIBERO-datasets, and are licensed under CC BY 4.0. scripts/download_datasets.py fetches them; docs/datasets.md records sizes and verification.

Citation

If you use this work, please cite the original LIBERO paper and MuJoCo Warp:

@inproceedings{wang2024libero,
  title={LIBERO: Lifelong Robot Learning},
  author={Wang, Haoyu and Wang, Junlin and Mayne, Matthew and Bao, Cheng and Li, Zichen and Ma, Wenlong and Konidaris, George},
  booktitle={ICLR 2024},
}

@misc{mujoco_warp,
  title={MuJoCo Warp: GPU-parallel simulation for MuJoCo},
  author={DeepMind},
  url={https://github.com/google-deepmind/mujoco_warp},
}

License

MIT — see LICENSE for the full text including third-party attributions.