# 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""" """ TEST_XML_NO_SENSORS = r""" """ TEST_XML_NO_ACTUATORS = r""" """ TEST_XML_MOCAP = r""" """ TEST_XML_EMPTY = r""" """ TEST_XML_DIVERGE = r""" """ 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()