rollout fixups and correctly initialize unspecified controls

python/mujoco/rollout.cc

@@ -73,26 +73,28 @@ void _unsafe_rollout(std::vector<const mjModel*>& m, mjData* d, int nroll, int n
   if (!(control_spec & mjSTATE_XFRC_APPLIED)) {
     mju_zero(d->xfrc_applied, 6*nbody);
   }
-  if (!(control_spec & mjSTATE_MOCAP_POS)) {
-    for (int i = 0; i < nbody; i++) {
-      int id = m[0]->body_mocapid[i];
-      if (id >= 0) mju_copy3(d->mocap_pos+3*id, m[0]->body_pos+3*i);
-    }
-  }
-  if (!(control_spec & mjSTATE_MOCAP_QUAT)) {
-    for (int i = 0; i < nbody; i++) {
-      int id = m[0]->body_mocapid[i];
-      if (id >= 0) mju_copy4(d->mocap_quat+4*id, m[0]->body_quat+4*i);
-    }
-  }
-  if (!(control_spec & mjSTATE_EQ_ACTIVE)) {
-    for (int i = 0; i < neq; i++) {
-      d->eq_active[i] = m[0]->eq_active0[i];
-    }
-  }
 
   // loop over rollouts
   for (int r = 0; r < nroll; r++) {
+    // clear user inputs if unspecified
+    if (!(control_spec & mjSTATE_MOCAP_POS)) {
+      for (int i = 0; i < nbody; i++) {
+        int id = m[r]->body_mocapid[i];
+        if (id >= 0) mju_copy3(d->mocap_pos+3*id, m[r]->body_pos+3*i);
+      }
+    }
+    if (!(control_spec & mjSTATE_MOCAP_QUAT)) {
+      for (int i = 0; i < nbody; i++) {
+        int id = m[r]->body_mocapid[i];
+        if (id >= 0) mju_copy4(d->mocap_quat+4*id, m[r]->body_quat+4*i);
+      }
+    }
+    if (!(control_spec & mjSTATE_EQ_ACTIVE)) {
+      for (int i = 0; i < neq; i++) {
+        d->eq_active[i] = m[r]->eq_active0[i];
+      }
+    }
+
     // set initial state
     mj_setState(m[r], d, state0 + r*nstate, mjSTATE_FULLPHYSICS);
 

python/mujoco/rollout.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Roll out open-loop trajectories from initial states, get subsequent states and sensor values."""
 
+from collections.abc import Sequence
 from typing import Optional, Union
 
 import mujoco
@@ -22,7 +23,7 @@ import numpy as np
 from numpy import typing as npt
 
 
-def rollout(model: Union[mujoco.MjModel, list[mujoco.MjModel]],
+def rollout(model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
             data: mujoco.MjData,
             initial_state: npt.ArrayLike,
             control: Optional[npt.ArrayLike] = None,
@@ -41,7 +42,7 @@ def rollout(model: Union[mujoco.MjModel, list[mujoco.MjModel]],
   Allocates outputs if none are given.
 
   Args:
-    model: An mjModel or a list of MjModel with the same size signature.
+    model: An mjModel or a sequence of MjModel with the same size signature.
     data: An associated mjData instance.
     initial_state: Array of initial states from which to roll out trajectories.
       ([nroll or 1] x nstate)
@@ -76,6 +77,9 @@ def rollout(model: Union[mujoco.MjModel, list[mujoco.MjModel]],
                      initial_warmstart, control, state, sensordata)
     return state, sensordata
 
+  if not isinstance(model, mujoco.MjModel):
+    model = list(model)
+
   # check control_spec
   if control_spec & ~mujoco.mjtState.mjSTATE_USER.value:
     raise ValueError('control_spec can only contain bits in mjSTATE_USER')
@@ -151,7 +155,7 @@ def rollout(model: Union[mujoco.MjModel, list[mujoco.MjModel]],
   _check_trailing_dimension(nsensordata, sensordata=sensordata)
 
   # tile input arrays/lists if required (singleton expansion)
-  model = model*nroll if len(model) == 1 else model
+  model = model * nroll if len(model) == 1 else model
   initial_state = _tile_if_required(initial_state, nroll)
   initial_warmstart = _tile_if_required(initial_warmstart, nroll)
   control = _tile_if_required(control, nroll, nstep)

python/mujoco/rollout_test.py

@@ -458,7 +458,7 @@ class MuJoCoRolloutTest(parameterized.TestCase):
     def thread_initializer():
       thread_local.data = mujoco.MjData(model)
 
-    model_list = [model]*nroll
+    model_list = [model] * nroll
     def call_rollout(initial_state, control, state, sensordata):
       rollout.rollout(model_list, thread_local.data, initial_state, control,
                       skip_checks=True,