Add a floating point precision option to to_xml(_string).

PiperOrigin-RevId: 467664520
Change-Id: Ibac426918ae8466dc09e65635afbf961157cb62a

Author: saran-t

dm_control/mjcf/attribute.py

@@ -105,7 +105,7 @@ def _before_clear(self):
   def _assign_from_string(self, string):
     self._assign(string)
 
-  def to_xml_string(self, prefix_root):  # pylint: disable=unused-argument
+  def to_xml_string(self, prefix_root, **kwargs):  # pylint: disable=unused-argument
     if self._value is None:
       return None
     else:
@@ -157,6 +157,15 @@ def _assign(self, value):
       raise ValueError('Expect a float value: got {}'.format(value)) from None
     self._value = float_value
 
+  def to_xml_string(self, prefix_root=None,
+                    *, precision=constants.XML_DEFAULT_PRECISION, **kwargs):
+    if self._value is None:
+      return None
+    else:
+      out = io.BytesIO()
+      np.savetxt(out, [self._value], fmt=f'%.{precision:d}g', newline=' ')
+      return util.to_native_string(out.getvalue())[:-1]  # Strip trailing space.
+
 
 class Keyword(_Attribute):
   """A keyword MJCF attribute."""
@@ -199,15 +208,13 @@ def _assign(self, value):
   def _assign_from_string(self, string):
     self._assign(np.fromstring(string, dtype=self._dtype, sep=' '))
 
-  def to_xml_string(self, prefix_root=None):  # pylint: disable=unused-argument
+  def to_xml_string(self, prefix_root=None,
+                    *, precision=constants.XML_DEFAULT_PRECISION, **kwargs):
     if self._value is None:
       return None
     else:
       out = io.BytesIO()
-      # 17 decimal digits is sufficient to represent a double float without loss
-      # of precision.
-      # https://en.wikipedia.org/wiki/IEEE_754#Character_representation
-      np.savetxt(out, self._value, fmt='%.17g', newline=' ')
+      np.savetxt(out, self._value, fmt=f'%.{precision:d}g', newline=' ')
       return util.to_native_string(out.getvalue())[:-1]  # Strip trailing space.
 
   def _check_shape(self, array):
@@ -255,7 +262,7 @@ def _defaults_string(self, prefix_root):
     prefix.append(self._value or '')
     return constants.PREFIX_SEPARATOR.join(prefix) or constants.PREFIX_SEPARATOR
 
-  def to_xml_string(self, prefix_root=None):
+  def to_xml_string(self, prefix_root=None, **kwargs):
     if self._parent.tag == constants.DEFAULT:
       return self._defaults_string(prefix_root)
     elif self._value:
@@ -352,7 +359,7 @@ def _defaults_string(self, prefix_root):
       out_string = prefix + self._value
     return out_string
 
-  def to_xml_string(self, prefix_root):
+  def to_xml_string(self, prefix_root, **kwargs):
     self._check_dead_reference()
     if isinstance(self._value, base.Element):
       return self._value.prefixed_identifier(prefix_root)
@@ -388,7 +395,7 @@ def _before_clear(self):
     if self._value:
       self._parent.namescope.remove(constants.BASEPATH, self._path_namespace)
 
-  def to_xml_string(self, prefix_root=None):
+  def to_xml_string(self, prefix_root=None, **kwargs):
     return None
 
 
@@ -527,7 +534,7 @@ def get_contents(self):
                          'querying the contents.')
     return self._value.contents
 
-  def to_xml_string(self, prefix_root=None):
+  def to_xml_string(self, prefix_root=None, **kwargs):
     """Returns the asset filename as it will appear in the generated XML."""
     del prefix_root  # Unused
     if self._value is not None:

dm_control/mjcf/attribute_test.py

@@ -131,15 +131,37 @@ def assertCanNotBeCleared(self, mjcf_element, attribute_name):
 
   def testFloatScalar(self):
     mujoco = self._mujoco
-    mujoco.optional.float = 5
-    self.assertEqual(mujoco.optional.float, 5)
+    mujoco.optional.float = 0.357357
+    self.assertEqual(mujoco.optional.float, 0.357357)
     self.assertEqual(type(mujoco.optional.float), float)
     with self.assertRaisesRegex(ValueError, 'Expect a float value'):
       mujoco.optional.float = 'five'
     # failed assignment should not change the value
-    self.assertEqual(mujoco.optional.float, 5)
-    self.assertXMLStringEqual(mujoco.optional, 'float', '5.0')
-    self.assertCanBeCleared(mujoco.optional, 'float')
+    self.assertEqual(mujoco.optional.float, 0.357357)
+    self.assertEqual(
+        mujoco.optional.get_attribute_xml_string('float', precision=1),
+        '0.4')
+    self.assertEqual(
+        mujoco.optional.get_attribute_xml_string('float', precision=2),
+        '0.36')
+    self.assertEqual(
+        mujoco.optional.get_attribute_xml_string('float', precision=3),
+        '0.357')
+    self.assertEqual(
+        mujoco.optional.get_attribute_xml_string('float', precision=4),
+        '0.3574')
+    self.assertEqual(
+        mujoco.optional.get_attribute_xml_string('float', precision=5),
+        '0.35736')
+    self.assertEqual(
+        mujoco.optional.get_attribute_xml_string('float', precision=6),
+        '0.357357')
+    self.assertEqual(
+        mujoco.optional.get_attribute_xml_string('float', precision=7),
+        '0.357357')
+    self.assertEqual(
+        mujoco.optional.get_attribute_xml_string('float', precision=8),
+        '0.357357')
 
   def testIntScalar(self):
     mujoco = self._mujoco
@@ -178,6 +200,9 @@ def testFloatArray(self):
     mujoco.optional.float_array = [np.pi, 2, 1e-16]
     self.assertXMLStringEqual(mujoco.optional, 'float_array',
                               '3.1415926535897931 2 9.9999999999999998e-17')
+    self.assertEqual(
+        mujoco.optional.get_attribute_xml_string('float_array', precision=5),
+        '3.1416 2 1e-16')
     self.assertCanBeCleared(mujoco.optional, 'float_array')
 
   def testFormatVeryLargeArray(self):

dm_control/mjcf/base.py

@@ -17,6 +17,8 @@
 
 import abc
 
+from dm_control.mjcf import constants
+
 
 class Element(metaclass=abc.ABCMeta):
   """Abstract base class for an MJCF element.
@@ -220,7 +222,8 @@ def all_children(self):
     pass
 
   @abc.abstractmethod
-  def to_xml(self, prefix_root=None, debug_context=None):
+  def to_xml(self, prefix_root=None, debug_context=None,
+             *, precision=constants.XML_DEFAULT_PRECISION):
     """Generates an etree._Element corresponding to this MJCF element.
 
     Args:
@@ -231,14 +234,17 @@ def to_xml(self, prefix_root=None, debug_context=None):
         the debugging information associated with the generated XML is written.
         This is intended for internal use within PyMJCF; users should never need
         manually pass this argument.
+      precision: (optional) Number of digits to output for floating point
+        quantities.
 
     Returns:
       An etree._Element object.
     """
 
   @abc.abstractmethod
   def to_xml_string(self, prefix_root=None,
-                    self_only=False, pretty_print=True, debug_context=None):
+                    self_only=False, pretty_print=True, debug_context=None,
+                    *, precision=constants.XML_DEFAULT_PRECISION):
     """Generates an XML string corresponding to this MJCF element.
 
     Args:
@@ -253,6 +259,8 @@ def to_xml_string(self, prefix_root=None,
         the debugging information associated with the generated XML is written.
         This is intended for internal use within PyMJCF; users should never need
         manually pass this argument.
+      precision: (optional) Number of digits to output for floating point
+        quantities.
 
     Returns:
       A string.

dm_control/mjcf/constants.py

@@ -72,3 +72,8 @@
 INDIRECT_REFERENCE_ATTRIB = {
     'xbody': 'body',
 }
+
+# 17 decimal digits is sufficient to represent a double float without loss
+# of precision.
+# https://en.wikipedia.org/wiki/IEEE_754#Character_representation
+XML_DEFAULT_PRECISION = 17

dm_control/mjcf/element.py

@@ -75,7 +75,7 @@ def _mjcf_property(self):
       err_with_next_tb = err.with_traceback(tb.tb_next)
       if isinstance(err, AttributeError):
         self._last_attribute_error = err_with_next_tb  # pylint: disable=protected-access
-      raise err_with_next_tb
+      raise err_with_next_tb  # pylint: disable=raise-missing-from
   return _raw_property(_mjcf_property)
 
 
@@ -182,7 +182,7 @@ def __init__(self, spec, parent, attributes=None):
             attribute_obj._force_clear()  # pylint: disable=protected-access
           # Then raise a meaningful error
           err_type, err, tb = sys.exc_info()
-          raise err_type(
+          raise err_type(  # pylint: disable=raise-missing-from
               f'during initialization of attribute {attribute_spec.name!r} of '
               f'element <{self._spec.name}>: {err}').with_traceback(tb)
 
@@ -507,9 +507,14 @@ def _get_attribute(self, attribute_name):
     self._check_valid_attribute(attribute_name)
     return self._attributes[attribute_name].value
 
-  def get_attribute_xml_string(self, attribute_name, prefix_root=None):
+  def get_attribute_xml_string(self,
+                               attribute_name,
+                               prefix_root=None,
+                               *,
+                               precision=constants.XML_DEFAULT_PRECISION):
     self._check_valid_attribute(attribute_name)
-    return self._attributes[attribute_name].to_xml_string(prefix_root)
+    return self._attributes[attribute_name].to_xml_string(
+        prefix_root, precision=precision)
 
   def get_attributes(self):
     fix_attribute_name = (
@@ -541,7 +546,7 @@ def set_attributes(self, **kwargs):
             self._set_attribute(name, old_value)
           # Then raise a meaningful error.
           err_type, err, tb = sys.exc_info()
-          raise err_type(
+          raise err_type(  # pylint: disable=raise-missing-from
               f'during assignment to attribute {attribute_name!r} of '
               f'element <{self._spec.name}>: {err}').with_traceback(tb)
 
@@ -554,7 +559,7 @@ def _check_valid_child(self, element_name):
     try:
       return self._spec.children[element_name]
     except KeyError:
-      raise AttributeError(
+      raise AttributeError(  # pylint: disable=raise-missing-from
           '<{}> is not a valid child of <{}>'
           .format(element_name, self._spec.name))
 
@@ -690,7 +695,8 @@ def all_children(self):
                        if child.spec.repeated]
     return all_children
 
-  def to_xml(self, prefix_root=None, debug_context=None):
+  def to_xml(self, prefix_root=None, debug_context=None,
+             *, precision=constants.XML_DEFAULT_PRECISION):
     """Generates an etree._Element corresponding to this MJCF element.
 
     Args:
@@ -701,30 +707,37 @@ def to_xml(self, prefix_root=None, debug_context=None):
         the debugging information associated with the generated XML is written.
         This is intended for internal use within PyMJCF; users should never need
         manually pass this argument.
+      precision: (optional) Number of digits to output for floating point
+        quantities.
 
     Returns:
       An etree._Element object.
     """
     prefix_root = prefix_root or self.namescope
     xml_element = etree.Element(self._spec.name)
-    self._attributes_to_xml(xml_element, prefix_root, debug_context)
-    self._children_to_xml(xml_element, prefix_root, debug_context)
+    self._attributes_to_xml(xml_element, prefix_root, debug_context,
+                            precision=precision)
+    self._children_to_xml(xml_element, prefix_root, debug_context,
+                          precision=precision)
     return xml_element
 
-  def _attributes_to_xml(self, xml_element, prefix_root, debug_context=None):
+  def _attributes_to_xml(self, xml_element, prefix_root, debug_context=None,
+                         *, precision):
     del debug_context  # Unused.
     for attribute_name, attribute in self._attributes.items():
-      attribute_value = attribute.to_xml_string(prefix_root)
+      attribute_value = attribute.to_xml_string(prefix_root,
+                                                precision=precision)
       if attribute_name == self._spec.identifier and attribute_value is None:
         xml_element.set(attribute_name, self.full_identifier)
       elif attribute_value is None:
         continue
       else:
         xml_element.set(attribute_name, attribute_value)
 
-  def _children_to_xml(self, xml_element, prefix_root, debug_context=None):
+  def _children_to_xml(self, xml_element, prefix_root, debug_context=None,
+                       *, precision):
     for child in self.all_children():
-      child_xml = child.to_xml(prefix_root, debug_context)
+      child_xml = child.to_xml(prefix_root, debug_context, precision=precision)
       if (child_xml.attrib or len(child_xml)  # pylint: disable=g-explicit-length-test
           or child.spec.repeated or child.spec.on_demand):
         xml_element.append(child_xml)
@@ -735,7 +748,8 @@ def _children_to_xml(self, xml_element, prefix_root, debug_context=None):
             child_xml.insert(0, copy.deepcopy(debug_comment))
 
   def to_xml_string(self, prefix_root=None,
-                    self_only=False, pretty_print=True, debug_context=None):
+                    self_only=False, pretty_print=True, debug_context=None,
+                    *, precision=constants.XML_DEFAULT_PRECISION):
     """Generates an XML string corresponding to this MJCF element.
 
     Args:
@@ -750,16 +764,19 @@ def to_xml_string(self, prefix_root=None,
         the debugging information associated with the generated XML is written.
         This is intended for internal use within PyMJCF; users should never need
         manually pass this argument.
+      precision: (optional) Number of digits to output for floating point
+        quantities.
 
     Returns:
       A string.
     """
-    xml_element = self.to_xml(prefix_root, debug_context)
+    xml_element = self.to_xml(prefix_root, debug_context, precision=precision)
     if self_only and len(xml_element) > 0:  # pylint: disable=g-explicit-length-test
       etree.strip_elements(xml_element, '*')
       xml_element.text = '...'
     if (self_only and self._spec.identifier and
-        not self._attributes[self._spec.identifier].to_xml_string(prefix_root)):
+        not self._attributes[self._spec.identifier].to_xml_string(
+            prefix_root, precision=precision)):
       del xml_element.attrib[self._spec.identifier]
     xml_string = util.to_native_string(
         etree.tostring(xml_element, pretty_print=pretty_print))
@@ -987,8 +1004,10 @@ def prefixed_identifier(self, prefix_root=None):
     prefix = self.namescope.full_prefix(prefix_root)
     return prefix + self._attachment.namescope.name + constants.PREFIX_SEPARATOR
 
-  def to_xml(self, prefix_root=None, debug_context=None):
-    xml_element = (super().to_xml(prefix_root, debug_context))
+  def to_xml(self, prefix_root=None, debug_context=None,
+             *, precision=constants.XML_DEFAULT_PRECISION):
+    xml_element = (super().to_xml(prefix_root, debug_context,
+                                  precision=precision))
     xml_element.set('name', self.prefixed_identifier(prefix_root))
     return xml_element
 
@@ -1013,8 +1032,10 @@ class _AttachmentFrameChild(_ElementImpl):
   """
   __slots__ = []
 
-  def to_xml(self, prefix_root=None, debug_context=None):
-    xml_element = (super().to_xml(prefix_root, debug_context))
+  def to_xml(self, prefix_root=None, debug_context=None,
+             *, precision=constants.XML_DEFAULT_PRECISION):
+    xml_element = (super().to_xml(prefix_root, debug_context,
+                                  precision=precision))
     if self.spec.namespace is not None:
       if self.name:
         name = (self._parent.prefixed_identifier(prefix_root) +
@@ -1051,14 +1072,17 @@ def _attach(self, other, exclude_worldbody=False, dry_run=False):
   def all_children(self):
     return [child for child in self._children]
 
-  def to_xml(self, prefix_root=None, debug_context=None):
+  def to_xml(self, prefix_root=None, debug_context=None,
+             *, precision=constants.XML_DEFAULT_PRECISION):
     prefix_root = prefix_root or self.namescope
-    xml_element = (super().to_xml(prefix_root, debug_context))
+    xml_element = (super().to_xml(prefix_root, debug_context,
+                                  precision=precision))
     if isinstance(self._parent, RootElement):
       root_default = etree.Element(self._spec.name)
       root_default.append(xml_element)
       for attachment in self._attachments.values():
-        attachment_xml = attachment.to_xml(prefix_root, debug_context)
+        attachment_xml = attachment.to_xml(prefix_root, debug_context,
+                                           precision=precision)
         for attachment_child_xml in attachment_xml:
           root_default.append(attachment_child_xml)
       xml_element = root_default
@@ -1082,12 +1106,13 @@ def _is_third_order_actuator(self, child):
     else:
       return False  # No other actuator shortcuts have internal dynamics.
 
-  def _children_to_xml(self, xml_element, prefix_root, debug_context=None):
+  def _children_to_xml(self, xml_element, prefix_root, debug_context=None,
+                       *, precision=constants.XML_DEFAULT_PRECISION):
     second_order = []
     third_order = []
     debug_comments = {}
     for child in self.all_children():
-      child_xml = child.to_xml(prefix_root, debug_context)
+      child_xml = child.to_xml(prefix_root, debug_context, precision=precision)
       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):
@@ -1297,7 +1322,7 @@ def __getitem__(self, index):
         return scoped_elements[index[(len(scope_name) + 1):]]
       except KeyError:
         # Re-raise so that the error shows the full, un-stripped index string
-        raise self._identifier_not_found_error(index)
+        raise self._identifier_not_found_error(index)  # pylint: disable=raise-missing-from
     elif isinstance(index, slice) or (isinstance(index, int) and index < 0):
       return self._full_list()[index]
     else: