Add cpp rollout and sim

judo/utils/mujoco_cpp.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
+
+from typing import Callable, Literal
+
+import numpy as np
+from mujoco import MjData, MjModel
+
+from judo_cpp import rollout, sim
+
+
+class RolloutBackend:
+    """The backend for conducting multithreaded rollouts."""
+
+    def __init__(self, num_threads: int, backend: Literal["mujoco"], task_to_sim_ctrl: Callable) -> None:
+        """Initialize the backend with a number of threads."""
+        self.backend = backend
+        if self.backend == "mujoco":
+            self.setup_mujoco_backend()
+        else:
+            raise ValueError(f"Unknown backend: {self.backend}")
+        self.task_to_sim_ctrl = task_to_sim_ctrl
+
+    def setup_mujoco_backend(self) -> None:
+        """Setup the mujoco backend."""
+        if self.backend == "mujoco":
+            self.rollout_func = rollout
+        else:
+            raise ValueError(f"Unknown backend: {self.backend}")
+
+    def rollout(
+        self,
+        model_data_pairs: list[tuple[MjModel, MjData]],
+        x0: np.ndarray,
+        controls: np.ndarray,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Conduct a rollout depending on the backend."""
+        # unpack models into a list of models and data
+        ms, ds = zip(*model_data_pairs, strict=True)
+        ms = list(ms)
+        ds = list(ds)
+
+        # getting shapes
+        nq = ms[0].nq
+        nv = ms[0].nv
+        nu = ms[0].nu
+
+        # the state passed into mujoco's rollout function includes the time
+        # shape = (num_rollouts, num_states + 1)
+        x0_batched = np.tile(x0, (len(ms), 1))
+        processed_controls = self.task_to_sim_ctrl(controls)
+        assert x0_batched.shape[-1] == nq + nv
+        assert x0_batched.ndim == 2
+        assert processed_controls.ndim == 3
+        assert processed_controls.shape[-1] == nu
+        assert processed_controls.shape[0] == x0_batched.shape[0]
+
+        # rollout
+        if self.backend == "mujoco":
+            _states, _out_sensors = self.rollout_func(ms, ds, x0_batched, processed_controls)
+        else:
+            raise ValueError(f"Unknown backend: {self.backend}")
+        out_states = np.array(_states)
+        out_sensors = np.array(_out_sensors)
+        return out_states, out_sensors
+
+    def update(self, num_threads: int) -> None:
+        """Update the backend with a new number of threads."""
+        if self.backend == "mujoco":
+            self.setup_mujoco_backend()
+        else:
+            raise ValueError(f"Unknown backend: {self.backend}")
+
+
+class SimBackend:
+    """The backend for conducting simulation."""
+
+    def __init__(self, task_to_sim_ctrl: Callable) -> None:
+        """Initialize the backend."""
+        self.task_to_sim_ctrl = task_to_sim_ctrl
+
+    def sim(self, sim_model: MjModel, sim_data: MjData, sim_controls: np.ndarray) -> None:
+        """Conduct a simulation step using cpp sim."""
+        processed_ctrl = self.task_to_sim_ctrl(sim_controls)
+        x0 = np.concatenate([sim_data.qpos, sim_data.qvel])
+        sim(sim_model, sim_data, x0, processed_ctrl)

judo_cpp/CMakeLists.txt

@@ -0,0 +1,86 @@
+cmake_minimum_required(VERSION 3.15)
+project(judo_cpp)
+
+set(CMAKE_CXX_STANDARD 17)
+find_package(OpenMP REQUIRED)
+find_package(Eigen3 REQUIRED)
+
+# 1. Locate Python, which is used to find the MuJoCo package
+find_package(Python
+  COMPONENTS
+    Interpreter
+    Development.Module
+  REQUIRED
+)
+execute_process(
+  COMMAND ${Python_EXECUTABLE} -c
+          "import mujoco, os; print(os.path.dirname(mujoco.__file__))"
+  OUTPUT_VARIABLE MJ_PKG_ROOT
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+)
+message(STATUS "MuJoCo package root: ${MJ_PKG_ROOT}")
+
+# 2. Glob for the actual libmujoco.so file(s) there
+file(GLOB MUJOCO_LIB_FILES
+     "${MJ_PKG_ROOT}/libmujoco.so*"  # Linux
+     "${MJ_PKG_ROOT}/libmujoco.dylib"  # macOS default
+     "${MJ_PKG_ROOT}/libmujoco.*.dylib"  # macOS versioned dylib
+)
+if(NOT MUJOCO_LIB_FILES)
+  message(FATAL_ERROR "Could not find any libmujoco.so* in ${MJ_PKG_ROOT}")
+endif()
+list(GET MUJOCO_LIB_FILES 0 MUJOCO_LIB)
+message(STATUS "Linking against MuJoCo library: ${MUJOCO_LIB}")
+
+# 3. Pull in mujoco headers - they're located under the include directory
+set(MJINC_DIR "${MJ_PKG_ROOT}/include")
+message(STATUS "MuJoCo include dir: ${MJINC_DIR}")
+include_directories(${MJINC_DIR})
+
+# 4. Build the pybind11 module
+find_package(pybind11 REQUIRED)
+pybind11_add_module(_judo_cpp
+  bindings.cpp
+  rollout.cpp
+)
+execute_process(
+  COMMAND ${Python_EXECUTABLE} -c "import site; print(site.getsitepackages()[0])"
+  OUTPUT_VARIABLE SITE_PACKAGES_DIR
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+)  # find the site-packages directory
+set_target_properties(_judo_cpp PROPERTIES
+  LIBRARY_OUTPUT_DIRECTORY "${SITE_PACKAGES_DIR}/judo_cpp"
+)  # rebuild the bound module to the site-packages directory
+
+# Configure RPATHs correctly for each platform
+if(APPLE)
+  set_target_properties(_judo_cpp PROPERTIES
+    MACOSX_RPATH ON
+    BUILD_WITH_INSTALL_RPATH TRUE
+    BUILD_RPATH "@loader_path/../mujoco"
+    INSTALL_RPATH "@loader_path/../mujoco"
+  )
+  # MuJoCo pip wheel on macOS ships a flat lib, but exposes a framework-style install_name.
+  # Rewrite the dependency on the built extension to the flat @rpath/libmujoco.*.dylib.
+  get_filename_component(MUJOCO_LIB_NAME "${MUJOCO_LIB}" NAME)
+  add_custom_command(TARGET _judo_cpp POST_BUILD
+    COMMAND install_name_tool -change
+      "@rpath/mujoco.framework/Versions/A/${MUJOCO_LIB_NAME}"
+      "@rpath/${MUJOCO_LIB_NAME}"
+      "$<TARGET_FILE:_judo_cpp>"
+    VERBATIM)
+else()
+  set_target_properties(_judo_cpp PROPERTIES
+    BUILD_RPATH "$ORIGIN/../mujoco"
+    INSTALL_RPATH "$ORIGIN/../mujoco"
+    BUILD_RPATH_USE_ORIGIN TRUE
+  )
+endif()
+target_link_libraries(_judo_cpp PRIVATE
+  Eigen3::Eigen
+  OpenMP::OpenMP_CXX
+  ${MUJOCO_LIB}
+)
+
+# 5. Install the module
+install(TARGETS _judo_cpp DESTINATION judo_cpp)

judo_cpp/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
+
+from judo_cpp._judo_cpp import (
+    rollout,
+    sim,
+)
+
+__all__ = [
+    "rollout",
+    "sim",
+]

judo_cpp/bindings.cpp

@@ -0,0 +1,103 @@
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/numpy.h>
+#include <mujoco/mujoco.h>
+
+#include "rollout.h"
+
+namespace py = pybind11;
+
+static std::vector<const mjModel*> getModelVector(const py::list& python_models) {
+    std::vector<const mjModel*> model_vector;
+    model_vector.reserve(python_models.size());
+    for (auto&& item : python_models) {
+        auto ptr = item.attr("_address").cast<std::uintptr_t>();
+        model_vector.push_back(reinterpret_cast<const mjModel*>(ptr));
+    }
+    return model_vector;
+}
+
+static std::vector<mjData*> getDataVector(const py::list& python_data) {
+    std::vector<mjData*> data_vector;
+    data_vector.reserve(python_data.size());
+    for (auto&& item : python_data) {
+        auto ptr = item.attr("_address").cast<std::uintptr_t>();
+        data_vector.push_back(reinterpret_cast<mjData*>(ptr));
+    }
+    return data_vector;
+}
+
+PYBIND11_MODULE(_judo_cpp, m) {
+    // Function to shutdown persistent thread pool
+    m.def("shutdown_thread_pool",
+          []() {
+              ThreadPoolManager::instance().shutdown();
+          },
+          R"doc(
+Shutdown the persistent thread pool.
+
+Call this function to clean up the persistent thread pool when done with rollouts.
+The pool will be automatically recreated on the next call to persistent_cpp_rollout.
+)doc");
+
+    // Rollout
+    m.def("rollout",
+          [](const py::list& models,
+             const py::list& data,
+             const py::array_t<double>& x0,
+             const py::array_t<double>& controls)
+          {
+              // turn Python lists into vectors of mjModel*/mjData*
+              auto models_cpp = getModelVector(models);
+              auto data_cpp   = getDataVector(data);
+
+              // call into your C++ implementation
+              return Rollout(models_cpp, data_cpp, x0, controls);
+          },
+          py::arg("models"),
+          py::arg("data"),
+          py::arg("x0"),
+          py::arg("controls"),
+          R"doc(
+Run parallel MuJoCo rollouts.
+
+    Args:
+    models:                 length-B list of mujoco._structs.MjModel
+    data:                   length-B list of mujoco._structs.MjData
+    x0:                     2D array of shape (B, nq+nv), batched initial [qpos;qvel]
+    controls:               3D array of shape (B, horizon, nu), batched control inputs
+
+Returns:
+    tuple of three np.ndarray:
+    states  -> shape (B, horizon+1, nq+nv) - MuJoCo states (includes initial state)
+    sensors -> shape (B, horizon, nsensordata) - sensor data
+)doc");
+
+    // Sim
+    m.def("sim",
+          [](py::object model,
+             py::object data,
+             const py::array_t<double>& x0,
+             const py::array_t<double>& controls)
+          {
+              auto model_ptr = reinterpret_cast<const mjModel*>(model.attr("_address").cast<std::uintptr_t>());
+              auto data_ptr  = reinterpret_cast<mjData*>(data.attr("_address").cast<std::uintptr_t>());
+              Sim(model_ptr, data_ptr, x0, controls);
+          },
+          py::arg("model"),
+          py::arg("data"),
+          py::arg("x0"),
+          py::arg("controls"),
+          R"doc(
+Run a single MuJoCo simulation step.
+
+    Args:
+    model:      mujoco._structs.MjModel
+    data:       mujoco._structs.MjData
+    x0:         1D array of shape (nq+nv), initial [qpos;qvel]
+    controls:   1D array of shape (nu), control input
+
+Returns:
+    None
+)doc");
+}

judo_cpp/pyproject.toml

@@ -0,0 +1,20 @@
+[build-system]
+requires = ["scikit-build-core[pyproject]", "pybind11", "mujoco"]
+build-backend = "scikit_build_core.build"
+
+[project]
+name = "judo-cpp"
+version = "0.0.1"
+
+[project.optional-dependencies]
+dev = [
+    "ninja",
+    "pybind11",
+]
+
+[tool.scikit-build]
+cmake.source-dir = "."
+wheel.packages = ["judo_cpp"]
+cmake.build-type = "Release"
+build-dir = "build"  # for manual rebuilding during dev
+cmake.args = ["-DEigen3_DIR=$ENV{CONDA_PREFIX}/lib/cmake/eigen3"]