diff --git a/doc/changelog.rst b/doc/changelog.rst index 3449b398..f148d26b 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -16,7 +16,7 @@ Python bindings - Added ``bind`` method and removed id attribute from :ref:`mjSpec` objects. Using ids is error prone in scenarios of repeated attachment and detachment. Python users are encouraged to use names for unique identification of model elements. - Removed ``nroll`` argument from :ref:`rollout` because its value can always be inferred. -- :ref:`rollout` can now accept lists of MjModel of length ``nroll``. ``nroll`` argument deprecated because +- :ref:`rollout` can now accept sequences of MjModel of length ``nroll``. ``nroll`` argument deprecated because its value can always be inferred. Bug fixes diff --git a/doc/python.rst b/doc/python.rst index 5b82292e..5a97f377 100644 --- a/doc/python.rst +++ b/doc/python.rst @@ -700,7 +700,7 @@ states and sensor values. The basic usage form is state, sensordata = rollout.rollout(model, data, initial_state, control) -``model`` is either a single instance of MjModel or a list of compatible MjModel of length ``nroll``. +``model`` is either a single instance of MjModel or a sequence of compatible MjModel of length ``nroll``. ``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 diff --git a/python/mujoco/rollout.cc b/python/mujoco/rollout.cc index 05576ba8..ffedda3f 100644 --- a/python/mujoco/rollout.cc +++ b/python/mujoco/rollout.cc @@ -73,26 +73,28 @@ void _unsafe_rollout(std::vector& m, mjData* d, int nroll, int n 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[0]->body_mocapid[i]; - if (id >= 0) mju_copy3(d->mocap_pos+3*id, m[0]->body_pos+3*i); - } - } - if (!(control_spec & mjSTATE_MOCAP_QUAT)) { - for (int i = 0; i < nbody; i++) { - int id = m[0]->body_mocapid[i]; - if (id >= 0) mju_copy4(d->mocap_quat+4*id, m[0]->body_quat+4*i); - } - } - if (!(control_spec & mjSTATE_EQ_ACTIVE)) { - for (int i = 0; i < neq; i++) { - d->eq_active[i] = m[0]->eq_active0[i]; - } - } // loop over rollouts for (int r = 0; r < nroll; r++) { + // clear user inputs if unspecified + if (!(control_spec & mjSTATE_MOCAP_POS)) { + for (int i = 0; i < nbody; i++) { + int id = m[r]->body_mocapid[i]; + if (id >= 0) mju_copy3(d->mocap_pos+3*id, m[r]->body_pos+3*i); + } + } + if (!(control_spec & mjSTATE_MOCAP_QUAT)) { + for (int i = 0; i < nbody; i++) { + int id = m[r]->body_mocapid[i]; + if (id >= 0) mju_copy4(d->mocap_quat+4*id, m[r]->body_quat+4*i); + } + } + if (!(control_spec & mjSTATE_EQ_ACTIVE)) { + for (int i = 0; i < neq; i++) { + d->eq_active[i] = m[r]->eq_active0[i]; + } + } + // set initial state mj_setState(m[r], d, state0 + r*nstate, mjSTATE_FULLPHYSICS); diff --git a/python/mujoco/rollout.py b/python/mujoco/rollout.py index 25852652..f314e795 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.""" +from collections.abc import Sequence from typing import Optional, Union import mujoco @@ -22,7 +23,7 @@ import numpy as np from numpy import typing as npt -def rollout(model: Union[mujoco.MjModel, list[mujoco.MjModel]], +def rollout(model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]], data: mujoco.MjData, initial_state: npt.ArrayLike, control: Optional[npt.ArrayLike] = None, @@ -41,7 +42,7 @@ def rollout(model: Union[mujoco.MjModel, list[mujoco.MjModel]], Allocates outputs if none are given. Args: - model: An mjModel or a list of MjModel with the same size signature. + model: An mjModel or a sequence of MjModel with the same size signature. data: An associated mjData instance. initial_state: Array of initial states from which to roll out trajectories. ([nroll or 1] x nstate) @@ -76,6 +77,9 @@ def rollout(model: Union[mujoco.MjModel, list[mujoco.MjModel]], initial_warmstart, control, state, sensordata) 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') @@ -151,7 +155,7 @@ def rollout(model: Union[mujoco.MjModel, list[mujoco.MjModel]], _check_trailing_dimension(nsensordata, sensordata=sensordata) # tile input arrays/lists if required (singleton expansion) - model = model*nroll if len(model) == 1 else model + 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) diff --git a/python/mujoco/rollout_test.py b/python/mujoco/rollout_test.py index 8f118d6b..32c670e6 100644 --- a/python/mujoco/rollout_test.py +++ b/python/mujoco/rollout_test.py @@ -458,7 +458,7 @@ class MuJoCoRolloutTest(parameterized.TestCase): def thread_initializer(): thread_local.data = mujoco.MjData(model) - model_list = [model]*nroll + model_list = [model] * nroll def call_rollout(initial_state, control, state, sensordata): rollout.rollout(model_list, thread_local.data, initial_state, control, skip_checks=True,