From 943eb6bc7e8b405029f0e6e7bd94e776f38b6bb3 Mon Sep 17 00:00:00 2001 From: Levi Burner Date: Sun, 24 Nov 2024 16:50:28 -0500 Subject: [PATCH 1/5] rollout accepts a list of models of length nroll --- doc/changelog.rst | 1 + doc/python.rst | 1 + python/mujoco/rollout.cc | 55 ++++++++++++++++++----------------- python/mujoco/rollout.py | 47 +++++++++++++++++++++--------- python/mujoco/rollout_test.py | 41 +++++++++++++++++++++++--- 5 files changed, 102 insertions(+), 43 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 17bdf077..36865344 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -16,6 +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``. Bug fixes ^^^^^^^^^ diff --git a/doc/python.rst b/doc/python.rst index 3af46b6d..5b82292e 100644 --- a/doc/python.rst +++ b/doc/python.rst @@ -700,6 +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``. ``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 c250751f..a738c87a 100644 --- a/python/mujoco/rollout.cc +++ b/python/mujoco/rollout.cc @@ -37,7 +37,7 @@ const auto rollout_doc = R"( Roll out open-loop trajectories from initial states, get resulting states and sensor values. input arguments (required): - model instance of MjModel + model list of MjModel instances of length nroll data associated instance of MjData nstep integer, number of steps to be taken for each trajectory control_spec specification of controls, ncontrol = mj_stateSize(m, control_spec) @@ -54,18 +54,18 @@ Roll out open-loop trajectories from initial states, get resulting states and se // 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 nroll, int nstep, unsigned int control_spec, +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) { // 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; + int nstate = mj_stateSize(m[0], mjSTATE_FULLPHYSICS); + int ncontrol = mj_stateSize(m[0], control_spec); + int nv = m[0]->nv, nbody = m[0]->nbody, neq = m[0]->neq; + int nsensordata = m[0]->nsensordata; // clear user inputs if unspecified if (!(control_spec & mjSTATE_CTRL)) { - mju_zero(d->ctrl, m->nu); + mju_zero(d->ctrl, m[0]->nu); } if (!(control_spec & mjSTATE_QFRC_APPLIED)) { mju_zero(d->qfrc_applied, nv); @@ -75,26 +75,26 @@ void _unsafe_rollout(const mjModel* m, mjData* d, int nroll, int nstep, unsigned } 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); + 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->body_mocapid[i]; - if (id >= 0) mju_copy4(d->mocap_quat+4*id, m->body_quat+4*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->eq_active0[i]; + d->eq_active[i] = m[0]->eq_active0[i]; } } // loop over rollouts for (int r = 0; r < nroll; r++) { // set initial state - mj_setState(m, d, state0 + r*nstate, mjSTATE_FULLPHYSICS); + mj_setState(m[r], d, state0 + r*nstate, mjSTATE_FULLPHYSICS); // set warmstart accelerations if (warmstart0) { @@ -124,7 +124,7 @@ void _unsafe_rollout(const mjModel* m, mjData* d, int nroll, int nstep, unsigned for (; t < nstep; t++) { int step = r*nstep + t; if (state) { - mj_getState(m, d, state + step*nstate, mjSTATE_FULLPHYSICS); + mj_getState(m[r], d, state + step*nstate, mjSTATE_FULLPHYSICS); } if (sensordata) { mju_copy(sensordata + step*nsensordata, d->sensordata, nsensordata); @@ -137,15 +137,15 @@ void _unsafe_rollout(const mjModel* m, mjData* d, int nroll, int nstep, unsigned // controls if (control) { - mj_setState(m, d, control + step*ncontrol, control_spec); + mj_setState(m[r], d, control + step*ncontrol, control_spec); } // step - mj_step(m, d); + mj_step(m[r], d); // copy out new state if (state) { - mj_getState(m, d, state + step*nstate, mjSTATE_FULLPHYSICS); + mj_getState(m[r], d, state + step*nstate, mjSTATE_FULLPHYSICS); } // copy out sensor values @@ -188,7 +188,7 @@ PYBIND11_MODULE(_rollout, pymodule) { // get subsequent states and corresponding sensor values pymodule.def( "rollout", - [](const MjModelWrapper& m, MjDataWrapper& d, + [](py::list m, MjDataWrapper& d, int nstep, unsigned int control_spec, const PyCArray state0, std::optional warmstart0, @@ -196,7 +196,12 @@ PYBIND11_MODULE(_rollout, pymodule) { std::optional state, std::optional sensordata ) { - const raw::MjModel* model = m.get(); + // get raw pointers + int nroll = state0.shape(0); + const raw::MjModel* 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 @@ -205,19 +210,17 @@ PYBIND11_MODULE(_rollout, pymodule) { } // get sizes - int nstate = mj_stateSize(model, mjSTATE_FULLPHYSICS); - int ncontrol = mj_stateSize(model, control_spec); - int nroll = state0.shape(0); + int nstate = mj_stateSize(model_ptrs[0], mjSTATE_FULLPHYSICS); + int ncontrol = mj_stateSize(model_ptrs[0], control_spec); - // get raw pointers mjtNum* state0_ptr = get_array_ptr(state0, "state0", nroll, 1, nstate); mjtNum* warmstart0_ptr = get_array_ptr(warmstart0, "warmstart0", nroll, - 1, model->nv); + 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->nsensordata); + nstep, model_ptrs[0]->nsensordata); // perform rollouts { @@ -226,7 +229,7 @@ PYBIND11_MODULE(_rollout, pymodule) { // call unsafe rollout function InterceptMjErrors(_unsafe_rollout)( - model, data, nroll, nstep, control_spec, state0_ptr, + model_ptrs, data, nroll, nstep, control_spec, state0_ptr, warmstart0_ptr, control_ptr, state_ptr, sensordata_ptr); } }, diff --git a/python/mujoco/rollout.py b/python/mujoco/rollout.py index 9321f62f..25852652 100644 --- a/python/mujoco/rollout.py +++ b/python/mujoco/rollout.py @@ -14,7 +14,7 @@ # ============================================================================== """Roll out open-loop trajectories from initial states, get subsequent states and sensor values.""" -from typing import Optional +from typing import Optional, Union import mujoco from mujoco import _rollout @@ -22,7 +22,7 @@ import numpy as np from numpy import typing as npt -def rollout(model: mujoco.MjModel, +def rollout(model: Union[mujoco.MjModel, list[mujoco.MjModel]], data: mujoco.MjData, initial_state: npt.ArrayLike, control: Optional[npt.ArrayLike] = None, @@ -41,7 +41,7 @@ def rollout(model: mujoco.MjModel, Allocates outputs if none are given. Args: - model: An mjModel instance. + model: An mjModel or a list 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) @@ -90,6 +90,7 @@ def rollout(model: mujoco.MjModel, state=state, sensordata=sensordata) + # check number of dimensions _check_number_of_dimensions(2, initial_state=initial_state, @@ -108,14 +109,6 @@ def rollout(model: mujoco.MjModel, state = _ensure_3d(state) sensordata = _ensure_3d(sensordata) - # check trailing dimensions - 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 nroll, check for incompatibilities nroll = _infer_dimension(0, 1, initial_state=initial_state, @@ -123,6 +116,14 @@ def rollout(model: mujoco.MjModel, 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} ' + f'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, @@ -130,7 +131,27 @@ def rollout(model: mujoco.MjModel, state=state, sensordata=sensordata) - # tile input arrays if required (singleton expansion) + # 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) @@ -139,7 +160,7 @@ def rollout(model: mujoco.MjModel, if state is None: state = np.empty((nroll, nstep, nstate)) if sensordata is None: - sensordata = np.empty((nroll, nstep, model.nsensordata)) + sensordata = np.empty((nroll, nstep, nsensordata)) # call rollout _rollout.rollout(model, data, nstep, control_spec, initial_state, diff --git a/python/mujoco/rollout_test.py b/python/mujoco/rollout_test.py index 4c2563e9..8f118d6b 100644 --- a/python/mujoco/rollout_test.py +++ b/python/mujoco/rollout_test.py @@ -334,6 +334,34 @@ class MuJoCoRolloutTest(parameterized.TestCase): np.testing.assert_array_equal(state, py_state) np.testing.assert_array_equal(sensordata, py_sensordata) + @parameterized.parameters(ALL_MODELS.keys()) + def test_multi_model(self, model_name): + nroll = 3 # number of initial states and models + nstep = 3 # number of timesteps + + spec = mujoco.MjSpec.from_string(ALL_MODELS[model_name]) + + if len(spec.bodies) > 1: + model = [] + for i in range(nroll): + body = spec.bodies[1] + assert body.name != 'world' + body.pos = body.pos + i + model.append(spec.compile()) + else: + model = [spec.compile() for i in range(nroll)] + + nstate = mujoco.mj_stateSize(model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS) + data = mujoco.MjData(model[0]) + + initial_state = np.random.randn(nroll, nstate) + control = np.random.randn(nroll, nstep, model[0].nu) + state, sensordata = rollout.rollout(model, data, initial_state, control) + + 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]) @@ -430,8 +458,9 @@ class MuJoCoRolloutTest(parameterized.TestCase): def thread_initializer(): thread_local.data = mujoco.MjData(model) + model_list = [model]*nroll def call_rollout(initial_state, control, state, sensordata): - rollout.rollout(model, thread_local.data, initial_state, control, + rollout.rollout(model_list, thread_local.data, initial_state, control, skip_checks=True, nstep=nstep, state=state, sensordata=sensordata) @@ -677,13 +706,17 @@ def py_rollout(model, data, initial_state, control, control = ensure_3d(control) nroll = initial_state.shape[0] nstep = control.shape[1] - nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) + + if isinstance(model, mujoco.MjModel): + model = [model]*nroll + + nstate = mujoco.mj_stateSize(model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS) state = np.empty((nroll, nstep, nstate)) - sensordata = np.empty((nroll, nstep, model.nsensordata)) + sensordata = np.empty((nroll, nstep, model[0].nsensordata)) for r in range(nroll): state_r, sensordata_r = one_rollout( - model, data, initial_state[r], control[r], control_spec + model[r], data, initial_state[r], control[r], control_spec ) state[r] = state_r sensordata[r] = sensordata_r From bd361447bb1cdb41bb5615ee5dd8a55d5efe770d Mon Sep 17 00:00:00 2001 From: Levi Burner Date: Tue, 26 Nov 2024 08:22:28 -0500 Subject: [PATCH 2/5] merge items in changelog --- doc/changelog.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 36865344..3449b398 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -16,7 +16,8 @@ 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``. +- :ref:`rollout` can now accept lists of MjModel of length ``nroll``. ``nroll`` argument deprecated because + its value can always be inferred. Bug fixes ^^^^^^^^^ From 5ba4d5849716389d7ac7e0a2095435c521e06007 Mon Sep 17 00:00:00 2001 From: Levi Burner Date: Tue, 26 Nov 2024 13:16:05 -0500 Subject: [PATCH 3/5] replace array on stack with std::vector --- python/mujoco/rollout.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/mujoco/rollout.cc b/python/mujoco/rollout.cc index a738c87a..838ecc5b 100644 --- a/python/mujoco/rollout.cc +++ b/python/mujoco/rollout.cc @@ -54,7 +54,7 @@ Roll out open-loop trajectories from initial states, get resulting states and se // 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 nroll, int nstep, unsigned int control_spec, +void _unsafe_rollout(std::vector& m, mjData* d, int nroll, int nstep, unsigned int control_spec, const mjtNum* state0, const mjtNum* warmstart0, const mjtNum* control, mjtNum* state, mjtNum* sensordata) { // sizes @@ -198,7 +198,8 @@ PYBIND11_MODULE(_rollout, pymodule) { ) { // get raw pointers int nroll = state0.shape(0); - const raw::MjModel* model_ptrs[nroll]; + std::vector model_ptrs; + model_ptrs.reserve(nroll); for (int r = 0; r < nroll; r++) { model_ptrs[r] = m[r].cast()->get(); } From 41b325785a630d45e80101667262f6be8d3a1f3d Mon Sep 17 00:00:00 2001 From: Levi Burner Date: Tue, 26 Nov 2024 16:42:08 -0500 Subject: [PATCH 4/5] fix creation of std::vector --- python/mujoco/rollout.cc | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/mujoco/rollout.cc b/python/mujoco/rollout.cc index 838ecc5b..05576ba8 100644 --- a/python/mujoco/rollout.cc +++ b/python/mujoco/rollout.cc @@ -198,8 +198,7 @@ PYBIND11_MODULE(_rollout, pymodule) { ) { // get raw pointers int nroll = state0.shape(0); - std::vector model_ptrs; - model_ptrs.reserve(nroll); + std::vector model_ptrs(nroll); for (int r = 0; r < nroll; r++) { model_ptrs[r] = m[r].cast()->get(); } From e6eab04dd092554af7d37ac2f03faa40116ec862 Mon Sep 17 00:00:00 2001 From: Levi Burner Date: Wed, 27 Nov 2024 07:02:16 -0500 Subject: [PATCH 5/5] rollout fixups and correctly initialize unspecified controls --- doc/changelog.rst | 2 +- doc/python.rst | 2 +- python/mujoco/rollout.cc | 36 ++++++++++++++++++----------------- python/mujoco/rollout.py | 10 +++++++--- python/mujoco/rollout_test.py | 2 +- 5 files changed, 29 insertions(+), 23 deletions(-) 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,