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

Merge pull request #2246 from aftersomemath:remove-nroll

Browse Source 
PiperOrigin-RevId: 700128273
Change-Id: If9c5b5b0a18483730c8b4cf0bd6788ef1e021a29
This commit is contained in:
Copybara-Service
2024-11-25 16:05:29 -08:00
parent 7e46e21ef3 96a8f002bc
commit 6e129a5ab2
4 changed files with 112 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+1
View File
@@ -15,6 +15,7 @@ Python bindings
- Provide prebuilt wheels for Python 3.13.
- 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<PyRollout>` because its value can always be inferred.
Bug fixes
^^^^^^^^^
python/mujoco/rollout.cc
+3 -4
View File
@@ -39,7 +39,6 @@ Roll out open-loop trajectories from initial states, get resulting states and se
input arguments (required):
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,
@@ -190,7 +189,7 @@ PYBIND11_MODULE(_rollout, pymodule) {
pymodule.def(
"rollout",
[](const MjModelWrapper& m, MjDataWrapper& d,
int nroll, int nstep, unsigned int control_spec,
int nstep, unsigned int control_spec,
const PyCArray state0,
std::optional<const PyCArray> warmstart0,
std::optional<const PyCArray> control,
@@ -201,13 +200,14 @@ PYBIND11_MODULE(_rollout, pymodule) {
raw::MjData* data = d.get();
// check that some steps need to be taken, return if not
if (nroll < 1 || nstep < 1) {
if (nstep < 1) {
return;
}
// get sizes
int nstate = mj_stateSize(model, mjSTATE_FULLPHYSICS);
int ncontrol = mj_stateSize(model, control_spec);
int nroll = state0.shape(0);
// get raw pointers
mjtNum* state0_ptr = get_array_ptr(state0, "state0", nroll, 1, nstate);
@@ -232,7 +232,6 @@ PYBIND11_MODULE(_rollout, pymodule) {
},
py::arg("model"),
py::arg("data"),
py::arg("nroll"),
py::arg("nstep"),
py::arg("control_spec"),
py::arg("state0"),
python/mujoco/rollout.py
+3 -7
View File
@@ -29,7 +29,6 @@ def rollout(model: mujoco.MjModel,
*, # 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,
@@ -50,7 +49,6 @@ def rollout(model: mujoco.MjModel,
([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)
@@ -74,7 +72,7 @@ def rollout(model: mujoco.MjModel,
# don't allocate output arrays
# just call rollout and return
if skip_checks:
_rollout.rollout(model, data, nroll, nstep, control_spec, initial_state,
_rollout.rollout(model, data, nstep, control_spec, initial_state,
initial_warmstart, control, state, sensordata)
return state, sensordata
@@ -83,8 +81,6 @@ def rollout(model: mujoco.MjModel,
raise ValueError('control_spec can only contain bits in mjSTATE_USER')
# check types
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(
@@ -121,7 +117,7 @@ def rollout(model: mujoco.MjModel,
_check_trailing_dimension(model.nsensordata, sensordata=sensordata)
# infer nroll, check for incompatibilities
nroll = _infer_dimension(0, nroll or 1,
nroll = _infer_dimension(0, 1,
initial_state=initial_state,
initial_warmstart=initial_warmstart,
control=control,
@@ -146,7 +142,7 @@ def rollout(model: mujoco.MjModel,
sensordata = np.empty((nroll, nstep, model.nsensordata))
# call rollout
_rollout.rollout(model, data, nroll, nstep, control_spec, initial_state,
_rollout.rollout(model, data, nstep, control_spec, initial_state,
initial_warmstart, control, state, sensordata)
# return outputs
python/mujoco/rollout_test.py
+105 -1
View File
@@ -192,6 +192,110 @@ 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_infer_nroll_initial_state(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)
nroll = 5 # number of rollouts
nstep = 1 # number of steps
initial_state = np.random.randn(nroll, nstate)
control = np.random.randn(nstep, model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, control)
mujoco.mj_resetData(model, data)
control = np.tile(control, (nroll, 1, 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_infer_nroll_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)
nroll = 5 # number of rollouts
nstep = 1 # number of steps
initial_state = np.random.randn(nstate)
control = np.random.randn(nroll, nstep, model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, control)
mujoco.mj_resetData(model, data)
initial_state = np.tile(initial_state, (nroll, 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_infer_nroll_warmstart(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)
nroll = 5 # number of rollouts
nstep = 1 # number of steps
initial_state = np.random.randn(nstate)
control = np.random.randn(nstep, model.nu)
initial_warmstart = np.tile(data.qacc_warmstart.copy(), (nroll, 1))
state, sensordata = rollout.rollout(model, data, initial_state, control,
initial_warmstart=initial_warmstart)
mujoco.mj_resetData(model, data)
initial_state = np.tile(initial_state, (nroll, 1))
control = np.tile(control, (nroll, 1, 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_infer_nroll_state(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)
nroll = 5 # number of rollouts
nstep = 1 # number of steps
initial_state = np.random.randn(nstate)
control = np.random.randn(nstep, model.nu)
state = np.empty((nroll, nstep, nstate))
state, sensordata = rollout.rollout(model, data, initial_state, control,
state=state)
mujoco.mj_resetData(model, data)
initial_state = np.tile(initial_state, (nroll, 1))
control = np.tile(control, (nroll, 1, 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_infer_nroll_sensordata(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)
nroll = 5 # number of rollouts
nstep = 1 # number of steps
initial_state = np.random.randn(nstate)
control = np.random.randn(nstep, model.nu)
sensordata = np.empty((nroll, nstep, model.nsensordata))
state, sensordata = rollout.rollout(model, data, initial_state, control,
sensordata=sensordata)
mujoco.mj_resetData(model, data)
initial_state = np.tile(initial_state, (nroll, 1))
control = np.tile(control, (nroll, 1, 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_one_rollout_fixed_ctrl(self, model_name):
model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
@@ -328,7 +432,7 @@ class MuJoCoRolloutTest(parameterized.TestCase):
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],
skip_checks=True,
nstep=nstep, state=state, sensordata=sensordata)
n = nroll // num_workers # integer division