rollout add chunk_divisor parameter

Author: aftersomemath

doc/changelog.rst

@@ -9,6 +9,11 @@ Bug fixes
 ^^^^^^^^^
 - Fixed a bug in the box-sphere collider, depth was incorrect for deep penetrations (:github:issue:`2206`).
 
+Python bindings
+^^^^^^^^^^^^^^^
+- :ref:`rollout<PyRollout>` can now accept sequences of MjData of length ``nthread``. If passed, :ref:`rollout<PyRollout>`
+  will automatically create a persistent threadpool and parallelize rollouts.
+
 Version 3.2.6 (Dec 2, 2024)
 ---------------------------
 

python/mujoco/rollout.cc

@@ -39,7 +39,7 @@ Roll out open-loop trajectories from initial states, get resulting states and se
 
   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,
@@ -50,6 +50,8 @@ 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_divisor      integer, determines threadpool chunk size according to
+                                chunk_size = max(1, nroll / (nthread * chunk_divisor)
 )";
 
 // C-style rollout function, assumes all arguments are valid
@@ -238,7 +240,8 @@ PYBIND11_MODULE(_rollout, pymodule) {
          std::optional<const PyCArray> warmstart0,
          std::optional<const PyCArray> control,
          std::optional<const PyCArray> state,
-         std::optional<const PyCArray> sensordata
+         std::optional<const PyCArray> sensordata,
+         int chunk_divisor
          ) {
         // get raw pointers
         int nroll = state0.shape(0);
@@ -278,7 +281,7 @@ PYBIND11_MODULE(_rollout, pymodule) {
 
           // call unsafe rollout function
           if (nthread > 1 && nroll > 1) {
-            int chunk_size = std::max(1, nroll / (10 * nthread));
+            int chunk_size = std::max(1, nroll / (chunk_divisor * nthread));
             InterceptMjErrors(_unsafe_rollout_threaded)(
                 model_ptrs, data_ptrs, nroll, nstep, control_spec, state0_ptr,
                 warmstart0_ptr, control_ptr, state_ptr, sensordata_ptr,
@@ -300,6 +303,7 @@ PYBIND11_MODULE(_rollout, pymodule) {
       py::arg("control")    = py::none(),
       py::arg("state")      = py::none(),
       py::arg("sensordata") = py::none(),
+      py::arg("chunk_divisor") = 10,
       py::doc(rollout_doc)
   );
 }

python/mujoco/rollout.py

@@ -35,6 +35,7 @@ def rollout(
     initial_warmstart: Optional[npt.ArrayLike] = None,
     state: Optional[npt.ArrayLike] = None,
     sensordata: Optional[npt.ArrayLike] = None,
+    chunk_divisor: int = 10,
 ):
   """Rolls out open-loop trajectories from initial states, get subsequent states and sensor values.
 
@@ -59,6 +60,8 @@ def rollout(
       (nroll x nstep x nstate)
     sensordata: Sensor data output array (optional).
       (nroll x nstep x nsensordata)
+    chunk_divisor: Determines threadpool chunk size according to
+                   chunk_size = max(1, nroll / (nthread * chunk_divisor)
 
   Returns:
     state:
@@ -85,6 +88,7 @@ def rollout(
         control,
         state,
         sensordata,
+        chunk_divisor,
     )
     return state, sensordata
 
@@ -198,6 +202,7 @@ def rollout(
       control,
       state,
       sensordata,
+      chunk_divisor,
   )
 
   # return outputs