rebase onto the latest master

Author: bd-jahn

mjx/mujoco/mjx/BUILD.bazel

@@ -268,6 +268,7 @@ py_library(
         ":scan",
         ":types",
         "//:mujoco-py",
+        "//python/mujoco/introspect:mjxmacro",
         requirement("jax"),
         requirement("numpy"),
     ],
@@ -478,6 +479,7 @@ py_test(
         requirement("absl-py"),
         requirement("numpy"),
         requirement("trimesh"),
+        requirement("typing_extensions"),
     ],
 )
 

python/mujoco/BUILD.bazel

@@ -251,7 +251,11 @@ pybind_extension(
 
 pybind_extension(
     name = "_rollout",
-    srcs = ["rollout.cc"],
+    srcs = [
+        "rollout.cc",
+        "threadpool.h",
+        "threadpool.cc"
+    ],
     dynamic_deps = [
         "//:mujoco_shared",
     ],
@@ -358,6 +362,7 @@ py_test(
         ":mujoco-py",
         requirement("absl-py"),
         requirement("numpy"),
+        requirement("typing-extensions"),
     ],
 )
 
@@ -393,22 +398,26 @@ py_test(
         ":mujoco-py",
         requirement("absl-py"),
         requirement("numpy"),
+        requirement("typing-extensions"),
     ],
 )
 
-# This fails due to the missing pystub
-# py_test(
-#     name = "specs_test",
-#     srcs = ["specs_test.py"],
-#     data = [
-#         "//:test_data",
-#     ],
-#     deps = [
-#         ":mujoco-py",
-#         requirement("absl-py"),
-#         requirement("numpy"),
-#     ],
-# )
+py_test(
+    name = "specs_test",
+    srcs = ["specs_test.py"],
+    data = [
+        "//python/mujoco/testdata:model",
+        ":symlinked_plugins",
+    ],
+    deps = [
+        ":mujoco-py",
+        requirement("absl-py"),
+        requirement("numpy"),
+        requirement("etils"),
+        requirement("typing-extensions"),
+        requirement("importlib-resources"),
+    ],
+)
 
 py_test(
     name = "minimize_test",
@@ -418,6 +427,7 @@ py_test(
         ":mujoco-py",
         requirement("absl-py"),
         requirement("numpy"),
+        requirement("typing-extensions"),
     ],
 )
 
@@ -456,5 +466,6 @@ py_test(
     deps = [
         ":mujoco-py",
         requirement("absl-py"),
+        requirement("typing-extensions"),
     ],
 )

python/mujoco/codegen/BUILD.bazel

@@ -6,8 +6,8 @@ py_binary(
     name = "generate_enum_traits",
     srcs = ["generate_enum_traits.py"],
     deps = [
-        "//introspect:ast_nodes",
-        "//introspect:enums",
+        "//python/mujoco/introspect:ast_nodes",
+        "//python/mujoco/introspect:enums",
         requirement("absl-py"),
     ],
 )
@@ -16,8 +16,8 @@ py_binary(
     name = "generate_function_traits",
     srcs = ["generate_function_traits.py"],
     deps = [
-        "//introspect:ast_nodes",
-        "//introspect:functions",
+        "//python/mujoco/introspect:ast_nodes",
+        "//python/mujoco/introspect:functions",
         requirement("absl-py"),
     ],
 )
@@ -26,8 +26,8 @@ py_binary(
     name = "generate_spec_bindings",
     srcs = ["generate_spec_bindings.py"],
     deps = [
-        "//introspect:ast_nodes",
-        "//introspect:structs",
+        "//python/mujoco/introspect:ast_nodes",
+        "//python/mujoco/introspect:structs",
         requirement("absl-py"),
     ],
 )

python/mujoco/enums.cc

@@ -44,7 +44,7 @@ inline T FloorDiv(T a, T b) {
 
 template <typename Trait>
 void DefEnum(py::module_& m) {
-  py::enum_<typename Trait::type> e(m, Trait::name, py::module_local());
+  py::enum_<typename Trait::type> e(m, Trait::name);
   for (const auto& [name, enumerator] : Trait::values) {
     e.value(name, enumerator);
   }

python/mujoco/introspect/BUILD.bazel

@@ -29,3 +29,9 @@ py_library(
     srcs = ["structs.py"],
     deps = [":ast_nodes"],
 )
+
+py_library(
+    name = "mjxmacro",
+    srcs = ["mjxmacro.py"],
+    deps = [":structs"],
+)

python/mujoco/introspect/codegen/BUILD.bazel

@@ -24,7 +24,7 @@ py_binary(
     srcs = ["generate_enums.py"],
     deps = [
         ":formatter",
-        "//introspect:ast_nodes",
+        "//python/mujoco/introspect:ast_nodes",
         requirement("absl-py"),
         "@abseil-cpp//absl/flags:flag",
     ],
@@ -50,8 +50,8 @@ py_binary(
     srcs = ["generate_functions.py"],
     deps = [
         ":formatter",
-        "//introspect:ast_nodes",
-        "//introspect:type_parsing",
+        "//python/mujoco/introspect:ast_nodes",
+        "//python/mujoco/introspect:type_parsing",
         requirement("absl-py"),
         "@abseil-cpp//absl/flags:flag",
     ],
@@ -78,8 +78,8 @@ py_binary(
     srcs = ["generate_structs.py"],
     deps = [
         ":formatter",
-        "//introspect:ast_nodes",
-        "//introspect:type_parsing",
+        "//python/mujoco/introspect:ast_nodes",
+        "//python/mujoco/introspect:type_parsing",
         requirement("absl-py"),
         "@abseil-cpp//absl/flags:flag",
     ],

test/benchmark/BUILD.bazel

@@ -4,6 +4,11 @@ names = [
     "engine_util_spatial_benchmark_test",
     "engine_core_smooth_benchmark_test",
     "engine_util_sparse_benchmark_test",
+    "ccd_benchmark_test",
+    "factorI_benchmark_test",
+    "inertia_benchmark_test",
+    "solveLD_benchmark_test",
+    "thread_performance_test",
 ]
 
 [

test/benchmark/ccd_benchmark_test.cc

@@ -19,7 +19,7 @@
 #include <string>
 
 #include <benchmark/benchmark.h>
-#include <absl/base/attributes.h>
+#include "absl/base/attributes.h"
 #include <mujoco/mjmodel.h>
 #include <mujoco/mujoco.h>
 #include "test/fixture.h"

test/engine/BUILD.bazel

@@ -22,6 +22,7 @@ cc_test(
         "engine_ray_test.cc",
         "engine_sensor_test.cc",
         "engine_solver_test.cc",
+        "engine_sort_test.cc",
         "engine_support_test.cc",
         "engine_thread_test.cc",
         "engine_util_blas_test.cc",
@@ -32,6 +33,7 @@ cc_test(
         "engine_util_sparse_test.cc",
         "engine_util_spatial_test.cc",
         "engine_vis_state_test.cc",
+        "engine_vis_visualize_test.cc",
     ],
     copts = [
         "-Wno-dangling-else",
@@ -54,6 +56,7 @@ cc_test(
     deps = [
         "//:mujoco",
         "//src/engine:engine_util_container",
+        "//src/engine:engine_util_sparse",
         "//test:fixture",
         "@googletest//:gtest_main",
     ],

test/user/BUILD.bazel

@@ -35,6 +35,7 @@ cc_test(
         "//src/user:user_resource",
         "//src/user:user_util",
         "//src/user:user_vfs",
+        "//src/cc:array_safety",
         "//test:fixture",
         "@googletest//:gtest_main",
     ],