Stop requiring stateful actuators to come after stateless ones.

PiperOrigin-RevId: 487201455
Change-Id: Iec3076769486b763053f32b5941678b4d6583eba

Author: saran-t

dm_control/mjcf/element.py

@@ -1134,47 +1134,24 @@ def to_xml(self, prefix_root=None, debug_context=None,
 
 
 class _ActuatorElement(_ElementImpl):
-  """Specialized object representing an <actuator> element.
+  """Specialized object representing an <actuator> element."""
 
-  This is necessary because MuJoCo requires that all 3rd-order actuators (i.e.
-  those with internal dynamics) come after all 2nd-order actuators in the
-  generated XML.
-  """
   __slots__ = ()
 
-  def _is_third_order_actuator(self, child):
-    if child.tag == 'general':
-      return child.dyntype and child.dyntype != 'none'
-    elif child.tag == 'cylinder':
-      return True  # The `<cylinder>` shortcut has 'filter' dynamics.
-    else:
-      return False  # No other actuator shortcuts have internal dynamics.
-
   def _children_to_xml(self, xml_element, prefix_root, debug_context=None,
                        *,
                        precision=constants.XML_DEFAULT_PRECISION,
                        zero_threshold=0):
-    second_order = []
-    third_order = []
     debug_comments = {}
     for child in self.all_children():
       child_xml = child.to_xml(prefix_root, debug_context,
                                precision=precision,
                                zero_threshold=zero_threshold)
-      if (child_xml.attrib or len(child_xml)  # pylint: disable=g-explicit-length-test
-          or child.spec.repeated or child.spec.on_demand):
-        if self._is_third_order_actuator(child):
-          third_order.append(child_xml)
-        else:
-          second_order.append(child_xml)
-        if debugging.debug_mode() and debug_context:
-          debug_comment = debug_context.register_element_for_debugging(child)
-          debug_comments[child_xml] = debug_comment
-          if len(child_xml) > 0:  # pylint: disable=g-explicit-length-test
-            child_xml.insert(0, copy.deepcopy(debug_comment))
-    # Ensure that all second-order actuators come before third-order actuators
-    # in the XML.
-    for child_xml in second_order + third_order:
+      if debugging.debug_mode() and debug_context:
+        debug_comment = debug_context.register_element_for_debugging(child)
+        debug_comments[child_xml] = debug_comment
+        if len(child_xml) > 0:  # pylint: disable=g-explicit-length-test
+          child_xml.insert(0, copy.deepcopy(debug_comment))
       xml_element.append(child_xml)
       if debugging.debug_mode() and debug_context:
         xml_element.append(debug_comments[child_xml])

dm_control/mjcf/physics.py

@@ -91,36 +91,6 @@ def _get_attributes(size_names, strip_prefixes):
   return out
 
 
-# Fields related to the internal states of actuators (i.e. with a leading
-# dimension of 'na') require special treatment.
-def _get_actuator_state_fields():
-  actuator_state_fields = []
-  for sizes_dict in sizes.array_sizes.values():
-    for field_name, dimensions in sizes_dict.items():
-      if dimensions[0] == 'na':
-        actuator_state_fields.append(field_name)
-  return frozenset(actuator_state_fields)
-
-_ACTUATOR_STATE_FIELDS = _get_actuator_state_fields()
-
-
-def _filter_stateful_actuators(physics, actuator_names):
-  """Removes any stateless actuators from the list of actuator names."""
-  if isinstance(actuator_names, str):
-    actuator_names = [actuator_names]
-
-  if physics.model.na:
-    # MuJoCo requires that stateful actuators always come after stateless
-    # actuators in the model, so we keep actuator names only if their
-    # corresponding IDs are >= to the total number of stateless actuators.
-    num_stateless_actuators = physics.model.nu - physics.model.na
-    return [
-        name for name in actuator_names
-        if physics.model.name2id(name, 'actuator') >= num_stateless_actuators]
-  else:
-    return []
-
-
 _ATTRIBUTES = {
     'actuator': _get_attributes(['na', 'nu'], strip_prefixes=['actuator']),
     'body': _get_attributes(['nbody'], strip_prefixes=['body']),
@@ -320,14 +290,7 @@ def _get_cached_array_and_index(self, name):
     try:
       index = self._array_index_cache[name]
     except KeyError:
-      # If we are indexing into a field relating to actuator internal states
-      # then we must first remove the names of any stateless actuators.
-      if name in _ACTUATOR_STATE_FIELDS:
-        named_index = _filter_stateful_actuators(
-            self._physics, self._named_index)
-      else:
-        named_index = self._named_index
-      index = named_indexer._convert_key(named_index)  # pylint: disable=protected-access
+      index = named_indexer._convert_key(self._named_index)  # pylint: disable=protected-access
       self._array_index_cache[name] = index
     return array, index
 
@@ -376,7 +339,7 @@ def __getattr__(self, name):
 
         if self._physics.is_dirty and not triggers_dirty:
           self._physics.forward()
-        if isinstance(index, int) and array.ndim == 1:
+        if np.issubdtype(type(index), np.integer) and array.ndim == 1:
           # Case where indexing results in a scalar.
           out = array[index]
         else:

dm_control/mujoco/index.py

@@ -94,6 +94,7 @@
 _RAGGED_ADDRS = {
     'nq': 'jnt_qposadr',
     'nv': 'jnt_dofadr',
+    'na': 'actuator_actadr',
     'nsensordata': 'sensor_adr',
     'nnumericdata': 'numeric_adr',
 }
@@ -218,6 +219,8 @@ def _get_size_name_to_element_names(model):
   # For example, the element names for "nv" axis come from "njnt".
   for size_name, address_field_name in _RAGGED_ADDRS.items():
     donor = 'n' + address_field_name.split('_')[0]
+    if donor == 'nactuator':
+      donor = 'nu'
     if donor in size_name_to_element_names:
       size_name_to_element_names[size_name] = size_name_to_element_names[donor]
 
@@ -230,13 +233,6 @@ def _get_size_name_to_element_names(model):
   assert None not in mocap_body_names
   size_name_to_element_names['nmocap'] = mocap_body_names
 
-  # Arrays with dimension `na` correspond to stateful actuators. MuJoCo's
-  # compiler requires that these are always defined after stateless actuators,
-  # so we only need the final `na` elements in the list of all actuator names.
-  if model.na:
-    all_actuator_names = size_name_to_element_names['nu']
-    size_name_to_element_names['na'] = all_actuator_names[-model.na:]
-
   return size_name_to_element_names
 
 
@@ -255,8 +251,11 @@ def _get_size_name_to_element_sizes(model):
 
   for size_name, address_field_name in _RAGGED_ADDRS.items():
     addresses = getattr(model, address_field_name).ravel()
-    total_length = getattr(model, size_name)
-    element_sizes = np.diff(np.r_[addresses, total_length])
+    if size_name == 'na':
+      element_sizes = np.where(addresses == -1, 0, 1)
+    else:
+      total_length = getattr(model, size_name)
+      element_sizes = np.diff(np.r_[addresses, total_length])
     size_name_to_element_sizes[size_name] = element_sizes
 
   return size_name_to_element_sizes
@@ -282,7 +281,9 @@ def make_axis_indexers(model):
     element_names = size_name_to_element_names[size_name]
     if size_name in _RAGGED_ADDRS:
       element_sizes = size_name_to_element_sizes[size_name]
-      indexer = RaggedNamedAxis(element_names, element_sizes)
+      singleton = (size_name == 'na')
+      indexer = RaggedNamedAxis(element_names, element_sizes,
+                                singleton=singleton)
     else:
       indexer = RegularNamedAxis(element_names)
     axis_indexers[size_name] = indexer
@@ -377,12 +378,13 @@ def names(self):
 class RaggedNamedAxis(Axis):
   """Represents an axis where the named elements may vary in size."""
 
-  def __init__(self, element_names, element_sizes):
+  def __init__(self, element_names, element_sizes, singleton=False):
     """Initializes a new `RaggedNamedAxis` instance.
 
     Args:
       element_names: A list or array containing the element names.
       element_sizes: A list or array containing the size of each element.
+      singleton: Whether to reduce singleton slices to scalars.
     """
     names_to_slices = {}
     names_to_indices = {}
@@ -391,7 +393,10 @@ def __init__(self, element_names, element_sizes):
     for name, size in zip(element_names, element_sizes):
       # Don't add unnamed elements to the dicts.
       if name:
-        names_to_slices[name] = slice(offset, offset + size)
+        if size == 1 and singleton:
+          names_to_slices[name] = offset
+        else:
+          names_to_slices[name] = slice(offset, offset + size)
         names_to_indices[name] = range(offset, offset + size)
       offset += size
 
@@ -400,29 +405,29 @@ def __init__(self, element_names, element_sizes):
     self._names_to_slices = names_to_slices
     self._names_to_indices = names_to_indices
 
-  def convert_key_item(self, key):
+  def convert_key_item(self, key_item):
     """Converts a named indexing expression to a numpy-friendly index."""
 
-    _validate_key_item(key)
+    _validate_key_item(key_item)
 
-    if isinstance(key, str):
-      key = self._names_to_slices[util.to_native_string(key)]
+    if isinstance(key_item, str):
+      key_item = self._names_to_slices[util.to_native_string(key_item)]
 
-    elif isinstance(key, (list, np.ndarray)):
+    elif isinstance(key_item, (list, np.ndarray)):
       # We assume that either all or none of the items in the sequence are
       # strings representing names. If there is a mix, we will let NumPy throw
       # an error when trying to index with the returned key.
-      if isinstance(key[0], str):
+      if isinstance(key_item[0], str):
         new_key = []
-        for k in key:
+        for k in key_item:
           idx = self._names_to_indices[util.to_native_string(k)]
           if isinstance(idx, int):
             new_key.append(idx)
           else:
             new_key.extend(idx)
-        key = new_key
+        key_item = new_key
 
-    return key
+    return key_item
 
   @property
   def names(self):

dm_control/mujoco/testing/assets/model_with_third_order_actuators.xml

@@ -6,11 +6,9 @@
     </body>
   </worldbody>
   <actuator>
-    <!-- Second-order actuators -->
-    <motor name="motor" joint="slide_joint"/>
-    <velocity name="velocity" joint="slide_joint"/>
-    <!-- Third-order actuators -->
-    <cylinder name="cylinder" joint="slide_joint"/>
-    <general name="general" joint="slide_joint" dyntype="integrator" biastype="affine" dynprm="1 0 0"/>
+    <motor name="motor" joint="slide_joint"/> <!-- Second-order -->
+    <cylinder name="cylinder" joint="slide_joint"/> <!-- Third-order -->
+    <velocity name="velocity" joint="slide_joint"/> <!-- Second-order -->
+    <general name="general" joint="slide_joint" dyntype="integrator" biastype="affine" dynprm="1 0 0"/> <!-- Third-order -->
   </actuator>
 </mujoco>