Add plugin support to Python bindings.

PiperOrigin-RevId: 480675078
Change-Id: I2967654b6662efced8ae62cae9b03c6c8e65c67c

python/mujoco/__init__.py

@@ -54,5 +54,15 @@ if _MUJOCO_GL not in ('disable', 'disabled', 'off', 'false', '0'):
     from mujoco.glfw import GLContext
 
 HEADERS_DIR = os.path.join(os.path.dirname(__file__), 'include/mujoco')
+PLUGINS_DIR = os.path.join(os.path.dirname(__file__), 'plugin')
+
+PLUGIN_HANDLES = []
+
+def _load_all_bundled_plugins():
+  for directory, _, filenames in os.walk(PLUGINS_DIR):
+    for filename in filenames:
+      PLUGIN_HANDLES.append(ctypes.CDLL(os.path.join(directory, filename)))
+
+_load_all_bundled_plugins()
 
 __version__ = mj_versionString()  # pylint: disable=undefined-variable

python/mujoco/bindings_test.py

@@ -69,6 +69,14 @@ TEST_XML_SENSOR = r"""
 </mujoco>
 """
 
+TEST_XML_PLUGIN = r"""
+<mujoco model="test">
+  <extension>
+    <required plugin="mujoco.elasticity.cable"/>
+  </extension>
+</mujoco>
+"""
+
 
 @contextlib.contextmanager
 def temporary_callback(setter, callback):
@@ -1111,5 +1119,8 @@ Euler integrator, semi-implicit in velocity.
         self.fail("Attribute '{}' differs from expected value: {}".format(
             name, str(e)))
 
+  def test_load_plugin(self):
+    mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
+
 if __name__ == '__main__':
   absltest.main()

python/setup.py

@@ -35,6 +35,7 @@ __version__ = '2.2.2'
 MUJOCO_CMAKE = 'MUJOCO_CMAKE'
 MUJOCO_CMAKE_ARGS = 'MUJOCO_CMAKE_ARGS'
 MUJOCO_PATH = 'MUJOCO_PATH'
+MUJOCO_PLUGIN_PATH = 'MUJOCO_PLUGIN_PATH'
 
 EXT_PREFIX = 'mujoco.'
 
@@ -70,6 +71,15 @@ def get_external_lib_patterns():
     return ['libmujoco.so.*']
 
 
+def get_plugin_lib_patterns():
+  if platform.system() == 'Windows':
+    return ['*.dll']
+  elif platform.system() == 'Darwin':
+    return ['lib*.dylib']
+  else:
+    return ['lib*']
+
+
 def start_and_end(iterable):
   it = iter(iterable)
   while True:
@@ -140,6 +150,7 @@ class BuildCMakeExtension(build_ext.build_ext):
     self._is_apple = (platform.system() == 'Darwin')
     (self._mujoco_library_path,
      self._mujoco_include_path,
+     self._mujoco_plugins_path,
      self._mujoco_framework_path) = self._find_mujoco()
     self._configure_cmake()
     for ext in self.extensions:
@@ -148,23 +159,30 @@ class BuildCMakeExtension(build_ext.build_ext):
       self.build_extension(ext)
     self._copy_external_libraries()
     self._copy_mujoco_headers()
+    self._copy_plugin_libraries()
 
   def _find_mujoco(self):
     if MUJOCO_PATH not in os.environ:
-      raise RuntimeError(f'{MUJOCO_PATH} environment variable is not set')
+      raise RuntimeError(
+          f'{MUJOCO_PATH} environment variable is not set')
+    if MUJOCO_PLUGIN_PATH not in os.environ:
+      raise RuntimeError(
+          f'{MUJOCO_PLUGIN_PATH} environment variable is not set')
     library_path = None
     include_path = None
+    plugin_path = os.environ['MUJOCO_PLUGIN_PATH']
     for directory, subdirs, filenames in os.walk(os.environ['MUJOCO_PATH']):
       if self._is_apple and 'mujoco.framework' in subdirs:
         return (os.path.join(directory, 'mujoco.framework/Versions/A'),
                 os.path.join(directory, 'mujoco.framework/Headers'),
+                plugin_path,
                 directory)
       if fnmatch.filter(filenames, get_mujoco_lib_pattern()):
         library_path = directory
       if os.path.exists(os.path.join(directory, 'mujoco/mujoco.h')):
         include_path = directory
       if library_path and include_path:
-        return library_path, include_path, None
+        return library_path, include_path, plugin_path, None
     raise RuntimeError('Cannot find MuJoCo library and/or include paths')
 
   def _copy_external_libraries(self):
@@ -175,6 +193,17 @@ class BuildCMakeExtension(build_ext.build_ext):
           shutil.copyfile(os.path.join(directory, filename),
                           os.path.join(dst, filename))
 
+  def _copy_plugin_libraries(self):
+    dst = os.path.join(
+        os.path.dirname(self.get_ext_fullpath(self.extensions[0].name)),
+        'plugin')
+    os.makedirs(dst)
+    for directory, _, filenames in os.walk(self._mujoco_plugins_path):
+      for pattern in get_plugin_lib_patterns():
+        for filename in fnmatch.filter(filenames, pattern):
+          shutil.copyfile(os.path.join(directory, filename),
+                          os.path.join(dst, filename))
+
   def _copy_mujoco_headers(self):
     dst = os.path.join(
         os.path.dirname(self.get_ext_fullpath(self.extensions[0].name)),