From a7eb6efd4e3181d1f15428cdcb3f42458613e115 Mon Sep 17 00:00:00 2001 From: Levi Burner Date: Mon, 16 Dec 2024 09:56:12 -0800 Subject: [PATCH] Copybara import of the project: -- 3a95b62f59e81bfef0f076afb173ecc14b27943d by Levi Burner : rollout prototype native threadpool for comparing to python threads -- efd8be1124ac839b902de45973a3ca8b9f2215e6 by Levi Burner : copy mjpcs threadpool into python bindings -- 75603eea3e8362e354a9675e8a6cd14e56ec3d28 by Levi Burner : rollout use threadpool as translation unit -- 06b90febd021663f6cc81fd7895e4d6e2008ed97 by Levi Burner : rollout add chunk_divisor parameter -- 298ab2f3c0d6e12530832c3cdbf784dd92d54806 by Levi Burner : rollout add native threading test -- 169cf9978e7abad6edd1392b8e6aab995e4f8f10 by Levi Burner : rollout exchange chunk_divisor arg for chunk_size -- 265af851d74432d261277d3dbda11cdef1841bc8 by Levi Burner : rollout fix cosmetics -- 1e8bffa88bf36190501b334bef31147e23db39f7 by Levi Burner : make native rollout a class instead of a function -- ba788214b047577f58c41ce0ab6c62c277cd8b0d by Levi Burner : rollout update docs and changelog -- e4cb7732319e04cba2ab2c2ad848c659f6309808 by Levi Burner : rollout don't register atexit handler for Rollout objects -- 5a08d2efdbbbb01d4b1231ff9a36a1dc44f4d9ee by Levi Burner : rollout nthread kwarg, rename shutdown_pool to close, fixups -- f622378543596a208339af0208fa3a70bf2a8007 by Levi Burner : rollout add missing .close() calls -- 50f3ebca43c53eac03f03943c34bb1e46967bd4f by Levi Burner : 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 | 8 + doc/python.rst | 44 +++- python/mujoco/CMakeLists.txt | 2 +- python/mujoco/rollout.cc | 226 +++++++++++++------ python/mujoco/rollout.py | 401 +++++++++++++++++++++++----------- python/mujoco/rollout_test.py | 116 +++++++++- python/mujoco/threadpool.cc | 87 ++++++++ python/mujoco/threadpool.h | 80 +++++++ 8 files changed, 761 insertions(+), 203 deletions(-) create mode 100644 python/mujoco/threadpool.cc create mode 100644 python/mujoco/threadpool.h diff --git a/doc/changelog.rst b/doc/changelog.rst index 5afaf2f0..33456f82 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -5,6 +5,14 @@ Changelog Upcoming version (not yet released) ----------------------------------- +Python bindings +^^^^^^^^^^^^^^^ +- :ref:`rollout` 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`). diff --git a/doc/python.rst b/doc/python.rst index ddaeb0a1..35c27d73 100644 --- a/doc/python.rst +++ b/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 -implemented in `rollout.cc `__ +``mujoco.rollout`` and ``mujoco.rollout.Rollout`` shows how to add additional C/C++ functionality, exposed as a Python module +via pybind11. It is implemented in `rollout.cc `__ and wrapped in `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`. ``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 -the ``test_threading`` function in +By default ``rollout.rollout`` creates a new thread pool every call if ``len(data) > 1``. To reuse the thread pool +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 `__ for an example -of threaded operation (and more generally for usage examples). +of threaded operation (and for more general usage examples). .. _PyMinimize: diff --git a/python/mujoco/CMakeLists.txt b/python/mujoco/CMakeLists.txt index aa97b9e6..b6d6c078 100644 --- a/python/mujoco/CMakeLists.txt +++ b/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( diff --git a/python/mujoco/rollout.cc b/python/mujoco/rollout.cc index ffedda3f..3f27f1d6 100644 --- a/python/mujoco/rollout.cc +++ b/python/mujoco/rollout.cc @@ -13,6 +13,7 @@ // limitations under the License. #include +#include #include #include @@ -20,6 +21,7 @@ #include "errors.h" #include "raw.h" #include "structs.h" +#include "threadpool.h" #include #include #include @@ -31,14 +33,24 @@ namespace { namespace py = ::pybind11; +using PyCArray = py::array_t; + // 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& m, mjData* d, int nroll, int nstep, unsigned int control_spec, +void _unsafe_rollout(std::vector& 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& 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& m, mjData* d, int nroll, int n } } +// C-style threaded version of _unsafe_rollout +void _unsafe_rollout_threaded(std::vector& m, std::vector& 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> arg, return static_cast(info.ptr); } +class Rollout { + public: + Rollout(int nthread) : nthread_(nthread) { + if (this->nthread_ > 0) { + this->pool_ = std::make_shared(this->nthread_); + } + } + + void rollout(py::list m, py::list d, int nstep, unsigned int control_spec, + const PyCArray state0, std::optional warmstart0, + std::optional control, + std::optional state, + std::optional sensordata, + std::optional chunk_size) { + // get raw pointers + int nroll = state0.shape(0); + std::vector model_ptrs(nroll); + for (int r = 0; r < nroll; r++) { + model_ptrs[r] = m[r].cast()->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 data_ptrs(py::len(d)); + for (int t = 0; t < py::len(d); t++) { + data_ptrs[t] = d[t].cast()->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 pool_; +}; PYBIND11_MODULE(_rollout, pymodule) { namespace py = ::pybind11; - using PyCArray = py::array_t; - // roll out open loop trajectories from multiple initial states - // get subsequent states and corresponding sensor values - pymodule.def( - "rollout", - [](py::list m, MjDataWrapper& d, - int nstep, unsigned int control_spec, - const PyCArray state0, - std::optional warmstart0, - std::optional control, - std::optional state, - std::optional sensordata - ) { - // get raw pointers - int nroll = state0.shape(0); - std::vector model_ptrs(nroll); - for (int r = 0; r < nroll; r++) { - model_ptrs[r] = m[r].cast()->get(); - } - raw::MjData* data = d.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 - 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) - ); + py::class_(pymodule, "Rollout") + .def( + py::init([](int nthread) { + return std::make_unique(nthread); + }), + py::kw_only(), + py::arg("nthread"), + py::doc(rollout_init_doc)) + .def( + "rollout", + &Rollout::rollout, + 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::arg("chunk_size") = py::none(), + py::doc(rollout_doc)); } } // namespace diff --git a/python/mujoco/rollout.py b/python/mujoco/rollout.py index 98eaa3f2..95b6ad3f 100644 --- a/python/mujoco/rollout.py +++ b/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. - data: An associated mjData instance. + 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. @@ -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: - # don't infer nroll/nstep - # don't support singleton expansion - # don't allocate output arrays - # just call rollout and return - if skip_checks: - _rollout.rollout( + if not isinstance(data, list): + data = [data] # Use a length 1 list to simplify code below + + nthread = len(data) if len(data) > 1 else 0 + + # Use a persistent thread pool if requested + 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, - sensordata, + control_spec=control_spec, + skip_checks=skip_checks, + nstep=nstep, + initial_warmstart=initial_warmstart, + state=state, + sensordata=sensordata, + chunk_size=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') - _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 + finally: + if not persistent_pool: + rollout_.close() def _check_must_be_numeric(**kwargs): diff --git a/python/mujoco/rollout_test.py b/python/mujoco/rollout_test.py index 3cc0d062..af8a5d3a 100644 --- a/python/mujoco/rollout_test.py +++ b/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): diff --git a/python/mujoco/threadpool.cc b/python/mujoco/threadpool.cc new file mode 100644 index 00000000..cd131b18 --- /dev/null +++ b/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 +#include +#include +#include +#include + +#include + +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 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 task) { + std::unique_lock 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 lock(m_); + cv_in_.wait(lock, [&]() { return !queue_.empty(); }); + std::function task = std::move(queue_.front()); + queue_.pop(); + cv_in_.notify_one(); + return task; + }(); + if (task == nullptr) { + { + std::unique_lock lock(m_); + ++ctr_; + cv_ext_.notify_one(); + } + break; + } + task(); + + { + std::unique_lock lock(m_); + ++ctr_; + cv_ext_.notify_one(); + } + } +} + +} // namespace mujoco::python diff --git a/python/mujoco/threadpool.h b/python/mujoco/threadpool.h new file mode 100644 index 00000000..5a142ad0 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include + +#include + +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 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 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 threads_; + std::mutex m_; + std::condition_variable cv_in_; + std::condition_variable cv_ext_; + std::queue> queue_; + std::uint64_t ctr_; +}; + +} // namespace mujoco::python + +#endif // MUJOCO_PYTHON_THREADPOOL_H_