Merge pull request #2377 from aftersomemath:rollout-fix-ndata-check

PiperOrigin-RevId: 718181763
Change-Id: Ie88bf1d26c77478ce3a5ca8fdf29ae90e54b5834

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<raw::MjData*> data_ptrs(py::len(d));

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)}'
       )

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 ----------------