mujoco 3.13.0__cp315-cp315-win_amd64.whl

mujoco/__init__.py

@@ -0,0 +1,256 @@
+# Copyright 2022 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Python bindings for MuJoCo."""
+
+import ctypes
+import ctypes.util
+import os
+import platform
+import subprocess
+from typing import Any, IO, Union, Sequence
+from typing_extensions import TypeAlias
+import warnings
+import zipfile
+
+# Extend the path to enable multiple directories to contribute to the same
+# package. Without this line, the `mujoco-mjx` package would not be able to
+# be discovered by import. For more information, see: https://packaging.python.org/guides/packaging-namespace-packages/#pkgutil-style-namespace-packages
+# NOTE: As per the Python Packaging User Guide linked above, the preferable way
+# of declaring the namespace package is to use the native namespace packages.
+# This seems non-trivial at the current state of the project, however.
+# For more information, see: https://github.com/google-deepmind/mujoco/issues/2119
+__path__ = __import__('pkgutil').extend_path(__path__, __name__)
+
+_SYSTEM = platform.system()
+if _SYSTEM == 'Windows':
+  ctypes.WinDLL(os.path.join(os.path.dirname(__file__), 'mujoco.dll'))
+elif _SYSTEM == 'Darwin':
+  proc_translated = subprocess.run(
+      ['sysctl', '-n', 'sysctl.proc_translated'], capture_output=True).stdout
+  try:
+    is_rosetta = bool(int(proc_translated))
+  except ValueError:
+    is_rosetta = False
+  if is_rosetta and platform.machine() == 'x86_64':
+    raise ImportError(
+        'You are running an x86_64 build of Python on an Apple Silicon '
+        'machine. This is not supported by MuJoCo. Please install and run a '
+        'native, arm64 build of Python.')
+
+from mujoco import _specs
+from mujoco import _structs
+from mujoco._callbacks import *
+from mujoco._constants import *
+from mujoco._enums import *
+from mujoco._errors import *
+from mujoco._functions import *
+from mujoco._specs import *
+from mujoco._structs import *
+
+import numpy as np
+MJTNUM_DTYPE = np.float32 if MJTNUM_BYTES == 4 else np.float64  # pylint: disable=undefined-variable
+
+try:
+  from mujoco._render import *  # pytype: disable=import-error
+except ImportError:
+  pass
+
+try:
+  from mujoco.rendering.classic.renderer import Renderer  # pytype: disable=import-error
+except ImportError:
+  pass
+
+try:
+  from mujoco.rendering.classic.gl_context import *  # pytype: disable=import-error
+except ImportError:
+  pass
+
+MjStruct: TypeAlias = Union[
+    _specs.MjsBody,
+    _specs.MjsFrame,
+    _specs.MjsGeom,
+    _specs.MjsJoint,
+    _specs.MjsLight,
+    _specs.MjsMaterial,
+    _specs.MjsSite,
+    _specs.MjsMesh,
+    _specs.MjsSkin,
+    _specs.MjsTexture,
+    _specs.MjsText,
+    _specs.MjsTuple,
+    _specs.MjsCamera,
+    _specs.MjsFlex,
+    _specs.MjsHField,
+    _specs.MjsKey,
+    _specs.MjsNumeric,
+    _specs.MjsPair,
+    _specs.MjsExclude,
+    _specs.MjsEquality,
+    _specs.MjsTendon,
+    _specs.MjsSensor,
+    _specs.MjsActuator,
+    _specs.MjsPlugin,
+]
+
+
+def to_zip(spec: _specs.MjSpec, file: Union[str, IO[bytes]]) -> None:
+  """Converts an MjSpec to a zip file.
+
+  Args:
+    spec: The mjSpec to save to a file.
+    file: The path to the file to save to or the file object to write to.
+  """
+  files_to_zip = spec.assets
+  files_to_zip[spec.modelname + '.xml'] = spec.to_xml()
+  if isinstance(file, str):
+    directory = os.path.dirname(file)
+    os.makedirs(directory, exist_ok=True)
+    file = open(file, 'wb')
+  with zipfile.ZipFile(file, 'w') as zip_file:
+    for filename, contents in files_to_zip.items():
+      zip_info = zipfile.ZipInfo(filename)
+      zip_file.writestr(zip_info, contents)
+
+
+def from_zip(file: Union[str, IO[bytes]]) -> _specs.MjSpec:
+  """Reads a zip file and returns an MjSpec.
+
+  Args:
+    file: The path to the file to read from or the file object to read from.
+  Returns:
+    An MjSpec object.
+  """
+  assets = {}
+  xml_string = None
+  if isinstance(file, str):
+    file = open(file, 'rb')
+  if not zipfile.is_zipfile(file):
+    raise ValueError(f'File {file} is not a zip file.')
+  with zipfile.ZipFile(file, 'r') as zip_file:
+    xml_dir = None
+    for zip_info in zip_file.infolist():
+      if not zip_info.filename.endswith(os.path.sep):
+        with zip_file.open(zip_info.filename) as f:
+          if zip_info.filename.endswith('.xml'):
+            xml_string = f.read()
+            xml_dir = os.path.dirname(zip_info.filename)
+          else:
+            assets[zip_info.filename] = f.read()
+
+  if not xml_string:
+    raise ValueError('No XML file found in zip file.')
+
+  relative_assets = {}
+  for key, value in assets.items():
+    new_key = os.path.relpath(key, xml_dir)
+    relative_assets[new_key] = value
+  assets = relative_assets
+
+  return _specs.MjSpec.from_string(xml_string, assets=assets)
+
+
+class _MjBindModel:
+  """Wrapper for MjModel that allows binding multiple specs."""
+
+  def __init__(self, elements: Sequence[Any]):
+    object.__setattr__(self, 'elements', elements)
+
+  def __getattr__(self, key: str):
+    items = []
+    for e in self.elements:
+      items.extend(getattr(e, key))
+    return items
+
+  def __setattr__(self, key: str, value: Any):
+    raise AttributeError(f'Cannot set {key} on MjModel.')
+
+
+class _MjBindData:
+  """Wrapper for MjData that allows binding multiple specs."""
+
+  def __init__(self, elements: Sequence[Any]):
+    object.__setattr__(self, 'elements', elements)
+
+  def __getattr__(self, key: str):
+    items = []
+    for e in self.elements:
+      items.extend(getattr(e, key))
+    return items
+
+  def __setattr__(self, key: str, value: Any):
+    value_it = iter(value)
+    for element in self.elements:
+      setattr(element, key, next(value_it))
+
+
+def _bind_model(
+    model: _structs.MjModel, specs: Union[Sequence[MjStruct], MjStruct]
+):
+  """Bind a Mujoco spec to a mjModel.
+
+  Args:
+    model: The mjModel to bind to.
+    specs: The mjSpec elements to use for binding, can be a single element or a
+      sequence.
+  Returns:
+    A MjModelGroupedViews object or a list of the same type.
+  """
+  if isinstance(specs, Sequence):
+    return _MjBindModel([model.bind_scalar(s) for s in specs])
+  else:
+    return model.bind_scalar(specs)
+
+
+def _bind_data(
+    data: _structs.MjData, specs: Union[Sequence[MjStruct], MjStruct]
+):
+  """Bind a Mujoco spec to a mjData.
+
+  Args:
+    data: The mjData to bind to.
+    specs: The mjSpec elements to use for binding, can be a single element or a
+      sequence.
+  Returns:
+    A MjDataGroupedViews object or a list of the same type.
+  """
+  if isinstance(specs, Sequence):
+    return _MjBindData([data.bind_scalar(s) for s in specs])
+  else:
+    return data.bind_scalar(specs)
+
+_specs.MjSpec.from_zip = from_zip
+_specs.MjSpec.to_zip = to_zip
+_structs.MjData.bind = _bind_data  # pyrefly: ignore[bad-assignment]
+_structs.MjModel.bind = _bind_model  # pyrefly: ignore[bad-assignment]
+
+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:
+      if os.path.splitext(filename)[-1] in [".dll", ".dylib", ".so"]:
+        PLUGIN_HANDLES.append(ctypes.CDLL(os.path.join(directory, filename)))
+      elif filename == "__init__.py":
+        pass
+      else:
+        warnings.warn('Ignoring non-library in plugin directory: '
+                      f'{os.path.join(directory, filename)}', ImportWarning)
+
+_load_all_bundled_plugins()
+
+__version__ = mj_versionString()  # pylint: disable=undefined-variable