Add size.memory to MJCF schema.

PiperOrigin-RevId: 509904430
Change-Id: I1265c3e1d57613146ae4382e49f069e475295abf

Author: sbohez

dm_control/mjcf/element.py

@@ -34,8 +34,40 @@
 
 _raw_property = property  # pylint: disable=invalid-name
 
+_UNITS = ('K', 'M', 'G', 'T', 'P', 'E')
 
-_CONFLICT_BEHAVIOR_FUNC = {'min': min, 'max': max}
+
+def _to_bytes(value_str):
+  """Converts a `str` value representing a size in bytes to `int`.
+
+  Args:
+    value_str: `str` value to be converted.
+
+  Returns:
+    `int` corresponding size in bytes.
+
+  Raises:
+    ValueError: if the `str` value passed has an unsupported unit.
+  """
+  if value_str.isdigit():
+    value_int = int(value_str)
+  else:
+    value_int = int(value_str[:-1])
+    unit = value_str[-1].upper()
+    if unit not in _UNITS:
+      raise ValueError(
+          f'unit of `size.memory` should be one of [{", ".join(_UNITS)}], got'
+          f' {unit}')
+    power = 10 * (_UNITS.index(unit) + 1)
+    value_int *= 2**power
+  return value_int
+
+
+def _max_bytes(first, second):
+  return str(max(_to_bytes(first), _to_bytes(second)))
+
+
+_CONFLICT_BEHAVIOR_FUNC = {'min': min, 'max': max, 'max_bytes': _max_bytes}
 
 
 def property(method):  # pylint: disable=redefined-builtin

dm_control/mjcf/element_test.py

@@ -1055,6 +1055,16 @@ def testMaxConflictingValues(self):
     self.assertEqual(model_1.size.nconmax, 345)
     self.assertEqual(model_1.size.njmax, 456)
 
+  def testMaxBytesConflictingValues(self):
+    model_1 = mjcf.RootElement()
+    model_1.size.memory = '10000'
+
+    model_2 = mjcf.RootElement()
+    model_2.size.memory = '1M'
+
+    model_1.attach(model_2)
+    self.assertEqual(model_1.size.memory, '1048576')
+
 
 if __name__ == '__main__':
   absltest.main()

dm_control/mjcf/schema.xml

@@ -95,6 +95,7 @@
     </element>
     <element name="size">
       <attributes>
+        <attribute name="memory" type="string" conflict_allowed="true" conflict_behavior="max_bytes"/>
         <attribute name="njmax" type="int" conflict_allowed="true" conflict_behavior="max"/>
         <attribute name="nconmax" type="int" conflict_allowed="true" conflict_behavior="max"/>
         <attribute name="nstack" type="int"/>