diff --git a/python/mujoco/rollout.cc b/python/mujoco/rollout.cc index 3f27f1d6..42519189 100644 --- a/python/mujoco/rollout.cc +++ b/python/mujoco/rollout.cc @@ -254,16 +254,20 @@ class Rollout { } // check length d and nthread are consistent - if (this->nthread_ == 0 && py::len(d) > 1) { + if (py::len(d) == 0) { + std::ostringstream msg; + msg << "The list of data instances is empty"; + throw py::value_error(msg.str()); + } else if (this->nthread_ == 0 && py::len(d) > 1) { std::ostringstream msg; msg << "More than one data instance passed but " << "rollout is configured to run on main thread"; - py::value_error(msg.str()); - } else if (this->nthread_ != py::len(d)) { + throw py::value_error(msg.str()); + } else if (this->nthread_ > 0 && this->nthread_ != py::len(d)) { std::ostringstream msg; msg << "Length of data: " << py::len(d) << " not equal to nthread: " << this->nthread_; - py::value_error(msg.str()); + throw py::value_error(msg.str()); } std::vector data_ptrs(py::len(d)); diff --git a/python/mujoco/rollout.py b/python/mujoco/rollout.py index 95b6ad3f..68a9be23 100644 --- a/python/mujoco/rollout.py +++ b/python/mujoco/rollout.py @@ -172,7 +172,7 @@ class Rollout: if isinstance(model, list) and nroll == 1: nroll = len(model) - if isinstance(model, list) and len(model) != nroll: + if isinstance(model, list) and len(model) > 1 and len(model) != nroll: raise ValueError( f'nroll inferred as {nroll} but model is length {len(model)}' ) diff --git a/python/mujoco/rollout_test.py b/python/mujoco/rollout_test.py index 23139ba8..5c8c5fbe 100644 --- a/python/mujoco/rollout_test.py +++ b/python/mujoco/rollout_test.py @@ -20,9 +20,10 @@ import threading from absl.testing import absltest from absl.testing import parameterized +import numpy as np + import mujoco from mujoco import rollout -import numpy as np # -------------------------- models used for testing --------------------------- @@ -796,6 +797,84 @@ class MuJoCoRolloutTest(parameterized.TestCase): 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 = np.random.randn(nstate) + control = np.random.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 = np.random.randn(nstate) + control = np.random.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 i 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 i in range(3)], initial_state, control + ) + # -------------- Python implementation of rollout functionality ----------------