Yuval Tassa
a0f014593b
- Update rollout.py to use mujoco.MJTNUM_DTYPE instead of np.float64. - Update rollout_test.py to use DTYPE for all state/control/sensor arrays to avoid copies and type mismatches. PiperOrigin-RevId: 928539946 Change-Id: Iee2250c049de799cb0e23ffe5558fe6eea8b3fd8
1006 lines
32 KiB
Python
1006 lines
32 KiB
Python
# Copyright 2022 DeepMind Technologies Limited
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""tests for rollout function."""
|
|
|
|
import concurrent.futures
|
|
import copy
|
|
import os
|
|
import threading
|
|
|
|
from absl.testing import absltest
|
|
from absl.testing import parameterized
|
|
|
|
import mujoco
|
|
from mujoco import rollout
|
|
|
|
import numpy as np
|
|
|
|
|
|
DTYPE = mujoco.MJTNUM_DTYPE
|
|
|
|
|
|
def randn(*shape):
|
|
return np.random.standard_normal(shape).astype(DTYPE)
|
|
|
|
|
|
def empty(*shape):
|
|
return np.empty(shape, dtype=DTYPE)
|
|
|
|
|
|
def zeros(*shape):
|
|
return np.zeros(shape, dtype=DTYPE)
|
|
|
|
# -------------------------- models used for testing ---------------------------
|
|
|
|
TEST_XML = r"""
|
|
<mujoco>
|
|
<worldbody>
|
|
<light pos="0 0 2"/>
|
|
<geom type="plane" size="5 5 .1"/>
|
|
<body pos="0 0 .1">
|
|
<joint name="yaw" axis="0 0 1"/>
|
|
<joint name="pitch" axis="0 1 0"/>
|
|
<geom type="capsule" size=".02" fromto="0 0 0 1 0 0"/>
|
|
<geom type="box" pos="1 0 0" size=".1 .1 .1"/>
|
|
<site name="site" pos="1 0 0"/>
|
|
</body>
|
|
</worldbody>
|
|
<actuator>
|
|
<general joint="pitch" gainprm="100"/>
|
|
<general joint="yaw" dyntype="filter" dynprm="1" gainprm="100"/>
|
|
</actuator>
|
|
<sensor>
|
|
<accelerometer site="site"/>
|
|
</sensor>
|
|
</mujoco>
|
|
"""
|
|
|
|
TEST_XML_NO_SENSORS = r"""
|
|
<mujoco>
|
|
<worldbody>
|
|
<light pos="0 0 2"/>
|
|
<geom type="plane" size="5 5 .1"/>
|
|
<body pos="0 0 .1">
|
|
<joint name="yaw" axis="0 0 1"/>
|
|
<joint name="pitch" axis="0 1 0"/>
|
|
<geom type="capsule" size=".02" fromto="0 0 0 1 0 0"/>
|
|
<geom type="box" pos="1 0 0" size=".1 .1 .1"/>
|
|
<site name="site" pos="1 0 0"/>
|
|
</body>
|
|
</worldbody>
|
|
<actuator>
|
|
<general joint="pitch" gainprm="100"/>
|
|
<general joint="yaw" dyntype="filter" dynprm="1" gainprm="100"/>
|
|
</actuator>
|
|
</mujoco>
|
|
"""
|
|
|
|
TEST_XML_NO_ACTUATORS = r"""
|
|
<mujoco>
|
|
<worldbody>
|
|
<light pos="0 0 2"/>
|
|
<geom type="plane" size="5 5 .1"/>
|
|
<body pos="0 0 .1">
|
|
<joint name="yaw" axis="0 0 1"/>
|
|
<joint name="pitch" axis="0 1 0"/>
|
|
<geom type="capsule" size=".02" fromto="0 0 0 1 0 0"/>
|
|
<geom type="box" pos="1 0 0" size=".1 .1 .1"/>
|
|
<site name="site" pos="1 0 0"/>
|
|
</body>
|
|
</worldbody>
|
|
<sensor>
|
|
<accelerometer site="site"/>
|
|
</sensor>
|
|
</mujoco>
|
|
"""
|
|
|
|
TEST_XML_MOCAP = r"""
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="1" mocap="true">
|
|
</body>
|
|
<body name="2" mocap="true">
|
|
</body>
|
|
</worldbody>
|
|
<sensor>
|
|
<framepos objtype="xbody" objname="1"/>
|
|
<framequat objtype="xbody" objname="1"/>
|
|
</sensor>
|
|
</mujoco>
|
|
"""
|
|
|
|
TEST_XML_EMPTY = r"""
|
|
<mujoco>
|
|
</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 -----------------------------------------
|
|
|
|
|
|
class MuJoCoRolloutTest(parameterized.TestCase):
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
np.random.seed(42)
|
|
|
|
# ----------------------------- 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 = randn(nstate)
|
|
control = randn(model.nu)
|
|
state, sensordata = rollout.rollout(model, data, initial_state, control)
|
|
|
|
mujoco.mj_resetData(model, data)
|
|
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(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 = randn(nstate)
|
|
control = randn(nstep, model.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_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)
|
|
|
|
nbatch = 5 # number of rollouts
|
|
nstep = 1 # number of steps
|
|
|
|
initial_state = randn(nbatch, nstate)
|
|
control = randn(nbatch, nstep, model.nu)
|
|
state, sensordata = rollout.rollout(model, data, initial_state, control)
|
|
|
|
mujoco.mj_resetData(model, data)
|
|
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_nbatch_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)
|
|
|
|
nbatch = 5 # number of rollouts
|
|
nstep = 1 # number of steps
|
|
|
|
initial_state = randn(nbatch, nstate)
|
|
control = randn(nstep, model.nu)
|
|
state, sensordata = rollout.rollout(model, data, initial_state, control)
|
|
|
|
mujoco.mj_resetData(model, data)
|
|
control = np.tile(control, (nbatch, 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_nbatch_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)
|
|
|
|
nbatch = 5 # number of rollouts
|
|
nstep = 1 # number of steps
|
|
|
|
initial_state = randn(nstate)
|
|
control = randn(nbatch, nstep, model.nu)
|
|
state, sensordata = rollout.rollout(model, data, initial_state, control)
|
|
|
|
mujoco.mj_resetData(model, data)
|
|
initial_state = np.tile(initial_state, (nbatch, 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_nbatch_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)
|
|
|
|
nbatch = 5 # number of rollouts
|
|
nstep = 1 # number of steps
|
|
|
|
initial_state = randn(nstate)
|
|
control = randn(nstep, model.nu)
|
|
initial_warmstart = np.tile(data.qacc_warmstart.copy(), (nbatch, 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, (nbatch, 1))
|
|
control = np.tile(control, (nbatch, 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_nbatch_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)
|
|
|
|
nbatch = 5 # number of rollouts
|
|
nstep = 1 # number of steps
|
|
|
|
initial_state = randn(nstate)
|
|
control = randn(nstep, model.nu)
|
|
state = empty(nbatch, nstep, nstate)
|
|
state, sensordata = rollout.rollout(
|
|
model, data, initial_state, control, state=state
|
|
)
|
|
|
|
mujoco.mj_resetData(model, data)
|
|
initial_state = np.tile(initial_state, (nbatch, 1))
|
|
control = np.tile(control, (nbatch, 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_nbatch_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)
|
|
|
|
nbatch = 5 # number of rollouts
|
|
nstep = 1 # number of steps
|
|
|
|
initial_state = randn(nstate)
|
|
control = randn(nstep, model.nu)
|
|
sensordata = empty(nbatch, 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, (nbatch, 1))
|
|
control = np.tile(control, (nbatch, 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])
|
|
nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
|
|
data = mujoco.MjData(model)
|
|
|
|
nbatch = 1 # number of rollouts
|
|
nstep = 3 # number of steps
|
|
|
|
initial_state = randn(nstate)
|
|
control = randn(model.nu)
|
|
state = empty(nbatch, nstep, nstate)
|
|
sensordata = empty(nbatch, nstep, model.nsensordata)
|
|
rollout.rollout(
|
|
model, data, initial_state, control, state=state, sensordata=sensordata
|
|
)
|
|
|
|
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)
|
|
|
|
nbatch = 2 # number of initial states
|
|
nstep = 3 # number of timesteps
|
|
|
|
initial_state = randn(nbatch, nstate)
|
|
control = randn(nbatch, nstep, model.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_model(self, model_name):
|
|
nbatch = 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(nbatch):
|
|
body = spec.bodies[1]
|
|
assert body.name != 'world'
|
|
body.pos = body.pos + i
|
|
model.append(spec.compile())
|
|
else:
|
|
model = [spec.compile() for _ in range(nbatch)]
|
|
|
|
nstate = mujoco.mj_stateSize(model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS)
|
|
data = mujoco.MjData(model[0])
|
|
|
|
initial_state = randn(nbatch, nstate)
|
|
control = randn(nbatch, 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])
|
|
nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
|
|
data = mujoco.MjData(model)
|
|
|
|
nbatch = 2 # number of rollouts
|
|
nstep = 3 # number of timesteps
|
|
|
|
initial_state = randn(nbatch, nstate)
|
|
control = randn(nbatch, 1, model.nu)
|
|
state = empty(nbatch, nstep, nstate)
|
|
state, sensordata = rollout.rollout(
|
|
model, data, initial_state, control, state=state
|
|
)
|
|
|
|
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.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)
|
|
|
|
nbatch = 4 # number of rollouts
|
|
nstep = 3 # number of timesteps
|
|
|
|
initial_state = randn(nbatch, nstate)
|
|
|
|
control_spec = (
|
|
mujoco.mjtState.mjSTATE_CTRL
|
|
| mujoco.mjtState.mjSTATE_QFRC_APPLIED
|
|
| mujoco.mjtState.mjSTATE_XFRC_APPLIED
|
|
)
|
|
ncontrol = mujoco.mj_stateSize(model, control_spec)
|
|
control = randn(nbatch, nstep, ncontrol)
|
|
|
|
state, sensordata = rollout.rollout(
|
|
model, data, initial_state, control, control_spec=control_spec
|
|
)
|
|
|
|
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)
|
|
|
|
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)
|
|
|
|
nbatch = 4 # number of rollouts
|
|
initial_state = empty(nbatch, 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 = randn(nbatch, 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
|
|
nbatch = 100
|
|
nstep = 5
|
|
initial_state = randn(nbatch, nstate)
|
|
state = empty(nbatch, nstep, nstate)
|
|
sensordata = empty(nbatch, nstep, model.nsensordata)
|
|
control = randn(nbatch, nstep, model.nu)
|
|
|
|
thread_local = threading.local()
|
|
|
|
def thread_initializer():
|
|
thread_local.data = mujoco.MjData(model)
|
|
|
|
model_list = [copy.copy(model) for _ in range(nbatch)]
|
|
|
|
def call_rollout(initial_state, control, state, sensordata):
|
|
rollout.rollout(
|
|
model_list,
|
|
[thread_local.data],
|
|
initial_state,
|
|
control,
|
|
skip_checks=True,
|
|
nstep=nstep,
|
|
state=state,
|
|
sensordata=sensordata,
|
|
)
|
|
|
|
n = nbatch // 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],
|
|
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 :],
|
|
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:
|
|
futures = []
|
|
for chunk in chunks:
|
|
futures.append(executor.submit(call_rollout, *chunk))
|
|
for future in concurrent.futures.as_completed(futures):
|
|
future.result()
|
|
|
|
data = mujoco.MjData(model)
|
|
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)
|
|
|
|
def test_threading_native(self):
|
|
model = mujoco.MjModel.from_xml_string(TEST_XML)
|
|
nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
|
|
num_workers = 32
|
|
nbatch = 100
|
|
nstep = 5
|
|
initial_state = randn(nbatch, nstate)
|
|
state = empty(nbatch, nstep, nstate)
|
|
sensordata = empty(nbatch, nstep, model.nsensordata)
|
|
control = randn(nbatch, nstep, model.nu)
|
|
|
|
model_list = [copy.copy(model) for _ in range(nbatch)]
|
|
data_list = [mujoco.MjData(model) for _ in range(num_workers)]
|
|
|
|
rollout.rollout(
|
|
model_list,
|
|
data_list,
|
|
initial_state,
|
|
control,
|
|
nstep=nstep,
|
|
state=state,
|
|
sensordata=sensordata,
|
|
)
|
|
|
|
data = mujoco.MjData(model)
|
|
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)
|
|
|
|
def test_threading_native_persistent_object(self):
|
|
model = mujoco.MjModel.from_xml_string(TEST_XML)
|
|
nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
|
|
num_workers = 32
|
|
nbatch = 100
|
|
nstep = 5
|
|
initial_state = randn(nbatch, nstate)
|
|
state = empty(nbatch, nstep, nstate)
|
|
sensordata = empty(nbatch, nstep, model.nsensordata)
|
|
control = randn(nbatch, nstep, model.nu)
|
|
|
|
model_list = [copy.copy(model) for _ in range(nbatch)]
|
|
data_list = [mujoco.MjData(model) for _ in range(num_workers)]
|
|
|
|
with rollout.Rollout(nthread=num_workers) as rollout_:
|
|
for _ in range(2):
|
|
rollout_.rollout(
|
|
model_list,
|
|
data_list,
|
|
initial_state,
|
|
control,
|
|
nstep=nstep,
|
|
state=state,
|
|
sensordata=sensordata,
|
|
)
|
|
|
|
data = mujoco.MjData(model)
|
|
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)
|
|
|
|
rollout_ = rollout.Rollout(nthread=num_workers)
|
|
for _ in range(2):
|
|
rollout_.rollout(
|
|
model_list,
|
|
data_list,
|
|
initial_state,
|
|
control,
|
|
nstep=nstep,
|
|
state=state,
|
|
sensordata=sensordata,
|
|
)
|
|
|
|
data = mujoco.MjData(model)
|
|
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)
|
|
rollout_.close()
|
|
|
|
def test_threading_native_persistent_function(self):
|
|
model = mujoco.MjModel.from_xml_string(TEST_XML)
|
|
nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
|
|
num_workers = 32
|
|
nbatch = 100
|
|
nstep = 5
|
|
initial_state = randn(nbatch, nstate)
|
|
state = empty(nbatch, nstep, nstate)
|
|
sensordata = empty(nbatch, nstep, model.nsensordata)
|
|
control = randn(nbatch, nstep, model.nu)
|
|
|
|
model_list = [copy.copy(model) for _ in range(nbatch)]
|
|
data_list = [mujoco.MjData(model) for _ in range(num_workers)]
|
|
|
|
for _ in range(2):
|
|
rollout.rollout(
|
|
model_list,
|
|
data_list,
|
|
initial_state,
|
|
control,
|
|
nstep=nstep,
|
|
state=state,
|
|
sensordata=sensordata,
|
|
persistent_pool=True,
|
|
)
|
|
|
|
data = mujoco.MjData(model)
|
|
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)
|
|
rollout.shutdown_persistent_pool()
|
|
|
|
# ---------------------------- 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)
|
|
|
|
# take one step, save the state
|
|
state0 = zeros(nstate)
|
|
control = zeros(model.nu)
|
|
state1, _ = step(model, data, state0, control)
|
|
|
|
# save qacc_warmstart
|
|
initial_warmstart = data.qacc_warmstart.copy()
|
|
|
|
# take one more step (uses correct warmstart)
|
|
state2, _ = step(model, data, state1[0], control)
|
|
|
|
# take step using rollout, don't take warmstart into account
|
|
state, _ = rollout.rollout(model, data, state1[0], control)
|
|
|
|
# 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
|
|
)
|
|
|
|
# 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 = zeros(nstate)
|
|
|
|
control_spec = (
|
|
mujoco.mjtState.mjSTATE_MOCAP_POS | mujoco.mjtState.mjSTATE_MOCAP_QUAT
|
|
)
|
|
|
|
pos1 = np.array((1.0, 2.0, 3.0), dtype=DTYPE)
|
|
quat1 = np.array((1.0, 2.0, 3.0, 4.0), dtype=DTYPE)
|
|
quat1 /= np.linalg.norm(quat1)
|
|
pos2 = np.array((2.0, 3.0, 4.0), dtype=DTYPE)
|
|
quat2 = np.array((2.0, 3.0, 4.0, 5.0), dtype=DTYPE)
|
|
quat2 /= np.linalg.norm(quat2)
|
|
control = np.hstack((pos1, pos2, quat1, quat2))
|
|
|
|
_, sensordata = rollout.rollout(
|
|
model, data, initial_state, control, control_spec=control_spec
|
|
)
|
|
|
|
np.testing.assert_array_almost_equal(sensordata[0][0][:3], pos1)
|
|
np.testing.assert_array_almost_equal(sensordata[0][0][3:], quat1)
|
|
|
|
# ---------------------------- 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)
|
|
|
|
nbatch = 1
|
|
nstep = 3
|
|
|
|
initial_state = zeros(nbatch, nstate)
|
|
ctrl = zeros(nbatch, nstep, model.nu)
|
|
|
|
model.opt.solver = 10 # invalid solver type
|
|
with self.assertRaisesWithLiteralMatch(
|
|
mujoco.FatalError, 'mj_fwdConstraint: unknown solver type 10'
|
|
):
|
|
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)
|
|
|
|
nbatch = 1
|
|
|
|
initial_state = zeros(nbatch, nstate)
|
|
|
|
control = 'string'
|
|
with self.assertRaisesWithLiteralMatch(
|
|
ValueError, 'control must be a numpy array or float'
|
|
):
|
|
rollout.rollout(model, data, initial_state, control)
|
|
|
|
control = 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)
|
|
|
|
nbatch = 1
|
|
nstep = 3
|
|
|
|
initial_state = randn(nbatch, nstate + 1)
|
|
with self.assertRaisesWithLiteralMatch(
|
|
ValueError, 'trailing dimension of initial_state must be 6, got 7'
|
|
):
|
|
rollout.rollout(model, data, initial_state)
|
|
|
|
initial_state = randn(nbatch, nstate)
|
|
control = randn(1, nstep, model.nu + 1)
|
|
with self.assertRaisesWithLiteralMatch(
|
|
ValueError, 'trailing dimension of control must be 2, got 3'
|
|
):
|
|
rollout.rollout(model, data, initial_state, control)
|
|
|
|
control = randn(nbatch, nstep, model.nu)
|
|
state = randn(nbatch, 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 = randn(nbatch, nstate)
|
|
control = randn(nbatch, 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)
|
|
nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
|
|
data = mujoco.MjData(model)
|
|
|
|
# step with a clean mjData
|
|
initial_state = randn(nstate)
|
|
control = randn(3, 3, model.nu)
|
|
state, sensordata = rollout.rollout(model, data, initial_state, control)
|
|
|
|
# fill user fields with random values
|
|
for attr in [
|
|
'ctrl',
|
|
'qfrc_applied',
|
|
'xfrc_applied',
|
|
'mocap_pos',
|
|
'mocap_quat',
|
|
]:
|
|
setattr(data, attr, randn(*getattr(data, attr).shape))
|
|
|
|
# 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)
|
|
|
|
def test_length_one_model_list(self):
|
|
model = mujoco.MjModel.from_xml_string(TEST_XML)
|
|
nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
|
|
data = mujoco.MjData(model)
|
|
|
|
initial_state = randn(nstate)
|
|
control = randn(3, 3, model.nu)
|
|
|
|
state, sensordata = rollout.rollout(model, data, initial_state, control)
|
|
state2, sensordata2 = rollout.rollout([model], data, initial_state, control)
|
|
|
|
# assert that we get same outputs
|
|
np.testing.assert_array_equal(state, state2)
|
|
np.testing.assert_array_equal(sensordata, sensordata2)
|
|
|
|
def test_data_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 = randn(nstate)
|
|
control = randn(3, 3, model.nu)
|
|
|
|
# Test passing empty lists for data
|
|
with self.assertRaisesWithLiteralMatch(
|
|
ValueError, 'The list of data instances is empty'
|
|
):
|
|
rollout.rollout(model, [], initial_state, control)
|
|
|
|
with self.assertRaisesWithLiteralMatch(
|
|
ValueError, 'The list of data instances is empty'
|
|
):
|
|
with rollout.Rollout(nthread=0) as rollout_:
|
|
rollout_.rollout(model, [], initial_state, control)
|
|
|
|
with self.assertRaisesWithLiteralMatch(
|
|
ValueError, 'The list of data instances is empty'
|
|
):
|
|
with rollout.Rollout(nthread=1) as rollout_:
|
|
rollout_.rollout(model, [], initial_state, control)
|
|
|
|
with self.assertRaisesWithLiteralMatch(
|
|
ValueError, 'The list of data instances is empty'
|
|
):
|
|
with rollout.Rollout(nthread=2) as rollout_:
|
|
rollout_.rollout(model, [], initial_state, control)
|
|
|
|
# Test checking that len(data) equals nthread
|
|
with self.assertRaisesWithLiteralMatch(
|
|
ValueError,
|
|
'More than one data instance passed but rollout is configured to run on'
|
|
' main thread',
|
|
):
|
|
with rollout.Rollout(nthread=0) as rollout_:
|
|
rollout_.rollout(
|
|
model, [copy.copy(data) for _ in range(2)], initial_state, control
|
|
)
|
|
|
|
with self.assertRaisesWithLiteralMatch(
|
|
ValueError, 'Length of data: 1 not equal to nthread: 2'
|
|
):
|
|
with rollout.Rollout(nthread=2) as rollout_:
|
|
rollout_.rollout(model, data, initial_state, control)
|
|
|
|
with self.assertRaisesWithLiteralMatch(
|
|
ValueError, 'Length of data: 1 not equal to nthread: 2'
|
|
):
|
|
with rollout.Rollout(nthread=2) as rollout_:
|
|
rollout_.rollout(model, [data], initial_state, control)
|
|
|
|
with self.assertRaisesWithLiteralMatch(
|
|
ValueError, 'Length of data: 3 not equal to nthread: 2'
|
|
):
|
|
with rollout.Rollout(nthread=2) as rollout_:
|
|
rollout_.rollout(
|
|
model, [copy.copy(data) for _ in range(3)], initial_state, control
|
|
)
|
|
|
|
@absltest.skip(reason='Takes a long time to run')
|
|
def test_large_state(self):
|
|
model = mujoco.MjModel.from_xml_string(TEST_XML)
|
|
nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
|
|
data = mujoco.MjData(model)
|
|
|
|
nthread = os.cpu_count()
|
|
nbatch = nthread
|
|
|
|
nstep = ((2**31) // (nstate * nbatch)) + 2
|
|
assert nstep * nstate * nbatch > 2**31
|
|
|
|
initial_state = randn(nbatch, nstate)
|
|
rollout.rollout(
|
|
model,
|
|
[copy.copy(data) for _ in range(nthread)],
|
|
initial_state,
|
|
nstep=nstep,
|
|
)
|
|
|
|
|
|
# -------------- Python implementation of rollout functionality ----------------
|
|
|
|
|
|
def get_state(model, data):
|
|
nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
|
|
state = empty(nstate)
|
|
mujoco.mj_getState(model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS)
|
|
return state.reshape((1, nstate))
|
|
|
|
|
|
def step(
|
|
model, data, state, control, control_spec=mujoco.mjtState.mjSTATE_CTRL
|
|
):
|
|
if state is not None:
|
|
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(model, data), data.sensordata)
|
|
|
|
|
|
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 = empty(nstep, nstate)
|
|
sensordata = empty(nstep, model.nsensordata)
|
|
|
|
mujoco.mj_resetData(model, data)
|
|
for t in range(nstep):
|
|
state[t], sensordata[t] = step(
|
|
model, data, initial_state if t == 0 else None, control[t], control_spec
|
|
)
|
|
return state, sensordata
|
|
|
|
|
|
def ensure_2d(arg):
|
|
if arg is None:
|
|
return None
|
|
else:
|
|
return np.ascontiguousarray(np.atleast_2d(arg), dtype=DTYPE)
|
|
|
|
|
|
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=DTYPE)
|
|
|
|
|
|
def py_rollout(
|
|
model,
|
|
data,
|
|
initial_state,
|
|
control,
|
|
control_spec=mujoco.mjtState.mjSTATE_CTRL,
|
|
):
|
|
initial_state = ensure_2d(initial_state)
|
|
control = ensure_3d(control)
|
|
nbatch = initial_state.shape[0]
|
|
nstep = control.shape[1]
|
|
|
|
if isinstance(model, mujoco.MjModel):
|
|
model = [copy.copy(model) for _ in range(nbatch)]
|
|
|
|
nstate = mujoco.mj_stateSize(model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS)
|
|
|
|
state = empty(nbatch, nstep, nstate)
|
|
sensordata = empty(nbatch, nstep, model[0].nsensordata)
|
|
for r in range(nbatch):
|
|
state_r, sensordata_r = one_rollout(
|
|
model[r], data, initial_state[r], control[r], control_spec
|
|
)
|
|
state[r] = state_r
|
|
sensordata[r] = sensordata_r
|
|
return state, sensordata
|
|
|
|
|
|
if __name__ == '__main__':
|
|
absltest.main()
|