chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improvements to mujoco.rollout:

Browse Source 
- `mjSTATE_FULLPHYSICS` as state spec, enabling divergence detection by inspecting time.
- User-defined control spec.
- Stop squeezing: outputs always have dim=3.

PiperOrigin-RevId: 600445256
Change-Id: I4466e88929cb7081e1c94968a5cfe10485bb7475
This commit is contained in:
Yuval Tassa
2024-01-22 07:01:01 -08:00
committed by Copybara-Service
parent 5cbaa23388
commit aceb52bd09
5 changed files with 564 additions and 427 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/rollout.cc
+116 -127
View File
@@ -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<const py::array_t<mjtNum>> 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<const py::array_t<mjtNum>> 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<mjtNum*>(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<const PyCArray> init_state,
std::optional<const PyCArray> init_time,
std::optional<const PyCArray> init_warmstart,
std::optional<const PyCArray> ctrl,
std::optional<const PyCArray> qfrc,
std::optional<const PyCArray> xfrc,
std::optional<const PyCArray> mocap,
[](const MjModelWrapper& m, MjDataWrapper& d,
int nroll, int nstep, unsigned int control_spec,
const PyCArray state0,
std::optional<const PyCArray> warmstart0,
std::optional<const PyCArray> control,
std::optional<const PyCArray> state,
std::optional<const PyCArray> 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
python/mujoco/rollout.py
+112 -77
View File
@@ -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
python/mujoco/rollout_test.py
+301 -199
View File
@@ -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"""
<mujoco>
@@ -96,7 +97,7 @@ TEST_XML_MOCAP = r"""
</worldbody>
<sensor>
<framepos objtype="xbody" objname="1"/>
<framequat objtype="xbody" objname="2"/>
<framequat objtype="xbody" objname="1"/>
</sensor>
</mujoco>
"""
@@ -106,12 +107,33 @@ TEST_XML_EMPTY = r"""
</mujoco>
"""
TEST_XML_DIVERGE = r"""
<mujoco>
<option>
<flag gravity="disable"/>
</option>
<worldbody>
<geom type="plane" size="5 5 .1"/>
<body pos="0 0 -.3" euler="30 45 90">
<freejoint/>
<geom type="box" size=".1 .2 .4"/>
</body>
</worldbody>
<keyframe>
<key name="non-diverging" qpos="0 0 .5 1 0 0 0"/>
</keyframe>
</mujoco>
"""
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()