diff --git a/doc/changelog.rst b/doc/changelog.rst index 857edacd..cb27d8b4 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -27,12 +27,21 @@ MJX 6. Updated MJX colab tutorial with more stable quadruped environment. 7. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, boxes, and meshes. +Python bindings +^^^^^^^^^^^^^^^ +8. Improved the implmentation of the :ref:`rollout` module. Note the changes below are breaking, dependent + code will require modification. + + - Uses :ref:`mjSTATE_FULLPHYSICS` as state spec, enabling divergence detection by inspecting time. + - Allows user-defined control spec for any combination of :ref:`user input` fields as controls. + - Outputs are no longer squeezed and always have dim=3. + Bug fixes ^^^^^^^^^ -8. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes +9. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes :github:issue:`1270`. -9. Fixed a bug in the :ref:`muscle model` that led to non-zero values outside the lower - bound of the length range. Fixes :github:issue:`1342`. +10. Fixed a bug in the :ref:`muscle model` that led to non-zero values outside the lower + bound of the length range. Fixes :github:issue:`1342`. Version 3.1.1 (December 18, 2023) @@ -40,7 +49,8 @@ Version 3.1.1 (December 18, 2023) Bug fixes ^^^^^^^^^ -1. Fixed a bug (introduced in 3.1.0) where box-box collisions produced no contacts if one box was deeply embedded in the other. +1. Fixed a bug (introduced in 3.1.0) where box-box collisions produced no contacts if one box was deeply embedded in the + other. 2. Fixed a bug in :ref:`simulate` where the "LOADING..." message was not showing correctly. 3. Fixed a crash in the Python :ref:`passive viewer`, when used with models containing Flex objects. 4. Fixed a bug in MJX where ``site_xmat`` was ignored in ``get_data`` and ``put_data`` @@ -53,8 +63,8 @@ Version 3.1.0 (December 12, 2023) General ^^^^^^^ 1. Improved convergence of Signed Distance Function (SDF) collisions by using line search and a new objective function - for the optimization. This allows to decrease the number of initial points needed for finding the contacts and is more - robust for very small or large geom sizes. + for the optimization. This allows to decrease the number of initial points needed for finding the contacts and is + more robust for very small or large geom sizes. 2. Added :ref:`frame` to MJCF, a :ref:`meta-element` which defines a pure coordinate transformation on its direct children, without requiring a :ref:`body`. 3. Added the :at:`kv` attribute to the :ref:`position` and :ref:`intvelocity` @@ -64,17 +74,14 @@ General Plugins ^^^^^^^ - 4. Allow actuator plugins to use activation variables in ``mjData.act`` as their internal state, rather than ``mjData.plugin_state``. Actuator plugins can now specify :ref:`callbacks` that compute activation variables, and they can be used with built-in :ref:`dyntype` actuator dynamics. - 5. Added the `pid `__ actuator plugin, a configurable PID controller that implements the Integral term, which is not available with native MuJoCo actuators. MJX ^^^ - 6. Added ``site_xpos`` and ``site_xmat`` to MJX. 7. Added ``put_data``, ``put_model``, ``get_data`` to replace ``device_put`` and ``device_get_into``, which will be deprecated. These new functions correctly translate fields that are the result of intermediate calculations such as @@ -86,8 +93,8 @@ Bug fixes Before this fix such actuators could lead to non-conservation of momentum. 9. Fix bug that prevented using flex with :ref:`simulate`. 10. Fix bug that prevented the use of elasticity plugins in combination with pinned flex vertices. -11. Release Python wheels targeting macOS 10.16 to support x86_64 systems where SYSTEM_VERSION_COMPAT is set. The minimum - supported version is still 11.0, but we release these wheels to fix compatibility for those users. See +11. Release Python wheels targeting macOS 10.16 to support x86_64 systems where ``SYSTEM_VERSION_COMPAT`` is set. + The minimum supported version is still 11.0, but we release these wheels to fix compatibility for those users. See :github:issue:`1213`. Version 3.0.1 (November 15, 2023) diff --git a/doc/python.rst b/doc/python.rst index e18149cb..6f1ce280 100644 --- a/doc/python.rst +++ b/doc/python.rst @@ -465,28 +465,32 @@ Open-loop rollouts ================== We include a code sample showing how to add additional C/C++ functionality, exposed as a Python module via pybind11. The -sample, implemented in ``rollout.cc`` and wrapped in ``rollout.py``, implements a common use case where tight loops -implemented outside of Python are beneficial: rolling out a trajectory (i.e., calling ``mj_step()`` in a loop), given an -intial state and sequence of controls, and returning subsequent states and sensor values. The canonical usage form is +sample, implemented in `rollout.cc `__ +and wrapped in `rollout.py `__, +implements a common use case where tight loops implemented outside of Python are beneficial: rolling out a trajectory +(i.e., calling ``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 .. code-block:: python - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + state, sensordata = rollout.rollout(model, data, initial_state, control) -``initial_state`` is a ``nstate x nqva`` array, with ``nstate`` initial states of length ``nqva``, where ``nqva = -model.nq + model.nv + model.na`` is the size of the full MuJoCo mechanical state: positions (``data.qpos``), velocities -(``data.qvel``) and actuator activations (``data.act``). ``ctrl`` is a ``nstate x nstep x nu`` array of control -sequences. +``initial_state`` is a ``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 +by default the ``mjModel.nu`` standard actuators, but any combination of :ref:`user input` arrays can be +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 -values already present in the given ``MjData`` instance will have no effect on the output. In order to facilitate this, -all inputs including ``time`` and ``qacc_warmstart`` are set to default values, as are auxillary controls -(``qfrc_applied``, ``xfrc_applied`` and ``mocap_{pos,quat}``). These can also be optionally set by the user. +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 -`rollout_test.py `_ for an example of -threaded operation. +`rollout_test.py `__ for an example +of threaded operation (and more generally for usage examples). .. _PyMjpy_migration: diff --git a/python/mujoco/rollout.cc b/python/mujoco/rollout.cc index eacdc8fe..6cfbb470 100644 --- a/python/mujoco/rollout.cc +++ b/python/mujoco/rollout.cc @@ -31,111 +31,114 @@ namespace { namespace py = ::pybind11; +// NOLINTBEGIN(whitespace/line_length) + const auto rollout_doc = R"( -Roll out open-loop trajectories from initial states, get subsequent states and sensor values. +Roll out open-loop trajectories from initial states, get resulting states and sensor values. input arguments (required): - model an instance of MjModel - data an associated instance of MjData - nstate an integer, number of initial states from which to roll out trajectories - nstep an integer, number of steps to be taken for each trajectory + model instance of MjModel + data associated instance of MjData + nroll integer, number of initial states from which to roll out trajectories + 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, + nstate = mj_stateSize(m, mjSTATE_FULLPHYSICS) input arguments (optional): - initial_state (nstate x nqva) nstate initial state vectors, nqva=nq+nv+na - initial_time (nstate x 1) nstate initial times - initial_warmstart (nstate x nv) nstate qacc_warmstart vectors - ctrl (nstate x nstep x nu) nstate length-nstep controls - qfrc_applied (nstate x nstep x nv) nstate length-nstep generalized forces - xfrc_applied (nstate x nstep x nbody*6) nstate length-nstep Cartesian wrenches - mocap (nstate x nstep x nmocap*7) nstate length-nstep mocap body poses + warmstart0 (nroll x nv) nroll qacc_warmstart vectors + control (nroll x nstep x ncontrol) nroll trajectories of nstep controls output arguments (optional): - state (nstate x nstep x nqva) nstate length-nstep states - sensordata (nstate x nstep x nsendordata) nstate length-nstep sensordatas + state (nroll x nstep x nstate) nroll nstep states + sensordata (nroll x nstep x nsendordata) nroll trajectories of nstep sensordata vectors )"; // 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(const mjModel* m, mjData* d, int nstate, int nstep, - const mjtNum* state0, const mjtNum* ctrl, - const mjtNum* qfrc, const mjtNum* xfrc, - const mjtNum* mocap, const mjtNum* time0, - const mjtNum* warmstart0, +void _unsafe_rollout(const mjModel* m, mjData* d, int nroll, int nstep, unsigned int control_spec, + const mjtNum* state0, const mjtNum* warmstart0, const mjtNum* control, mjtNum* state, mjtNum* sensordata) { - // model sizes - int nq = m->nq; - int nv = m->nv; - int na = m->na; - int nqva = nq + nv + na; - int nu = m->nu; - int nbody = m->nbody; - int nmocap = m->nmocap; + // sizes + int nstate = mj_stateSize(m, mjSTATE_FULLPHYSICS); + int ncontrol = mj_stateSize(m, control_spec); + int nv = m->nv, nbody = m->nbody, neq = m->neq; int nsensordata = m->nsensordata; - // loop over initial states - for (int s=0; s < nstate; s++) { - // set initial state - if (state0) { - mju_copy(d->qpos, state0 + s*nqva, nq); - mju_copy(d->qvel, state0 + s*nqva + nq, nv); - mju_copy(d->act, state0 + s*nqva + nq + nv, na); - } else { - mju_copy(d->qpos, m->qpos0, nq); - mju_zero(d->qvel, nv); - mju_zero(d->act, na); + // clear user inputs if unspecified + if (!(control_spec & mjSTATE_CTRL)) { + mju_zero(d->ctrl, m->nu); + } + if (!(control_spec & mjSTATE_QFRC_APPLIED)) { + mju_zero(d->qfrc_applied, nv); + } + if (!(control_spec & mjSTATE_XFRC_APPLIED)) { + mju_zero(d->xfrc_applied, 6*nbody); + } + if (!(control_spec & mjSTATE_MOCAP_POS)) { + for (int i = 0; i < nbody; i++) { + int id = m->body_mocapid[i]; + if (id >= 0) mju_copy3(d->mocap_pos+3*id, m->body_pos+3*i); } + } + if (!(control_spec & mjSTATE_MOCAP_QUAT)) { + for (int i = 0; i < nbody; i++) { + int id = m->body_mocapid[i]; + if (id >= 0) mju_copy4(d->mocap_quat+4*id, m->body_quat+4*i); + } + } + if (!(control_spec & mjSTATE_EQ_ACTIVE)) { + for (int i = 0; i < neq; i++) { + d->eq_active[i] = m->eq_active0[i]; + } + } - // set initial time - d->time = time0 ? time0[s] : 0; + // loop over rollouts + for (int r = 0; r < nroll; r++) { + // set initial state + mj_setState(m, d, state0 + r*nstate, mjSTATE_FULLPHYSICS); // set warmstart accelerations if (warmstart0) { - mju_copy(d->qacc_warmstart, warmstart0 + s*nv, nv); + mju_copy(d->qacc_warmstart, warmstart0 + r*nv, nv); } else { mju_zero(d->qacc_warmstart, nv); } - // clear control inputs if unspecified - if (s == 0) { - if (!ctrl) { - mju_zero(d->ctrl, nu); - } - if (!qfrc) { - mju_zero(d->qfrc_applied, nv); - } - if (!xfrc) { - mju_zero(d->xfrc_applied, 6*nbody); - } - if (!mocap) { - for (int j=0; j < nbody; j++) { - int id = m->body_mocapid[j]; - if (id >= 0) { - mju_copy3(d->mocap_pos+3*id, m->body_pos+3*j); - mju_copy4(d->mocap_quat+4*id, m->body_quat+4*j); - } - } - } + // clear warning counters + for (int i = 0; i < mjNWARNING; i++) { + d->warning[i].number = 0; } - // roll out trajectories + // roll out trajectory for (int t = 0; t < nstep; t++) { + // check for warnings + bool nwarning = false; + for (int i = 0; i < mjNWARNING; i++) { + if (d->warning[i].number) { + nwarning = true; + break; + } + } + + // if any warnings, fill remaining outputs with current outputs, break + if (nwarning) { + for (; t < nstep; t++) { + int step = r*nstep + t; + if (state) { + mj_getState(m, d, state + step*nstate, mjSTATE_FULLPHYSICS); + } + if (sensordata) { + mju_copy(sensordata + step*nsensordata, d->sensordata, nsensordata); + } + } + break; + } + + int step = r*nstep + t; + // controls - if (ctrl) { - mju_copy(d->ctrl, ctrl + s*nstep*nu + t*nu, nu); - } - // generalized forces - if (qfrc) { - mju_copy(d->qfrc_applied, qfrc + s*nstep*nv + t*nv, nv); - } - // Cartesian wrenches - if (xfrc) { - mju_copy(d->xfrc_applied, xfrc + s*nstep*6*nbody + t*6*nbody, 6*nbody); - } - // mocap bodies - if (mocap) { - mju_copy(d->mocap_pos, - mocap + s*nstep*7*nmocap + t*7*nmocap, 3*nmocap); - mju_copy(d->mocap_quat, - mocap + s*nstep*7*nmocap + t*7*nmocap + 3*nmocap, 4*nmocap); + if (control) { + mj_setState(m, d, control + step*ncontrol, control_spec); } // step @@ -143,23 +146,22 @@ void _unsafe_rollout(const mjModel* m, mjData* d, int nstate, int nstep, // copy out new state if (state) { - mju_copy(state + s*nstep*nqva + t*nqva, d->qpos, nq); - mju_copy(state + s*nstep*nqva + t*nqva + nq, d->qvel, nv); - mju_copy(state + s*nstep*nqva + t*nqva + nq + nv, d->act, na); + mj_getState(m, d, state + step*nstate, mjSTATE_FULLPHYSICS); } + // copy out sensor values if (sensordata) { - mju_copy(sensordata + s*nstep*nsensordata + t*nsensordata, - d->sensordata, nsensordata); + mju_copy(sensordata + step*nsensordata, d->sensordata, nsensordata); } } } } +// NOLINTEND(whitespace/line_length) // check size of optional argument to rollout(), return raw pointer mjtNum* get_array_ptr(std::optional> arg, - const char* name, int nstate, int nstep, int dim) { + const char* name, int nroll, int nstep, int dim) { // if empty return nullptr if (!arg.has_value()) { return nullptr; @@ -169,11 +171,10 @@ mjtNum* get_array_ptr(std::optional> arg, py::buffer_info info = arg->request(); // check size - int expected_size = nstate * nstep * dim; + int expected_size = nroll * nstep * dim; if (info.size != expected_size) { std::ostringstream msg; - msg << name << ".size should be " << expected_size << - ", got " << info.size; + msg << name << ".size should be " << expected_size << ", got " << info.size; throw py::value_error(msg.str()); } return static_cast(info.ptr); @@ -188,14 +189,11 @@ PYBIND11_MODULE(_rollout, pymodule) { // get subsequent states and corresponding sensor values pymodule.def( "rollout", - [](const MjModelWrapper& m, MjDataWrapper& d, int nstate, int nstep, - std::optional init_state, - std::optional init_time, - std::optional init_warmstart, - std::optional ctrl, - std::optional qfrc, - std::optional xfrc, - std::optional mocap, + [](const MjModelWrapper& m, MjDataWrapper& d, + int nroll, int nstep, unsigned int control_spec, + const PyCArray state0, + std::optional warmstart0, + std::optional control, std::optional state, std::optional sensordata ) { @@ -203,28 +201,22 @@ PYBIND11_MODULE(_rollout, pymodule) { raw::MjData* data = d.get(); // check that some steps need to be taken, return if not - if (nstate < 1 || nstep < 1) { + if (nroll < 1 || nstep < 1) { return; } + // get sizes + int nstate = mj_stateSize(model, mjSTATE_FULLPHYSICS); + int ncontrol = mj_stateSize(model, control_spec); + // get raw pointers - int nqva = model->nq + model->nv + model->na; - mjtNum* init_state_ptr = - get_array_ptr(init_state, "initial_state", nstate, 1, nqva); - mjtNum* ctrl_ptr = - get_array_ptr(ctrl, "ctrl", nstate, nstep, model->nu); - mjtNum* qfrc_ptr = - get_array_ptr(qfrc, "qfrc_applied", nstate, nstep, model->nv); - mjtNum* xfrc_ptr = - get_array_ptr(xfrc, "xfrc_applied", nstate, nstep, 6*model->nbody); - mjtNum* mocap_ptr = - get_array_ptr(mocap, "mocap", nstate, nstep, 7*model->nmocap); - mjtNum* init_time_ptr = - get_array_ptr(init_time, "init_time", nstate, 1, 1); - mjtNum* init_warmstart_ptr = get_array_ptr( - init_warmstart, "init_warmstart", nstate, 1, model->nv); - mjtNum* state_ptr = get_array_ptr(state, "state", nstate, nstep, nqva); - mjtNum* sensordata_ptr = get_array_ptr(sensordata, "sensordata", nstate, + mjtNum* state0_ptr = get_array_ptr(state0, "state0", nroll, 1, nstate); + mjtNum* warmstart0_ptr = get_array_ptr(warmstart0, "warmstart0", nroll, + 1, model->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->nsensordata); // perform rollouts @@ -234,24 +226,20 @@ PYBIND11_MODULE(_rollout, pymodule) { // call unsafe rollout function InterceptMjErrors(_unsafe_rollout)( - model, data, nstate, nstep, init_state_ptr, ctrl_ptr, qfrc_ptr, - xfrc_ptr, mocap_ptr, init_time_ptr, init_warmstart_ptr, state_ptr, - sensordata_ptr); + model, data, nroll, nstep, control_spec, state0_ptr, + warmstart0_ptr, control_ptr, state_ptr, sensordata_ptr); } }, py::arg("model"), py::arg("data"), - py::arg("nstate"), + py::arg("nroll"), py::arg("nstep"), - py::arg("initial_state") = py::none(), - py::arg("initial_time") = py::none(), - py::arg("initial_warmstart") = py::none(), - py::arg("ctrl") = py::none(), - py::arg("qfrc_applied") = py::none(), - py::arg("xfrc_applied") = py::none(), - py::arg("mocap") = py::none(), - py::arg("state") = py::none(), - py::arg("sensordata") = py::none(), + 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) ); } @@ -259,3 +247,4 @@ PYBIND11_MODULE(_rollout, pymodule) { } // namespace } // namespace mujoco::python + diff --git a/python/mujoco/rollout.py b/python/mujoco/rollout.py index 10ad328f..8904e30c 100644 --- a/python/mujoco/rollout.py +++ b/python/mujoco/rollout.py @@ -14,132 +14,144 @@ # ============================================================================== """Roll out open-loop trajectories from initial states, get subsequent states and sensor values.""" +from typing import Optional + +import mujoco from mujoco import _rollout import numpy as np +from numpy import typing as npt -def rollout(model, data, initial_state=None, ctrl=None, - *, # require following arguments to be named - skip_checks=False, - nstate=None, - nstep=None, - initial_time=None, - initial_warmstart=None, - qfrc_applied=None, - xfrc_applied=None, - mocap=None, - state=None, - sensordata=None): - """Roll out open-loop trajectories from initial states, get subsequent states and sensor values. +def rollout(model: mujoco.MjModel, + data: 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, + nroll: Optional[int] = None, + nstep: Optional[int] = None, + initial_warmstart: Optional[npt.ArrayLike] = None, + state: Optional[npt.ArrayLike] = None, + sensordata: Optional[npt.ArrayLike] = None): + """Rolls out open-loop trajectories from initial states, get subsequent states and sensor values. - This function serves as a Python wrapper for the C++ functionality in - `rollout.cc`, please see documentation therein. This python funtion will - infer `nstate` and `nstep`, tile input arguments with singleton dimensions, - and allocate output arguments if none are given. + 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 mjModel instance. + data: An associated mjData instance. + 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. + nroll: Number of rollouts (inferred if unspecified). + 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) + + Returns: + state: + State output array, (nroll x nstep x nstate). + sensordata: + Sensor data output array, (nroll x nstep x nsensordata). + + Raises: + ValueError: bad shapes or sizes. """ - # don't infer nstate/nstep, don't support singleton expansion, don't allocate - # output arrays, just call rollout + # 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(model, data, nstate, nstep, initial_state, initial_time, - initial_warmstart, ctrl, qfrc_applied, xfrc_applied, mocap, - state, sensordata) + _rollout.rollout(model, data, nroll, nstep, control_spec, initial_state, + initial_warmstart, control, state, sensordata) return state, sensordata + # 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 nstate and not isinstance(nstate, int): - raise ValueError('nstate must be an integer') + if nroll and not isinstance(nroll, int): + raise ValueError('nroll must be an integer') if nstep and not isinstance(nstep, int): raise ValueError('nstep must be an integer') _check_must_be_numeric( initial_state=initial_state, - initial_time=initial_time, initial_warmstart=initial_warmstart, - ctrl=ctrl, - qfrc_applied=qfrc_applied, - xfrc_applied=xfrc_applied, - mocap=mocap, + control=control, state=state, sensordata=sensordata) # check number of dimensions _check_number_of_dimensions(2, initial_state=initial_state, - initial_time=initial_time, initial_warmstart=initial_warmstart) _check_number_of_dimensions(3, - ctrl=ctrl, - qfrc_applied=qfrc_applied, - xfrc_applied=xfrc_applied, - mocap=mocap, + control=control, state=state, sensordata=sensordata) # ensure 2D, make contiguous, row-major (C ordering) initial_state = _ensure_2d(initial_state) - initial_time = _ensure_2d(initial_time) initial_warmstart = _ensure_2d(initial_warmstart) # ensure 3D, make contiguous, row-major (C ordering) - ctrl = _ensure_3d(ctrl) - qfrc_applied = _ensure_3d(qfrc_applied) - xfrc_applied = _ensure_3d(xfrc_applied) - mocap = _ensure_3d(mocap) + control = _ensure_3d(control) state = _ensure_3d(state) sensordata = _ensure_3d(sensordata) # check trailing dimensions - _check_trailing_dimension(model.nq + model.nv + model.na, - initial_state=initial_state, state=state) - _check_trailing_dimension(1, initial_time=initial_time) - _check_trailing_dimension(model.nu, ctrl=ctrl) - _check_trailing_dimension(model.nv, qfrc_applied=qfrc_applied) - _check_trailing_dimension(model.nbody*6, xfrc_applied=xfrc_applied) - _check_trailing_dimension(model.nmocap*7, mocap=mocap) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS.value) + _check_trailing_dimension(nstate, initial_state=initial_state, state=state) + ncontrol = mujoco.mj_stateSize(model, control_spec) + _check_trailing_dimension(ncontrol, control=control) + _check_trailing_dimension(model.nv, initial_warmstart=initial_warmstart) _check_trailing_dimension(model.nsensordata, sensordata=sensordata) - # infer nstate, check for incompatibilities - nstate = _infer_dimension(0, nstate or 1, - initial_state=initial_state, - initial_time=initial_time, - initial_warmstart=initial_warmstart, - ctrl=ctrl, - qfrc_applied=qfrc_applied, - xfrc_applied=xfrc_applied, - mocap=mocap, - state=state, - sensordata=sensordata) + # infer nroll, check for incompatibilities + nroll = _infer_dimension(0, nroll or 1, + initial_state=initial_state, + initial_warmstart=initial_warmstart, + control=control, + state=state, + sensordata=sensordata) # infer nstep, check for incompatibilities nstep = _infer_dimension(1, nstep or 1, - ctrl=ctrl, - qfrc_applied=qfrc_applied, - xfrc_applied=xfrc_applied, - mocap=mocap, + control=control, state=state, sensordata=sensordata) # tile input arrays if required (singleton expansion) - initial_state = _tile_if_required(initial_state, nstate) - initial_time = _tile_if_required(initial_time, nstate) - initial_warmstart = _tile_if_required(initial_warmstart, nstate) - ctrl = _tile_if_required(ctrl, nstate, nstep) - qfrc_applied = _tile_if_required(qfrc_applied, nstate, nstep) - xfrc_applied = _tile_if_required(xfrc_applied, nstate, nstep) - mocap = _tile_if_required(mocap, nstate, nstep) + 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((nstate, nstep, model.nq + model.nv + model.na)) + state = np.empty((nroll, nstep, nstate)) if sensordata is None: - sensordata = np.empty((nstate, nstep, model.nsensordata)) + sensordata = np.empty((nroll, nstep, model.nsensordata)) # call rollout - _rollout.rollout(model, data, nstate, nstep, initial_state, initial_time, - initial_warmstart, ctrl, qfrc_applied, xfrc_applied, mocap, - state, sensordata) + _rollout.rollout(model, data, nroll, nstep, control_spec, initial_state, + initial_warmstart, control, state, sensordata) + + # return outputs + return state, sensordata - # return squeezed outputs - return state.squeeze(), sensordata.squeeze() def _check_must_be_numeric(**kwargs): for key, value in kwargs.items(): @@ -148,6 +160,7 @@ def _check_must_be_numeric(**kwargs): if not isinstance(value, np.ndarray) and not isinstance(value, float): raise ValueError(f'{key} must be a numpy array or float') + def _check_number_of_dimensions(ndim, **kwargs): for key, value in kwargs.items(): if value is None: @@ -155,12 +168,16 @@ def _check_number_of_dimensions(ndim, **kwargs): if value.ndim > ndim: raise ValueError(f'{key} can have at most {ndim} dimensions') + def _check_trailing_dimension(dim, **kwargs): for key, value in kwargs.items(): if value is None: continue if value.shape[-1] != dim: - raise ValueError(f'trailing dimension of {key} must be {dim}, got {value.shape[-1]}') + raise ValueError( + f'trailing dimension of {key} must be {dim}, got {value.shape[-1]}' + ) + def _ensure_2d(arg): if arg is None: @@ -168,6 +185,7 @@ def _ensure_2d(arg): else: return np.ascontiguousarray(np.atleast_2d(arg), dtype=np.float64) + def _ensure_3d(arg): if arg is None: return None @@ -181,7 +199,22 @@ def _ensure_3d(arg): arg = arg[np.newaxis, ...] return np.ascontiguousarray(arg, dtype=np.float64) + def _infer_dimension(dim, value, **kwargs): + """Infers dimension `dim` given guess `value` from set of arrays. + + Args: + dim: Dimension to be inferred. + value: Initial guess of inferred value (1: unknown). + **kwargs: List of arrays which should all have the same size (or 1) + along dimension dim. + + Returns: + Inferred dimension. + + Raises: + ValueError: If mismatch between array shapes or initial guess. + """ for name, array in kwargs.items(): if array is None: continue @@ -190,10 +223,12 @@ def _infer_dimension(dim, value, **kwargs): value = array.shape[dim] elif array.shape[dim] != 1: raise ValueError( - f'dimension {dim} inferred as {value} but {name} has {array.shape[dim]}' + f'dimension {dim} inferred as {value} ' + f'but {name} has {array.shape[dim]}' ) return value + def _tile_if_required(array, dim0, dim1=None): if array is None: return diff --git a/python/mujoco/rollout_test.py b/python/mujoco/rollout_test.py index 50ca4568..e1915592 100644 --- a/python/mujoco/rollout_test.py +++ b/python/mujoco/rollout_test.py @@ -14,15 +14,16 @@ # ============================================================================== """tests for rollout function.""" +import concurrent.futures +import threading + from absl.testing import absltest from absl.testing import parameterized import mujoco -import numpy as np -import concurrent.futures -import threading from mujoco import rollout +import numpy as np -#--------------------------- models used for testing --------------------------- +# -------------------------- models used for testing --------------------------- TEST_XML = r""" @@ -96,7 +97,7 @@ TEST_XML_MOCAP = r""" - + """ @@ -106,12 +107,33 @@ TEST_XML_EMPTY = r""" """ +TEST_XML_DIVERGE = r""" + + + + + + + + + + + + + + + +""" + ALL_MODELS = {'TEST_XML': TEST_XML, 'TEST_XML_NO_SENSORS': TEST_XML_NO_SENSORS, 'TEST_XML_NO_ACTUATORS': TEST_XML_NO_ACTUATORS, 'TEST_XML_EMPTY': TEST_XML_EMPTY} -#------------------------------- tests ----------------------------------------- +# ------------------------------ tests ----------------------------------------- + class MuJoCoRolloutTest(parameterized.TestCase): @@ -119,179 +141,209 @@ class MuJoCoRolloutTest(parameterized.TestCase): super().setUp() np.random.seed(42) - #----------------------------- test basic operation + # ----------------------------- test basic operation @parameterized.parameters(ALL_MODELS.keys()) def test_single_step(self, model_name): model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name]) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - initial_state = np.random.randn(model.nq + model.nv + model.na) - ctrl = np.random.randn(model.nu) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + initial_state = np.random.randn(nstate) + control = np.random.randn(model.nu) + state, sensordata = rollout.rollout(model, data, initial_state, control) mujoco.mj_resetData(model, data) - py_state, py_sensordata = step(model, data, initial_state, ctrl=ctrl) + 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) - - @parameterized.parameters(ALL_MODELS.keys()) - def test_single_rollout(self, model_name): + def test_one_rollout(self, model_name): nstep = 3 # number of timesteps model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name]) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - initial_state = np.random.randn(model.nq + model.nv + model.na) - ctrl = np.random.randn(nstep, model.nu) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + initial_state = np.random.randn(nstate) + control = np.random.randn(nstep, model.nu) + state, sensordata = rollout.rollout(model, data, initial_state, control) - py_state, py_sensordata = single_rollout(model, data, initial_state, - ctrl=ctrl) - np.testing.assert_array_equal(state, np.asarray(py_state)) - np.testing.assert_array_equal(sensordata, np.asarray(py_sensordata)) + 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) @parameterized.parameters(ALL_MODELS.keys()) def test_multi_step(self, model_name): model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name]) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - nstate = 5 # number of initial states + nroll = 5 # number of rollouts + nstep = 1 # number of steps - initial_state = np.random.randn(nstate, model.nq + model.nv + model.na) - ctrl = np.random.randn(nstate, 1, model.nu) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + initial_state = np.random.randn(nroll, nstate) + control = np.random.randn(nroll, nstep, model.nu) + state, sensordata = rollout.rollout(model, data, initial_state, control) mujoco.mj_resetData(model, data) - py_state, py_sensordata = multi_rollout(model, data, initial_state, - ctrl=ctrl) + 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) @parameterized.parameters(ALL_MODELS.keys()) - def test_single_rollout_fixed_ctrl(self, model_name): - nstep = 3 + def test_one_rollout_fixed_ctrl(self, model_name): model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name]) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - initial_state = np.random.randn(model.nq + model.nv + model.na) - ctrl = np.random.randn(model.nu) - state = np.empty((nstep, model.nq + model.nv + model.na)) - sensordata = np.empty((nstep, model.nsensordata)) - rollout.rollout(model, data, initial_state, ctrl, + nroll = 1 # number of rollouts + nstep = 3 # number of steps + + initial_state = np.random.randn(nstate) + control = np.random.randn(model.nu) + state = np.empty((nroll, nstep, nstate)) + sensordata = np.empty((nroll, nstep, model.nsensordata)) + rollout.rollout(model, data, initial_state, control, state=state, sensordata=sensordata) - ctrl = np.tile(ctrl, (nstep, 1)) # repeat?? - py_state, py_sensordata = single_rollout(model, data, initial_state, - ctrl=ctrl) + control = np.tile(control, (nstep, 1)) + 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) @parameterized.parameters(ALL_MODELS.keys()) def test_multi_rollout(self, model_name): model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name]) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - nstate = 2 # number of initial states + nroll = 2 # number of initial states nstep = 3 # number of timesteps - initial_state = np.random.randn(nstate, model.nq + model.nv + model.na) - ctrl = np.random.randn(nstate, nstep, model.nu) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + initial_state = np.random.randn(nroll, nstate) + control = np.random.randn(nroll, nstep, model.nu) + state, sensordata = rollout.rollout(model, data, initial_state, control) - py_state, py_sensordata = multi_rollout(model, data, initial_state, - ctrl=ctrl) - np.testing.assert_array_equal(py_state, py_state) - np.testing.assert_array_equal(py_sensordata, py_sensordata) + 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) @parameterized.parameters(ALL_MODELS.keys()) def test_multi_rollout_fixed_ctrl_infer_from_output(self, model_name): model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name]) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - nstate = 2 # number of initial states + nroll = 2 # number of rollouts nstep = 3 # number of timesteps - initial_state = np.random.randn(nstate, model.nq + model.nv + model.na) - ctrl = np.random.randn(nstate, 1, model.nu) # 1 control in the time dimension - state = np.empty((nstate, nstep, model.nq + model.nv + model.na)) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl, + initial_state = np.random.randn(nroll, nstate) + control = np.random.randn(nroll, 1, model.nu) + state = np.empty((nroll, nstep, nstate)) + state, sensordata = rollout.rollout(model, data, initial_state, control, state=state) - ctrl = np.repeat(ctrl, nstep, axis=1) - py_state, py_sensordata = multi_rollout(model, data, initial_state, - ctrl=ctrl) + control = np.repeat(control, nstep, axis=1) + 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) - @parameterized.product(arg_nstep=[[3, 1, 1], [3, 3, 1], [3, 1, 3]], - model_name=list(ALL_MODELS.keys())) - def test_multi_rollout_multiple_inputs(self, arg_nstep, model_name): + @parameterized.parameters(ALL_MODELS.keys()) + def test_py_rollout_generalized_control(self, model_name): model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name]) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - nstate = 4 # number of initial states + nroll = 4 # number of rollouts + nstep = 3 # number of timesteps - initial_state = np.random.randn(nstate, model.nq + model.nv + model.na) + initial_state = np.random.randn(nroll, nstate) - # arg_nstep is the horizon for {ctrl, qfrc_applied, xfrc_applied}, respectively - ctrl = np.random.randn(nstate, arg_nstep[0], model.nu) - qfrc_applied = np.random.randn(nstate, arg_nstep[1], model.nv) - xfrc_applied = np.random.randn(nstate, arg_nstep[2], model.nbody*6) + control_spec = (mujoco.mjtState.mjSTATE_CTRL | + mujoco.mjtState.mjSTATE_QFRC_APPLIED | + mujoco.mjtState.mjSTATE_XFRC_APPLIED) + ncontrol = mujoco.mj_stateSize(model, control_spec) + control = np.random.randn(nroll, nstep, ncontrol) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl, - qfrc_applied=qfrc_applied, - xfrc_applied=xfrc_applied) + state, sensordata = rollout.rollout(model, data, initial_state, control, + control_spec=control_spec) - # tile singleton arguments - nstep = max(arg_nstep) - if arg_nstep[0] == 1: - ctrl = np.repeat(ctrl, nstep, axis=1) - if arg_nstep[1] == 1: - qfrc_applied = np.repeat(qfrc_applied, nstep, axis=1) - if arg_nstep[2] == 1: - xfrc_applied = np.repeat(xfrc_applied, nstep, axis=1) - - py_state, py_sensordata = multi_rollout(model, data, initial_state, - ctrl=ctrl, - qfrc_applied=qfrc_applied, - xfrc_applied=xfrc_applied) + py_state, py_sensordata = py_rollout(model, data, initial_state, control, + control_spec=control_spec) np.testing.assert_array_equal(state, py_state) np.testing.assert_array_equal(sensordata, py_sensordata) - #----------------------------- test threaded operation + def test_detect_divergence(self): + model = mujoco.MjModel.from_xml_string(TEST_XML_DIVERGE) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) + data = mujoco.MjData(model) + + nroll = 4 # number of rollouts + initial_state = np.empty((nroll, nstate)) + + # get diverging (0, 2) and non-diverging (1, 3) states + mujoco.mj_getState(model, data, initial_state[0], + mujoco.mjtState.mjSTATE_FULLPHYSICS) + mujoco.mj_getState(model, data, initial_state[2], + mujoco.mjtState.mjSTATE_FULLPHYSICS) + mujoco.mj_resetDataKeyframe(model, data, 0) # keyframe 0 does not diverge + mujoco.mj_getState(model, data, initial_state[1], + mujoco.mjtState.mjSTATE_FULLPHYSICS) + mujoco.mj_getState(model, data, initial_state[3], + mujoco.mjtState.mjSTATE_FULLPHYSICS) + + nstep = 10000 # divergence after ~15s, timestep = 2e-3 + + state = np.random.randn(nroll, nstep, nstate) + + rollout.rollout(model, data, initial_state, state=state) + + # initial_state[0,2] diverged, final timesteps are identical + assert state[0][-1][0] == state[0][-2][0] + assert state[2][-1][0] == state[2][-2][0] + + # initial_state[1,3] did not diverge, final timesteps are different + assert state[1][-1][0] != state[1][-2][0] + assert state[3][-1][0] != state[3][-2][0] + + # ----------------------------- test threaded operation def test_threading(self): model = mujoco.MjModel.from_xml_string(TEST_XML) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) num_workers = 32 - nstate = 10000 + nroll = 10000 nstep = 5 - initial_state = np.random.randn(nstate, model.nq+model.nv+model.na) - state = np.zeros((nstate, nstep, model.nq+model.nv+model.na)) - sensordata = np.zeros((nstate, nstep, model.nsensordata)) - ctrl = np.random.randn(nstate, nstep, model.nu) + 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) thread_local = threading.local() def thread_initializer(): thread_local.data = mujoco.MjData(model) - def call_rollout(initial_state, ctrl, state): - rollout.rollout(model, thread_local.data, skip_checks=True, - nstate=initial_state.shape[0], nstep=nstep, - initial_state=initial_state, ctrl=ctrl, state=state) + def call_rollout(initial_state, control, state, sensordata): + rollout.rollout(model, thread_local.data, initial_state, control, + skip_checks=True, nroll=initial_state.shape[0], + nstep=nstep, state=state, sensordata=sensordata) - n = initial_state.shape[0] // num_workers # integer division + n = nroll // num_workers # integer division chunks = [] # a list of tuples, one per worker for i in range(num_workers-1): - chunks.append( - (initial_state[i*n:(i+1)*n], ctrl[i*n:(i+1)*n], state[i*n:(i+1)*n])) + chunks.append((initial_state[i*n:(i+1)*n], + control[i*n:(i+1)*n], + state[i*n:(i+1)*n], + sensordata[i*n:(i+1)*n])) + # last chunk, absorbing the remainder: - chunks.append( - (initial_state[(num_workers-1)*n:], ctrl[(num_workers-1)*n:], - state[(num_workers-1)*n:])) + chunks.append((initial_state[(num_workers-1)*n:], + control[(num_workers-1)*n:], + state[(num_workers-1)*n:], + sensordata[(num_workers-1)*n:])) with concurrent.futures.ThreadPoolExecutor( max_workers=num_workers, initializer=thread_initializer) as executor: @@ -302,187 +354,237 @@ class MuJoCoRolloutTest(parameterized.TestCase): future.result() data = mujoco.MjData(model) - py_state, py_sensordata = multi_rollout(model, data, initial_state, - ctrl=ctrl) + 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) - #----------------------------- test advanced operation - - def test_time(self): - model = mujoco.MjModel.from_xml_string(TEST_XML) - data = mujoco.MjData(model) - - nstate = 1 - nstep = 3 - - initial_time = np.array([[2.]]) - initial_state = np.random.randn(nstate, model.nq + model.nv + model.na) - ctrl = np.random.randn(nstate, nstep, model.nu) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl, - initial_time=initial_time) - - self.assertAlmostEqual(data.time, 2 + nstep*model.opt.timestep) + # ---------------------------- test advanced operation def test_warmstart(self): model = mujoco.MjModel.from_xml_string(TEST_XML) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - state0 = np.zeros(model.nq + model.nv + model.na) - ctrl = np.zeros(model.nu) - state1, _ = step(model, data, state0, ctrl=ctrl) + # take one step, save the state + state0 = np.zeros(nstate) + control = np.zeros(model.nu) + state1, _ = step(model, data, state0, control) + + # save qacc_warmstart initial_warmstart = data.qacc_warmstart.copy() - state2, _ = step(model, data, state1, ctrl=ctrl) + # take one more step (uses correct warmstart) + state2, _ = step(model, data, state1[0], control) - state, _ = rollout.rollout(model, data, state1, ctrl) - assert np.linalg.norm(state-state2) > 0 + # take step using rollout, don't take warmstart into account + state, _ = rollout.rollout(model, data, state1[0], control) - state, _ = rollout.rollout(model, data, state1, ctrl, + # assert that stepping without warmstarts is not exact + np.testing.assert_raises(AssertionError, + np.testing.assert_array_equal, state, state2) + + # take step using rollout, take warmstart into account + state, _ = rollout.rollout(model, data, state1, control, initial_warmstart=initial_warmstart) - np.testing.assert_array_equal(state, state2) + + # assert exact equality + np.testing.assert_array_equal(state, np.expand_dims(state2, axis=0)) def test_mocap(self): model = mujoco.MjModel.from_xml_string(TEST_XML_MOCAP) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - initial_state = np.zeros(model.nq + model.nv + model.na) + initial_state = np.zeros(nstate) + + control_spec = (mujoco.mjtState.mjSTATE_MOCAP_POS | + mujoco.mjtState.mjSTATE_MOCAP_QUAT) + pos1 = np.array((1., 2., 3.)) quat1 = np.array((1., 2., 3., 4.)) quat1 /= np.linalg.norm(quat1) pos2 = np.array((2., 3., 4.)) quat2 = np.array((2., 3., 4., 5.)) quat2 /= np.linalg.norm(quat2) - mocap = np.hstack((pos1, quat1, pos2, quat2)) + control = np.hstack((pos1, pos2, quat1, quat2)) - state, sensordata = rollout.rollout(model, data, initial_state, mocap=mocap) + _, sensordata = rollout.rollout(model, data, initial_state, control, + control_spec=control_spec) - np.testing.assert_array_almost_equal(sensordata[:3], pos1) - np.testing.assert_array_almost_equal(sensordata[3:], quat2) + np.testing.assert_array_almost_equal(sensordata[0][0][:3], pos1) + np.testing.assert_array_almost_equal(sensordata[0][0][3:], quat1) - #----------------------------- test correctness + # ---------------------------- test correctness def test_intercept_mj_errors(self): model = mujoco.MjModel.from_xml_string(TEST_XML) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - initial_state = np.zeros(model.nq + model.nv + model.na) - ctrl = np.zeros((3, model.nu)) + nroll = 1 + nstep = 3 + + initial_state = np.zeros((nroll, nstate)) + ctrl = np.zeros((nroll, nstep, model.nu)) model.opt.solver = 10 # invalid solver type with self.assertRaisesWithLiteralMatch( mujoco.FatalError, 'mj_fwdConstraint: unknown solver type 10'): - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + rollout.rollout(model, data, initial_state, ctrl) def test_invalid(self): model = mujoco.MjModel.from_xml_string(TEST_XML) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - initial_state = np.zeros(model.nq + model.nv + model.na) + nroll = 1 - ctrl = 'string' - with self.assertRaisesWithLiteralMatch( - ValueError, 'ctrl must be a numpy array or float'): - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + initial_state = np.zeros((nroll, nstate)) - qfrc_applied = np.zeros((2, 3, 4, 5)) + control = 'string' with self.assertRaisesWithLiteralMatch( - ValueError, 'qfrc_applied can have at most 3 dimensions'): - state, sensordata = rollout.rollout(model, data, initial_state, - qfrc_applied=qfrc_applied) + ValueError, 'control must be a numpy array or float'): + rollout.rollout(model, data, initial_state, control) + + control = np.zeros((2, 3, 4, 5)) + with self.assertRaisesWithLiteralMatch( + ValueError, 'control can have at most 3 dimensions'): + rollout.rollout(model, data, initial_state, control) def test_bad_sizes(self): model = mujoco.MjModel.from_xml_string(TEST_XML) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - initial_state = np.random.randn(model.nq + model.nv + model.na+1) - with self.assertRaisesWithLiteralMatch( - ValueError, 'trailing dimension of initial_state must be 5, got 6'): - state, sensordata = rollout.rollout(model, data, initial_state) + nroll = 1 + nstep = 3 - initial_state = np.random.randn(model.nq + model.nv + model.na) - ctrl = np.random.randn(model.nu+1) + initial_state = np.random.randn(nroll, nstate + 1) with self.assertRaisesWithLiteralMatch( - ValueError, 'trailing dimension of ctrl must be 2, got 3'): - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + ValueError, 'trailing dimension of initial_state must be 6, got 7'): + rollout.rollout(model, data, initial_state) - ctrl = np.random.randn(2, model.nu) - qfrc_applied = np.random.randn(3, model.nv) # incompatible horizon + initial_state = np.random.randn(nroll, nstate) + control = np.random.randn(1, nstep, model.nu + 1) with self.assertRaisesWithLiteralMatch( - ValueError, 'dimension 1 inferred as 2 but qfrc_applied has 3'): - state, sensordata = rollout.rollout(model, data, initial_state, ctrl, - qfrc_applied=qfrc_applied) + ValueError, 'trailing dimension of control must be 2, got 3'): + rollout.rollout(model, data, initial_state, control) + + control = np.random.randn(nroll, nstep, model.nu) + state = np.random.randn(nroll, nstep+1, nstate) # incompatible nstep + with self.assertRaisesWithLiteralMatch( + ValueError, 'dimension 1 inferred as 3 but state has 4'): + rollout.rollout(model, data, initial_state, control, state=state) + + initial_state = np.random.randn(nroll, nstate) + control = np.random.randn(nroll, nstep, model.nu) + bad_spec = mujoco.mjtState.mjSTATE_ACT + with self.assertRaisesWithLiteralMatch( + ValueError, 'control_spec can only contain bits in mjSTATE_USER'): + rollout.rollout(model, data, initial_state, control, + control_spec=bad_spec) def test_stateless(self): model = mujoco.MjModel.from_xml_string(TEST_XML) - model.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_WARMSTART.value + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) data = mujoco.MjData(model) - # call step with a clean mjData - initial_state = np.random.randn(model.nq + model.nv + model.na) - ctrl = np.random.randn(model.nu) - state, sensordata = rollout.rollout(model, data, initial_state, ctrl) + # step with a clean mjData + initial_state = np.random.randn(nstate) + control = np.random.randn(3, 3, model.nu) + state, sensordata = rollout.rollout(model, data, initial_state, control) - # fill mjData with some debug value, see that we still get the same outputs - mujoco.mj_resetDataDebug(model, data, 255) - debug_state, debug_sensordata = rollout.rollout(model, data, initial_state, - ctrl) + # fill user fields with random values + for attr in [ + 'ctrl', + 'qfrc_applied', + 'xfrc_applied', + 'mocap_pos', + 'mocap_quat', + ]: + setattr(data, attr, np.random.randn(*getattr(data, attr).shape)) - np.testing.assert_array_equal(state, debug_state) - np.testing.assert_array_equal(sensordata, debug_sensordata) + # roll out again + state2, sensordata2 = rollout.rollout(model, data, initial_state, control) + + # assert that we still get the same outputs + np.testing.assert_array_equal(state, state2) + np.testing.assert_array_equal(sensordata, sensordata2) -#--------------- Python implementation of rollout functionality ---------------- +# -------------- Python implementation of rollout functionality ---------------- -def get_state(data): - return np.hstack((data.qpos, data.qvel, data.act)) -def set_state(model, data, state): - data.qpos = state[:model.nq] - data.qvel = state[model.nq:model.nq+model.nv] - data.act = state[model.nq+model.nv:model.nq+model.nv+model.na] +def get_state(model, data): + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) + state = np.empty(nstate) + mujoco.mj_getState(model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS) + return state.reshape((1, nstate)) -def step(model, data, state, **kwargs): + +def step(model, data, state, control, + control_spec=mujoco.mjtState.mjSTATE_CTRL): if state is not None: - set_state(model, data, state) - for key, value in kwargs.items(): - if value is not None: - setattr(data, key, np.reshape(value, getattr(data, key).shape)) + mujoco.mj_setState(model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS) + mujoco.mj_setState(model, data, control, control_spec) mujoco.mj_step(model, data) - return (get_state(data), data.sensordata) + return (get_state(model, data), data.sensordata) -def single_rollout(model, data, initial_state, **kwargs): - arg_nstep = set([a.shape[0] for a in kwargs.values()]) - assert len(arg_nstep) == 1 # nstep dimensions must match - nstep = arg_nstep.pop() - state = np.empty((nstep, model.nq + model.nv + model.na)) +def one_rollout(model, data, initial_state, control, + control_spec=mujoco.mjtState.mjSTATE_CTRL): + nstep = control.shape[0] + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) + state = np.empty((nstep, nstate)) sensordata = np.empty((nstep, model.nsensordata)) mujoco.mj_resetData(model, data) for t in range(nstep): - kwargs_t = {} - for key, value in kwargs.items(): - kwargs_t[key] = value[0 if value.ndim == 1 else t] state[t], sensordata[t] = step(model, data, - initial_state if t==0 else None, - **kwargs_t) + initial_state if t == 0 else None, + control[t], control_spec) return state, sensordata -def multi_rollout(model, data, initial_state, **kwargs): - nstate = initial_state.shape[0] - arg_nstep = set([a.shape[1] for a in kwargs.values()]) - assert len(arg_nstep) == 1 # nstep dimensions must match - nstep = arg_nstep.pop() - state = np.empty((nstate, nstep, model.nq + model.nv + model.na)) - sensordata = np.empty((nstate, nstep, model.nsensordata)) - for s in range(nstate): - kwargs_s = {key : value[s] for key, value in kwargs.items()} - state_s, sensordata_s = single_rollout(model, data, initial_state[s], - **kwargs_s) - state[s] = state_s - sensordata[s] = sensordata_s - return state.squeeze(), sensordata.squeeze() +def ensure_2d(arg): + if arg is None: + return None + else: + return np.ascontiguousarray(np.atleast_2d(arg), dtype=np.float64) + + +def ensure_3d(arg): + if arg is None: + return None + else: + # np.atleast_3d adds both leading and trailing dims, we want only leading + if arg.ndim == 0: + arg = arg[np.newaxis, np.newaxis, np.newaxis, ...] + elif arg.ndim == 1: + arg = arg[np.newaxis, np.newaxis, ...] + elif arg.ndim == 2: + arg = arg[np.newaxis, ...] + return np.ascontiguousarray(arg, dtype=np.float64) + + +def py_rollout(model, data, initial_state, control, + control_spec=mujoco.mjtState.mjSTATE_CTRL): + initial_state = ensure_2d(initial_state) + control = ensure_3d(control) + nroll = initial_state.shape[0] + nstep = control.shape[1] + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) + + state = np.empty((nroll, nstep, nstate)) + sensordata = np.empty((nroll, nstep, model.nsensordata)) + for r in range(nroll): + state_r, sensordata_r = one_rollout( + model, data, initial_state[r], control[r], control_spec + ) + state[r] = state_r + sensordata[r] = sensordata_r + return state, sensordata + if __name__ == '__main__': absltest.main()