diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml
index 368a9c06..a264dca3 100644
--- a/mjx/pyproject.toml
+++ b/mjx/pyproject.toml
@@ -32,6 +32,10 @@ dependencies = [
     "mujoco>=3.3.5.dev0",
     "scipy",
     "trimesh",
+]
+
+[project.optional-dependencies]
+warp = [
     "warp-lang==1.8.1",
 ]