Change MuJoCo lookup paths to work with version 2.1.1.

saran-t · saran-t · commit 84fc2faa95ca · 2021-12-16T19:22:32.000Z
PiperOrigin-RevId: 416848645
Change-Id: I5e7c7ea60559d5afb40a94d71b59390b8c5760a2
diff --git a/README.md b/README.md
@@ -57,29 +57,35 @@ versions:
 *   3.9
 
 Various people have been successful in getting `dm_control` to work on other
-Linux distros, OS X, and Windows. We do not provide active support for these,
+Linux distros, macOS, and Windows. We do not provide active support for these,
 but will endeavour to answer questions on a best-effort basis.
 
 Follow these steps to install `dm_control`:
 
-1.  Download MuJoCo 2.1.0 from the Download page on the [MuJoCo website]. MuJoCo
-    must be installed before `dm_control`, since `dm_control`'s install script
-    generates Python [`ctypes`] bindings based on MuJoCo's header files. By
-    default, `dm_control` assumes that the MuJoCo archive is extracted into
-    `~/.mujoco`.
+1.  Download MuJoCo 2.1.1 from the
+    [Releases page on the MuJoCo GitHub repository]. MuJoCo must be installed
+    before `dm_control`, since `dm_control`'s install script generates Python
+    [`ctypes`] bindings based on MuJoCo's header files. By default, `dm_control`
+    assumes that MuJoCo is installed via the following instructions:
+
+    -   On Linux, extract the tarball into `~/.mujoco`.
+    -   On Windows, extract the zip archive into either `%HOMEPATH%\MuJoCo` or
+        `%PUBLIC%\MuJoCo`.
+    -   On macOS, either place `MuJoCo.app` into `/Applications`, or place
+        `MuJoCo.Framework` into `~/.mujoco`.
 
 2.  Install the `dm_control` Python package by running `pip install dm_control`.
     We recommend `pip install`ing into a `virtualenv`, or with the `--user` flag
     to avoid interfering with system packages. At installation time,
-    `dm_control` looks for the MuJoCo headers from Step 1 in
-    `~/.mujoco/mujoco210/include`, however this path can be configured with the
-    `headers-dir` command line argument.
+    `dm_control` looks for the MuJoCo headers at the paths described in Step 1
+    by default, however this path can be configured with the `headers-dir`
+    command line argument.
 
-3.  If the shared library provided by MuJoCo (e.g. `libmujoco210.so` or
-    `libmujoco210.dylib`) is installed at a non-default path, specify its
-    location using the `MJLIB_PATH` environment variable. This environment
-    variable should be set to the full path to the library file itself, e.g.
-    `export MJLIB_PATH=/path/to/libmujoco210.so`.
+3.  If the shared library provided by MuJoCo (i.e. `libmujoco.so.2.1.1` or
+    `libmujoco.2.1.1.dylib` or `mujoco.dll`) is installed at a non-default path,
+    specify its location using the `MJLIB_PATH` environment variable. This
+    environment variable should be set to the full path to the library file
+    itself, e.g. `export MJLIB_PATH=/path/to/libmujoco.so.2.1.1`.
 
 ## Versioning
 
@@ -133,6 +139,7 @@ setting the environment variable `EGL_DEVICE_ID=` to the target GPU ID.
     `export DYLD_LIBRARY_PATH=$(brew --prefix)/lib:$DYLD_LIBRARY_PATH`.
 
 [EXT_platform_device]: https://www.khronos.org/registry/EGL/extensions/EXT/EGL_EXT_platform_device.txt
+[Releases page on the MuJoCo GitHub repository]: https://github.com/deepmind/mujoco/releases
 [MuJoCo website]: https://mujoco.org/
 [tech report]: https://arxiv.org/abs/2006.12983
 [`ctypes`]: https://docs.python.org/3/library/ctypes.html
diff --git a/dm_control/mujoco/wrapper/util.py b/dm_control/mujoco/wrapper/util.py
@@ -27,14 +27,6 @@
 from dm_control.utils import io as resources
 
 _PLATFORM = platform.system()
-try:
-  _PLATFORM_SUFFIX = {
-      "Linux": "linux",
-      "Darwin": "macos",
-      "Windows": "win64"
-  }[_PLATFORM]
-except KeyError:
-  raise OSError("Unsupported platform: {}".format(_PLATFORM))
 
 # Environment variable that can be used to override the default path to the
 # MuJoCo shared library.
@@ -58,15 +50,19 @@ def _get_shared_library_filename():
 def _get_default_library_paths():  # pylint: disable=missing-function-docstring
   candidate_paths = []
   if _PLATFORM == "Linux":
+    candidate_paths.append(os.path.expanduser("~/.mujoco/mujoco-2.1.1/lib"))
     candidate_paths.append(os.path.expanduser("~/.mujoco/lib"))
   elif _PLATFORM == "Darwin":
     framework_path = "MuJoCo.Framework/Versions/A"
     candidate_paths.append(
         os.path.join(os.path.expanduser("~/.mujoco"), framework_path))
     candidate_paths.append(
-        os.path.join("/Applications/MuJoCo.App/Frameworks", framework_path))
+        os.path.join(
+            "/Applications/MuJoCo.App/Contents/Frameworks", framework_path))
   elif _PLATFORM == "Windows":
-    candidate_paths.append(os.path.expanduser("~/.mujoco/bin"))
+    candidate_paths.append(os.path.join(
+        os.environ["HOMEDRIVE"], os.environ["HOMEPATH"], "MuJoCo\\bin"))
+    candidate_paths.append(os.path.join(os.environ["PUBLIC"], "MuJoCo\\bin"))
   else:
     raise OSError("Unsupported platform: {}".format(_PLATFORM))
   return [os.path.join(path, _get_shared_library_filename())
diff --git a/setup.py b/setup.py
@@ -19,6 +19,7 @@
 from distutils import log
 import fnmatch
 import os
+import platform
 import subprocess
 import sys
 
@@ -27,12 +28,7 @@
 from setuptools.command import install
 from setuptools.command import test
 
-PLATFORM_SUFFIXES = {
-    'Linux': 'linux',
-    'Windows': 'win64',
-    'Darwin': 'macos',
-}
-DEFAULT_HEADERS_DIR = '~/.mujoco/include'
+PLATFORM = platform.system()
 
 # Relative paths to the binding generator script and the output directory.
 AUTOWRAP_PATH = 'dm_control/autowrap/autowrap.py'
@@ -56,7 +52,28 @@ def _initialize_mjbindings_options(cmd_instance):
   """Set default values for options relating to `build_mjbindings`."""
   # A default value must be assigned to each user option here.
   cmd_instance.inplace = 0
-  cmd_instance.headers_dir = DEFAULT_HEADERS_DIR
+
+  candidate_paths = []
+  if PLATFORM == 'Linux':
+    candidate_paths.append(os.path.expanduser('~/.mujoco/mujoco-2.1.1/include'))
+    candidate_paths.append(os.path.expanduser('~/.mujoco/include'))
+  elif PLATFORM == 'Darwin':
+    framework_path = 'MuJoCo.Framework/Headers'
+    candidate_paths.append(
+        os.path.join(os.path.expanduser('~/.mujoco'), framework_path))
+    candidate_paths.append(os.path.join(
+        '/Applications/MuJoCo.App/Contents/Frameworks', framework_path))
+  elif PLATFORM == 'Windows':
+    candidate_paths.append(os.path.join(
+        os.environ['HOMEDRIVE'], os.environ['HOMEPATH'], 'MuJoCo\\include'))
+    candidate_paths.append(os.path.join(
+        os.environ['PUBLIC'], 'MuJoCo\\include'))
+
+  cmd_instance.headers_dir = ''
+  for path in candidate_paths:
+    if os.path.isdir(path):
+      cmd_instance.headers_dir = path
+      break
 
 
 def _finalize_mjbindings_options(cmd_instance):
@@ -175,7 +192,7 @@ def is_excluded(s):
 
 setup(
     name='dm_control',
-    version='0.0.416463896',
+    version='0.0.416848645',
     description='Continuous control environments and MuJoCo Python bindings.',
     author='DeepMind',
     license='Apache License, Version 2.0',
