rollout don't register atexit handler for Rollout objects

Author: aftersomemath

doc/python.rst

@@ -759,9 +759,11 @@ or
 
 .. code-block:: python
 
-  # pool shutdown on object deletion, interpreter shutdown, or call to rollout_.shutdown_pool
+  # pool shutdown on object deletion or call to rollout_.shutdown_pool
+  # to ensure clean shutdown of threads, call shutdown_pool before interpreter exit
   rollout_ = rollout.Rollout(nthread)
   rollout_.rollout(model, data, initial_state)
+  rollout_.shutdown_pool()
 
 Since the Global Interpreter Lock is released, this function can also be threaded using Python threads. However, this
 is less efficient than using native threads. See the ``test_threading`` function in

python/mujoco/rollout.cc

@@ -176,7 +176,7 @@ void _unsafe_rollout_threaded(std::vector<const mjModel*>& m, std::vector<mjData
                               int nroll, int nstep, unsigned int control_spec,
                               const mjtNum* state0, const mjtNum* warmstart0,
                               const mjtNum* control, mjtNum* state, mjtNum* sensordata,
-                              std::shared_ptr<ThreadPool> pool, int chunk_size) {
+                              std::shared_ptr<ThreadPool>& pool, int chunk_size) {
   int nfulljobs = nroll / chunk_size;
   int chunk_remainder = nroll % chunk_size;
   int njobs = (chunk_remainder > 0) ? nfulljobs + 1 : nfulljobs;

python/mujoco/rollout.py

@@ -33,7 +33,6 @@ def __init__(self, nthread: int = None):
     """  # fmt: skip
     self.nthread = 0 if nthread is None else nthread
     self.rollout_ = _rollout.Rollout(self.nthread)
-    atexit.register(self.shutdown_pool)
 
   def __enter__(self):
     return self

python/mujoco/rollout_test.py

@@ -579,6 +579,24 @@ def test_threading_native_persistent_object(self):
       np.testing.assert_array_equal(state, py_state)
       np.testing.assert_array_equal(sensordata, py_sensordata)
 
+    rollout_ = rollout.Rollout(num_workers)
+    for i 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_.shutdown_pool()
+
   def test_threading_native_persistent_function(self):
     model = mujoco.MjModel.from_xml_string(TEST_XML)
     nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)