rollout add native threading test

Author: aftersomemath

python/mujoco/rollout_test.py

@@ -519,6 +519,35 @@ def call_rollout(initial_state, control, state, sensordata):
     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
+    nroll = 10000
+    nstep = 5
+    initial_state = np.random.randn(nroll, nstate)
+    state = np.empty((nroll, nstep, nstate))
+    sensordata = np.empty((nroll, nstep, model.nsensordata))
+    control = np.random.randn(nroll, nstep, model.nu)
+
+    model_list = [model] * nroll
+    data_list = [mujoco.MjData(model) for i 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)
+
   # ---------------------------- test advanced operation
 
   def test_warmstart(self):