Copybara import of the project:

--
3a95b62f59e81bfef0f076afb173ecc14b27943d by Levi Burner <leviburner@gmail.com>:

rollout prototype native threadpool for comparing to python threads

--
efd8be1124ac839b902de45973a3ca8b9f2215e6 by Levi Burner <leviburner@gmail.com>:

copy mjpcs threadpool into python bindings

--
75603eea3e8362e354a9675e8a6cd14e56ec3d28 by Levi Burner <leviburner@gmail.com>:

rollout use threadpool as translation unit

--
06b90febd021663f6cc81fd7895e4d6e2008ed97 by Levi Burner <leviburner@gmail.com>:

rollout add chunk_divisor parameter

--
298ab2f3c0d6e12530832c3cdbf784dd92d54806 by Levi Burner <leviburner@gmail.com>:

rollout add native threading test

--
169cf9978e7abad6edd1392b8e6aab995e4f8f10 by Levi Burner <leviburner@gmail.com>:

rollout exchange chunk_divisor arg for chunk_size

--
265af851d74432d261277d3dbda11cdef1841bc8 by Levi Burner <leviburner@gmail.com>:

rollout fix cosmetics

--
1e8bffa88bf36190501b334bef31147e23db39f7 by Levi Burner <leviburner@gmail.com>:

make native rollout a class instead of a function

--
ba788214b047577f58c41ce0ab6c62c277cd8b0d by Levi Burner <leviburner@gmail.com>:

rollout update docs and changelog

--
e4cb7732319e04cba2ab2c2ad848c659f6309808 by Levi Burner <leviburner@gmail.com>:

rollout don't register atexit handler for Rollout objects

--
5a08d2efdbbbb01d4b1231ff9a36a1dc44f4d9ee by Levi Burner <leviburner@gmail.com>:

rollout nthread kwarg, rename shutdown_pool to close, fixups

--
f622378543596a208339af0208fa3a70bf2a8007 by Levi Burner <leviburner@gmail.com>:

rollout add missing .close() calls

--
50f3ebca43c53eac03f03943c34bb1e46967bd4f by Levi Burner <leviburner@gmail.com>:

rollout return immediately

COPYBARA_INTEGRATE_REVIEW=https://github.com/google-deepmind/mujoco/pull/2282 from aftersomemath:rollout-threaded 50f3ebca43c53eac03f03943c34bb1e46967bd4f
PiperOrigin-RevId: 706744277
Change-Id: I1ab2263b7d6ce30cf1908aec8fd5f2eb976a19e6

doc/changelog.rst

@@ -5,6 +5,14 @@ Changelog
 Upcoming version (not yet released)
 -----------------------------------
 
+Python bindings
+^^^^^^^^^^^^^^^
+- :ref:`rollout<PyRollout>` now features native multi-threading. If a sequence of ``MjData`` instances
+  of length ``nthread`` is passed in, ``rollout`` will automatically create a thread pool and parallelize
+  the computation. The thread pool can be resused across calls, but then the function cannot be called simultaneously
+  from multiple threads. To run multiple threaded rollouts simultaneously, use the new class ``Rollout`` which
+  encapsulates the thread pool. Contribution by :github:user:`aftersomemath`.
+
 Bug fixes
 ^^^^^^^^^
 - Fixed a bug in the box-sphere collider, depth was incorrect for deep penetrations (:github:issue:`2206`).

doc/python.rst

@@ -711,18 +711,20 @@ The ``mujoco`` package contains two sub-modules: ``mujoco.rollout`` and ``mujoco
 rollout
 -------
 
-``mujoco.rollout`` shows how to add additional C/C++ functionality, exposed as a Python module via pybind11. It is
+``mujoco.rollout`` and ``mujoco.rollout.Rollout`` shows how to add additional C/C++ functionality, exposed as a Python module
-implemented in `rollout.cc <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.cc>`__
+via pybind11. It is implemented in `rollout.cc <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.cc>`__
 and wrapped in `rollout.py <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.py>`__. The module
 performs a common functionality where tight loops implemented outside of Python are beneficial: rolling out a trajectory
 (i.e., calling :ref:`mj_step` in a loop), given an intial state and sequence of controls, and returning subsequent
-states and sensor values. The basic usage form is
+states and sensor values. The rollouts are run in parallel with an internally managed thread pool if multiple MjData instances
+(one per thread) are passed as an argument. The basic usage form is
 
 .. code-block:: python
 
    state, sensordata = rollout.rollout(model, data, initial_state, control)
 
 ``model`` is either a single instance of MjModel or a sequence of compatible MjModel of length ``nroll``.
+``data`` is either a single instance of MjData or a sequence of compatible MjData of length ``nthread``.
 ``initial_state`` is an ``nroll x nstate`` array, with ``nroll`` initial states of size ``nstate``, where
 ``nstate = mj_stateSize(model, mjtState.mjSTATE_FULLPHYSICS)`` is the size of the
 :ref:`full physics state<geFullPhysics>`. ``control`` is a ``nroll x nstep x ncontrol`` array of controls. Controls are
@@ -732,13 +734,41 @@ specified by passing an optional ``control_spec`` bitflag.
 If a rollout diverges, the current state and sensor values are used to fill the remainder of the trajectory.
 Therefore, non-increasing time values can be used to detect diverged rollouts.
 
-The ``rollout`` function is designed to be completely stateless, so all inputs of the stepping pipeline are set and any
+The ``rollout`` function is designed to be computationally stateless, so all inputs of the stepping pipeline are set and any
 values already present in the given ``MjData`` instance will have no effect on the output.
 
-Since the Global Interpreter Lock can be released, this function can be efficiently threaded using Python threads. See
+By default ``rollout.rollout`` creates a new thread pool every call if ``len(data) > 1``. To reuse the thread pool
-the ``test_threading`` function in
+over multiple calls use the ``persistent_pool`` argument. ``rollout.rollout`` is not thread safe when using
+a persistent pool. The basic usage form is
+
+.. code-block:: python
+
+   state, sensordata = rollout.rollout(model, data, initial_state, persistent_pool=True)
+
+The pool is shutdown on interpreter shutdown or by a call to ``rollout.shutdown_persistent_pool``.
+
+To use multiple thread pools from multiple threads, use ``Rollout`` objects. The basic usage form is
+
+.. code-block:: python
+
+   # Pool shutdown upon exiting block.
+   with rollout.Rollout(nthread=nthread) as rollout_:
+    rollout_.rollout(model, data, initial_state)
+
+or
+
+.. code-block:: python
+
+   # Pool shutdown on object deletion or call to rollout_.close().
+   # To ensure clean shutdown of threads, call close() before interpreter exit.
+   rollout_ = rollout.Rollout(nthread=nthread)
+   rollout_.rollout(model, data, initial_state)
+   rollout_.close()
+
+Since the Global Interpreter Lock is released, this function can also be threaded using Python threads. However, this
+is less efficient than using native threads. See the ``test_threading`` function in
 `rollout_test.py <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout_test.py>`__ for an example
-of threaded operation (and more generally for usage examples).
+of threaded operation (and for more general usage examples).
 
 .. _PyMinimize:
 

python/mujoco/CMakeLists.txt

@@ -383,7 +383,7 @@ target_link_libraries(
           structs_header
 )
 
-mujoco_pybind11_module(_rollout rollout.cc)
+mujoco_pybind11_module(_rollout rollout.cc threadpool.cc)
 target_link_libraries(_rollout PRIVATE functions_header mujoco raw)
 
 mujoco_pybind11_module(

python/mujoco/rollout.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include <iostream>
+#include <memory>
 #include <optional>
 #include <sstream>
 
@@ -20,6 +21,7 @@
 #include "errors.h"
 #include "raw.h"
 #include "structs.h"
+#include "threadpool.h"
 #include <pybind11/buffer_info.h>
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
@@ -31,14 +33,24 @@ namespace {
 
 namespace py = ::pybind11;
 
+using PyCArray = py::array_t<mjtNum, py::array::c_style>;
+
 // NOLINTBEGIN(whitespace/line_length)
 
+const auto rollout_init_doc = R"(
+Construct a rollout object containing a thread pool for parallel rollouts.
+
+  input arguments (optional):
+    nthread            integer, number of threads in pool
+                       if zero, this pool is not started and rollouts run on the calling thread
+)";
+
 const auto rollout_doc = R"(
 Roll out open-loop trajectories from initial states, get resulting states and sensor values.
 
   input arguments (required):
     model              list of MjModel instances of length nroll
-    data               associated instance of MjData
+    data               list of associated MjData instances of length nthread
     nstep              integer, number of steps to be taken for each trajectory
     control_spec       specification of controls, ncontrol = mj_stateSize(m, control_spec)
     state0             (nroll x nstate) nroll initial state vectors,
@@ -49,12 +61,14 @@ Roll out open-loop trajectories from initial states, get resulting states and se
   output arguments (optional):
     state              (nroll x nstep x nstate)       nroll nstep states
     sensordata         (nroll x nstep x nsendordata)  nroll trajectories of nstep sensordata vectors
+    chunk_size         integer, determines threadpool chunk size. If unspecified
+                                chunk_size = max(1, nroll / (nthread * 10))
 )";
 
 // C-style rollout function, assumes all arguments are valid
 // all input fields of d are initialised, contents at call time do not matter
 // after returning, d will contain the last step of the last rollout
-void _unsafe_rollout(std::vector<const mjModel*>& m, mjData* d, int nroll, int nstep, unsigned int control_spec,
+void _unsafe_rollout(std::vector<const mjModel*>& m, mjData* d, int start_roll, int end_roll, int nstep, unsigned int control_spec,
                      const mjtNum* state0, const mjtNum* warmstart0, const mjtNum* control,
                      mjtNum* state, mjtNum* sensordata) {
   // sizes
@@ -75,7 +89,7 @@ void _unsafe_rollout(std::vector<const mjModel*>& m, mjData* d, int nroll, int n
   }
 
   // loop over rollouts
-  for (int r = 0; r < nroll; r++) {
+  for (int r = start_roll; r < end_roll; r++) {
     // clear user inputs if unspecified
     if (!(control_spec & mjSTATE_MOCAP_POS)) {
       for (int i = 0; i < nbody; i++) {
@@ -158,6 +172,43 @@ void _unsafe_rollout(std::vector<const mjModel*>& m, mjData* d, int nroll, int n
   }
 }
 
+// C-style threaded version of _unsafe_rollout
+void _unsafe_rollout_threaded(std::vector<const mjModel*>& m, std::vector<mjData*>& d,
+                              int nroll, int nstep, unsigned int control_spec,
+                              const mjtNum* state0, const mjtNum* warmstart0,
+                              const mjtNum* control, mjtNum* state, mjtNum* sensordata,
+                              ThreadPool* pool, int chunk_size) {
+  int nfulljobs = nroll / chunk_size;
+  int chunk_remainder = nroll % chunk_size;
+  int njobs = (chunk_remainder > 0) ? nfulljobs + 1 : nfulljobs;
+
+  // Reset the pool counter
+  pool->ResetCount();
+
+  // schedule all jobs of full (chunk) size
+  for (int j = 0; j < nfulljobs; j++) {
+    auto task = [=, &m, &d](void) {
+      int id = pool->WorkerId();
+      _unsafe_rollout(m, d[id], j*chunk_size, (j+1)*chunk_size,
+        nstep, control_spec, state0, warmstart0, control, state, sensordata);
+    };
+    pool->Schedule(task);
+  }
+
+  // schedule any remaining jobs of size < chunk_size
+  if (chunk_remainder > 0) {
+    auto task = [=, &m, &d](void) {
+      _unsafe_rollout(m, d[pool->WorkerId()], nfulljobs*chunk_size,
+        nfulljobs*chunk_size+chunk_remainder,
+        nstep, control_spec, state0, warmstart0, control, state, sensordata);
+    };
+    pool->Schedule(task);
+  }
+
+  // wait for job counter to incremented up to the number of jobs submitted by this thread
+  pool->WaitCount(njobs);
+}
+
 // NOLINTEND(whitespace/line_length)
 
 // check size of optional argument to rollout(), return raw pointer
@@ -181,71 +232,118 @@ mjtNum* get_array_ptr(std::optional<const py::array_t<mjtNum>> arg,
   return static_cast<mjtNum*>(info.ptr);
 }
 
+class Rollout {
+ public:
+  Rollout(int nthread) : nthread_(nthread) {
+    if (this->nthread_ > 0) {
+      this->pool_ = std::make_shared<ThreadPool>(this->nthread_);
+    }
+  }
+
+  void rollout(py::list m, py::list d, int nstep, unsigned int control_spec,
+               const PyCArray state0, std::optional<const PyCArray> warmstart0,
+               std::optional<const PyCArray> control,
+               std::optional<const PyCArray> state,
+               std::optional<const PyCArray> sensordata,
+               std::optional<int> chunk_size) {
+    // get raw pointers
+    int nroll = state0.shape(0);
+    std::vector<const raw::MjModel*> model_ptrs(nroll);
+    for (int r = 0; r < nroll; r++) {
+      model_ptrs[r] = m[r].cast<const MjModelWrapper*>()->get();
+    }
+
+    // check length d and nthread are consistent
+    if (this->nthread_ == 0 && py::len(d) > 1) {
+      std::ostringstream msg;
+      msg << "More than one data instance passed but "
+          << "rollout is configured to run on main thread";
+      py::value_error(msg.str());
+    } else if (this->nthread_ != py::len(d)) {
+      std::ostringstream msg;
+      msg << "Length of data: " << py::len(d)
+          << " not equal to nthread: " << this->nthread_;
+      py::value_error(msg.str());
+    }
+
+    std::vector<raw::MjData*> data_ptrs(py::len(d));
+    for (int t = 0; t < py::len(d); t++) {
+      data_ptrs[t] = d[t].cast<MjDataWrapper*>()->get();
+    }
+
+    // check that some steps need to be taken, return if not
+    if (nstep < 1) {
+      return;
+    }
+
+    // get sizes
+    int nstate = mj_stateSize(model_ptrs[0], mjSTATE_FULLPHYSICS);
+    int ncontrol = mj_stateSize(model_ptrs[0], control_spec);
+
+    mjtNum* state0_ptr = get_array_ptr(state0, "state0", nroll, 1, nstate);
+    mjtNum* warmstart0_ptr =
+        get_array_ptr(warmstart0, "warmstart0", nroll, 1, model_ptrs[0]->nv);
+    mjtNum* control_ptr =
+        get_array_ptr(control, "control", nroll, nstep, ncontrol);
+    mjtNum* state_ptr = get_array_ptr(state, "state", nroll, nstep, nstate);
+    mjtNum* sensordata_ptr = get_array_ptr(sensordata, "sensordata", nroll,
+                                           nstep, model_ptrs[0]->nsensordata);
+
+    // perform rollouts
+    {
+      // release the GIL
+      py::gil_scoped_release no_gil;
+
+      // call unsafe rollout function, multi or single threaded
+      if (this->nthread_ > 0 && nroll > 1) {
+        int chunk_size_final = 1;
+        if (!chunk_size.has_value()) {
+          chunk_size_final = std::max(1, nroll / (10 * this->nthread_));
+        } else {
+          chunk_size_final = *chunk_size;
+        }
+        InterceptMjErrors(_unsafe_rollout_threaded)(
+            model_ptrs, data_ptrs, nroll, nstep, control_spec, state0_ptr,
+            warmstart0_ptr, control_ptr, state_ptr, sensordata_ptr,
+            this->pool_.get(), chunk_size_final);
+      } else {
+        InterceptMjErrors(_unsafe_rollout)(
+            model_ptrs, data_ptrs[0], 0, nroll, nstep, control_spec, state0_ptr,
+            warmstart0_ptr, control_ptr, state_ptr, sensordata_ptr);
+      }
+    }
+  }
+
+ private:
+  int nthread_;
+  std::shared_ptr<ThreadPool> pool_;
+};
 
 PYBIND11_MODULE(_rollout, pymodule) {
   namespace py = ::pybind11;
-  using PyCArray = py::array_t<mjtNum, py::array::c_style>;
 
-  // roll out open loop trajectories from multiple initial states
+  py::class_<Rollout>(pymodule, "Rollout")
-  // get subsequent states and corresponding sensor values
+      .def(
-  pymodule.def(
+        py::init([](int nthread) {
-      "rollout",
+          return std::make_unique<Rollout>(nthread);
-      [](py::list m, MjDataWrapper& d,
+        }),
-         int nstep, unsigned int control_spec,
+        py::kw_only(),
-         const PyCArray state0,
+        py::arg("nthread"),
-         std::optional<const PyCArray> warmstart0,
+        py::doc(rollout_init_doc))
-         std::optional<const PyCArray> control,
+      .def(
-         std::optional<const PyCArray> state,
+        "rollout",
-         std::optional<const PyCArray> sensordata
+        &Rollout::rollout,
-         ) {
+        py::arg("model"),
-        // get raw pointers
+        py::arg("data"),
-        int nroll = state0.shape(0);
+        py::arg("nstep"),
-        std::vector<const raw::MjModel*> model_ptrs(nroll);
+        py::arg("control_spec"),
-        for (int r = 0; r < nroll; r++) {
+        py::arg("state0"),
-          model_ptrs[r] = m[r].cast<const MjModelWrapper*>()->get();
+        py::arg("warmstart0") = py::none(),
-        }
+        py::arg("control")    = py::none(),
-        raw::MjData* data = d.get();
+        py::arg("state")      = py::none(),
-
+        py::arg("sensordata") = py::none(),
-        // check that some steps need to be taken, return if not
+        py::arg("chunk_size") = py::none(),
-        if (nstep < 1) {
+        py::doc(rollout_doc));
-          return;
-        }
-
-        // get sizes
-        int nstate = mj_stateSize(model_ptrs[0], mjSTATE_FULLPHYSICS);
-        int ncontrol = mj_stateSize(model_ptrs[0], control_spec);
-
-        mjtNum* state0_ptr = get_array_ptr(state0, "state0", nroll, 1, nstate);
-        mjtNum* warmstart0_ptr = get_array_ptr(warmstart0, "warmstart0", nroll,
-                                               1, model_ptrs[0]->nv);
-        mjtNum* control_ptr = get_array_ptr(control, "control", nroll,
-                                            nstep, ncontrol);
-        mjtNum* state_ptr = get_array_ptr(state, "state", nroll, nstep, nstate);
-        mjtNum* sensordata_ptr = get_array_ptr(sensordata, "sensordata", nroll,
-                                               nstep, model_ptrs[0]->nsensordata);
-
-        // perform rollouts
-        {
-          // release the GIL
-          py::gil_scoped_release no_gil;
-
-          // call unsafe rollout function
-          InterceptMjErrors(_unsafe_rollout)(
-              model_ptrs, data, nroll, nstep, control_spec, state0_ptr,
-              warmstart0_ptr, control_ptr, state_ptr, sensordata_ptr);
-        }
-      },
-      py::arg("model"),
-      py::arg("data"),
-      py::arg("nstep"),
-      py::arg("control_spec"),
-      py::arg("state0"),
-      py::arg("warmstart0") = py::none(),
-      py::arg("control")    = py::none(),
-      py::arg("state")      = py::none(),
-      py::arg("sensordata") = py::none(),
-      py::doc(rollout_doc)
-  );
 }
 
 }  // namespace

python/mujoco/rollout.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Roll out open-loop trajectories from initial states, get subsequent states and sensor values."""
 
+import atexit
 from collections.abc import Sequence
 from typing import Optional, Union
 
@@ -23,9 +24,243 @@ import numpy as np
 from numpy import typing as npt
 
 
+class Rollout:
+  """Rollout object containing a thread pool for parallel rollouts."""
+
+  def __init__(self, *, nthread: Optional[int] = None):
+    """Construct a rollout object containing a thread pool for parallel rollouts.
+
+    Args:
+      nthread: Number of threads in pool.
+        If zero, this pool is not started and rollouts run on the calling thread.
+    """  # fmt: skip
+    self.nthread = 0 if nthread is None else nthread
+    self.rollout_ = _rollout.Rollout(nthread=self.nthread)
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    self.close()
+
+  def close(self):
+    del self.rollout_
+    self.rollout_ = None
+
+  def rollout(
+      self,
+      model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
+      data: Union[mujoco.MjData, Sequence[mujoco.MjData]],
+      initial_state: npt.ArrayLike,
+      control: Optional[npt.ArrayLike] = None,
+      *,  # require subsequent arguments to be named
+      control_spec: int = mujoco.mjtState.mjSTATE_CTRL.value,
+      skip_checks: bool = False,
+      nstep: Optional[int] = None,
+      initial_warmstart: Optional[npt.ArrayLike] = None,
+      state: Optional[npt.ArrayLike] = None,
+      sensordata: Optional[npt.ArrayLike] = None,
+      chunk_size: Optional[int] = None,
+  ):
+    """Rolls out open-loop trajectories from initial states, get subsequent state and sensor values.
+
+    Python wrapper for rollout.cc, see documentation therein.
+    Infers nroll and nstep.
+    Tiles inputs with singleton dimensions.
+    Allocates outputs if none are given.
+
+    Args:
+      model: An instance or length nroll sequence of MjModel with the same size signature.
+      data: Associated mjData instance or sequence of instances with length nthread.
+      initial_state: Array of initial states from which to roll out trajectories.
+        ([nroll or 1] x nstate)
+      control: Open-loop controls array to apply during the rollouts.
+        ([nroll or 1] x [nstep or 1] x ncontrol)
+      control_spec: mjtState specification of control vectors.
+      skip_checks: Whether to skip internal shape and type checks.
+      nstep: Number of steps in rollouts (inferred if unspecified).
+      initial_warmstart: Initial qfrc_warmstart array (optional).
+        ([nroll or 1] x nv)
+      state: State output array (optional).
+        (nroll x nstep x nstate)
+      sensordata: Sensor data output array (optional).
+        (nroll x nstep x nsensordata)
+      chunk_size: Determines threadpool chunk size. If unspecified,
+                  chunk_size = max(1, nroll / (nthread * 10))
+
+    Returns:
+      state:
+        State output array, (nroll x nstep x nstate).
+      sensordata:
+        Sensor data output array, (nroll x nstep x nsensordata).
+
+    Raises:
+      RuntimeError: rollout requested after thread pool shutdown.
+      ValueError: bad shapes or sizes.
+    """  # fmt: skip
+
+    if self.rollout_ is None:
+      raise RuntimeError('rollout requested after thread pool shutdown')
+
+    # skip_checks shortcut:
+    #   don't infer nroll/nstep
+    #   don't support singleton expansion
+    #   don't allocate output arrays
+    #   just call rollout and return
+    if skip_checks:
+      self.rollout_.rollout(
+          model,
+          data,
+          nstep,
+          control_spec,
+          initial_state,
+          initial_warmstart,
+          control,
+          state,
+          sensordata,
+          chunk_size,
+      )
+      return state, sensordata
+
+    if not isinstance(model, mujoco.MjModel):
+      model = list(model)
+
+    # check control_spec
+    if control_spec & ~mujoco.mjtState.mjSTATE_USER.value:
+      raise ValueError('control_spec can only contain bits in mjSTATE_USER')
+
+    # check types
+    if nstep and not isinstance(nstep, int):
+      raise ValueError('nstep must be an integer')
+    if chunk_size and not isinstance(chunk_size, int):
+      raise ValueError('chunk_size must be an integer')
+    _check_must_be_numeric(
+        initial_state=initial_state,
+        initial_warmstart=initial_warmstart,
+        control=control,
+        state=state,
+        sensordata=sensordata,
+    )
+
+    # check number of dimensions
+    _check_number_of_dimensions(
+        2, initial_state=initial_state, initial_warmstart=initial_warmstart
+    )
+    _check_number_of_dimensions(
+        3, control=control, state=state, sensordata=sensordata
+    )
+
+    # ensure 2D, make contiguous, row-major (C ordering)
+    initial_state = _ensure_2d(initial_state)
+    initial_warmstart = _ensure_2d(initial_warmstart)
+
+    # ensure 3D, make contiguous, row-major (C ordering)
+    control = _ensure_3d(control)
+    state = _ensure_3d(state)
+    sensordata = _ensure_3d(sensordata)
+
+    # infer nroll, check for incompatibilities
+    nroll = _infer_dimension(
+        0,
+        1,
+        initial_state=initial_state,
+        initial_warmstart=initial_warmstart,
+        control=control,
+        state=state,
+        sensordata=sensordata,
+    )
+    if isinstance(model, list) and nroll == 1:
+      nroll = len(model)
+
+    if isinstance(model, list) and len(model) != nroll:
+      raise ValueError(
+          f'nroll inferred as {nroll} but model is length {len(model)}'
+      )
+    elif not isinstance(model, list):
+      model = [model]  # Use a length 1 list to simplify code below
+
+    if not isinstance(data, list):
+      data = [data]  # Use a length 1 list to simplify code below
+
+    # infer nstep, check for incompatibilities
+    nstep = _infer_dimension(
+        1, nstep or 1, control=control, state=state, sensordata=sensordata
+    )
+
+    # get nstate/ncontrol/nv/nsensordata
+    # check that they are equal across models
+    nstate = mujoco.mj_stateSize(
+        model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS.value
+    )
+    ncontrol = mujoco.mj_stateSize(model[0], control_spec)
+    nv = model[0].nv
+    nsensordata = model[0].nsensordata
+    for m in model[1:]:
+      if (
+          nstate
+          != mujoco.mj_stateSize(m, mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
+          or ncontrol != mujoco.mj_stateSize(m, control_spec)
+          or nv != m.nv
+          or nsensordata != m.nsensordata
+      ):
+        raise ValueError('models are not compatible')
+
+    # check trailing dimensions
+    _check_trailing_dimension(nstate, initial_state=initial_state, state=state)
+    _check_trailing_dimension(ncontrol, control=control)
+    _check_trailing_dimension(nv, initial_warmstart=initial_warmstart)
+    _check_trailing_dimension(nsensordata, sensordata=sensordata)
+
+    # tile input arrays/lists if required (singleton expansion)
+    model = model * nroll if len(model) == 1 else model
+    initial_state = _tile_if_required(initial_state, nroll)
+    initial_warmstart = _tile_if_required(initial_warmstart, nroll)
+    control = _tile_if_required(control, nroll, nstep)
+
+    # allocate output if not provided
+    if state is None:
+      state = np.empty((nroll, nstep, nstate))
+    if sensordata is None:
+      sensordata = np.empty((nroll, nstep, nsensordata))
+
+    # call rollout
+    self.rollout_.rollout(
+        model,
+        data,
+        nstep,
+        control_spec,
+        initial_state,
+        initial_warmstart,
+        control,
+        state,
+        sensordata,
+        chunk_size,
+    )
+
+    # return outputs
+    return state, sensordata
+
+
+persistent_rollout = None
+
+
+def shutdown_persistent_pool():
+  """Shutdown the persistent thread pool that is optionally created by rollout.
+
+  This is called automatically interpreter shutdown, but can also be called manually.
+  """  # fmt: skip
+  global persistent_rollout
+  if persistent_rollout is not None:
+    persistent_rollout.close()
+  persistent_rollout = None
+
+
+atexit.register(shutdown_persistent_pool)
+
+
 def rollout(
     model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
-    data: mujoco.MjData,
+    data: Union[mujoco.MjData, Sequence[mujoco.MjData]],
     initial_state: npt.ArrayLike,
     control: Optional[npt.ArrayLike] = None,
     *,  # require subsequent arguments to be named
@@ -35,6 +270,8 @@ def rollout(
     initial_warmstart: Optional[npt.ArrayLike] = None,
     state: Optional[npt.ArrayLike] = None,
     sensordata: Optional[npt.ArrayLike] = None,
+    chunk_size: Optional[int] = None,
+    persistent_pool: bool = False,
 ):
   """Rolls out open-loop trajectories from initial states, get subsequent states and sensor values.
 
@@ -44,8 +281,8 @@ def rollout(
   Allocates outputs if none are given.
 
   Args:
-    model: An mjModel or a sequence of MjModel with the same size signature.
+    model: An instance or length nroll sequence of MjModel with the same size signature.
-    data: An associated mjData instance.
+    data: Associated mjData instance or sequence of instances with length nthread.
     initial_state: Array of initial states from which to roll out trajectories.
       ([nroll or 1] x nstate)
     control: Open-loop controls array to apply during the rollouts.
@@ -59,6 +296,9 @@ def rollout(
       (nroll x nstep x nstate)
     sensordata: Sensor data output array (optional).
       (nroll x nstep x nsensordata)
+    chunk_size: Determines threadpool chunk size. If unspecified,
+                chunk_size = max(1, nroll / (nthread * 10))
+    persistent_pool: Determines if a persistent thread pool is created or reused.
 
   Returns:
     state:
@@ -69,136 +309,41 @@ def rollout(
   Raises:
     ValueError: bad shapes or sizes.
   """  # fmt: skip
-  # skip_checks shortcut:
+  if not isinstance(data, list):
-  #   don't infer nroll/nstep
+    data = [data]  # Use a length 1 list to simplify code below
-  #   don't support singleton expansion
+
-  #   don't allocate output arrays
+  nthread = len(data) if len(data) > 1 else 0
-  #   just call rollout and return
+
-  if skip_checks:
+  # Use a persistent thread pool if requested
-    _rollout.rollout(
+  if persistent_pool:
+    # Create or restart persistent threadpool
+    global persistent_rollout
+    if persistent_rollout is None:
+      persistent_rollout = Rollout(nthread=nthread)
+    if persistent_rollout.nthread != nthread:
+      persistent_rollout.close()
+      persistent_rollout = Rollout(nthread=nthread)
+    rollout_ = persistent_rollout
+  else:
+    rollout_ = Rollout(nthread=nthread)
+
+  try:
+    return rollout_.rollout(
         model,
         data,
-        nstep,
-        control_spec,
         initial_state,
-        initial_warmstart,
         control,
-        state,
+        control_spec=control_spec,
-        sensordata,
+        skip_checks=skip_checks,
+        nstep=nstep,
+        initial_warmstart=initial_warmstart,
+        state=state,
+        sensordata=sensordata,
+        chunk_size=chunk_size,
     )
-    return state, sensordata
+  finally:
-
+    if not persistent_pool:
-  if not isinstance(model, mujoco.MjModel):
+      rollout_.close()
-    model = list(model)
-
-  # check control_spec
-  if control_spec & ~mujoco.mjtState.mjSTATE_USER.value:
-    raise ValueError('control_spec can only contain bits in mjSTATE_USER')
-
-  # check types
-  if nstep and not isinstance(nstep, int):
-    raise ValueError('nstep must be an integer')
-  _check_must_be_numeric(
-      initial_state=initial_state,
-      initial_warmstart=initial_warmstart,
-      control=control,
-      state=state,
-      sensordata=sensordata,
-  )
-
-  # check number of dimensions
-  _check_number_of_dimensions(
-      2, initial_state=initial_state, initial_warmstart=initial_warmstart
-  )
-  _check_number_of_dimensions(
-      3, control=control, state=state, sensordata=sensordata
-  )
-
-  # ensure 2D, make contiguous, row-major (C ordering)
-  initial_state = _ensure_2d(initial_state)
-  initial_warmstart = _ensure_2d(initial_warmstart)
-
-  # ensure 3D, make contiguous, row-major (C ordering)
-  control = _ensure_3d(control)
-  state = _ensure_3d(state)
-  sensordata = _ensure_3d(sensordata)
-
-  # infer nroll, check for incompatibilities
-  nroll = _infer_dimension(
-      0,
-      1,
-      initial_state=initial_state,
-      initial_warmstart=initial_warmstart,
-      control=control,
-      state=state,
-      sensordata=sensordata,
-  )
-  if isinstance(model, list) and nroll == 1:
-    nroll = len(model)
-
-  if isinstance(model, list) and len(model) != nroll:
-    raise ValueError(
-        f'nroll inferred as {nroll} but model is length {len(model)}'
-    )
-  elif not isinstance(model, list):
-    model = [model]  # Use a length 1 list to simplify code below
-
-  # infer nstep, check for incompatibilities
-  nstep = _infer_dimension(
-      1, nstep or 1, control=control, state=state, sensordata=sensordata
-  )
-
-  # get nstate/ncontrol/nv/nsensordata
-  # check that they are equal across models
-  nstate = mujoco.mj_stateSize(
-      model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS.value
-  )
-  ncontrol = mujoco.mj_stateSize(model[0], control_spec)
-  nv = model[0].nv
-  nsensordata = model[0].nsensordata
-  for m in model[1:]:
-    if (
-        nstate
-        != mujoco.mj_stateSize(m, mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
-        or ncontrol != mujoco.mj_stateSize(m, control_spec)
-        or nv != m.nv
-        or nsensordata != m.nsensordata
-    ):
-      raise ValueError('models are not compatible')
-
-  # check trailing dimensions
-  _check_trailing_dimension(nstate, initial_state=initial_state, state=state)
-  _check_trailing_dimension(ncontrol, control=control)
-  _check_trailing_dimension(nv, initial_warmstart=initial_warmstart)
-  _check_trailing_dimension(nsensordata, sensordata=sensordata)
-
-  # tile input arrays/lists if required (singleton expansion)
-  model = model * nroll if len(model) == 1 else model
-  initial_state = _tile_if_required(initial_state, nroll)
-  initial_warmstart = _tile_if_required(initial_warmstart, nroll)
-  control = _tile_if_required(control, nroll, nstep)
-
-  # allocate output if not provided
-  if state is None:
-    state = np.empty((nroll, nstep, nstate))
-  if sensordata is None:
-    sensordata = np.empty((nroll, nstep, nsensordata))
-
-  # call rollout
-  _rollout.rollout(
-      model,
-      data,
-      nstep,
-      control_spec,
-      initial_state,
-      initial_warmstart,
-      control,
-      state,
-      sensordata,
-  )
-
-  # return outputs
-  return state, sensordata
 
 
 def _check_must_be_numeric(**kwargs):

python/mujoco/rollout_test.py

@@ -355,7 +355,7 @@ class MuJoCoRolloutTest(parameterized.TestCase):
         body.pos = body.pos + i
         model.append(spec.compile())
     else:
-      model = [spec.compile() for i in range(nroll)]
+      model = [spec.compile() for _ in range(nroll)]
 
     nstate = mujoco.mj_stateSize(model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS)
     data = mujoco.MjData(model[0])
@@ -461,7 +461,7 @@ class MuJoCoRolloutTest(parameterized.TestCase):
     model = mujoco.MjModel.from_xml_string(TEST_XML)
     nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     num_workers = 32
-    nroll = 10000
+    nroll = 100
     nstep = 5
     initial_state = np.random.randn(nroll, nstate)
     state = np.empty((nroll, nstep, nstate))
@@ -478,7 +478,7 @@ class MuJoCoRolloutTest(parameterized.TestCase):
     def call_rollout(initial_state, control, state, sensordata):
       rollout.rollout(
           model_list,
-          thread_local.data,
+          [thread_local.data],
           initial_state,
           control,
           skip_checks=True,
@@ -519,6 +519,116 @@ class MuJoCoRolloutTest(parameterized.TestCase):
     np.testing.assert_array_equal(state, py_state)
     np.testing.assert_array_equal(sensordata, py_sensordata)
 
+  def test_threading_native(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML)
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
+    num_workers = 32
+    nroll = 100
+    nstep = 5
+    initial_state = np.random.randn(nroll, nstate)
+    state = np.empty((nroll, nstep, nstate))
+    sensordata = np.empty((nroll, nstep, model.nsensordata))
+    control = np.random.randn(nroll, nstep, model.nu)
+
+    model_list = [model] * nroll
+    data_list = [mujoco.MjData(model) for _ in range(num_workers)]
+
+    rollout.rollout(
+        model_list,
+        data_list,
+        initial_state,
+        control,
+        nstep=nstep,
+        state=state,
+        sensordata=sensordata,
+    )
+
+    data = mujoco.MjData(model)
+    py_state, py_sensordata = py_rollout(model, data, initial_state, control)
+    np.testing.assert_array_equal(state, py_state)
+    np.testing.assert_array_equal(sensordata, py_sensordata)
+
+  def test_threading_native_persistent_object(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML)
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
+    num_workers = 32
+    nroll = 100
+    nstep = 5
+    initial_state = np.random.randn(nroll, nstate)
+    state = np.empty((nroll, nstep, nstate))
+    sensordata = np.empty((nroll, nstep, model.nsensordata))
+    control = np.random.randn(nroll, nstep, model.nu)
+
+    model_list = [model] * nroll
+    data_list = [mujoco.MjData(model) for _ in range(num_workers)]
+
+    with rollout.Rollout(nthread=num_workers) as rollout_:
+      for _ in range(2):
+        rollout_.rollout(
+            model_list,
+            data_list,
+            initial_state,
+            control,
+            nstep=nstep,
+            state=state,
+            sensordata=sensordata,
+        )
+
+      data = mujoco.MjData(model)
+      py_state, py_sensordata = py_rollout(model, data, initial_state, control)
+      np.testing.assert_array_equal(state, py_state)
+      np.testing.assert_array_equal(sensordata, py_sensordata)
+
+    rollout_ = rollout.Rollout(nthread=num_workers)
+    for _ in range(2):
+      rollout_.rollout(
+          model_list,
+          data_list,
+          initial_state,
+          control,
+          nstep=nstep,
+          state=state,
+          sensordata=sensordata,
+      )
+
+      data = mujoco.MjData(model)
+      py_state, py_sensordata = py_rollout(model, data, initial_state, control)
+      np.testing.assert_array_equal(state, py_state)
+      np.testing.assert_array_equal(sensordata, py_sensordata)
+    rollout_.close()
+
+  def test_threading_native_persistent_function(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML)
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
+    num_workers = 32
+    nroll = 100
+    nstep = 5
+    initial_state = np.random.randn(nroll, nstate)
+    state = np.empty((nroll, nstep, nstate))
+    sensordata = np.empty((nroll, nstep, model.nsensordata))
+    control = np.random.randn(nroll, nstep, model.nu)
+
+    model_list = [model] * nroll
+    data_list = [mujoco.MjData(model) for _ in range(num_workers)]
+
+    for _ in range(2):
+      rollout.rollout(
+          model_list,
+          data_list,
+          initial_state,
+          control,
+          nstep=nstep,
+          state=state,
+          sensordata=sensordata,
+          persistent_pool=True,
+      )
+
+      data = mujoco.MjData(model)
+      py_state, py_sensordata = py_rollout(model, data, initial_state, control)
+      np.testing.assert_array_equal(state, py_state)
+      np.testing.assert_array_equal(sensordata, py_sensordata)
+    rollout.shutdown_persistent_pool()
+
   # ---------------------------- test advanced operation
 
   def test_warmstart(self):

python/mujoco/threadpool.cc

@@ -0,0 +1,87 @@
+// Copyright 2024 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "threadpool.h"
+
+#include <condition_variable>
+#include <functional>
+#include <mutex>
+#include <thread>
+#include <utility>
+
+#include <absl/base/attributes.h>
+
+namespace mujoco::python {
+
+ABSL_CONST_INIT thread_local int ThreadPool::worker_id_ = -1;
+
+// ThreadPool constructor
+ThreadPool::ThreadPool(int num_threads) : ctr_(0) {
+  for (int i = 0; i < num_threads; i++) {
+    threads_.push_back(std::thread(&ThreadPool::WorkerThread, this, i));
+  }
+}
+
+// ThreadPool destructor
+ThreadPool::~ThreadPool() {
+  {
+    std::unique_lock<std::mutex> lock(m_);
+    for (int i = 0; i < threads_.size(); i++) {
+      queue_.push(nullptr);
+    }
+    cv_in_.notify_all();
+  }
+  for (auto& thread : threads_) {
+    thread.join();
+  }
+}
+
+// ThreadPool scheduler
+void ThreadPool::Schedule(std::function<void()> task) {
+  std::unique_lock<std::mutex> lock(m_);
+  queue_.push(std::move(task));
+  cv_in_.notify_one();
+}
+
+// ThreadPool worker
+void ThreadPool::WorkerThread(int i) {
+  worker_id_ = i;
+  while (true) {
+    auto task = [&]() {
+      std::unique_lock<std::mutex> lock(m_);
+      cv_in_.wait(lock, [&]() { return !queue_.empty(); });
+      std::function<void()> task = std::move(queue_.front());
+      queue_.pop();
+      cv_in_.notify_one();
+      return task;
+    }();
+    if (task == nullptr) {
+      {
+        std::unique_lock<std::mutex> lock(m_);
+        ++ctr_;
+        cv_ext_.notify_one();
+      }
+      break;
+    }
+    task();
+
+    {
+      std::unique_lock<std::mutex> lock(m_);
+      ++ctr_;
+      cv_ext_.notify_one();
+    }
+  }
+}
+
+}  // namespace mujoco::python

python/mujoco/threadpool.h

@@ -0,0 +1,80 @@
+// Copyright 2024 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_PYTHON_THREADPOOL_H_
+#define MUJOCO_PYTHON_THREADPOOL_H_
+
+#include <condition_variable>
+#include <cstdint>
+#include <functional>
+#include <mutex>
+#include <queue>
+#include <thread>
+#include <vector>
+
+#include <absl/base/attributes.h>
+
+namespace mujoco::python {
+
+// ThreadPool class
+class ThreadPool {
+ public:
+  // constructor
+  explicit ThreadPool(int num_threads);
+
+  // destructor
+  ~ThreadPool();
+
+  int NumThreads() const { return threads_.size(); }
+
+  // returns an ID between 0 and NumThreads() - 1. must be called within
+  // worker thread (returns -1 if not).
+  static int WorkerId() { return worker_id_; }
+
+  // ----- methods ----- //
+  // set task for threadpool
+  void Schedule(std::function<void()> task);
+
+  // return number of tasks completed
+  std::uint64_t GetCount() { return ctr_; }
+
+  // reset count to zero
+  void ResetCount() { ctr_ = 0; }
+
+  // wait for count, then return
+  void WaitCount(int value) {
+    std::unique_lock<std::mutex> lock(m_);
+    cv_ext_.wait(lock, [&]() { return this->GetCount() >= value; });
+  }
+
+ private:
+  // ----- methods ----- //
+
+  // execute task with available thread
+  void WorkerThread(int i);
+
+  ABSL_CONST_INIT static thread_local int worker_id_;
+
+  // ----- members ----- //
+  std::vector<std::thread> threads_;
+  std::mutex m_;
+  std::condition_variable cv_in_;
+  std::condition_variable cv_ext_;
+  std::queue<std::function<void()>> queue_;
+  std::uint64_t ctr_;
+};
+
+}  // namespace mujoco::python
+
+#endif  // MUJOCO_PYTHON_THREADPOOL_H_