dm_control: Import of refs/pull/496/head

PiperOrigin-RevId: 686468082
Change-Id: I5f16943624183b578c1672b408430f66083ba61c

Author: barikata1984

dm_control/mjcf/attribute.py

@@ -424,7 +424,7 @@ def __init__(self, extension, prefix=''):
   def __eq__(self, other):
     return self.get_vfs_filename() == other.get_vfs_filename()
 
-  def get_vfs_filename(self):
+  def get_vfs_filename(self, filename_with_hash=True):
     """Returns the name of the asset file as registered in MuJoCo's VFS."""
     # Hash the contents of the asset to get a unique identifier.
     hash_string = hashlib.sha1(util.to_binary_string(self.contents)).hexdigest()
@@ -435,7 +435,10 @@ def get_vfs_filename(self):
       if raw_length > constants.MAX_VFS_FILENAME_LENGTH:
         trim_amount = raw_length - constants.MAX_VFS_FILENAME_LENGTH
         prefix = prefix[:-trim_amount]
-      filename = '-'.join([prefix, hash_string])
+      if filename_with_hash:
+        filename = '-'.join([prefix, hash_string])
+      else:
+        filename = prefix
     else:
       filename = hash_string
 
@@ -568,6 +571,7 @@ 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:
-      return self._value.get_vfs_filename()
+      filename_with_hash = kwargs.get('filename_with_hash', True)
+      return self._value.get_vfs_filename(filename_with_hash)
     else:
       return None

dm_control/mjcf/element.py

@@ -32,7 +32,7 @@
 import numpy as np
 
 
-_raw_property = property  # pylint: disable=invalid-name
+_raw_property = property  # pylint: disable=invalid-name,used-before-assignment
 
 _UNITS = ('K', 'M', 'G', 'T', 'P', 'E')
 
@@ -751,7 +751,7 @@ def all_children(self):
   def to_xml(self, prefix_root=None, debug_context=None,
              *,
              precision=constants.XML_DEFAULT_PRECISION,
-             zero_threshold=0):
+             zero_threshold=0, filename_with_hash=True):
     """Generates an etree._Element corresponding to this MJCF element.
 
     Args:
@@ -766,25 +766,31 @@ def to_xml(self, prefix_root=None, debug_context=None,
         quantities.
       zero_threshold: (optional) When outputting XML, floating point quantities
         whose absolute value falls below this threshold will be treated as zero.
+      filename_with_hash: (optional) A boolean, whether to append a hash string
+        to the name of a file when it is registered in MuJoCo's VFS.
 
     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,
-                            precision=precision, zero_threshold=zero_threshold)
+                            precision=precision, zero_threshold=zero_threshold,
+                            filename_with_hash=filename_with_hash)
     self._children_to_xml(xml_element, prefix_root, debug_context,
-                          precision=precision, zero_threshold=zero_threshold)
+                          precision=precision, zero_threshold=zero_threshold,
+                          filename_with_hash=filename_with_hash)
     return xml_element
 
   def _attributes_to_xml(self, xml_element, prefix_root, debug_context=None,
-                         *, precision, zero_threshold):
+                         *, precision, zero_threshold, filename_with_hash):
     del debug_context  # Unused.
     for attribute_name, attribute in self._attributes.items():
-      attribute_value = attribute.to_xml_string(prefix_root,
-                                                precision=precision,
-                                                zero_threshold=zero_threshold)
+      attribute_value = attribute.to_xml_string(
+          prefix_root,
+          precision=precision,
+          zero_threshold=zero_threshold,
+          filename_with_hash=filename_with_hash)
       if attribute_name == self._spec.identifier and attribute_value is None:
         xml_element.set(attribute_name, self.full_identifier)
       elif attribute_value is None:
@@ -793,11 +799,12 @@ def _attributes_to_xml(self, xml_element, prefix_root, debug_context=None,
         xml_element.set(attribute_name, attribute_value)
 
   def _children_to_xml(self, xml_element, prefix_root, debug_context=None,
-                       *, precision, zero_threshold):
+                       *, precision, zero_threshold, filename_with_hash):
     for child in self.all_children():
       child_xml = child.to_xml(prefix_root, debug_context,
                                precision=precision,
-                               zero_threshold=zero_threshold)
+                               zero_threshold=zero_threshold,
+                               filename_with_hash=filename_with_hash)
       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)
@@ -811,7 +818,7 @@ def to_xml_string(self, prefix_root=None,
                     self_only=False, pretty_print=True, debug_context=None,
                     *,
                     precision=constants.XML_DEFAULT_PRECISION,
-                    zero_threshold=0):
+                    zero_threshold=0, filename_with_hash=True):
     """Generates an XML string corresponding to this MJCF element.
 
     Args:
@@ -830,13 +837,16 @@ def to_xml_string(self, prefix_root=None,
         quantities.
       zero_threshold: (optional) When outputting XML, floating point quantities
         whose absolute value falls below this threshold will be treated as zero.
+      filename_with_hash: (optional) A boolean, whether to append a hash string
+        to the name of a file when it is registered in MuJoCo's VFS.
 
     Returns:
       A string.
     """
     xml_element = self.to_xml(prefix_root, debug_context,
                               precision=precision,
-                              zero_threshold=zero_threshold)
+                              zero_threshold=zero_threshold,
+                              filename_with_hash=filename_with_hash)
     if self_only and len(xml_element) > 0:  # pylint: disable=g-explicit-length-test
       etree.strip_elements(xml_element, '*')
       xml_element.text = '...'
@@ -1073,10 +1083,12 @@ def prefixed_identifier(self, prefix_root=None):
   def to_xml(self, prefix_root=None, debug_context=None,
              *,
              precision=constants.XML_DEFAULT_PRECISION,
-             zero_threshold=0):
+             zero_threshold=0,
+             filename_with_hash=True):
     xml_element = (super().to_xml(prefix_root, debug_context,
                                   precision=precision,
-                                  zero_threshold=zero_threshold))
+                                  zero_threshold=zero_threshold,
+                                  filename_with_hash=filename_with_hash))
     xml_element.set('name', self.prefixed_identifier(prefix_root))
     return xml_element
 
@@ -1104,10 +1116,12 @@ class _AttachmentFrameChild(_ElementImpl):
   def to_xml(self, prefix_root=None, debug_context=None,
              *,
              precision=constants.XML_DEFAULT_PRECISION,
-             zero_threshold=0):
+             zero_threshold=0,
+             filename_with_hash=True):
     xml_element = (super().to_xml(prefix_root, debug_context,
                                   precision=precision,
-                                  zero_threshold=zero_threshold))
+                                  zero_threshold=zero_threshold,
+                                  filename_with_hash=filename_with_hash))
     if self.spec.namespace is not None:
       if self.name:
         name = (self._parent.prefixed_identifier(prefix_root) +
@@ -1147,18 +1161,22 @@ def all_children(self):
   def to_xml(self, prefix_root=None, debug_context=None,
              *,
              precision=constants.XML_DEFAULT_PRECISION,
-             zero_threshold=0):
+             zero_threshold=0,
+             filename_with_hash=True):
     prefix_root = prefix_root or self.namescope
     xml_element = (super().to_xml(prefix_root, debug_context,
                                   precision=precision,
-                                  zero_threshold=zero_threshold))
+                                  zero_threshold=zero_threshold,
+                                  filename_with_hash=filename_with_hash))
     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,
-                                           precision=precision,
-                                           zero_threshold=zero_threshold)
+        attachment_xml = attachment.to_xml(
+            prefix_root, debug_context,
+            precision=precision,
+            zero_threshold=zero_threshold,
+            filename_with_hash=filename_with_hash)
         for attachment_child_xml in attachment_xml:
           root_default.append(attachment_child_xml)
       xml_element = root_default
@@ -1173,12 +1191,14 @@ class _ActuatorElement(_ElementImpl):
   def _children_to_xml(self, xml_element, prefix_root, debug_context=None,
                        *,
                        precision=constants.XML_DEFAULT_PRECISION,
-                       zero_threshold=0):
+                       zero_threshold=0,
+                       filename_with_hash=True):
     debug_comments = {}
     for child in self.all_children():
       child_xml = child.to_xml(prefix_root, debug_context,
                                precision=precision,
-                               zero_threshold=zero_threshold)
+                               zero_threshold=zero_threshold,
+                               filename_with_hash=filename_with_hash)
       if debugging.debug_mode() and debug_context:
         debug_comment = debug_context.register_element_for_debugging(child)
         debug_comments[child_xml] = debug_comment