chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Version 2.1.2: Python bindings, OBJ assets support, bugfixes.

Browse Source 
PiperOrigin-RevId: 434731612
Change-Id: I0cfda3e7a3d1c72036764986efc252ffa1b8c6b0
This commit is contained in:
Saran Tunyasuvunakool
2022-03-15 14:04:44 +00:00
parent 175d25cd9f
commit 3577e2cf8b
304 changed files with 23906 additions and 1013 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/MANIFEST.in
+3
View File
@@ -0,0 +1,3 @@
include LICENSE *.md
recursive-include mujoco *.h *.cc CMakeLists.txt
recursive-include cmake *.cmake
python/README.md
+62
View File
@@ -0,0 +1,62 @@
# MuJoCo Python Bindings
This package is the canonical Python bindings for the
[MuJoCo physics engine](https://github.com/deepmind/mujoco).
These bindings are developed and maintained by DeepMind, and is kept up-to-date
with the latest developments in MuJoCo itself.
The `mujoco` package provides direct access to raw MuJoCo C API functions,
structs, constants, and enumerations. Structs are provided as Python classes,
with Pythonic initialization and deletion semantics.
It is not the aim of this package to provide fully fledged
scene/environment/game authoring API, as there are already a number of existing
packages that do this well. However, this package does provide a number of
lower-level components outside of MuJoCo itself that are likely to be useful to
most users who access MuJoCo through Python. For example, the `egl`, `glfw`, and
`osmesa` subpackages contain utilities for setting up OpenGL rendering contexts.
## Installation
The package can be installed from [PyPI](https://pypi.org/project/mujoco/) via
```sh
pip install mujoco
```
A copy of the MuJoCo library is provided as part of the package and does **not**
need to be downloaded or installed separately.
If you wish to modify and build the bindings from source, you should clone the
entire `mujoco` repository from GitHub, then run the `make_sdist.sh` script to
generate a [source distribution](https://packaging.python.org/en/latest/glossary/#term-Source-Distribution-or-sdist)
tarball, then run `pip wheel name_of_sdist.tar.gz`. The `make_sdist.sh` script
generates additional C++ header files that are needed to build the bindings,
and also pulls in required files from elsewhere in the repository outside the
`python` directory into the sdist.
CMake and a C++17 compiler are needed to build the bindings from source.
## Usage
Once installed, the package can be imported via `import mujoco`. Please consult
our [documentation](https://mujoco.readthedocs.io/en/latest/python.html) for
further detail on the package's API.
## Versioning
The `major.minor.micro` portion of the version number matches the version of
MuJoCo that the bindings provide. Optionally, if we release updates to the
Python bindings themselves that target the same version of MuJoCo, a `-rN`
suffix is added, for example `2.1.2-r2` represents the second update to the
bindings for MuJoCo 2.1.2.
## License and Disclaimer
Copyright 2022 DeepMind Technologies Limited
MuJoCo and its Python bindings are licensed under the Apache License,
Version 2.0. You may obtain a copy of the License at
https://www.apache.org/licenses/LICENSE-2.0.
This is not an officially supported Google product.
python/make_sdist.sh
Executable
+61
View File
@@ -0,0 +1,61 @@
#!/bin/bash -xe
# 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
#
# https://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.
if [[ -z ${VIRTUAL_ENV} ]]; then
echo "This script must be run from within a Python virtual environment"
exit 1
fi
# Figure out the path to this script (https://stackoverflow.com/a/246128).
package_dir="$(cd -- "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"
if [[ "$(uname)" == CYGWIN* ]]; then
package_dir="$(cygpath -m ${package_dir})"
readonly tmp_dir="$(TMPDIR="${LOCALAPPDATA//\\/$'/'}/Temp" mktemp -d)"
else
readonly tmp_dir="$(mktemp -d)"
fi
python -m pip install --upgrade pip
python -m pip install absl-py
pushd ${tmp_dir}
cp -r "${package_dir}"/* .
# Generate header files.
old_pythonpath="${PYTHONPATH}"
if [[ "$(uname)" == CYGWIN* ]]; then
export PYTHONPATH="${old_pythonpath};${package_dir}/.."
else
export PYTHONPATH="${old_pythonpath}:${package_dir}/.."
fi
python "${package_dir}"/mujoco/codegen/generate_enum_traits.py > \
mujoco/enum_traits.h
python "${package_dir}"/mujoco/codegen/generate_function_traits.py > \
mujoco/function_traits.h
export PYTHONPATH="${old_pythonpath}"
# Copy over the LICENSE file.
cp "${package_dir}"/../../LICENSE .
# Copy over CMake scripts.
mkdir cmake
cp "${package_dir}"/../../cmake/*.cmake cmake
python setup.py sdist --formats=gztar
tar -tf dist/mujoco-*.tar.gz
popd
mkdir -p "${package_dir}"/dist
mv "${tmp_dir}"/dist/* "${package_dir}"/dist
python/mujoco/CMakeLists.txt
+381
View File
@@ -0,0 +1,381 @@
# 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
#
# https://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.
project(mujoco_python)
cmake_minimum_required(VERSION 3.15)
# Support new IN_LIST if() operator.
set(CMAKE_POLICY_DEFAULT_CMP0057 NEW)
# INTERPROCEDURAL_OPTIMIZATION is enforced when enabled.
set(CMAKE_POLICY_DEFAULT_CMP0069 NEW)
enable_language(C)
enable_language(CXX)
if(MSVC AND MSVC_VERSION GREATER_EQUAL 1927)
set(CMAKE_CXX_STANDARD 20) # For forceinline lambdas.
else()
set(CMAKE_CXX_STANDARD 17)
endif()
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
set(CMAKE_C_VISIBILITY_PRESET hidden)
set(CMAKE_CXX_VISIBILITY_PRESET hidden)
set(CMAKE_VISIBILITY_INLINES_HIDDEN ON)
# TODO(stunya) Figure out why this is required.
separate_arguments(CMDLINE_LINK_OPTIONS UNIX_COMMAND ${CMAKE_SHARED_LINKER_FLAGS})
add_link_options(${CMDLINE_LINK_OPTIONS})
include(MujocoLinkOptions)
get_mujoco_extra_link_options(EXTRA_LINK_OPTIONS)
add_link_options(${EXTRA_LINK_OPTIONS})
if(APPLE)
add_compile_options(-Werror=partial-availability -Werror=unguarded-availability)
add_link_options(-Wl,-no_weak_imports)
endif()
find_package(Python3 COMPONENTS Interpreter Development)
include(FindOrFetch)
# ==================== MUJOCO LIBRARY ==========================================
if(NOT TARGET mujoco)
find_library(MUJOCO_LIBRARY mujoco mujoco.2.1.2 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED)
find_path(MUJOCO_INCLUDE mujoco.h HINTS ${MUJOCO_INCLUDE_DIR} REQUIRED)
message("MuJoCo is at ${MUJOCO_LIBRARY}")
message("MuJoCo headers are at ${MUJOCO_INCLUDE}")
add_library(mujoco SHARED IMPORTED)
if(WIN32)
set_target_properties(mujoco PROPERTIES IMPORTED_IMPLIB "${MUJOCO_LIBRARY}")
else()
set_target_properties(mujoco PROPERTIES IMPORTED_LOCATION "${MUJOCO_LIBRARY}")
endif()
target_include_directories(mujoco INTERFACE "${MUJOCO_INCLUDE}")
if(APPLE)
execute_process(
COMMAND otool -XD ${MUJOCO_LIBRARY}
COMMAND head -n 1
COMMAND xargs dirname
COMMAND xargs echo -n
OUTPUT_VARIABLE MUJOCO_INSTALL_NAME_DIR
)
set_target_properties(mujoco PROPERTIES INSTALL_NAME_DIR "${MUJOCO_INSTALL_NAME_DIR}")
elseif(UNIX)
execute_process(
COMMAND objdump -p ${MUJOCO_LIBRARY}
COMMAND grep SONAME
COMMAND grep -Po [^\\s]+$
COMMAND xargs echo -n
OUTPUT_VARIABLE MUJOCO_SONAME
)
set_target_properties(mujoco PROPERTIES IMPORTED_SONAME "${MUJOCO_SONAME}")
endif()
endif()
# ==================== ABSEIL ==================================================
if(APPLE)
set(ABSL_EXTRA_FETCH_ARGS
PATCH_COMMAND
"sed"
"-i"
" "
"s/-march=armv8-a+crypto/-mcpu=apple-m1+crypto/g"
"${CMAKE_BINARY_DIR}/_deps/abseil-cpp-src/absl/copts/GENERATED_AbseilCopts.cmake"
)
else()
set(ABSL_EXTRA_FETCH_ARGS "")
endif()
findorfetch(
USE_SYSTEM_PACKAGE
OFF
PACKAGE_NAME
absl
LIBRARY_NAME
abseil-cpp
GIT_REPO
https://github.com/abseil/abseil-cpp
GIT_TAG
215105818dfde3174fe799600bb0f3cae233d0bf # 20211102.0
TARGETS
absl::core_headers
absl::flat_hash_map
absl::span
${ABSL_EXTRA_FETCH_ARGS}
EXCLUDE_FROM_ALL
)
# ==================== EIGEN ===================================================
add_compile_definitions(EIGEN_MPL2_ONLY)
findorfetch(
USE_SYSTEM_PACKAGE
OFF
PACKAGE_NAME
Eigen3
LIBRARY_NAME
eigen
GIT_REPO
https://gitlab.com/libeigen/eigen
GIT_TAG
3147391d946bb4b6c68edd901f2add6ac1f31f8c # 3.4.0
TARGETS
Eigen3::Eigen
EXCLUDE_FROM_ALL
)
# ==================== PYBIND11 ================================================
if(MUJOCO_PYBIND11_DIR)
FetchContent_Declare(pybind11 SOURCE_DIR ${MUJOCO_PYBIND11_DIR} EXCLUDE_FROM_ALL)
FetchContent_MakeAvailable(pybind11)
else()
findorfetch(
USE_SYSTEM_PACKAGE
OFF
PACKAGE_NAME
pybind11
LIBRARY_NAME
pybind11
GIT_REPO
https://github.com/pybind/pybind11
GIT_TAG
4c4b33f14215f8992e2c67ace2a7a4c114c48bb9
TARGETS
pybind11::pybind11_headers
EXCLUDE_FROM_ALL
)
endif()
# ==================== MUJOCO PYTHON BINDINGS ==================================
add_subdirectory(util)
if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/enum_traits.h)
add_library(enum_traits INTERFACE enum_traits.h)
else()
add_custom_command(
OUTPUT enum_traits.h
COMMAND ${CMAKE_COMMAND} -E env PYTHONPATH=${mujoco_SOURCE_DIR}/mujoco ${Python3_EXECUTABLE}
${CMAKE_CURRENT_SOURCE_DIR}/codegen/generate_enum_traits.py > enum_traits.h
DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/codegen/generate_enum_traits.py
)
add_library(enum_traits INTERFACE enum_traits.h)
target_include_directories(
enum_traits INTERFACE ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR}
)
endif()
target_link_libraries(enum_traits INTERFACE mujoco absl::core_headers)
if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/function_traits.h)
add_library(function_traits INTERFACE function_traits.h)
else()
add_custom_command(
OUTPUT function_traits.h
COMMAND ${CMAKE_COMMAND} -E env PYTHONPATH=${mujoco_SOURCE_DIR}/mujoco ${Python3_EXECUTABLE}
${CMAKE_CURRENT_SOURCE_DIR}/codegen/generate_function_traits.py > function_traits.h
DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/codegen/generate_function_traits.py
)
add_library(function_traits INTERFACE function_traits.h)
target_include_directories(
function_traits INTERFACE ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR}
)
endif()
target_link_libraries(function_traits INTERFACE mujoco absl::core_headers)
add_library(errors_header INTERFACE errors.h)
set_target_properties(errors_header PROPERTIES PUBLIC_HEADER errors.h)
target_link_libraries(errors_header INTERFACE crossplatform func_wrap mujoco)
add_library(raw INTERFACE raw.h)
set_target_properties(raw PROPERTIES PUBLIC_HEADER raw.h)
target_link_libraries(raw INTERFACE mujoco)
add_library(
structs_header
INTERFACE
indexer_xmacro.h
indexers.h
mjdata_meta.h
structs.h
)
set_target_properties(structs_header PROPERTIES PUBLIC_HEADER structs.h)
target_link_libraries(
structs_header
INTERFACE absl::flat_hash_map
absl::span
mujoco
raw
)
add_library(functions_header INTERFACE functions.h)
set_target_properties(functions_header PROPERTIES PUBLIC_HEADER functions.h)
target_link_libraries(
functions_header
INTERFACE array_traits
crossplatform
Eigen3::Eigen
errors_header
func_wrap
structs_header
tuple_tools
)
include(CheckAvxSupport)
get_avx_compile_options(AVX_COMPILE_OPTIONS)
macro(mujoco_pybind11_module name)
pybind11_add_module(${name} ${ARGN})
target_compile_options(${name} PRIVATE ${AVX_COMPILE_OPTIONS})
set_target_properties(${name} PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
if(APPLE)
add_custom_command(
TARGET ${name}
POST_BUILD
COMMAND
install_name_tool -change
$<TARGET_PROPERTY:mujoco,INSTALL_NAME_DIR>/$<TARGET_FILE_NAME:mujoco>
@rpath/$<TARGET_FILE_NAME:mujoco> -add_rpath @loader_path $<TARGET_FILE:${name}>
)
elseif(NOT WIN32)
add_custom_command(
TARGET ${name}
POST_BUILD
COMMAND patchelf --remove-needed $<TARGET_SONAME_FILE_NAME:mujoco>
$<TARGET_FILE_NAME:${name}>
)
endif()
endmacro()
mujoco_pybind11_module(_callbacks callbacks.cc)
target_link_libraries(
_callbacks
PRIVATE errors_header
mujoco
raw
structs_header
)
mujoco_pybind11_module(_constants constants.cc)
target_link_libraries(_constants PRIVATE mujoco)
mujoco_pybind11_module(_enums enums.cc)
target_link_libraries(
_enums
PRIVATE crossplatform
enum_traits
mujoco
tuple_tools
)
mujoco_pybind11_module(_errors errors.cc)
target_link_libraries(_errors PRIVATE errors_header)
mujoco_pybind11_module(_functions functions.cc)
target_link_libraries(
_functions
PRIVATE Eigen3::Eigen
functions_header
function_traits
mujoco
raw
)
if(APPLE)
# C++17 aligned allocation is not available until macOS 10.14.
target_compile_options(_functions PRIVATE -fno-aligned-allocation)
endif()
mujoco_pybind11_module(_render render.cc)
target_link_libraries(
_render
PRIVATE Eigen3::Eigen
errors_header
functions_header
function_traits
mujoco
raw
structs_header
)
mujoco_pybind11_module(_rollout rollout.cc)
target_link_libraries(_rollout PRIVATE functions_header mujoco raw)
mujoco_pybind11_module(
_structs
indexers.cc
serialization.h
structs.cc
)
target_link_libraries(
_structs
PRIVATE absl::flat_hash_map
mujoco
raw
errors_header
func_wrap
function_traits
structs_header
)
set(LIBRARIES_FOR_WHEEL
"$<TARGET_FILE:_callbacks>"
"$<TARGET_FILE:_constants>"
"$<TARGET_FILE:_enums>"
"$<TARGET_FILE:_errors>"
"$<TARGET_FILE:_functions>"
"$<TARGET_FILE:_render>"
"$<TARGET_FILE:_rollout>"
"$<TARGET_FILE:_structs>"
"$<TARGET_FILE:mujoco>"
)
if(NOT APPLE)
set(LIBRARIES_FOR_WHEEL ${LIBRARIES_FOR_WHEEL} $<TARGET_FILE:mujoco_nogl>)
endif()
if(NOT (APPLE OR WIN32))
set(LIBRARIES_FOR_WHEEL
${LIBRARIES_FOR_WHEEL}
"$<TARGET_FILE:glew>"
"$<TARGET_FILE:glewegl>"
"$<TARGET_FILE:glewosmesa>"
)
endif()
if(MUJOCO_PYTHON_MAKE_WHEEL)
add_custom_target(
wheel ALL
COMMAND "${CMAKE_COMMAND}" -E rm -rf "${CMAKE_CURRENT_BINARY_DIR}/dist"
COMMAND "${CMAKE_COMMAND}" -E copy_directory "${CMAKE_CURRENT_SOURCE_DIR}/.."
"${CMAKE_CURRENT_BINARY_DIR}/dist"
COMMAND "${CMAKE_COMMAND}" -E copy "${CMAKE_CURRENT_SOURCE_DIR}/../../LICENSE"
"${CMAKE_CURRENT_BINARY_DIR}/dist/LICENSE"
COMMAND "${CMAKE_COMMAND}" -E copy ${LIBRARIES_FOR_WHEEL}
"${CMAKE_CURRENT_BINARY_DIR}/dist/mujoco"
COMMAND "${Python3_EXECUTABLE}" -m pip wheel --wheel-dir "${CMAKE_BINARY_DIR}" --no-deps -vvv
"${CMAKE_CURRENT_BINARY_DIR}/dist"
)
add_dependencies(
wheel
_callbacks
_constants
_enums
_errors
_functions
_render
_rollout
_structs
mujoco
)
endif()
python/mujoco/__init__.py
+82
View File
@@ -0,0 +1,82 @@
# 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 os
import platform
import subprocess
HEADERS_DIR = os.path.join(os.path.dirname(__file__), 'include')
_MUJOCO_GL_ENABLE = ('enable', 'enabled', 'on', 'true', '1' , '')
_MUJOCO_GL_DISABLE = ('disable', 'disabled', 'off', 'false', '0')
_MUJOCO_GL = os.environ.get('MUJOCO_GL', '').lower().strip()
_MUJOCO_GL_IS_VALID = True
_SYSTEM = platform.system()
print(_MUJOCO_GL)
if _SYSTEM == 'Linux':
libglew_name = None
if _MUJOCO_GL in _MUJOCO_GL_ENABLE + ('glfw', 'glx'):
libglew_name = 'libglew.so'
elif _MUJOCO_GL == 'egl':
libglew_name = 'libglewegl.so'
elif _MUJOCO_GL == 'osmesa':
libglew_name = 'libglewosmesa.so'
elif _MUJOCO_GL not in _MUJOCO_GL_DISABLE:
_MUJOCO_GL_IS_VALID = False
if libglew_name is not None:
ctypes.CDLL(os.path.join(os.path.dirname(__file__), libglew_name),
ctypes.RTLD_GLOBAL)
ctypes.CDLL(
os.path.join(os.path.dirname(__file__), 'libmujoco.so.2.1.2'),
ctypes.RTLD_GLOBAL)
else:
ctypes.CDLL(
os.path.join(os.path.dirname(__file__), 'libmujoco_nogl.so.2.1.2'),
ctypes.RTLD_GLOBAL)
elif _SYSTEM == 'Windows':
if _MUJOCO_GL in _MUJOCO_GL_ENABLE + ('glfw', 'wgl'):
ctypes.WinDLL(os.path.join(os.path.dirname(__file__), 'mujoco.dll'))
elif _MUJOCO_GL in _MUJOCO_GL_DISABLE:
ctypes.WinDLL(os.path.join(os.path.dirname(__file__), 'mujoco_nogl.dll'))
else:
_MUJOCO_GL_IS_VALID = False
if not _MUJOCO_GL_IS_VALID:
raise RuntimeError(
f'invalid value for environment variable MUJOCO_GL: {_MUJOCO_GL}')
from mujoco._callbacks import *
from mujoco._constants import *
from mujoco._enums import *
from mujoco._errors import *
from mujoco._functions import *
from mujoco._structs import *
# pylint: disable=g-import-not-at-top
if _MUJOCO_GL not in _MUJOCO_GL_DISABLE:
from mujoco._render import *
if _SYSTEM != 'Linux':
from mujoco.glfw import GLContext
else:
_dl_handle = ctypes.CDLL(None)
if hasattr(_dl_handle, 'OSMesaCreateContextExt'):
from mujoco.osmesa import GLContext
elif hasattr(_dl_handle, 'eglCreateContext'):
from mujoco.egl import GLContext
else:
from mujoco.glfw import GLContext
python/mujoco/bindings_test.py
+847
View File
@@ -0,0 +1,847 @@
# 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.
# ==============================================================================
"""Tests for MuJoCo Python bindings."""
import contextlib
import copy
import pickle
import sys
from absl.testing import absltest
from absl.testing import parameterized
import mujoco
import numpy as np
TEST_XML = r"""
<mujoco model="test">
<compiler coordinate="local" angle="radian" eulerseq="xyz"/>
<option timestep="0.002" gravity="0 0 -9.81"/>
<visual>
<global fovy="50" />
<quality shadowsize="51" />
</visual>
<worldbody>
<geom name="myplane" type="plane" size="10 10 1"/>
<body pos="0 0 0.1">
<geom name="mybox" type="box" size="0.1 0.1 0.1" mass="0.25"/>
<freejoint name="myfree"/>
</body>
<body>
<inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
<site pos="0 0 -1" name="mysite" type="sphere"/>
<joint type="hinge" axis="0 1 0"/>
</body>
<body>
<inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
<joint name="myball" type="ball"/>
</body>
<body mocap="true" pos="42 0 42">
<geom type="sphere" size="0.1"/>
</body>
</worldbody>
</mujoco>
"""
TEST_XML_SENSOR = r"""
<mujoco model="test">
<worldbody>
<geom name="myplane" type="plane" size="10 10 1"/>
</worldbody>
<sensor>
<user objtype="geom" objname="myplane"
datatype="real" needstage="vel" dim="1"/>
</sensor>
</mujoco>
"""
@contextlib.contextmanager
def temporary_callback(setter, callback):
setter(callback)
yield
setter(None)
class MuJoCoBindingsTest(parameterized.TestCase):
def setUp(self):
super().setUp()
self.model: mujoco.MjModel = mujoco.MjModel.from_xml_string(TEST_XML)
self.data = mujoco.MjData(self.model)
def test_load_xml_can_handle_name_clash(self):
xml_1 = r"""
<mujoco>
<worldbody>
<geom name="plane" type="plane" size="1 1 1"/>
<include file="model_.xml"/>
<include file="model__.xml"/>
</worldbody>
</mujoco>"""
xml_2 = rb"""<mujoco><geom name="box" type="box" size="1 1 1"/></mujoco>"""
xml_3 = rb"""<mujoco><geom name="ball" type="sphere" size="1"/></mujoco>"""
model = mujoco.MjModel.from_xml_string(
xml_1, {'model_.xml': xml_2, 'model__.xml': xml_3})
self.assertEqual(
mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'plane'), 0)
self.assertEqual(
mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'box'), 1)
self.assertEqual(
mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'ball'), 2)
def test_can_read_array(self):
np.testing.assert_array_equal(
self.model.body_pos,
[[0, 0, 0], [0, 0, 0.1], [0, 0, 0], [0, 0, 0], [42.0, 0, 42.0]])
def test_can_set_array(self):
self.data.qpos = 0.12345
np.testing.assert_array_equal(
self.data.qpos, [0.12345]*len(self.data.qpos))
def test_array_is_a_view(self):
qpos_ref = self.data.qpos
self.data.qpos = 0.789
np.testing.assert_array_equal(
qpos_ref, [0.789]*len(self.data.qpos))
def test_array_keeps_struct_alive(self):
model = mujoco.MjModel.from_xml_string(TEST_XML)
qpos0 = model.qpos0
qpos_spring = model.qpos_spring
# This only fails reliably under ASAN, which detects heap-use-after-free.
# However, often the assertEqual is enough since the memory block is
# already reused between mjModel deallocation and the subsequent read.
qpos0[:] = 1
del model
self.assertEqual(qpos0[0], 1)
# When running under test coverage tools, the refcount of objects can be
# higher than normal. To take this into account, we first measure the
# refcount of a dummy object with no other referrer.
dummy = []
base_refcount = sys.getrefcount(dummy) - 1
# Here `base` is actually a PyCapsule that holds the raw mjModel* rather
# than the actual MjModel wrapper object itself.
capsule = qpos0.base
self.assertEqual(sys.getrefcount(capsule) - base_refcount, 3)
del qpos0
self.assertEqual(sys.getrefcount(capsule) - base_refcount, 2)
del qpos_spring
self.assertEqual(sys.getrefcount(capsule) - base_refcount, 1)
def test_named_indexing_geom_size(self):
box_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, 'mybox')
self.assertIs(self.model.geom('mybox'), self.model.geom(box_id))
self.assertIs(self.model.geom('mybox').size, self.model.geom(box_id).size)
self.assertEqual(self.model.geom('mybox').size.shape, (3,))
# Test that the indexer is returning a view into the underlying struct.
size_from_indexer = self.model.geom('mybox').size
self.model.geom_size[box_id] = [7, 11, 13]
np.testing.assert_array_equal(size_from_indexer, [7, 11, 13])
self.model.geom('mybox').size = [5, 3, 2]
np.testing.assert_array_equal(self.model.geom_size[box_id], [5, 3, 2])
def test_named_indexing_geom_quat(self):
box_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, 'mybox')
self.assertIs(self.model.geom('mybox'), self.model.geom(box_id))
self.assertIs(self.model.geom('mybox').quat, self.model.geom(box_id).quat)
self.assertEqual(self.model.geom('mybox').quat.shape, (4,))
# Test that the indexer is returning a view into the underlying struct.
quat_from_indexer = self.model.geom('mybox').quat
self.model.geom_quat[box_id] = [5, 10, 15, 20]
np.testing.assert_array_equal(quat_from_indexer, [5, 10, 15, 20])
self.model.geom('mybox').quat = [12, 9, 6, 3]
np.testing.assert_array_equal(self.model.geom_quat[box_id], [12, 9, 6, 3])
def test_named_indexing_ragged_qpos(self):
balljoint_id = mujoco.mj_name2id(
self.model, mujoco.mjtObj.mjOBJ_JOINT, 'myball')
self.assertIs(self.data.joint('myball'), self.data.joint(balljoint_id))
self.assertIs(self.data.joint('myball').qpos,
self.data.joint(balljoint_id).qpos)
self.assertEqual(self.data.joint('myball').qpos.shape, (4,))
# Test that the indexer is returning a view into the underlying struct.
qpos_from_indexer = self.data.joint('myball').qpos
qpos_idx = self.model.jnt_qposadr[balljoint_id]
self.data.qpos[qpos_idx:qpos_idx+4] = [4, 5, 6, 7]
np.testing.assert_array_equal(qpos_from_indexer, [4, 5, 6, 7])
self.data.joint('myball').qpos = [9, 8, 7, 6]
np.testing.assert_array_equal(self.data.qpos[qpos_idx:qpos_idx+4],
[9, 8, 7, 6])
def test_named_indexing_ragged2d_cdof(self):
freejoint_id = mujoco.mj_name2id(
self.model, mujoco.mjtObj.mjOBJ_JOINT, 'myfree')
self.assertIs(self.data.joint('myfree'), self.data.joint(freejoint_id))
self.assertIs(self.data.joint('myfree').cdof,
self.data.joint(freejoint_id).cdof)
self.assertEqual(self.data.joint('myfree').cdof.shape, (6, 6))
# Test that the indexer is returning a view into the underlying struct.
cdof_from_indexer = self.data.joint('myfree').cdof
dof_idx = self.model.jnt_dofadr[freejoint_id]
self.data.cdof[dof_idx:dof_idx+6, :] = np.reshape(range(36), (6, 6))
np.testing.assert_array_equal(cdof_from_indexer,
np.reshape(range(36), (6, 6)))
self.data.joint('myfree').cdof = 42
np.testing.assert_array_equal(self.data.cdof[dof_idx:dof_idx+6], [[42]*6]*6)
def test_addresses_differ_between_structs(self):
model2 = mujoco.MjModel.from_xml_string(TEST_XML)
data2 = mujoco.MjData(model2)
self.assertGreater(self.model._address, 0)
self.assertGreater(self.data._address, 0)
self.assertGreater(model2._address, 0)
self.assertGreater(data2._address, 0)
self.assertLen({self.model._address, self.data._address,
model2._address, data2._address}, 4)
def test_mjmodel_can_read_and_write_opt(self):
self.assertEqual(self.model.opt.timestep, 0.002)
np.testing.assert_array_equal(self.model.opt.gravity, [0, 0, -9.81])
opt = self.model.opt
self.model.opt.timestep = 0.001
self.assertEqual(opt.timestep, 0.001)
gravity = opt.gravity
self.model.opt.gravity[1] = 0.1
np.testing.assert_array_equal(gravity, [0, 0.1, -9.81])
self.model.opt.gravity = 0.2
np.testing.assert_array_equal(gravity, [0.2, 0.2, 0.2])
def test_mjmodel_can_read_and_write_stat(self):
self.assertNotEqual(self.model.stat.meanmass, 0)
stat = self.model.stat
self.model.stat.meanmass = 1.2
self.assertEqual(stat.meanmass, 1.2)
def test_mjmodel_can_read_and_write_vis(self):
self.assertEqual(self.model.vis.quality.shadowsize, 51)
self.model.vis.quality.shadowsize = 100
self.assertEqual(self.model.vis.quality.shadowsize, 100)
def test_mjmodel_can_access_names_directly(self):
# mjModel offers direct access to names array, to allow usecases other than
# id2name
model_name = str(self.model.names[0:self.model.names.find(b'\0')], 'utf-8')
self.assertEqual(model_name, 'test')
start_index = self.model.name_geomadr[0]
end_index = self.model.names.find(b'\0', start_index)
geom_name = str(self.model.names[start_index:end_index], 'utf-8')
self.assertEqual(geom_name, 'myplane')
def test_mjmodel_names_doesnt_copy(self):
names = self.model.names
self.assertIs(names, self.model.names)
def test_vis_global_exposed_as_global_(self):
self.assertEqual(self.model.vis.global_.fovy, 50)
self.model.vis.global_.fovy = 100
self.assertEqual(self.model.vis.global_.fovy, 100)
def test_mjoption_can_make_default(self):
opt = mujoco.MjOption()
self.assertEqual(opt.timestep, 0.002)
np.testing.assert_array_equal(opt.gravity, [0, 0, -9.81])
def test_mjoption_can_copy(self):
opt1 = mujoco.MjOption()
opt1.timestep = 0.001
opt1.gravity = 2
opt2 = copy.copy(opt1)
self.assertEqual(opt2.timestep, 0.001)
np.testing.assert_array_equal(opt2.gravity, [2, 2, 2])
# Make sure opt2 is actually a copy.
opt1.timestep = 0.005
opt1.gravity = 5
self.assertEqual(opt2.timestep, 0.001)
np.testing.assert_array_equal(opt2.gravity, [2, 2, 2])
def test_mjmodel_can_copy(self):
model_copy = copy.copy(self.model)
self.assertEqual(
mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_JOINT, 0),
'myfree')
self.assertEqual(
mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 0),
'myplane')
self.assertEqual(
mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 1),
'mybox')
# Make sure it's a copy.
self.model.geom_size[1] = 0.5
np.testing.assert_array_equal(self.model.geom_size[1], [0.5, 0.5, 0.5])
np.testing.assert_array_equal(model_copy.geom_size[1], [0.1, 0.1, 0.1])
def test_assets_array_filename_too_long(self):
# Longest allowed filename (excluding null byte)
limit = mujoco.mjMAXVFSNAME - 1
contents = b'<mujoco/>'
valid_filename = 'a' * limit
mujoco.MjModel.from_xml_path(valid_filename, {valid_filename: contents})
invalid_filename = 'a' * (limit + 1)
expected_message = (
f'Filename length 1000 exceeds 999 character limit: {invalid_filename}')
with self.assertRaisesWithLiteralMatch(ValueError, expected_message):
mujoco.MjModel.from_xml_path(invalid_filename,
{invalid_filename: contents})
def test_mjdata_can_copy(self):
self.data.qpos = [0, 0, 0.1*np.sqrt(2) - 0.001,
np.cos(np.pi/8), np.sin(np.pi/8), 0, 0, 0,
1, 0, 0, 0]
mujoco.mj_forward(self.model, self.data)
data_copy = copy.copy(self.data)
self.assertEqual(data_copy.ncon, 2)
# Make sure it's a copy.
mujoco.mj_resetData(self.model, self.data)
mujoco.mj_forward(self.model, self.data)
mujoco.mj_forward(self.model, data_copy)
self.assertEqual(self.data.ncon, 4)
self.assertEqual(data_copy.ncon, 2)
mujoco.mj_resetData(self.model, data_copy)
mujoco.mj_forward(self.model, data_copy)
self.assertEqual(data_copy.ncon, 4)
def test_mjdata_can_read_warning_array(self):
warnings = self.data.warning
self.assertLen(warnings, mujoco.mjtWarning.mjNWARNING)
self.data.qpos[0] = float('NaN')
mujoco.mj_checkPos(self.model, self.data)
self.assertEqual(warnings[mujoco.mjtWarning.mjWARN_BADQPOS].number, 1)
def test_mjcontact_can_copy(self):
mujoco.mj_forward(self.model, self.data)
contact_copy = []
for i in range(4):
contact_copy.append(copy.copy(self.data.contact[i]))
# Sort contacts in anticlockwise order
contact_copy = sorted(
contact_copy, key=lambda x: np.arctan2(x.pos[1], x.pos[0]))
np.testing.assert_allclose(contact_copy[0].pos[:2], [-0.1, -0.1])
np.testing.assert_allclose(contact_copy[1].pos[:2], [0.1, -0.1])
np.testing.assert_allclose(contact_copy[2].pos[:2], [0.1, 0.1])
np.testing.assert_allclose(contact_copy[3].pos[:2], [-0.1, 0.1])
# Make sure they're actually copies.
for i in range(4):
self.data.contact[i].pos[:2] = 55
np.testing.assert_allclose(self.data.contact[0].pos[:2], [55, 55])
np.testing.assert_allclose(self.data.contact[1].pos[:2], [55, 55])
np.testing.assert_allclose(self.data.contact[2].pos[:2], [55, 55])
np.testing.assert_allclose(self.data.contact[3].pos[:2], [55, 55])
np.testing.assert_allclose(contact_copy[0].pos[:2], [-0.1, -0.1])
np.testing.assert_allclose(contact_copy[1].pos[:2], [0.1, -0.1])
np.testing.assert_allclose(contact_copy[2].pos[:2], [0.1, 0.1])
np.testing.assert_allclose(contact_copy[3].pos[:2], [-0.1, 0.1])
def test_mj_step(self):
displacement = 0.25
self.data.qpos[2] += displacement
mujoco.mj_forward(self.model, self.data)
gravity = -self.model.opt.gravity[2]
expected_contact_time = np.sqrt(2 * displacement / gravity)
# Grab a reference to the contacts upfront so that we know that they're
# a view into mjData rather than a copy.
contact = self.data.contact[:4]
self.model.opt.timestep = 2**-9 # 0.001953125; allows exact comparisons
self.assertEqual(self.data.time, 0)
while self.data.time < expected_contact_time:
self.assertEqual(self.data.ncon, 0)
prev_time = self.data.time
mujoco.mj_step(self.model, self.data)
self.assertEqual(self.data.time, prev_time + self.model.opt.timestep)
mujoco.mj_forward(self.model, self.data)
self.assertEqual(self.data.ncon, 4)
# Sort contacts in anticlockwise order
sorted_contact = sorted(
contact, key=lambda x: np.arctan2(x.pos[1], x.pos[0]))
np.testing.assert_allclose(sorted_contact[0].pos[:2], [-0.1, -0.1])
np.testing.assert_allclose(sorted_contact[1].pos[:2], [0.1, -0.1])
np.testing.assert_allclose(sorted_contact[2].pos[:2], [0.1, 0.1])
np.testing.assert_allclose(sorted_contact[3].pos[:2], [-0.1, 0.1])
def test_mj_step_multiple(self):
self.model.opt.timestep = 2**-9 # 0.001953125; allows exact comparisons
self.assertEqual(self.data.time, 0)
for _ in range(10):
prev_time = self.data.time
mujoco.mj_step(self.model, self.data, nstep=7)
self.assertEqual(self.data.time, prev_time + 7 * self.model.opt.timestep)
self.assertIn('Optionally, repeat nstep times.', mujoco.mj_step.__doc__)
def test_mj_contact_list(self):
self.assertLen(self.data.contact, self.model.nconmax)
expected_pos = []
for contact in self.data.contact:
expected_pos.append(np.random.uniform(size=3))
contact.pos = expected_pos[-1]
np.testing.assert_array_equal(self.data.contact.pos, expected_pos)
expected_friction = []
for contact in self.data.contact:
expected_friction.append(np.random.uniform(size=5))
contact.friction = expected_friction[-1]
np.testing.assert_array_equal(self.data.contact.friction, expected_friction)
expected_H = [] # pylint: disable=invalid-name
for contact in self.data.contact:
expected_H.append(np.random.uniform(size=36))
contact.H = expected_H[-1]
np.testing.assert_array_equal(self.data.contact.H, expected_H)
def test_mj_struct_list_equality(self):
model2 = mujoco.MjModel.from_xml_string(TEST_XML)
data2 = mujoco.MjData(model2)
mujoco.mj_forward(self.model, self.data)
self.assertEqual(self.data.ncon, 4)
mujoco.mj_forward(model2, data2)
self.assertEqual(data2.ncon, 4)
self.assertEqual(data2.contact[:4], self.data.contact[:4])
self.data.qpos[3:7] = [np.cos(np.pi/8), np.sin(np.pi/8), 0, 0]
self.data.qpos[2] *= (np.sqrt(2) - 1) * 0.1 - 1e-6
mujoco.mj_forward(self.model, self.data)
self.assertEqual(self.data.ncon, 2)
self.assertNotEqual(data2.contact[:2], self.data.contact[:2])
# Check that we can compare slices of different lengths
self.assertNotEqual(data2.contact[:2], self.data.contact[:4])
# Check that comparing things of different types do not raise an error
self.assertNotEqual(self.data.contact, self.data.warning)
self.assertNotEqual(self.data.contact, 5)
@parameterized.named_parameters([
('MjOption', mujoco.MjOption, 'tolerance'),
('MjWarningStat', mujoco.MjWarningStat, 'number'),
('MjTimerStat', mujoco.MjTimerStat, 'number'),
('MjSolverStat', mujoco.MjSolverStat, 'neval'),
('MjContact', mujoco.MjContact, 'dist'),
('MjStatistic', mujoco.MjStatistic, 'extent'),
('MjLROpt', mujoco.MjLROpt, 'maxforce'),
('MjvPerturb', mujoco.MjvPerturb, 'select'),
('MjvCamera', mujoco.MjvCamera, 'fixedcamid'),
])
def test_mj_struct_equality(self, cls, attr):
struct = cls()
struct2 = cls()
setattr(struct, attr, 1)
self.assertNotEqual(struct, struct2)
setattr(struct2, attr, 1)
self.assertEqual(struct, struct2)
self.assertNotEqual(struct, 3)
self.assertNotEqual(struct, None)
# mutable structs shouldn't declare __hash__
with self.assertRaises(TypeError):
hash(struct)
def test_mj_struct_equality_array(self):
contact1 = mujoco.MjContact()
contact2 = mujoco.MjContact()
contact1.H[3] = 1
self.assertNotEqual(contact1, contact2)
contact2.H[3] = 1
self.assertEqual(contact1, contact2)
@parameterized.named_parameters([
('MjOption', mujoco.MjOption, 'tolerance'),
('MjWarningStat', mujoco.MjWarningStat, 'number'),
('MjTimerStat', mujoco.MjTimerStat, 'number'),
('MjSolverStat', mujoco.MjSolverStat, 'neval'),
('MjContact', mujoco.MjContact, 'dist'),
('MjStatistic', mujoco.MjStatistic, 'extent'),
('MjLROpt', mujoco.MjLROpt, 'maxforce'),
('MjvPerturb', mujoco.MjvPerturb, 'select'),
('MjvCamera', mujoco.MjvCamera, 'fixedcamid'),
])
def test_mj_struct_repr(self, cls, attr):
struct = cls()
setattr(struct, attr, 1)
representation = repr(struct)
self.assertStartsWith(representation, f'<{cls.__name__}')
self.assertIn(f'{attr}: 1', representation)
self.assertEqual(str(struct), repr(struct))
def test_mj_struct_repr_for_subclass(self):
class MjWarningStatSubclass(mujoco.MjWarningStat):
# ptr attribute could cause an infinite recursion, if the repr
# implementation simply looked at all attributes.
@property
def ptr(self):
return self
# repr should include name of subclass.
expected_repr = """<MjWarningStatSubclass
lastinfo: 0
number: 0
>"""
self.assertEqual(repr(MjWarningStatSubclass()), expected_repr)
def test_mju_rotVecQuat(self): # pylint: disable=invalid-name
vec = [1, 0, 0]
quat = [np.cos(np.pi/8), 0, 0, np.sin(np.pi/8)]
expected = np.array([1, 1, 0]) / np.sqrt(2)
# Check that the output argument works, and that the binding returns None.
res = np.empty(3, np.float64)
self.assertIsNone(mujoco.mju_rotVecQuat(res, vec, quat))
np.testing.assert_allclose(res, expected)
# Check that the function can be called via keyword arguments.
mujoco.mju_rotVecQuat(vec=vec, quat=quat, res=res)
np.testing.assert_allclose(res, expected)
# Check that the res argument must have the right size.
with self.assertRaises(TypeError):
mujoco.mju_rotVecQuat(np.empty(4, np.float64), vec, quat)
# Check that the vec argument must have the right size.
with self.assertRaises(TypeError):
mujoco.mju_rotVecQuat(res, [1, 2, 3, 4], quat)
# Check that the quat argument must have the right size.
with self.assertRaises(TypeError):
mujoco.mju_rotVecQuat(res, vec, [1, 2, 3])
# Check that the output argument must have the correct dtype.
with self.assertRaises(TypeError):
mujoco.mju_rotVecQuat(vec, quat, res=np.empty(3, int))
def test_mj_jacSite(self): # pylint: disable=invalid-name
mujoco.mj_forward(self.model, self.data)
site_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, 'mysite')
# Call mj_jacSite with only jacp.
jacp = np.empty((3, 10), np.float64)
mujoco.mj_jacSite(self.model, self.data, jacp, None, site_id)
expected_jacp = np.array(
[[0, 0, 0, 0, 0, 0, -1, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
np.testing.assert_array_equal(jacp, expected_jacp)
# Call mj_jacSite with only jacr.
jacr = np.empty((3, 10), np.float64)
mujoco.mj_jacSite(self.model, self.data, None, jacr, site_id)
expected_jacr = np.array(
[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
np.testing.assert_array_equal(jacr, expected_jacr)
# Call mj_jacSite with both jacp and jacr.
jacp[:] = 0
jacr[:] = 0
mujoco.mj_jacSite(self.model, self.data, jacp, jacr, site_id)
np.testing.assert_array_equal(jacp, expected_jacp)
np.testing.assert_array_equal(jacr, expected_jacr)
# Check that the jacp argument must have the right size.
with self.assertRaises(TypeError):
mujoco.mj_jacSite(
self.model, self.data, np.empty((3, 6), jacp.dtype), None, site_id)
# Check that the jacr argument must have the right size.
with self.assertRaises(TypeError):
mujoco.mj_jacSite(
self.model, self.data, None, np.empty((4, 7), jacr.dtype), site_id)
# Check that the jacp argument must have the right dtype.
with self.assertRaises(TypeError):
mujoco.mj_jacSite(
self.model, self.data, np.empty(jacp.shape, int), None, site_id)
# Check that the jacp argument must have the right dtype.
with self.assertRaises(TypeError):
mujoco.mj_jacSite(
self.model, self.data, None, np.empty(jacr.shape, int), site_id)
def test_docstrings(self): # pylint: disable=invalid-name
self.assertEqual(
mujoco.mj_versionString.__doc__,
"""mj_versionString() -> str
Return the current version of MuJoCo as a null-terminated string.
""")
self.assertEqual(
mujoco.mj_Euler.__doc__,
"""mj_Euler(m: mujoco._structs.MjModel, d: mujoco._structs.MjData) -> None
Euler integrator, semi-implicit in velocity.
""")
def test_int_constant(self):
self.assertEqual(mujoco.mjMAXVFSNAME, 1000)
def test_float_constant(self):
self.assertEqual(mujoco.mjMAXVAL, 1e10)
def test_string_constants(self):
self.assertLen(mujoco.mjDISABLESTRING, mujoco.mjtDisableBit.mjNDISABLE)
self.assertLen(mujoco.mjENABLESTRING, mujoco.mjtEnableBit.mjNENABLE)
self.assertLen(mujoco.mjTIMERSTRING, mujoco.mjtTimer.mjNTIMER)
self.assertLen(mujoco.mjLABELSTRING, mujoco.mjtLabel.mjNLABEL)
self.assertLen(mujoco.mjFRAMESTRING, mujoco.mjtFrame.mjNFRAME)
self.assertLen(mujoco.mjVISSTRING, mujoco.mjtVisFlag.mjNVISFLAG)
self.assertLen(mujoco.mjRNDSTRING, mujoco.mjtRndFlag.mjNRNDFLAG)
self.assertEqual(mujoco.mjDISABLESTRING[11], 'Refsafe')
self.assertEqual(mujoco.mjVISSTRING[mujoco.mjtVisFlag.mjVIS_INERTIA],
('&Inertia', '0', 'I'))
def test_enum_values(self):
self.assertEqual(mujoco.mjtJoint.mjJNT_FREE, 0)
self.assertEqual(mujoco.mjtJoint.mjJNT_BALL, 1)
self.assertEqual(mujoco.mjtJoint.mjJNT_SLIDE, 2)
self.assertEqual(mujoco.mjtJoint.mjJNT_HINGE, 3)
self.assertEqual(mujoco.mjtEnableBit.mjENBL_OVERRIDE, 1<<0)
self.assertEqual(mujoco.mjtEnableBit.mjENBL_ENERGY, 1<<1)
self.assertEqual(mujoco.mjtEnableBit.mjENBL_FWDINV, 1<<2)
self.assertEqual(mujoco.mjtEnableBit.mjENBL_SENSORNOISE, 1<<3)
self.assertEqual(mujoco.mjtEnableBit.mjNENABLE, 4)
self.assertEqual(mujoco.mjtGeom.mjGEOM_PLANE, 0)
self.assertEqual(mujoco.mjtGeom.mjGEOM_HFIELD, 1)
self.assertEqual(mujoco.mjtGeom.mjGEOM_SPHERE, 2)
self.assertEqual(mujoco.mjtGeom.mjGEOM_ARROW, 100)
self.assertEqual(mujoco.mjtGeom.mjGEOM_ARROW1, 101)
self.assertEqual(mujoco.mjtGeom.mjGEOM_ARROW2, 102)
self.assertEqual(mujoco.mjtGeom.mjGEOM_NONE, 1001)
def test_enum_from_int(self):
self.assertEqual(mujoco.mjtJoint.mjJNT_FREE, mujoco.mjtJoint(0))
self.assertEqual(mujoco.mjtGeom.mjGEOM_ARROW, mujoco.mjtGeom(value=100))
# mjENABLE_FWDINV and mjNENABLE have the same int value. Default to the
# first defined one.
self.assertEqual(
mujoco.mjtEnableBit.mjENBL_FWDINV,
mujoco.mjtEnableBit(mujoco.mjtEnableBit.mjNENABLE.value))
with self.assertRaises(ValueError):
mujoco.mjtJoint(1000)
with self.assertRaises(ValueError):
mujoco.mjtJoint(-1)
def test_can_raise_error(self):
self.data.pstack = self.data.nstack
with self.assertRaisesWithLiteralMatch(mujoco.FatalError, 'Stack overflow'):
mujoco.mj_forward(self.model, self.data)
def test_mjcb_time(self):
class CallCounter:
def __init__(self):
self.count = 0
def __call__(self):
self.count += 1
return self.count - 1
call_counter = CallCounter()
with temporary_callback(mujoco.set_mjcb_time, call_counter):
self.assertIs(mujoco.get_mjcb_time(), call_counter)
# Check that the callback setter and getter aren't callin g the function.
self.assertEqual(call_counter.count, 0)
mujoco.mj_forward(self.model, self.data)
self.assertGreater(call_counter.count, 0)
self.assertIsNone(mujoco.get_mjcb_time())
def test_mjcb_time_exception(self):
class TestError(RuntimeError):
pass
def raises_exception():
raise TestError('string', (1, 2, 3), {'a': 1, 'b': 2})
with temporary_callback(mujoco.set_mjcb_time, raises_exception):
with self.assertRaises(TestError) as e:
mujoco.mj_forward(self.model, self.data)
self.assertEqual(
e.exception.args, ('string', (1, 2, 3), {'a': 1, 'b': 2}))
# Should not raise now that we've cleared the callback.
mujoco.mj_forward(self.model, self.data)
def test_mjcb_time_wrong_return_type(self):
with temporary_callback(mujoco.set_mjcb_time, lambda: 'string'):
with self.assertRaisesWithLiteralMatch(
TypeError, 'mjcb_time callback did not return a number'):
mujoco.mj_forward(self.model, self.data)
def test_mjcb_time_not_callable(self):
with self.assertRaisesWithLiteralMatch(
TypeError, 'callback is not an Optional[Callable]'):
mujoco.set_mjcb_time(1)
def test_mjcb_sensor(self):
class SensorCallback:
def __init__(self, test, expected_model, expected_data):
self.test = test
self.expected_model = expected_model
self.expected_data = expected_data
self.count = 0
def __call__(self, m, d, stage):
self.test.assertIs(m, self.expected_model)
self.test.assertIs(d, self.expected_data)
self.test.assertEqual(stage, mujoco.mjtStage.mjSTAGE_VEL)
d.sensordata[0] = 17
self.count += 1
model_with_sensor = mujoco.MjModel.from_xml_string(TEST_XML_SENSOR)
data_with_sensor = mujoco.MjData(model_with_sensor)
sensor_callback = SensorCallback(self, model_with_sensor, data_with_sensor)
self.assertEqual(sensor_callback.count, 0)
with temporary_callback(mujoco.set_mjcb_sensor, sensor_callback):
mujoco.mj_forward(model_with_sensor, data_with_sensor)
self.assertEqual(sensor_callback.count, 1)
self.assertEqual(data_with_sensor.sensordata[0], 17)
def test_can_initialize_mjv_structs(self):
self.assertIsInstance(mujoco.MjvScene(), mujoco.MjvScene)
self.assertIsInstance(mujoco.MjvCamera(), mujoco.MjvCamera)
self.assertIsInstance(mujoco.MjvGLCamera(), mujoco.MjvGLCamera)
self.assertIsInstance(mujoco.MjvGeom(), mujoco.MjvGeom)
self.assertIsInstance(mujoco.MjvLight(), mujoco.MjvLight)
self.assertIsInstance(mujoco.MjvOption(), mujoco.MjvOption)
self.assertIsInstance(mujoco.MjvScene(), mujoco.MjvScene)
self.assertIsInstance(mujoco.MjvScene(self.model, 100), mujoco.MjvScene)
self.assertIsInstance(mujoco.MjvFigure(), mujoco.MjvFigure)
def test_mjv_camera(self):
camera = mujoco.MjvCamera()
camera.type = mujoco.mjtCamera.mjCAMERA_TRACKING
# IDs should be integers
camera.fixedcamid = 2**31 - 1
self.assertEqual(camera.fixedcamid, 2**31 - 1)
with self.assertRaises(TypeError):
camera.fixedcamid = 0.5
def test_mjv_scene(self):
scene = mujoco.MjvScene(model=self.model, maxgeom=100)
# scene.geoms is a fixed-length tuple of length maxgeom.
self.assertEqual(scene.ngeom, 0)
self.assertEqual(scene.maxgeom, 100)
self.assertLen(scene.geoms, scene.maxgeom)
# When the scene is updated, geoms are added to the scene
# (ngeom is incremented)
mujoco.mj_forward(self.model, self.data)
mujoco.mjv_updateScene(self.model, self.data, mujoco.MjvOption(),
None, mujoco.MjvCamera(),
mujoco.mjtCatBit.mjCAT_ALL, scene)
self.assertGreater(scene.ngeom, 0)
def test_mjv_scene_without_model(self):
scene = mujoco.MjvScene()
self.assertEqual(scene.scale, 1.0)
self.assertEqual(scene.maxgeom, 0)
def test_mj_ray(self):
# mj_ray has tricky argument types
geomid = np.zeros(1, np.int32)
mujoco.mj_ray(self.model, self.data, [0, 0, 0], [0, 0, 1], None, 0, 0,
geomid)
mujoco.mj_ray(self.model, self.data, [0, 0, 0], [0, 0, 1],
[0, 0, 0, 0, 0, 0], 0, 0, geomid)
# Check that named arguments work
mujoco.mj_ray(
m=self.model,
d=self.data,
pnt=[0, 0, 0],
vec=[0, 0, 1],
geomgroup=None,
flg_static=0,
bodyexclude=0,
geomid=geomid)
@parameterized.product(flg_html=(False, True), flg_pad=(False, True))
def test_mj_printSchema(self, flg_html, flg_pad): # pylint: disable=invalid-name
# Make sure that mj_printSchema doesn't raise an exception
# (e.g. because the internal output buffer is too small)
self.assertIn('mujoco', mujoco.mj_printSchema(flg_html, flg_pad))
def test_pickle_mjdata(self):
mujoco.mj_step(self.model, self.data)
data2 = pickle.loads(pickle.dumps(self.data))
attr_to_compare = (
'time', 'qpos', 'qvel', 'qacc', 'xpos', 'mocap_pos',
'warning', 'energy'
)
self._assert_attributes_equal(data2, self.data, attr_to_compare)
for _ in range(10):
mujoco.mj_step(self.model, self.data)
mujoco.mj_step(self.model, data2)
self._assert_attributes_equal(data2, self.data, attr_to_compare)
def test_pickle_mjmodel(self):
model2 = pickle.loads(pickle.dumps(self.model))
attr_to_compare = (
'nq', 'nmat', 'body_pos', 'names',
)
self._assert_attributes_equal(model2, self.model, attr_to_compare)
def _assert_attributes_equal(self, actual_obj, expected_obj, attr_to_compare):
for name in attr_to_compare:
actual_value = getattr(actual_obj, name)
expected_value = getattr(expected_obj, name)
try:
if isinstance(expected_value, np.ndarray):
np.testing.assert_array_equal(actual_value, expected_value)
else:
self.assertEqual(actual_value, expected_value)
except AssertionError as e:
self.fail("Attribute '{}' differs from expected value: {}".format(
name, str(e)))
if __name__ == '__main__':
absltest.main()
python/mujoco/callbacks.cc
+413
View File
@@ -0,0 +1,413 @@
// 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.
#include <cstddef>
#include <cstdint>
#include <exception>
#include <limits>
#include <sstream>
#include <type_traits>
#include <mujoco.h>
#include "errors.h"
#include "structs.h"
#include "raw.h"
#include <pybind11/eval.h>
#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
namespace mujoco::python {
namespace {
namespace py = ::pybind11;
[[noreturn]] static void EscapeWithPythonException() {
mju_error("Python exception raised");
std::terminate(); // not actually reachable, mju_error doesn't return
}
template <typename T, typename U>
using enable_if_not_const_t =
std::enable_if_t<std::is_same_v<std::remove_const_t<T>, T>, U>;
// MuJoCo passes raw mjModel* and mjData* as arguments to callbacks, but Python
// callables expect the corresponding MjWrapper objects. To avoid creating new
// wrappers each time we enter callbacks, we instead maintain a global lookup
// table that associates raw MuJoCo struct pointers back to the pointers to
// their corresponding wrappers.
template <typename Raw>
static enable_if_not_const_t<Raw, py::handle> MjWrapperLookup(Raw* ptr) {
using LookupFnType = MjWrapper<Raw>* (Raw*);
static LookupFnType* const lookup = []() -> LookupFnType* {
py::gil_scoped_acquire gil;
auto m = py::module_::import("mujoco._structs");
pybind11::handle builtins(PyEval_GetBuiltins());
if (!builtins.contains(MjWrapper<Raw>::kFromRawPointer)) {
return nullptr;
} else {
try {
return reinterpret_cast<LookupFnType*>(
builtins[MjWrapper<Raw>::kFromRawPointer]
.template cast<std::uintptr_t>());
} catch (const py::cast_error&) {
return nullptr;
}
}
}();
MjWrapper<Raw>* wrapper = nullptr;
if (lookup) {
wrapper = lookup(ptr);
} else {
{
py::gil_scoped_acquire gil;
PyErr_SetString(
UnexpectedError::GetPyExc(),
"_structs module did not register its raw pointer lookup functions");
}
}
if (!wrapper) {
{
py::gil_scoped_acquire gil;
PyErr_SetString(
UnexpectedError::GetPyExc(),
"cannot find the corresponding wrapper for the raw mjStruct");
}
}
// Now we find the existing Python instance of our wrapper.
// TODO(stunya): Figure out a way to do this without invoking py::detail.
{
py::gil_scoped_acquire gil;
const auto [src, type] =
py::detail::type_caster_base<MjWrapper<Raw>>::src_and_type(wrapper);
if (type) {
py::handle instance =
py::detail::find_registered_python_instance(wrapper, type);
if (!instance) {
if (!PyErr_Occurred()) {
PyErr_SetString(
UnexpectedError::GetPyExc(),
"cannot find the Python instance of the MjWrapper");
}
} else {
return instance;
}
} else {
if (!PyErr_Occurred()) {
PyErr_SetString(
UnexpectedError::GetPyExc(),
"MjWrapper type isn't registered with pybind11");
}
}
}
EscapeWithPythonException();
}
template <typename Raw>
static const py::handle MjWrapperLookup(const Raw* ptr) {
return MjWrapperLookup(const_cast<Raw*>(ptr));
}
template <typename Return, typename... Args>
static Return
CallPyCallback(const char* name, PyObject* py_callback, Args... args) {
{
py::gil_scoped_acquire gil;
if (!py_callback) {
std::ostringstream msg;
msg << "py_" << name << " is null";
PyErr_SetString(UnexpectedError::GetPyExc(), msg.str().c_str());
} else {
py::handle callback(py_callback);
try {
if constexpr (std::is_void_v<Return>) {
callback(args...);
return;
} else {
return callback(args...).template cast<Return>();
}
} catch (py::error_already_set& e) {
e.restore();
} catch (const py::cast_error&) {
std::ostringstream msg;
msg << name << " callback did not return ";
if constexpr (std::is_integral_v<Return>) {
msg << "an integer";
} else if constexpr (std::is_floating_point_v<Return>) {
msg << "a number";
} else {
msg << "the correct type";
}
PyErr_SetString(PyExc_TypeError, msg.str().c_str());
}
}
}
EscapeWithPythonException();
}
static PyObject* py_mju_user_warning = nullptr;
static void PyMjuUserWarning(const char* msg) {
CallPyCallback<void>("mju_user_warning", py_mju_user_warning, msg);
}
// We only support ctypes function pointers for these.
// The PyObject* are only here so that we can return the ctypes pointers back
// through the getters.
static PyObject* py_mju_user_malloc = nullptr;
static PyObject* py_mju_user_free = nullptr;
static PyObject* py_mjcb_passive = nullptr;
static void PyMjcbPassive(const raw::MjModel* m, raw::MjData* d) {
CallPyCallback<void>("mjcb_passive", py_mjcb_passive,
MjWrapperLookup(m), MjWrapperLookup(d));
}
static PyObject* py_mjcb_control = nullptr;
static void PyMjcbControl(const raw::MjModel* m, raw::MjData* d) {
CallPyCallback<void>("mjcb_control", py_mjcb_control,
MjWrapperLookup(m), MjWrapperLookup(d));
}
static PyObject* py_mjcb_contactfilter = nullptr;
static int PyMjcbContactfilter(
const raw::MjModel* m, raw::MjData* d, int geom1, int geom2) {
return CallPyCallback<int>("mjcb_contactfilter", py_mjcb_contactfilter,
MjWrapperLookup(m), MjWrapperLookup(d),
geom1, geom2);
}
static PyObject* py_mjcb_sensor = nullptr;
static void
PyMjcbSensor(const raw::MjModel* m, raw::MjData* d, int stage) {
CallPyCallback<void>("mjcb_sensor", py_mjcb_sensor,
MjWrapperLookup(m), MjWrapperLookup(d), stage);
}
static PyObject* py_mjcb_time = nullptr;
static mjtNum PyMjcbTime() {
return CallPyCallback<mjtNum>("mjcb_time", py_mjcb_time);
}
static PyObject* py_mjcb_act_dyn = nullptr;
static mjtNum
PyMjcbActDyn(const raw::MjModel* m, const raw::MjData* d, int id) {
return CallPyCallback<mjtNum>("mjcb_act_dyn", py_mjcb_act_dyn,
MjWrapperLookup(m), MjWrapperLookup(d), id);
}
static PyObject* py_mjcb_act_gain = nullptr;
static mjtNum
PyMjcbActGain(const raw::MjModel* m, const raw::MjData* d, int id) {
return CallPyCallback<mjtNum>("mjcb_act_gain", py_mjcb_act_gain,
MjWrapperLookup(m), MjWrapperLookup(d), id);
}
static PyObject* py_mjcb_act_bias = nullptr;
static mjtNum
PyMjcbActBias(const raw::MjModel* m, const raw::MjData* d, int id) {
return CallPyCallback<mjtNum>("mjcb_act_bias", py_mjcb_act_bias,
MjWrapperLookup(m), MjWrapperLookup(d), id);
}
// If the Python object is a ctypes function pointer, returns the corresponding
// C function pointer. Otherwise, returns a null pointer.
template <typename FuncPtr>
static FuncPtr GetCFuncPtr(py::handle h) {
struct CTypes { PyObject* cfuncptr; PyObject* cast; PyObject* c_void_p; };
static const CTypes ctypes = []() -> CTypes {
try {
auto m = py::module_::import("ctypes");
PyObject* cfuncptr = m.attr("_CFuncPtr").ptr();
PyObject* cast = m.attr("cast").ptr();
PyObject* c_void_p = m.attr("c_void_p").ptr();
Py_XINCREF(cfuncptr);
Py_XINCREF(cast);
Py_XINCREF(c_void_p);
return {cfuncptr, cast, c_void_p};
} catch (const py::error_already_set&) {
return {nullptr, nullptr, nullptr};
}
}();
if (!ctypes.cfuncptr) {
throw UnexpectedError("cannot find `ctypes._CFuncPtr`");
}
const int is_cfuncptr = PyObject_IsInstance(h.ptr(), ctypes.cfuncptr);
if (is_cfuncptr == -1) {
throw py::error_already_set();
} else if (is_cfuncptr) {
if (!ctypes.cast) {
throw UnexpectedError("cannot find `ctypes.cast`");
}
if (!ctypes.c_void_p) {
throw UnexpectedError("cannot find `ctypes.c_void_p`");
}
const uintptr_t func_address =
py::handle(ctypes.cast)(h, py::handle(ctypes.c_void_p))
.attr("value")
.template cast<std::uintptr_t>();
return reinterpret_cast<FuncPtr>(func_address);
} else {
return nullptr;
}
}
static bool IsCallable(py::handle h) {
static PyObject* const is_callable = []() -> PyObject* {
try{
PyObject* o = py::eval("callable").ptr();
Py_XINCREF(o);
return o;
} catch (const py::error_already_set&) {
return nullptr;
}
}();
if (!is_callable) {
throw UnexpectedError("cannot find `callable`");
}
return py::handle(is_callable)(h).cast<bool>();
}
template <typename CFuncPtr>
void SetCallback(py::handle h, CFuncPtr py_trampoline,
PyObject** py_callback, CFuncPtr* mj_callback) {
CFuncPtr cfuncptr = GetCFuncPtr<CFuncPtr>(h);
if (h.is_none()) {
Py_XDECREF(*py_callback);
*py_callback = nullptr;
*mj_callback = nullptr;
} else if (cfuncptr) {
Py_XDECREF(*py_callback);
Py_INCREF(h.ptr());
*py_callback = h.ptr();
*mj_callback = cfuncptr;
} else if (IsCallable(h)) {
Py_XDECREF(*py_callback);
Py_INCREF(h.ptr());
*py_callback = h.ptr();
*mj_callback = py_trampoline;
} else {
throw py::type_error("callback is not an Optional[Callable]");
}
}
py::object GetCallback(PyObject* py_callback) {
if (!py_callback) {
return py::none();
}
return py::reinterpret_borrow<py::object>(py_callback);
}
PYBIND11_MODULE(_callbacks, pymodule) {
// Setters
pymodule.def("set_mju_user_warning", [](py::handle h) {
SetCallback(h, PyMjuUserWarning, &py_mju_user_warning, &::mju_user_warning);
});
pymodule.def("set_mju_user_malloc", [](py::handle h) {
if (h.is_none()) {
Py_XDECREF(py_mju_user_malloc);
py_mju_user_malloc = nullptr;
} else {
auto* cfuncptr = GetCFuncPtr<decltype(::mju_user_malloc)>(h);
if (!cfuncptr) {
throw py::type_error("mju_user_malloc must be a C function pointer");
}
Py_XDECREF(py_mju_user_malloc);
Py_XINCREF(h.ptr());
py_mju_user_malloc = h.ptr();
::mju_user_malloc = cfuncptr;
}
});
pymodule.def("set_mju_user_free", [](py::handle h) {
if (h.is_none()) {
Py_XDECREF(py_mju_user_free);
py_mju_user_free = nullptr;
} else {
auto* cfuncptr = GetCFuncPtr<decltype(::mju_user_free)>(h);
if (!cfuncptr) {
throw py::type_error("mju_user_free must be a C function pointer");
}
Py_XDECREF(py_mju_user_free);
Py_XINCREF(h.ptr());
py_mju_user_free = h.ptr();
::mju_user_free = cfuncptr;
}
});
pymodule.def("set_mjcb_passive", [](py::handle h) {
SetCallback(h, PyMjcbPassive, &py_mjcb_passive, &::mjcb_passive);
});
pymodule.def("set_mjcb_control", [](py::handle h) {
SetCallback(h, PyMjcbControl, &py_mjcb_control, &::mjcb_control);
});
pymodule.def("set_mjcb_contactfilter", [](py::handle h) {
SetCallback(h, PyMjcbContactfilter,
&py_mjcb_contactfilter, &::mjcb_contactfilter);
});
pymodule.def("set_mjcb_sensor", [](py::handle h) {
SetCallback(h, PyMjcbSensor, &py_mjcb_sensor, &::mjcb_sensor);
});
pymodule.def("set_mjcb_time", [](py::handle h) {
SetCallback(h, PyMjcbTime, &py_mjcb_time, &::mjcb_time);
});
pymodule.def("set_mjcb_act_dyn", [](py::handle h) {
SetCallback(h, PyMjcbActDyn, &py_mjcb_act_dyn, &::mjcb_act_dyn);
});
pymodule.def("set_mjcb_act_gain", [](py::handle h) {
SetCallback(h, PyMjcbActGain, &py_mjcb_act_gain, &::mjcb_act_gain);
});
pymodule.def("set_mjcb_act_bias", [](py::handle h) {
SetCallback(h, PyMjcbActBias, &py_mjcb_act_bias, &::mjcb_act_bias);
});
// Getters
pymodule.def("get_mju_user_warning", []() {
return GetCallback(py_mju_user_warning);
});
pymodule.def("get_mju_user_malloc", []() {
return GetCallback(py_mju_user_malloc);
});
pymodule.def("get_mju_user_free", []() {
return GetCallback(py_mju_user_free);
});
pymodule.def("get_mjcb_passive", []() {
return GetCallback(py_mjcb_passive);
});
pymodule.def("get_mjcb_control", []() {
return GetCallback(py_mjcb_control);
});
pymodule.def("get_mjcb_contactfilter", []() {
return GetCallback(py_mjcb_contactfilter);
});
pymodule.def("get_mjcb_sensor", []() {
return GetCallback(py_mjcb_sensor);
});
pymodule.def("get_mjcb_time", []() {
return GetCallback(py_mjcb_time);
});
pymodule.def("get_mjcb_act_dyn", []() {
return GetCallback(py_mjcb_act_dyn);
});
pymodule.def("get_mjcb_act_gain", []() {
return GetCallback(py_mjcb_act_gain);
});
pymodule.def("get_mjcb_act_bias", []() {
return GetCallback(py_mjcb_act_bias);
});
} // PYBIND11_MODULE
} // namespace
} // namespace mujoco::python
python/mujoco/codegen/generate_enum_traits.py
+89
View File
@@ -0,0 +1,89 @@
# 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.
# ==============================================================================
"""Code generator for function_traits.h."""
from typing import Mapping, Sequence
from absl import app
from introspect import ast_nodes
from introspect import enums
ENUMS: Mapping[str, ast_nodes.EnumDecl] = enums.ENUMS
def main(argv: Sequence[str]) -> None:
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
struct_decls = []
for enum in ENUMS.values():
value_decls = []
for k in enum.values:
value_decls.append(f'std::make_pair("{k}", ::{enum.name}::{k})')
if len(value_decls) < 2:
value_decls = ''.join(value_decls)
else:
value_decls = '\n ' + ',\n '.join(value_decls)
struct_decls.append(f"""
struct {enum.name} {{
static constexpr char name[] = "{enum.name}";
using type = ::{enum.name};
static constexpr auto values = std::array{{{value_decls}}};
}};
""".strip())
all_structs = '\n\n'.join(struct_decls)
all_enum_inits = '\n ' + '{},\n '.join(ENUMS.keys()) + '{}'
print(f"""
// 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.
#ifndef MUJOCO_PYTHON_CODEGEN_ENUM_TRAITS_H_
#define MUJOCO_PYTHON_CODEGEN_ENUM_TRAITS_H_
#include <array>
#include <tuple>
#include <utility>
#include <mujoco.h>
namespace mujoco::python_traits {{
{all_structs}
static constexpr auto kAllEnums = std::make_tuple({all_enum_inits});
}} // namespace mujoco::python_traits
#endif // MUJOCO_PYTHON_CODEGEN_ENUM_TRAITS_H_
""".lstrip())
if __name__ == '__main__':
app.run(main)
python/mujoco/codegen/generate_function_traits.py
+108
View File
@@ -0,0 +1,108 @@
# 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.
# ==============================================================================
"""Code generator for function_traits.h."""
from typing import Mapping, Sequence
from absl import app
from introspect import ast_nodes
from introspect import functions
FUNCTIONS: Mapping[str, ast_nodes.FunctionDecl] = functions.FUNCTIONS
def main(argv: Sequence[str]) -> None:
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
struct_decls = []
for func in FUNCTIONS.values():
# Modify some parameter types.
parameters = []
modified = False
for p in func.parameters:
# Expose array parameters as pointer-to-arrays so that we can determine
# array extents in C++ templates.
if isinstance(p.type, ast_nodes.ArrayType):
parameters.append(ast_nodes.FunctionParameterDecl(
name=p.name, type=ast_nodes.PointerType(
ast_nodes.ArrayType(
inner_type=p.type.inner_type, extents=p.type.extents))))
modified = True
else:
parameters.append(p)
if modified:
func = ast_nodes.FunctionDecl(
name=func.name, return_type=func.return_type,
parameters=parameters, doc=func.doc)
getfunc = f'*reinterpret_cast<type*>(&::{func.name})'
else:
getfunc = f'::{func.name}'
param_names = ', '.join(f'"{p.name}"' for p in parameters)
struct_decls.append(f"""
struct {func.name} {{
static constexpr char name[] = "{func.name}";
static constexpr char doc[] = "{func.doc}";
using type = {func.decltype};
static constexpr auto param_names = std::make_tuple({param_names});
MUJOCO_ALWAYS_INLINE static type& GetFunc() {{
return {getfunc};
}}
}};
""".strip())
all_structs = '\n\n'.join(struct_decls)
print(f"""
// 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.
#ifndef MUJOCO_PYTHON_CODEGEN_FUNCTION_TRAITS_H_
#define MUJOCO_PYTHON_CODEGEN_FUNCTION_TRAITS_H_
#include <tuple>
#include <mujoco.h>
#include "util/crossplatform.h"
namespace mujoco::python_traits {{
{all_structs}
}} // namespace mujoco::python_traits
#endif // MUJOCO_PYTHON_CODEGEN_FUNCTION_TRAITS_H_
""".lstrip())
if __name__ == '__main__':
app.run(main)
python/mujoco/constants.cc
+91
View File
@@ -0,0 +1,91 @@
// 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.
#include <vector>
#include <mjmodel.h>
#include <mjvisualize.h>
#include <mujoco.h>
#include <pybind11/cast.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
namespace mujoco::python {
namespace {
namespace py = ::pybind11;
template <auto N>
py::tuple MakeTuple(
const char* (&strings)[N]) {
py::list result;
for (int i = 0; i < N; i++) {
result.append(py::str(strings[i]));
}
return result;
}
template <auto N>
py::tuple MakeTuple(const char* (&strings)[N][3]) {
py::list result;
for (int i = 0; i < N; i++) {
result.append(py::make_tuple(
py::str(strings[i][0]),
py::str(strings[i][1]),
py::str(strings[i][2])));
}
return result;
}
PYBIND11_MODULE(_constants, pymodule) {
#define X(var) pymodule.attr(#var) = var
// from mjmodel.h
X(mjPI);
X(mjMAXVAL);
X(mjMINMU);
X(mjMINIMP);
X(mjMAXIMP);
X(mjMAXCONPAIR);
X(mjMAXVFS);
X(mjMAXVFSNAME);
X(mjNEQDATA);
X(mjNDYN);
X(mjNGAIN);
X(mjNBIAS);
X(mjNREF);
X(mjNIMP);
X(mjNSOLVER);
// from mjvisualize.h
X(mjNGROUP);
X(mjMAXLIGHT);
X(mjMAXOVERLAY);
X(mjMAXLINE);
X(mjMAXLINEPNT);
X(mjMAXPLANEGRID);
// from mujoco.h
X(mjVERSION_HEADER);
#undef X
pymodule.attr("mjDISABLESTRING") = MakeTuple(mjDISABLESTRING);
pymodule.attr("mjENABLESTRING") = MakeTuple(mjENABLESTRING);
pymodule.attr("mjTIMERSTRING") = MakeTuple(mjTIMERSTRING);
pymodule.attr("mjLABELSTRING") = MakeTuple(mjLABELSTRING);
pymodule.attr("mjFRAMESTRING") = MakeTuple(mjFRAMESTRING);
pymodule.attr("mjVISSTRING") = MakeTuple(mjVISSTRING);
pymodule.attr("mjRNDSTRING") = MakeTuple(mjRNDSTRING);
}
} // namespace
} // namespace mujoco::python
python/mujoco/egl/__init__.py
+130
View File
@@ -0,0 +1,130 @@
# Copyright 2018 The dm_control Authors
#
# 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.
# ==============================================================================
"""An EGL context for headless accelerated OpenGL rendering on GPU devices."""
import atexit
import ctypes
import os
PYOPENGL_PLATFORM = os.environ.get('PYOPENGL_PLATFORM')
if not PYOPENGL_PLATFORM:
os.environ['PYOPENGL_PLATFORM'] = 'egl'
elif PYOPENGL_PLATFORM.lower() != 'egl':
raise ImportError(
'Cannot use EGL rendering platform. '
'The PYOPENGL_PLATFORM environment variable is set to {!r} '
'(should be either unset or \'egl\').')
from mujoco.egl import egl_ext as EGL
from OpenGL import error
def create_initialized_egl_device_display():
"""Creates an initialized EGL display directly on a device."""
all_devices = EGL.eglQueryDevicesEXT()
selected_device = os.environ.get('MUJOCO_EGL_DEVICE_ID', None)
if selected_device is None:
candidates = all_devices
else:
device_idx = int(selected_device)
if not 0 <= device_idx < len(all_devices):
raise RuntimeError(
f'The MUJOCO_EGL_DEVICE_ID environment variable must be an integer '
f'between 0 and {len(all_devices)-1} (inclusive), got {device_idx}.')
candidates = all_devices[device_idx:device_idx + 1]
for device in candidates:
display = EGL.eglGetPlatformDisplayEXT(
EGL.EGL_PLATFORM_DEVICE_EXT, device, None)
if display != EGL.EGL_NO_DISPLAY and EGL.eglGetError() == EGL.EGL_SUCCESS:
# `eglInitialize` may or may not raise an exception on failure depending
# on how PyOpenGL is configured. We therefore catch a `GLError` and also
# manually check the output of `eglGetError()` here.
try:
initialized = EGL.eglInitialize(display, None, None)
except error.GLError:
pass
else:
if initialized == EGL.EGL_TRUE and EGL.eglGetError() == EGL.EGL_SUCCESS:
return display
return EGL.EGL_NO_DISPLAY
EGL_DISPLAY = create_initialized_egl_device_display()
if EGL_DISPLAY == EGL.EGL_NO_DISPLAY:
raise ImportError(
'Cannot initialize a EGL device display. This likely means that your EGL '
'driver does not support the PLATFORM_DEVICE extension, which is '
'required for creating a headless rendering context.')
atexit.register(EGL.eglTerminate, EGL_DISPLAY)
EGL_ATTRIBUTES = (
EGL.EGL_RED_SIZE, 8,
EGL.EGL_GREEN_SIZE, 8,
EGL.EGL_BLUE_SIZE, 8,
EGL.EGL_ALPHA_SIZE, 8,
EGL.EGL_DEPTH_SIZE, 24,
EGL.EGL_STENCIL_SIZE, 8,
EGL.EGL_COLOR_BUFFER_TYPE, EGL.EGL_RGB_BUFFER,
EGL.EGL_SURFACE_TYPE, EGL.EGL_PBUFFER_BIT,
EGL.EGL_RENDERABLE_TYPE, EGL.EGL_OPENGL_BIT,
EGL.EGL_NONE
)
class GLContext:
"""An EGL context for headless accelerated OpenGL rendering on GPU devices."""
def __init__(self, max_width, max_height):
del max_width, max_height # unused
num_configs = ctypes.c_long()
config_size = 1
config = EGL.EGLConfig()
EGL.eglReleaseThread()
EGL.eglChooseConfig(
EGL_DISPLAY,
EGL_ATTRIBUTES,
ctypes.byref(config),
config_size,
num_configs)
if num_configs.value < 1:
raise RuntimeError(
'EGL failed to find a framebuffer configuration that matches the '
'desired attributes: {}'.format(EGL_ATTRIBUTES))
EGL.eglBindAPI(EGL.EGL_OPENGL_API)
self._context = EGL.eglCreateContext(
EGL_DISPLAY, config, EGL.EGL_NO_CONTEXT, None)
if not self._context:
raise RuntimeError('Cannot create an EGL context.')
def make_current(self):
if not EGL.eglMakeCurrent(
EGL_DISPLAY, EGL.EGL_NO_SURFACE, EGL.EGL_NO_SURFACE, self._context):
raise RuntimeError('Failed to make the EGL context current.')
def free(self):
"""Frees resources associated with this context."""
if self._context:
current_context = EGL.eglGetCurrentContext()
if current_context and self._context.address == current_context.address:
EGL.eglMakeCurrent(EGL_DISPLAY, EGL.EGL_NO_SURFACE,
EGL.EGL_NO_SURFACE, EGL.EGL_NO_CONTEXT)
EGL.eglDestroyContext(EGL_DISPLAY, self._context)
EGL.eglReleaseThread()
self._context = None
def __del__(self):
self.free()
python/mujoco/egl/egl_ext.py
+71
View File
@@ -0,0 +1,71 @@
# Copyright 2018 The dm_control Authors
#
# 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.
# ==============================================================================
"""Extends OpenGL.EGL with definitions necessary for headless rendering."""
import ctypes
from OpenGL.platform import ctypesloader # pylint: disable=g-bad-import-order
try:
# Nvidia driver seems to need libOpenGL.so (as opposed to libGL.so)
# for multithreading to work properly. We load this in before everything else.
ctypesloader.loadLibrary(ctypes.cdll, 'OpenGL', mode=ctypes.RTLD_GLOBAL)
except OSError:
pass
# pylint: disable=g-import-not-at-top
from OpenGL import EGL
from OpenGL import error
# From the EGL_EXT_device_enumeration extension.
PFNEGLQUERYDEVICESEXTPROC = ctypes.CFUNCTYPE(
EGL.EGLBoolean,
EGL.EGLint,
ctypes.POINTER(EGL.EGLDeviceEXT),
ctypes.POINTER(EGL.EGLint),
)
try:
_eglQueryDevicesEXT = PFNEGLQUERYDEVICESEXTPROC( # pylint: disable=invalid-name
EGL.eglGetProcAddress('eglQueryDevicesEXT'))
except TypeError as e:
raise ImportError('eglQueryDevicesEXT is not available.') from e
# From the EGL_EXT_platform_device extension.
EGL_PLATFORM_DEVICE_EXT = 0x313F
PFNEGLGETPLATFORMDISPLAYEXTPROC = ctypes.CFUNCTYPE(
EGL.EGLDisplay, EGL.EGLenum, ctypes.c_void_p, ctypes.POINTER(EGL.EGLint))
try:
eglGetPlatformDisplayEXT = PFNEGLGETPLATFORMDISPLAYEXTPROC( # pylint: disable=invalid-name
EGL.eglGetProcAddress('eglGetPlatformDisplayEXT'))
except TypeError as e:
raise ImportError('eglGetPlatformDisplayEXT is not available.') from e
# Wrap raw _eglQueryDevicesEXT function into something more Pythonic.
def eglQueryDevicesEXT(max_devices=10): # pylint: disable=invalid-name
devices = (EGL.EGLDeviceEXT * max_devices)()
num_devices = EGL.EGLint()
success = _eglQueryDevicesEXT(max_devices, devices, num_devices)
if success == EGL.EGL_TRUE:
return [devices[i] for i in range(num_devices.value)]
else:
raise error.GLError(err=EGL.eglGetError(),
baseOperation=eglQueryDevicesEXT,
result=success)
# Expose everything from upstream so that
# we can use this as a drop-in replacement for OpenGL.EGL.
# pylint: disable=wildcard-import,g-bad-import-order
from OpenGL.EGL import *
python/mujoco/enums.cc
+62
View File
@@ -0,0 +1,62 @@
// 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.
#include <sstream>
#include <type_traits>
#include "util/crossplatform.h"
#include "enum_traits.h"
#include "util/tuple_tools.h"
#include <pybind11/pybind11.h>
namespace mujoco::python {
namespace {
namespace py = ::pybind11;
template <typename Trait>
MUJOCO_ALWAYS_INLINE
void DefEnum(py::module_& m) {
py::enum_<typename Trait::type> e(m, Trait::name);
for (const auto& [name, enumerator] : Trait::values) {
e.value(name, enumerator);
}
e.def(py::init([](int value) {
for (const auto& [name, enumerator] : Trait::values) {
if (value == enumerator) {
return static_cast<typename Trait::type>(value);
}
}
std::ostringstream err;
err << "Invalid int value for " << Trait::name << ": " << value;
throw py::value_error(err.str());
}),
py::arg("value"), py::prepend());
}
template <typename Tuple>
MUJOCO_ALWAYS_INLINE
void DefAllEnums(py::module_& m, Tuple&& tuple) {
using TupleNoRef = std::remove_reference_t<Tuple>;
if constexpr (std::tuple_size_v<TupleNoRef> != 0) {
using Trait = std::remove_reference_t<std::tuple_element_t<0, TupleNoRef>>;
DefEnum<Trait>(m);
DefAllEnums(m, util::tuple_slice<1, void>(tuple));
}
}
PYBIND11_MODULE(_enums, pymodule) {
DefAllEnums(pymodule, python_traits::kAllEnums);
}
} // namespace
} // namespace mujoco::python
python/mujoco/errors.cc
+25
View File
@@ -0,0 +1,25 @@
// 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.
#include "errors.h"
#include <pybind11/pybind11.h>
namespace mujoco::python {
namespace {
PYBIND11_MODULE(_errors, m) {
m.attr("FatalError") = FatalError::GetPyExc();
m.attr("UnexpectedError") = UnexpectedError::GetPyExc();
}
} // namespace
} // namespace mujoco::python
python/mujoco/errors.h
+193
View File
@@ -0,0 +1,193 @@
// 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.
#ifndef MUJOCO_PYTHON_ERRORS_H_
#define MUJOCO_PYTHON_ERRORS_H_
#include <csetjmp>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <mjexport.h>
#include "util/crossplatform.h"
#include "util/func_wrap.h"
#include <pybind11/pybind11.h>
// DO NOT USE THIS FUNCTION ELSEWHERE.
// It should be regarded as part of MuJoCo's internal implementation detail.
extern "C" {
MJAPI void _mjPRIVATE__set_tls_error_fn(void (*h)(const char*));
}
// When building for Linux and statically linking against a "hermetic" libc++abi
// (i.e. where libc++/libc++abi symbols all have "hidden" visibility), exception
// types do not propagate correctly across shared library boundaries.
//
// This occurs even when PYBIND11_EXPORT_EXCEPTION is used, or when the virtual
// destructor is defined in errors.cc rather than than defined inline in the
// class definition (following the FAQ comment about "key functions" on
// https://libcxxabi.llvm.org/). For some reason, the type_info address is
// different across different shared objects (DSO), even when the symbols for
// typeinfo, type name, and vtable are left undefined in all but one DSO.
// Since libc++abi matches exception types by comparing type_info addresses,
// this breaks exception handling across DSO boundaries.
//
// Instead, we manually create and store exception types in PyEval_GetBuiltins,
// then effectively register separate pybind11 exception translators in each
// module that all translate to the same exception type (which is what's
// happening with pybind11's built-in exceptions under this setup).
//
// Effectively, we are doing almost the same thing as what a pybind11 is doing
// with its "internals" struct, but we store the exception types directly in
// the builtin context rather than in a PyCapsule.
namespace mujoco::python {
namespace _impl {
template <typename T>
class ErrorBase : public pybind11::builtin_exception {
public:
virtual ~ErrorBase() = default;
static PyObject* GetPyExc() {
static PyObject* const e = []() {
pybind11::gil_scoped_acquire gil;
std::string unique_identifier = "__MUJOCO_ERROR_";
unique_identifier += T::kName;
pybind11::str py_builtin_identifier(unique_identifier);
// We can end up here while handling another Python exception.
// Temporarily clear the Python error indicator since we need to interact
// with the interpreter.
struct PyErrCache {
PyErrCache() { PyErr_Fetch(&type, &value, &traceback); }
~PyErrCache() { if (type) PyErr_Restore(type, value, traceback); }
PyObject* type;
PyObject* value;
PyObject* traceback;
};
PyErrCache err_cache;
pybind11::handle builtins(PyEval_GetBuiltins());
if (!builtins.contains(py_builtin_identifier)) {
std::string full_name = "mujoco.";
full_name += T::kName;
auto ret =
PyErr_NewException(full_name.c_str(), PyExc_Exception, nullptr);
builtins[py_builtin_identifier] = pybind11::handle(ret);
return ret;
} else {
return builtins[py_builtin_identifier].ptr();
}
}();
return e;
}
void set_error() const override { PyErr_SetString(T::GetPyExc(), what()); }
protected:
using builtin_exception::builtin_exception;
};
// We shouldn't throw a C++ exception from a function that's a callback
// from C. (Usually it would work, but it would be undefined behavior since
// there'd be no guarantee that C++ can unwind the stack correctly through the
// C functions. On Windows, longjmp actually uses the same mechanism as
// C++ exceptions.)
// Instead, we call setjmp before entering MuJoCo, and do a longjmp from
// mju_user_error back to C++ to throw an exception.
static thread_local std::jmp_buf mju_error_jmp_buf;
static thread_local std::array<char, 1024> mju_error_msg{0};
static void MjErrorHandler(const char* msg) {
std::strncpy(mju_error_msg.data(), msg, mju_error_msg.size());
std::longjmp(mju_error_jmp_buf, 1);
}
template <typename InterceptAsType>
struct MjErrorIntercepter {
template <typename Return, typename... Args, typename Callable>
MUJOCO_ALWAYS_INLINE
static constexpr auto WrapFunc(Callable&& callable) {
return [callable](Args... args) MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE {
_mjPRIVATE__set_tls_error_fn(&MjErrorHandler);
// DON'T MIX RAII WITH SETJMP!
// From https://en.cppreference.com/w/cpp/utility/program/longjmp:
// If replacing of std::longjmp with `throw` and setjmp with `catch`
// would execute a non-trivial destructor for any automatic object, the
// behavior of such std::longjmp is undefined.
if (setjmp(mju_error_jmp_buf) == 0) {
if constexpr (std::is_void_v<decltype(callable(args...))>) {
callable(args...);
_mjPRIVATE__set_tls_error_fn(nullptr);
} else {
auto ret = callable(args...);
static_assert(std::is_trivially_destructible_v<decltype(ret)>);
_mjPRIVATE__set_tls_error_fn(nullptr);
return ret;
}
} else {
// This branch is entered via a longjmp back from our mju_error handler.
_mjPRIVATE__set_tls_error_fn(nullptr);
{
// Check if a Python callback has thrown an exception.
// We cannot use a py::gil_scoped_acquire here: on Windows its
// destructor isn't triggered before pybind returns control to the
// interpreter.
auto gil = PyGILState_Ensure();
if (PyErr_Occurred()) {
// We must hold the GIL when we create the py::error_already_set,
// since its constructor calls PyErr_Fetch.
pybind11::error_already_set err;
// But the GIL must be released before we throw, otherwise a
// deadlock ensues!
PyGILState_Release(gil);
throw err;
}
PyGILState_Release(gil);
}
throw InterceptAsType(std::string(mju_error_msg.data()));
}
};
}
};
} // namespace _impl
class FatalError : public _impl::ErrorBase<FatalError> {
public:
static constexpr char kName[] = "FatalError";
using ErrorBase<FatalError>::ErrorBase;
virtual ~FatalError() = default;
};
class UnexpectedError : public _impl::ErrorBase<UnexpectedError> {
public:
static constexpr char kName[] = "UnexpectedError";
UnexpectedError(const std::string& msg)
: ErrorBase(msg +
" (this error not expected to ever occur,"
" please report it to MuJoCo developers!)") {}
virtual ~UnexpectedError() = default;
};
template <typename Callable>
MUJOCO_ALWAYS_INLINE
static constexpr auto InterceptMjErrors(Callable&& callable) {
return util::_impl::WrapFunc<_impl::MjErrorIntercepter<FatalError>>(
std::forward<Callable>(callable));
}
} // namespace mujoco::python
#endif // MUJOCO_PYTHON_ERRORS_H_
python/mujoco/functions.cc
+1036
View File
File diff suppressed because it is too large Load Diff
python/mujoco/functions.h
+206
View File
@@ -0,0 +1,206 @@
// 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.
#ifndef MUJOCO_PYTHON_FUNCTIONS_H_
#define MUJOCO_PYTHON_FUNCTIONS_H_
#include <array>
#include <optional>
#include <string>
#include <string_view>
#include <tuple>
#include <type_traits>
#include <Eigen/Core>
#include <mujoco.h>
#include "errors.h"
#include "structs.h"
#include "util/array_traits.h"
#include "util/crossplatform.h"
#include "util/func_wrap.h"
#include "util/tuple_tools.h"
#include <pybind11/eigen.h>
#include <pybind11/pybind11.h>
// Performs a compile-time check that the omitted argument name list is a
// subset of the underlying function parameter names, and returns a helper that
// defines a pybind11 function whose argument names obtained by removing
// __VA_ARGS__ from the param_names tuple in the given MjTraits.
// (This has to be implemented as a a macro because we cannot perform a
// constexpr comparison of a tuple that is passed as a function argument.)
#define DEF_WITH_OMITTED_PY_ARGS(MJTRAITS, ...) \
static_assert( \
util::is_subset_strings(std::make_tuple(__VA_ARGS__), \
MJTRAITS::param_names), \
"omitted argument names is not a subset of function parameter names"); \
DefWithOmittedPyArgsImpl<MJTRAITS, decltype(std::make_tuple(__VA_ARGS__))> { \
std::make_tuple(__VA_ARGS__) \
}
namespace mujoco::util {
template <typename RawMj>
struct wrapped<RawMj*, python::enable_if_mj_struct_t<RawMj>> {
MUJOCO_ALWAYS_INLINE
static constexpr RawMj* unwrap(python::MjWrapper<RawMj>& wrapper) {
return wrapper.get();
}
};
// We use std::optional on pointer arguments to indicate that Python callers
// can pass None.
template <typename RawMj>
struct wrapped<std::optional<RawMj*>, python::enable_if_mj_struct_t<RawMj>> {
MUJOCO_ALWAYS_INLINE
static constexpr std::optional<RawMj*> unwrap(
std::optional<python::MjWrapper<RawMj>*> wrapper) {
if (wrapper.has_value()) {
return (*wrapper)->get();
}
return std::nullopt;
}
};
template <typename Arr>
using enable_if_arithmetic_array_t = std::enable_if_t<
std::is_array_v<Arr> &&
std::is_arithmetic_v<std::remove_all_extents_t<Arr>>>;
template <typename Arr>
struct wrapped<Arr*, enable_if_arithmetic_array_t<Arr>> {
MUJOCO_ALWAYS_INLINE
static constexpr Arr* unwrap(Eigen::Ref<array_eigen_t<Arr>> wrapper) {
return reinterpret_cast<Arr*>(wrapper.data());
}
};
} // namespace mujoco::util
namespace mujoco::python {
namespace _impl {
template <typename, typename>
struct py_arg_helper {};
template <typename... PyArg, typename... OmittedArg>
struct py_arg_helper<std::tuple<PyArg...>, std::tuple<OmittedArg...>> {
std::tuple<PyArg...> py_args;
std::tuple<OmittedArg...> omitted_args;
static constexpr int n_py_args = std::tuple_size_v<decltype(py_args)>;
static constexpr int n_omitted_args =
std::tuple_size_v<decltype(omitted_args)>;
template <typename... T>
MUJOCO_ALWAYS_INLINE
constexpr void def(::pybind11::module_& m, T&&... t) {
constexpr int NExtras = std::tuple_size_v<std::tuple<T...>>;
unpack_tuple_as_py_args<0, NExtras>(m, std::forward<T>(t)...);
}
template <int ArgIdx, int NExtras, typename... T>
MUJOCO_ALWAYS_INLINE
constexpr void unpack_tuple_as_py_args(::pybind11::module_& m, T&&... t) {
if constexpr (ArgIdx == n_py_args) {
if constexpr (std::tuple_size_v<std::tuple<T...>> ==
NExtras + n_py_args - n_omitted_args) {
m.def(std::forward<T>(t)...);
} else {
// This should ideally be a static_assert, but we need C++20 consteval
// to do that. When using the DEF_WITH_OMITTED_PY_ARGS macro, the
// static_assert in that macro would trigger first, rendering this
// branch unreachable.
throw UnexpectedError(
"omitted argument names do not match the underlying function "
"parameter names");
}
} else if (is_omitted<ArgIdx>()) {
unpack_tuple_as_py_args<ArgIdx+1, NExtras>(m, std::forward<T>(t)...);
} else {
unpack_tuple_as_py_args<ArgIdx+1, NExtras>(
m, std::forward<T>(t)..., ::pybind11::arg(std::get<ArgIdx>(py_args)));
}
}
template <int ArgIdx, int OmittedIdx = 0>
MUJOCO_ALWAYS_INLINE
constexpr bool is_omitted() {
if constexpr (OmittedIdx == n_omitted_args) {
return false;
// string_view comparison can be constexpr
} else if (std::string_view(std::get<ArgIdx>(py_args)) ==
std::string_view(std::get<OmittedIdx>(omitted_args))) {
return true;
} else {
return is_omitted<ArgIdx, OmittedIdx + 1>();
}
}
};
} // namespace _impl
template <typename Tuple>
MUJOCO_ALWAYS_INLINE
static constexpr auto WithNamedArgs(Tuple&& py_args) {
using ArgTuple = std::remove_cv_t<std::remove_reference_t<Tuple>>;
return _impl::py_arg_helper<ArgTuple, std::tuple<>>{
std::forward<Tuple>(py_args), std::tuple<>()};
}
template <typename Tuple1, typename Tuple2>
MUJOCO_ALWAYS_INLINE
static constexpr auto WithNamedArgs(Tuple1&& py_args, Tuple2&& omitted_args) {
using ArgTuple = std::remove_cv_t<std::remove_reference_t<Tuple1>>;
using OmittedTuple = std::remove_cv_t<std::remove_reference_t<Tuple2>>;
return _impl::py_arg_helper<ArgTuple, OmittedTuple>{
std::forward<Tuple1>(py_args), std::forward<Tuple2>(omitted_args)};
}
template <typename MjTraits>
MUJOCO_ALWAYS_INLINE
static constexpr void Def(::pybind11::module_& m) {
WithNamedArgs(MjTraits::param_names).def(
m, MjTraits::name,
util::UnwrapArgs(InterceptMjErrors(MjTraits::GetFunc())),
::pybind11::doc(MjTraits::doc),
::pybind11::call_guard<::pybind11::gil_scoped_release>());
}
template <typename MjTraits, typename Func>
MUJOCO_ALWAYS_INLINE
static constexpr void Def(::pybind11::module_& m, Func&& func) {
WithNamedArgs(MjTraits::param_names).def(
m, MjTraits::name, util::UnwrapArgs(std::forward<Func>(func)),
::pybind11::doc(MjTraits::doc),
::pybind11::call_guard<::pybind11::gil_scoped_release>());
}
template <typename MjTraits, typename OmittedArgs, typename Func>
MUJOCO_ALWAYS_INLINE
static constexpr void Def(
::pybind11::module_& m, OmittedArgs&& omitted_args, Func&& func) {
WithNamedArgs(MjTraits::param_names, std::forward<OmittedArgs>(omitted_args))
.def(m, MjTraits::name, util::UnwrapArgs(std::forward<Func>(func)),
::pybind11::doc(MjTraits::doc),
::pybind11::call_guard<::pybind11::gil_scoped_release>());
}
// Should only be invoked via the DEF_WITH_OMITTED_PY_ARGS macro.
template <typename MjTraits, typename OmittedArgsTuple>
struct DefWithOmittedPyArgsImpl {
OmittedArgsTuple omitted_args;
template <typename Func>
constexpr auto operator()(::pybind11::module_& m, Func&& func) {
return Def<MjTraits>(m, omitted_args, std::forward<Func>(func));
}
};
} // namespace mujoco::python
#endif // MUJOCO_PYTHON_FUNCTIONS_H_
python/mujoco/glfw/__init__.py
+41
View File
@@ -0,0 +1,41 @@
# Copyright 2017 The dm_control Authors
#
# 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.
# ==============================================================================
"""An OpenGL context created via GLFW."""
import glfw
class GLContext:
"""An OpenGL context created via GLFW."""
def __init__(self, max_width, max_height):
glfw.init()
glfw.window_hint(glfw.VISIBLE, 0)
self._context = glfw.create_window(width=max_width, height=max_height,
title='Invisible window', monitor=None,
share=None)
def make_current(self):
glfw.make_context_current(self._context)
def free(self):
if self._context:
if glfw.get_current_context() == self._context:
glfw.make_context_current(None)
glfw.destroy_window(self._context)
self._context = None
def __del__(self):
self.free()
python/mujoco/indexer_xmacro.h
+406
View File
@@ -0,0 +1,406 @@
// 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.
#ifndef MUJOCO_PYTHON_INDEXER_XMACRO_H_
#define MUJOCO_PYTHON_INDEXER_XMACRO_H_
#include <mjxmacro.h>
#define MJMODEL_ACTUATOR \
X( int, actuator_, trntype, nu, 1 ) \
X( int, actuator_, dyntype, nu, 1 ) \
X( int, actuator_, gaintype, nu, 1 ) \
X( int, actuator_, biastype, nu, 1 ) \
X( int, actuator_, trnid, nu, 2 ) \
X( int, actuator_, group, nu, 1 ) \
X( mjtByte, actuator_, ctrllimited, nu, 1 ) \
X( mjtByte, actuator_, forcelimited, nu, 1 ) \
X( mjtNum, actuator_, dynprm, nu, mjNDYN ) \
X( mjtNum, actuator_, gainprm, nu, mjNGAIN ) \
X( mjtNum, actuator_, biasprm, nu, mjNBIAS ) \
X( mjtNum, actuator_, ctrlrange, nu, 2 ) \
X( mjtNum, actuator_, forcerange, nu, 2 ) \
X( mjtNum, actuator_, gear, nu, 6 ) \
X( mjtNum, actuator_, cranklength, nu, 1 ) \
X( mjtNum, actuator_, acc0, nu, 1 ) \
X( mjtNum, actuator_, length0, nu, 1 ) \
X( mjtNum, actuator_, lengthrange, nu, 2 ) \
X( mjtNum, actuator_, user, nu, MJ_M(nuser_actuator) )
#define MJMODEL_BODY \
X( int, body_, parentid, nbody, 1 ) \
X( int, body_, rootid, nbody, 1 ) \
X( int, body_, weldid, nbody, 1 ) \
X( int, body_, mocapid, nbody, 1 ) \
X( int, body_, jntnum, nbody, 1 ) \
X( int, body_, jntadr, nbody, 1 ) \
X( int, body_, dofnum, nbody, 1 ) \
X( int, body_, dofadr, nbody, 1 ) \
X( int, body_, geomnum, nbody, 1 ) \
X( int, body_, geomadr, nbody, 1 ) \
X( mjtByte, body_, simple, nbody, 1 ) \
X( mjtByte, body_, sameframe, nbody, 1 ) \
X( mjtNum, body_, pos, nbody, 3 ) \
X( mjtNum, body_, quat, nbody, 4 ) \
X( mjtNum, body_, ipos, nbody, 3 ) \
X( mjtNum, body_, iquat, nbody, 4 ) \
X( mjtNum, body_, mass, nbody, 1 ) \
X( mjtNum, body_, subtreemass, nbody, 1 ) \
X( mjtNum, body_, inertia, nbody, 3 ) \
X( mjtNum, body_, invweight0, nbody, 2 ) \
X( mjtNum, body_, user, nbody, MJ_M(nuser_body) )
#define MJMODEL_CAMERA \
X( int, cam_, mode, ncam, 1 ) \
X( int, cam_, bodyid, ncam, 1 ) \
X( int, cam_, targetbodyid, ncam, 1 ) \
X( mjtNum, cam_, pos, ncam, 3 ) \
X( mjtNum, cam_, quat, ncam, 4 ) \
X( mjtNum, cam_, poscom0, ncam, 3 ) \
X( mjtNum, cam_, pos0, ncam, 3 ) \
X( mjtNum, cam_, mat0, ncam, 9 ) \
X( mjtNum, cam_, fovy, ncam, 1 ) \
X( mjtNum, cam_, ipd, ncam, 1 ) \
X( mjtNum, cam_, user, ncam, MJ_M(nuser_cam) )
#define MJMODEL_EQUALITY \
X( int, eq_, type, neq, 1 ) \
X( int, eq_, obj1id, neq, 1 ) \
X( int, eq_, obj2id, neq, 1 ) \
X( mjtByte, eq_, active, neq, 1 ) \
X( mjtNum, eq_, solref, neq, mjNREF ) \
X( mjtNum, eq_, solimp, neq, mjNIMP ) \
X( mjtNum, eq_, data, neq, mjNEQDATA )
#define MJMODEL_EXCLUDE \
X( int, exclude_, signature, nexclude, 1 )
#define MJMODEL_GEOM \
X( int, geom_, type, ngeom, 1 ) \
X( int, geom_, contype, ngeom, 1 ) \
X( int, geom_, conaffinity, ngeom, 1 ) \
X( int, geom_, condim, ngeom, 1 ) \
X( int, geom_, bodyid, ngeom, 1 ) \
X( int, geom_, dataid, ngeom, 1 ) \
X( int, geom_, matid, ngeom, 1 ) \
X( int, geom_, group, ngeom, 1 ) \
X( int, geom_, priority, ngeom, 1 ) \
X( mjtByte, geom_, sameframe, ngeom, 1 ) \
X( mjtNum, geom_, solmix, ngeom, 1 ) \
X( mjtNum, geom_, solref, ngeom, mjNREF ) \
X( mjtNum, geom_, solimp, ngeom, mjNIMP ) \
X( mjtNum, geom_, size, ngeom, 3 ) \
X( mjtNum, geom_, rbound, ngeom, 1 ) \
X( mjtNum, geom_, pos, ngeom, 3 ) \
X( mjtNum, geom_, quat, ngeom, 4 ) \
X( mjtNum, geom_, friction, ngeom, 3 ) \
X( mjtNum, geom_, margin, ngeom, 1 ) \
X( mjtNum, geom_, gap, ngeom, 1 ) \
X( mjtNum, geom_, user, ngeom, MJ_M(nuser_geom) ) \
X( float, geom_, rgba, ngeom, 4 )
#define MJMODEL_HFIELD \
X( mjtNum, hfield_, size, nhfield, 4 ) \
X( int, hfield_, nrow, nhfield, 1 ) \
X( int, hfield_, ncol, nhfield, 1 ) \
X( int, hfield_, adr, nhfield, 1 ) \
X( float, hfield_, data, nhfielddata, 1 )
#define MJMODEL_JOINT \
X( mjtNum, , qpos0, nq, 1 ) \
X( mjtNum, , qpos_spring, nq, 1 ) \
X( int, jnt_, type, njnt, 1 ) \
X( int, jnt_, qposadr, njnt, 1 ) \
X( int, jnt_, dofadr, njnt, 1 ) \
X( int, jnt_, group, njnt, 1 ) \
X( mjtByte, jnt_, limited, njnt, 1 ) \
X( mjtNum, jnt_, pos, njnt, 3 ) \
X( mjtNum, jnt_, axis, njnt, 3 ) \
X( mjtNum, jnt_, stiffness, njnt, 1 ) \
X( mjtNum, jnt_, range, njnt, 2 ) \
X( mjtNum, jnt_, margin, njnt, 1 ) \
X( mjtNum, jnt_, user, njnt, MJ_M(nuser_jnt) ) \
X( int, dof_, bodyid, nv, 1 ) \
X( int, dof_, jntid, nv, 1 ) \
X( int, dof_, parentid, nv, 1 ) \
X( int, dof_, Madr, nv, 1 ) \
X( int, dof_, simplenum, nv, 1 ) \
X( mjtNum, dof_, solref, nv, mjNREF ) \
X( mjtNum, dof_, solimp, nv, mjNIMP ) \
X( mjtNum, dof_, frictionloss, nv, 1 ) \
X( mjtNum, dof_, armature, nv, 1 ) \
X( mjtNum, dof_, damping, nv, 1 ) \
X( mjtNum, dof_, invweight0, nv, 1 ) \
X( mjtNum, dof_, M0, nv, 1 )
#define MJMODEL_LIGHT \
X( int, light_, mode, nlight, 1 ) \
X( int, light_, bodyid, nlight, 1 ) \
X( int, light_, targetbodyid, nlight, 1 ) \
X( mjtByte, light_, directional, nlight, 1 ) \
X( mjtByte, light_, castshadow, nlight, 1 ) \
X( mjtByte, light_, active, nlight, 1 ) \
X( mjtNum, light_, pos, nlight, 3 ) \
X( mjtNum, light_, dir, nlight, 3 ) \
X( mjtNum, light_, poscom0, nlight, 3 ) \
X( mjtNum, light_, pos0, nlight, 3 ) \
X( mjtNum, light_, dir0, nlight, 3 ) \
X( float, light_, attenuation, nlight, 3 ) \
X( float, light_, cutoff, nlight, 1 ) \
X( float, light_, exponent, nlight, 1 ) \
X( float, light_, ambient, nlight, 3 ) \
X( float, light_, diffuse, nlight, 3 ) \
X( float, light_, specular, nlight, 3 )
#define MJMODEL_MATERIAL \
X( int, mat_, texid, nmat, 1 ) \
X( mjtByte, mat_, texuniform, nmat, 1 ) \
X( float, mat_, texrepeat, nmat, 2 ) \
X( float, mat_, emission, nmat, 1 ) \
X( float, mat_, specular, nmat, 1 ) \
X( float, mat_, shininess, nmat, 1 ) \
X( float, mat_, reflectance, nmat, 1 ) \
X( float, mat_, rgba, nmat, 4 )
#define MJMODEL_MESH \
X( int, mesh_, vertadr, nmesh, 1 ) \
X( int, mesh_, vertnum, nmesh, 1 ) \
X( int, mesh_, texcoordadr, nmesh, 1 ) \
X( int, mesh_, faceadr, nmesh, 1 ) \
X( int, mesh_, facenum, nmesh, 1 ) \
X( int, mesh_, graphadr, nmesh, 1 )
#define MJMODEL_NUMERIC \
X( int, numeric_, adr, nnumeric, 1 ) \
X( int, numeric_, size, nnumeric, 1 ) \
X( mjtNum, numeric_, data, nnumericdata, 1 )
#define MJMODEL_PAIR \
X( int, pair_, dim, npair, 1 ) \
X( int, pair_, geom1, npair, 1 ) \
X( int, pair_, geom2, npair, 1 ) \
X( int, pair_, signature, npair, 1 ) \
X( mjtNum, pair_, solref, npair, mjNREF ) \
X( mjtNum, pair_, solimp, npair, mjNIMP ) \
X( mjtNum, pair_, margin, npair, 1 ) \
X( mjtNum, pair_, gap, npair, 1 ) \
X( mjtNum, pair_, friction, npair, 5 )
#define MJMODEL_SENSOR \
X( int, sensor_, type, nsensor, 1 ) \
X( int, sensor_, datatype, nsensor, 1 ) \
X( int, sensor_, needstage, nsensor, 1 ) \
X( int, sensor_, objtype, nsensor, 1 ) \
X( int, sensor_, objid, nsensor, 1 ) \
X( int, sensor_, reftype, nsensor, 1 ) \
X( int, sensor_, refid, nsensor, 1 ) \
X( int, sensor_, dim, nsensor, 1 ) \
X( int, sensor_, adr, nsensor, 1 ) \
X( mjtNum, sensor_, cutoff, nsensor, 1 ) \
X( mjtNum, sensor_, noise, nsensor, 1 ) \
X( mjtNum, sensor_, user, nsensor, MJ_M(nuser_sensor) )
#define MJMODEL_SITE \
X( int, site_, type, nsite, 1 ) \
X( int, site_, bodyid, nsite, 1 ) \
X( int, site_, matid, nsite, 1 ) \
X( int, site_, group, nsite, 1 ) \
X( mjtByte, site_, sameframe, nsite, 1 ) \
X( mjtNum, site_, size, nsite, 3 ) \
X( mjtNum, site_, pos, nsite, 3 ) \
X( mjtNum, site_, quat, nsite, 4 ) \
X( mjtNum, site_, user, nsite, MJ_M(nuser_site) ) \
X( float, site_, rgba, nsite, 4 )
#define MJMODEL_SKIN \
X( int, skin_, matid, nskin, 1 ) \
X( float, skin_, rgba, nskin, 4 ) \
X( float, skin_, inflate, nskin, 1 ) \
X( int, skin_, vertadr, nskin, 1 ) \
X( int, skin_, vertnum, nskin, 1 ) \
X( int, skin_, texcoordadr, nskin, 1 ) \
X( int, skin_, faceadr, nskin, 1 ) \
X( int, skin_, facenum, nskin, 1 ) \
X( int, skin_, boneadr, nskin, 1 ) \
X( int, skin_, bonenum, nskin, 1 )
#define MJMODEL_TENDON \
X( int, tendon, _adr, ntendon, 1 ) \
X( int, tendon, _num, ntendon, 1 ) \
X( int, tendon, _matid, ntendon, 1 ) \
X( int, tendon, _group, ntendon, 1 ) \
X( mjtByte, tendon, _limited, ntendon, 1 ) \
X( mjtNum, tendon, _width, ntendon, 1 ) \
X( mjtNum, tendon, _solref_lim, ntendon, mjNREF ) \
X( mjtNum, tendon, _solimp_lim, ntendon, mjNIMP ) \
X( mjtNum, tendon, _solref_fri, ntendon, mjNREF ) \
X( mjtNum, tendon, _solimp_fri, ntendon, mjNIMP ) \
X( mjtNum, tendon, _range, ntendon, 2 ) \
X( mjtNum, tendon, _margin, ntendon, 1 ) \
X( mjtNum, tendon, _stiffness, ntendon, 1 ) \
X( mjtNum, tendon, _damping, ntendon, 1 ) \
X( mjtNum, tendon, _frictionloss, ntendon, 1 ) \
X( mjtNum, tendon, _lengthspring, ntendon, 1 ) \
X( mjtNum, tendon, _length0, ntendon, 1 ) \
X( mjtNum, tendon, _invweight0, ntendon, 1 ) \
X( mjtNum, tendon, _user, ntendon, MJ_M(nuser_tendon) ) \
X( float, tendon, _rgba, ntendon, 4 )
#define MJMODEL_TEXTURE \
X( int, tex_, type, ntex, 1 ) \
X( int, tex_, height, ntex, 1 ) \
X( int, tex_, width, ntex, 1 ) \
X( int, tex_, adr, ntex, 1 ) \
X( mjtByte, tex_, rgb, ntexdata, 1 )
#define MJMODEL_TUPLE \
X( int, tuple_, adr, ntuple, 1 ) \
X( int, tuple_, size, ntuple, 1 ) \
X( int, tuple_, objtype, ntupledata, 1 ) \
X( int, tuple_, objid, ntupledata, 1 ) \
X( mjtNum, tuple_, objprm, ntupledata, 1 )
#define MJMODEL_KEYFRAME \
X( mjtNum, key_, time, nkey, 1 ) \
X( mjtNum, key_, qpos, nkey, MJ_M(nq) ) \
X( mjtNum, key_, qvel, nkey, MJ_M(nv) ) \
X( mjtNum, key_, act, nkey, MJ_M(na) ) \
X( mjtNum, key_, mpos, nkey, MJ_M(nmocap)*3 ) \
X( mjtNum, key_, mquat, nkey, MJ_M(nmocap)*4 )
#define MJMODEL_VIEW_GROUPS \
XGROUP( MjModelActuatorViews, actuator, nu, MJMODEL_ACTUATOR ) \
XGROUP( MjModelBodyViews, body, nbody, MJMODEL_BODY ) \
XGROUP( MjModelCameraViews, cam, ncam, MJMODEL_CAMERA ) \
XGROUP( MjModelEqualityViews, eq, neq, MJMODEL_EQUALITY ) \
XGROUP( MjModelExcludeViews, exclude, nexclude, MJMODEL_EXCLUDE ) \
XGROUP( MjModelGeomViews, geom, ngeom, MJMODEL_GEOM ) \
XGROUP( MjModelHfieldViews, hfield, nhfield, MJMODEL_HFIELD ) \
XGROUP( MjModelJointViews, jnt, njnt, MJMODEL_JOINT ) \
XGROUP( MjModelLightViews, light, nlight, MJMODEL_LIGHT ) \
XGROUP( MjModelMaterialViews, mat, nmat, MJMODEL_MATERIAL ) \
XGROUP( MjModelMeshViews, mesh, nmesh, MJMODEL_MESH ) \
XGROUP( MjModelNumericViews, numeric, nnumeric, MJMODEL_NUMERIC ) \
XGROUP( MjModelPairViews, pair, npair, MJMODEL_PAIR ) \
XGROUP( MjModelSensorViews, sensor, nsensor, MJMODEL_SENSOR ) \
XGROUP( MjModelSiteViews, site, nsite, MJMODEL_SITE ) \
XGROUP( MjModelSkinViews, skin, nskin, MJMODEL_SKIN ) \
XGROUP( MjModelTendonViews, tendon, ntendon, MJMODEL_TENDON ) \
XGROUP( MjModelTextureViews, tex, ntex, MJMODEL_TEXTURE ) \
XGROUP( MjModelTupleViews, tuple, ntuple, MJMODEL_TUPLE ) \
XGROUP( MjModelKeyframeViews, key, nkey, MJMODEL_KEYFRAME )
#define MJMODEL_VIEW_GROUPS_ALTNAMES \
XGROUP( cam, camera, MJMODEL_CAMERA ) \
XGROUP( eq, equality, MJMODEL_EQUALITY ) \
XGROUP( jnt, joint, MJMODEL_JOINT ) \
XGROUP( mat, material, MJMODEL_MATERIAL ) \
XGROUP( tex, texture, MJMODEL_TEXTURE ) \
XGROUP( key, keyframe, MJMODEL_KEYFRAME )
#define MJDATA_ACTUATOR \
X( mjtNum, , ctrl, nu, 1 ) \
X( mjtNum, actuator_, length, nu, 1 ) \
X( mjtNum, actuator_, moment, nu, MJ_M(nv) ) \
X( mjtNum, actuator_, velocity, nu, 1 ) \
X( mjtNum, actuator_, force, nu, 1 )
#define MJDATA_BODY \
X( mjtNum, , xfrc_applied, nbody, 6 ) \
X( mjtNum, , xpos, nbody, 3 ) \
X( mjtNum, , xquat, nbody, 4 ) \
X( mjtNum, , xmat, nbody, 9 ) \
X( mjtNum, , xipos, nbody, 3 ) \
X( mjtNum, , ximat, nbody, 9 ) \
X( mjtNum, , subtree_com, nbody, 3 ) \
X( mjtNum, , cinert, nbody, 10 ) \
X( mjtNum, , crb, nbody, 10 ) \
X( mjtNum, , cvel, nbody, 6 ) \
X( mjtNum, , subtree_linvel, nbody, 3 ) \
X( mjtNum, , subtree_angmom, nbody, 3 ) \
X( mjtNum, , cacc, nbody, 6 ) \
X( mjtNum, , cfrc_int, nbody, 6 ) \
X( mjtNum, , cfrc_ext, nbody, 6 )
#define MJDATA_CAMERA \
X( mjtNum, cam_, xpos, ncam, 3 ) \
X( mjtNum, cam_, xmat, ncam, 9 )
#define MJDATA_GEOM \
X( mjtNum, geom_, xpos, ngeom, 3 ) \
X( mjtNum, geom_, xmat, ngeom, 9 )
#define MJDATA_JOINT \
X( mjtNum, , qpos, nq, 1 ) \
X( mjtNum, , qvel, nv, 1 ) \
X( mjtNum, , qacc_warmstart, nv, 1 ) \
X( mjtNum, , qfrc_applied, nv, 1 ) \
X( mjtNum, , qacc, nv, 1 ) \
X( mjtNum, , xanchor, njnt, 3 ) \
X( mjtNum, , xaxis, njnt, 3 ) \
X( mjtNum, , cdof, nv, 6 ) \
X( mjtNum, , qLDiagInv, nv, 1 ) \
X( mjtNum, , qLDiagSqrtInv, nv, 1 ) \
X( int, , efc_JT_rownnz, nv, 1 ) \
X( int, , efc_JT_rowadr, nv, 1 ) \
X( int, , efc_JT_rowsuper, nv, 1 ) \
X( int, , efc_JT_colind, nv, MJ_M(njmax) ) \
X( mjtNum, , efc_JT, nv, MJ_M(njmax) ) \
X( mjtNum, , cdof_dot, nv, 6 ) \
X( mjtNum, , qfrc_bias, nv, 1 ) \
X( mjtNum, , qfrc_passive, nv, 1 ) \
X( mjtNum, , qfrc_actuator, nv, 1 ) \
X( mjtNum, , qfrc_unc, nv, 1 ) \
X( mjtNum, , qacc_unc, nv, 1 ) \
X( mjtNum, , qfrc_constraint, nv, 1 ) \
X( mjtNum, , qfrc_inverse, nv, 1 )
#define MJDATA_LIGHT \
X( mjtNum, light_, xpos, nlight, 3 ) \
X( mjtNum, light_, xdir, nlight, 3 )
#define MJDATA_SENSOR \
X( mjtNum, sensor, data, nsensordata, 1 )
#define MJDATA_SITE \
X( mjtNum, site_, xpos, nsite, 3 ) \
X( mjtNum, site_, xmat, nsite, 9 )
#define MJDATA_TENDON \
X( int, ten_, wrapadr , ntendon, 1 ) \
X( int, ten_, wrapnum , ntendon, 1 ) \
X( int, ten_, J_rownnz, ntendon, 1 ) \
X( int, ten_, J_rowadr, ntendon, 1 ) \
X( int, ten_, J_colind, ntendon, MJ_M(nv) ) \
X( mjtNum, ten_, length , ntendon, 1 ) \
X( mjtNum, ten_, J , ntendon, MJ_M(nv) ) \
X( mjtNum, ten_, velocity, ntendon, 1 )
#define MJDATA_VIEW_GROUPS \
XGROUP( MjDataActuatorViews, actuator, nu, MJDATA_ACTUATOR ) \
XGROUP( MjDataBodyViews, body, nbody, MJDATA_BODY ) \
XGROUP( MjDataCameraViews, cam, ncam, MJDATA_CAMERA ) \
XGROUP( MjDataGeomViews, geom, ngeom, MJDATA_GEOM ) \
XGROUP( MjDataJointViews, jnt, njnt, MJDATA_JOINT ) \
XGROUP( MjDataLightViews, light, nlight, MJDATA_LIGHT ) \
XGROUP( MjDataSensorViews, sensor, nsensor, MJDATA_SENSOR ) \
XGROUP( MjDataSiteViews, site, nsite, MJDATA_SITE ) \
XGROUP( MjDataTendonViews, tendon, ntendon, MJDATA_TENDON )
#define MJDATA_VIEW_GROUPS_ALTNAMES \
XGROUP( cam, camera, MJDATA_CAMERA ) \
XGROUP( jnt, joint, MJDATA_JOINT ) \
XGROUP( tendon, ten, MJDATA_TENDON )
#endif // MUJOCO_PYTHON_INDEXER_XMACRO_H_
python/mujoco/indexers.cc
+366
View File
@@ -0,0 +1,366 @@
// 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.
#include <memory>
#include <optional>
#include <string>
#include <string_view>
#include <utility>
#include <variant>
#include <vector>
#include "errors.h"
#include "indexers.h"
#include "mjdata_meta.h"
#include "raw.h"
#include <pybind11/pybind11.h>
namespace mujoco::python {
namespace py = ::pybind11;
namespace {
// Parses raw mjModel to create a mapping from names to array indices.
//
// Args:
// count: Number of names in the map.
// name_offsets: Array consisting of indices that correspond to the start of
// each name in the `names` array.
// names: Character arrays consisting of the concatenation of all names.
template <typename IntPtr, typename CharPtr>
NameToIDMap MakeMap(int count, IntPtr name_offsets, CharPtr names)
{
NameToIDMap name_to_id;
for (int index = 0; index < count; ++index) {
const char* name = &names[name_offsets[index]];
if (name[0] != '\0') {
name_to_id.insert({name, index});
}
}
return name_to_id;
}
// Makes an array view into an mjModel/mjData struct field at a given index.
//
// Template args:
// MjSize: An (int M::*) specifying the number of entities of the same type
// as the one whose array view is to be created. This is a way to indicate
// the MuJoCo category of the entity itself, e.g. nbody indicates that the
// field belongs to a body.
// T: Scalar data type of the field.
// M: Either raw::MjModel or MjDataMetadata.
//
// Args:
// base_ptr: Pointer to the first entry in the entire field.
// index: Entity ID at which a view is to be created.
// shape: The shape of the array as defined by MuJoCo. If MjSize corresponds
// to a ragged or indirect field, the shape may be prepended by additional
// dimensions as appropriate, e.g. if MjSize is nv then the view gains an
// additional dimension of the size len(qvel) of the particular joint.
// m: Used for dereferencing MjSize.
// owner: The base object whose lifetime is tied to the returned array.
template <auto MjSize, typename T, typename M>
py::array_t<T> MakeArray(T* base_ptr, int index, std::vector<int>&& shape,
const M& m, py::handle owner) {
int offset;
if (MjSize == &M::nq) {
offset = m.jnt_qposadr[index];
shape.insert(
shape.begin(),
((index < m.njnt-1) ? m.jnt_qposadr[index+1] : m.nq) - offset);
} else if (MjSize == &M::nv) {
offset = m.jnt_dofadr[index];
shape.insert(
shape.begin(),
((index < m.njnt-1) ? m.jnt_dofadr[index+1] : m.nv) - offset);
} else if (MjSize == &M::nhfielddata) {
offset = m.hfield_adr[index];
shape.insert(shape.begin(), m.hfield_ncol[index]);
shape.insert(shape.begin(), m.hfield_nrow[index]);
} else if (MjSize == &M::ntexdata) {
offset = m.tex_adr[index];
shape.insert(shape.begin(), m.tex_width[index]);
shape.insert(shape.begin(), m.tex_height[index]);
} else if (MjSize == &M::nsensordata) {
offset = m.sensor_adr[index];
shape.insert(shape.begin(), m.sensor_dim[index]);
} else if (MjSize == &M::nnumericdata) {
offset = m.numeric_adr[index];
shape.insert(shape.begin(), m.numeric_size[index]);
} else if (MjSize == &M::ntupledata) {
offset = m.tuple_adr[index];
shape.insert(shape.begin(), m.tuple_size[index]);
} else {
int size = 1;
for (int s : shape) {
size *= s;
}
offset = index * size;
}
return py::array_t<T>(std::move(shape), base_ptr + offset,
py::reinterpret_borrow<py::object>(owner));
}
} // namespace
// M is either a raw::MjModel or MjDataMetadata.
template <typename M>
NameToIDMaps::NameToIDMaps(const M& m)
: body(MakeMap(m.nbody, m.name_bodyadr, m.names)),
jnt(MakeMap(m.njnt, m.name_jntadr, m.names)),
geom(MakeMap(m.ngeom, m.name_geomadr, m.names)),
site(MakeMap(m.nsite, m.name_siteadr, m.names)),
cam(MakeMap(m.ncam, m.name_camadr, m.names)),
light(MakeMap(m.nlight, m.name_lightadr, m.names)),
mesh(MakeMap(m.nmesh, m.name_meshadr, m.names)),
skin(MakeMap(m.nskin, m.name_skinadr, m.names)),
hfield(MakeMap(m.nhfield, m.name_hfieldadr, m.names)),
tex(MakeMap(m.ntex, m.name_texadr, m.names)),
mat(MakeMap(m.nmat, m.name_matadr, m.names)),
pair(MakeMap(m.npair, m.name_pairadr, m.names)),
exclude(MakeMap(m.nexclude, m.name_excludeadr, m.names)),
eq(MakeMap(m.neq, m.name_eqadr, m.names)),
tendon(MakeMap(m.ntendon, m.name_tendonadr, m.names)),
actuator(MakeMap(m.nu, m.name_actuatoradr, m.names)),
sensor(MakeMap(m.nsensor, m.name_sensoradr, m.names)),
numeric(MakeMap(m.nnumeric, m.name_numericadr, m.names)),
text(MakeMap(m.ntext, m.name_textadr, m.names)),
tuple(MakeMap(m.ntuple, m.name_tupleadr, m.names)),
key(MakeMap(m.nkey, m.name_keyadr, m.names)) {}
MjModelIndexer::MjModelIndexer(raw::MjModel* m, py::handle owner)
: m_(m),
owner_(owner),
name_to_id_(*m)
#define XGROUP(MjModelFieldGroupedViews, field, nfield, FIELD_XMACROS) \
, field##_(m->nfield, std::nullopt)
MJMODEL_VIEW_GROUPS
#undef XGROUP
{}
#define XGROUP(MjModelFieldGroupedViews, field, nfield, FIELD_XMACROS) \
MjModelFieldGroupedViews& MjModelIndexer::field(int i) { \
if (i > field##_.size()) { \
throw py::index_error("index out of range"); \
} \
auto& indexer = field##_[i]; \
if (!indexer.has_value()) { \
indexer.emplace(i, m_, owner_); \
} \
return *indexer; \
}
MJMODEL_VIEW_GROUPS
#undef XGROUP
#define XGROUP(MjModelFieldGroupedViews, field, nfield, FIELD_XMACROS) \
MjModelFieldGroupedViews& MjModelIndexer::field##_by_name( \
std::string_view name) { \
try { \
return field(name_to_id_.field.at(name)); \
} catch (...) { \
throw py::key_error(std::string(name)); \
} \
}
MJMODEL_VIEW_GROUPS
#undef XGROUP
MjDataIndexer::MjDataIndexer(raw::MjData* d, const MjDataMetadata* m,
py::handle owner)
: d_(d),
m_(m),
owner_(owner),
name_to_id_(*m)
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
, field##_(m->nfield, std::nullopt)
MJDATA_VIEW_GROUPS
#undef XGROUP
{}
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
MjDataGroupedViews& MjDataIndexer::field(int i) { \
if (i > field##_.size()) { \
throw py::index_error("index out of range"); \
} \
auto& indexer = field##_[i]; \
if (!indexer.has_value()) { \
indexer.emplace(i, d_, m_, owner_); \
} \
return *indexer; \
}
MJDATA_VIEW_GROUPS
#undef XGROUP
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
MjDataGroupedViews& MjDataIndexer::field##_by_name( \
std::string_view name) { \
try { \
return field(name_to_id_.field.at(name)); \
} catch (...) { \
throw py::key_error(std::string(name)); \
} \
}
MJDATA_VIEW_GROUPS
#undef XGROUP
#define MAKE_SHAPE(dim) \
[n = (dim)]() -> std::vector<int> { \
if constexpr (std::string_view(#dim) == std::string_view("1")) { \
return {}; \
} else { \
return {n}; \
} \
}()
#undef MJ_M
#define MJ_M(n) m_->n
#define X(type, prefix, var, dim0, dim1) \
py::array_t<type> XGROUP::var() { \
if (!var##_.has_value()) { \
var##_.emplace(MakeArray<&raw::MjModel::dim0>( \
m_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \
} \
return *var##_; \
}
#define XGROUP MjModelActuatorViews
MJMODEL_ACTUATOR
#undef XGROUP
#define XGROUP MjModelBodyViews
MJMODEL_BODY
#undef XGROUP
#define XGROUP MjModelCameraViews
MJMODEL_CAMERA
#undef XGROUP
#define XGROUP MjModelEqualityViews
MJMODEL_EQUALITY
#undef XGROUP
#define XGROUP MjModelExcludeViews
MJMODEL_EXCLUDE
#undef XGROUP
#define XGROUP MjModelGeomViews
MJMODEL_GEOM
#undef XGROUP
#define XGROUP MjModelHfieldViews
MJMODEL_HFIELD
#undef XGROUP
#define XGROUP MjModelJointViews
MJMODEL_JOINT
#undef XGROUP
#define XGROUP MjModelLightViews
MJMODEL_LIGHT
#undef XGROUP
#define XGROUP MjModelMaterialViews
MJMODEL_MATERIAL
#undef XGROUP
#define XGROUP MjModelMeshViews
MJMODEL_MESH
#undef XGROUP
#define XGROUP MjModelNumericViews
MJMODEL_NUMERIC
#undef XGROUP
#define XGROUP MjModelPairViews
MJMODEL_PAIR
#undef XGROUP
#define XGROUP MjModelSensorViews
MJMODEL_SENSOR
#undef XGROUP
#define XGROUP MjModelSiteViews
MJMODEL_SITE
#undef XGROUP
#define XGROUP MjModelSkinViews
MJMODEL_SKIN
#undef XGROUP
#define XGROUP MjModelTendonViews
MJMODEL_TENDON
#undef XGROUP
#define XGROUP MjModelTextureViews
MJMODEL_TEXTURE
#undef XGROUP
#define XGROUP MjModelTupleViews
MJMODEL_TUPLE
#undef XGROUP
#define XGROUP MjModelKeyframeViews
MJMODEL_KEYFRAME
#undef XGROUP
#undef X
#define MJ_M(n) m_->n
#define X(type, prefix, var, dim0, dim1) \
py::array_t<type> XGROUP::var() { \
if (!var##_.has_value()) { \
var##_.emplace(MakeArray<&MjDataMetadata::dim0>( \
d_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \
} \
return *var##_; \
}
#define XGROUP MjDataActuatorViews
MJDATA_ACTUATOR
#undef XGROUP
#define XGROUP MjDataBodyViews
MJDATA_BODY
#undef XGROUP
#define XGROUP MjDataCameraViews
MJDATA_CAMERA
#undef XGROUP
#define XGROUP MjDataGeomViews
MJDATA_GEOM
#undef XGROUP
#define XGROUP MjDataJointViews
MJDATA_JOINT
#undef XGROUP
#define XGROUP MjDataLightViews
MJDATA_LIGHT
#undef XGROUP
#define XGROUP MjDataSensorViews
MJDATA_SENSOR
#undef XGROUP
#define XGROUP MjDataSiteViews
MJDATA_SITE
#undef XGROUP
#define XGROUP MjDataTendonViews
MJDATA_TENDON
#undef XGROUP
#undef X
#undef MJ_M
#define MJ_M(n) n
} // namespace mujoco::python
python/mujoco/indexers.h
+190
View File
@@ -0,0 +1,190 @@
// 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.
#ifndef MUJOCO_PYTHON_INDEXERS_H_
#define MUJOCO_PYTHON_INDEXERS_H_
#include <optional>
#include <string>
#include <string_view>
#include <utility>
#include <variant>
#include <vector>
#include <absl/container/flat_hash_map.h>
#include <mjxmacro.h>
#include "indexer_xmacro.h"
#include "mjdata_meta.h"
#include "raw.h"
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
namespace mujoco::python {
using NameToIDMap = absl::flat_hash_map<std::string, int>;
struct NameToIDMaps {
// M is either a raw::MjModel or MjDataMetadata.
template <typename M>
explicit NameToIDMaps(const M& m);
NameToIDMap body;
NameToIDMap jnt;
NameToIDMap geom;
NameToIDMap site;
NameToIDMap cam;
NameToIDMap light;
NameToIDMap mesh;
NameToIDMap skin;
NameToIDMap hfield;
NameToIDMap tex;
NameToIDMap mat;
NameToIDMap pair;
NameToIDMap exclude;
NameToIDMap eq;
NameToIDMap tendon;
NameToIDMap actuator;
NameToIDMap sensor;
NameToIDMap numeric;
NameToIDMap text;
NameToIDMap tuple;
NameToIDMap key;
};
// Base class for a collection of NumPy views into mjModel fields associated
// with the same "entity" (for example, fields corresponding to a particular
// geom, or a particular joint).
class MjModelGroupedViewsBase {
public:
MjModelGroupedViewsBase(int index, raw::MjModel* m,
pybind11::handle owner)
: index_(index), m_(m), owner_(owner) {}
protected:
int index_;
raw::MjModel* m_;
pybind11::handle owner_;
};
#define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS) \
class MjModelGroupedViews : public MjModelGroupedViewsBase { \
public: \
using MjModelGroupedViewsBase::MjModelGroupedViewsBase; \
FIELD_XMACROS \
};
// Lazily instantiate a NumPy array view object when accessed from Python, but
// cache it once made so that we can return the same one if requested again.
#define X(type, prefix, var, dim0, dim1) \
pybind11::array_t<type> var(); \
std::optional<pybind11::array_t<type>> var##_;
MJMODEL_VIEW_GROUPS
#undef XGROUP
#undef X
// A class for accessing fields corresponding to a single mjModel "entity"
// (e.g. a particular geom or joint) either by name or by ID.
class MjModelIndexer {
public:
MjModelIndexer(raw::MjModel* m, pybind11::handle owner);
#define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS) \
MjModelGroupedViews& field(int i); \
MjModelGroupedViews& field##_by_name(std::string_view name);
MJMODEL_VIEW_GROUPS
#undef XGROUP
private:
raw::MjModel* m_;
pybind11::handle owner_;
NameToIDMaps name_to_id_;
// Lazily instantiate a grouped views object when accessed from Python, but
// cache it once made so that we can return the same one if requested again.
#define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS) \
std::vector<std::optional<MjModelGroupedViews>> field##_;
MJMODEL_VIEW_GROUPS
#undef XGROUP
};
// Base class for a collection of NumPy views into mjData fields associated
// with the same "entity" (for example, fields corresponding to a particular
// geom, or a particular joint).
class MjDataGroupedViewsBase {
public:
MjDataGroupedViewsBase(int index, raw::MjData* d,
const MjDataMetadata* m,
pybind11::handle owner)
: index_(index), d_(d), m_(m), owner_(owner) {}
protected:
int index_;
raw::MjData* d_;
const MjDataMetadata* m_;
pybind11::handle owner_;
};
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
class MjDataGroupedViews : public MjDataGroupedViewsBase { \
public: \
using MjDataGroupedViewsBase::MjDataGroupedViewsBase; \
FIELD_XMACROS \
};
// Lazily instantiate a NumPy array view object when accessed from Python, but
// cache it once made so that we can return the same one if requested again.
#define X(type, prefix, var, dim0, dim1) \
pybind11::array_t<type> var(); \
std::optional<pybind11::array_t<type>> var##_;
MJDATA_VIEW_GROUPS
#undef XGROUP
#undef X
// A class for accessing fields corresponding to a single mjModel "entity"
// (e.g. a particular geom or joint) either by name or by ID.
class MjDataIndexer {
public:
MjDataIndexer(raw::MjData* d, const MjDataMetadata* m,
pybind11::handle owner);
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
MjDataGroupedViews& field(int i); \
MjDataGroupedViews& field##_by_name(std::string_view name);
MJDATA_VIEW_GROUPS
#undef XGROUP
private:
raw::MjData* d_;
const MjDataMetadata* m_;
pybind11::handle owner_;
NameToIDMaps name_to_id_;
// Lazily instantiate a grouped views object when accessed from Python, but
// cache it once made so that we can return the same one if requested again.
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
std::vector<std::optional<MjDataGroupedViews>> field##_;
MJDATA_VIEW_GROUPS
#undef XGROUP
};
} // namespace mujoco::python
#endif // MUJOCO_PYTHON_INDEXERS_H_
python/mujoco/mjdata_meta.h
+103
View File
@@ -0,0 +1,103 @@
// Copyright 2021 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.
#ifndef MUJOCO_PYTHON_MJDATA_META_H_
#define MUJOCO_PYTHON_MJDATA_META_H_
#include <mjxmacro.h>
#include "raw.h"
namespace mujoco::python {
namespace _impl {
template <typename T> class MjWrapper;
} // namespace _impl
#define MJDATA_METADATA \
X( int, jnt_qposadr, njnt ) \
X( int, jnt_dofadr, njnt ) \
X( int, hfield_nrow, nhfield ) \
X( int, hfield_ncol, nhfield ) \
X( int, hfield_adr, nhfield ) \
X( int, tex_height, ntex ) \
X( int, tex_width, ntex ) \
X( int, tex_adr, ntex ) \
X( int, sensor_dim, nsensor ) \
X( int, sensor_adr, nsensor ) \
X( int, numeric_adr, nnumeric ) \
X( int, numeric_size, nnumeric ) \
X( int, tuple_adr, ntuple ) \
X( int, tuple_size, ntuple ) \
X( int, name_bodyadr, nbody ) \
X( int, name_jntadr, njnt ) \
X( int, name_geomadr, ngeom ) \
X( int, name_siteadr, nsite ) \
X( int, name_camadr, ncam ) \
X( int, name_lightadr, nlight ) \
X( int, name_meshadr, nmesh ) \
X( int, name_skinadr, nskin ) \
X( int, name_hfieldadr, nhfield ) \
X( int, name_texadr, ntex ) \
X( int, name_matadr, nmat ) \
X( int, name_pairadr, npair ) \
X( int, name_excludeadr, nexclude ) \
X( int, name_eqadr, neq ) \
X( int, name_tendonadr, ntendon ) \
X( int, name_actuatoradr, nu ) \
X( int, name_sensoradr, nsensor ) \
X( int, name_numericadr, nnumeric ) \
X( int, name_textadr, ntext ) \
X( int, name_tupleadr, ntuple ) \
X( int, name_keyadr, nkey ) \
X( char, names, nnames )
// A subset of mjModel fields that are required to reconstruct an MjDataWrapper.
struct MjDataMetadata {
public:
friend class _impl::MjWrapper<raw::MjData>;
#define X(var) int var;
MJMODEL_INTS
#undef X
#define X(type, var, n) std::shared_ptr<type[]> var;
MJDATA_METADATA
#undef X
private:
MjDataMetadata() = default;
MjDataMetadata(const MjDataMetadata& other) = default;
MjDataMetadata(MjDataMetadata&& other) = default;
explicit MjDataMetadata(const raw::MjModel* m)
:
#define X(var) var(m->var),
MJMODEL_INTS
#undef X
#define X(dtype, var, n) \
var( \
[](dtype* src, int len) { \
dtype* dst = new dtype[len]; \
std::memcpy(dst, src, len * sizeof(dtype)); \
return dst; \
}(m->var, m->n)),
MJDATA_METADATA
#undef X
dummy_() {}
bool dummy_; // Dummy variable to terminate X macro sequences.
};
} // namespace mujoco::python
#endif // MUJOCO_PYTHON_MJDATA_META_H_
python/mujoco/osmesa/__init__.py
+81
View File
@@ -0,0 +1,81 @@
# Copyright 2018 The dm_control Authors
#
# 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.
# ==============================================================================
"""An OSMesa context for software-based OpenGL rendering."""
import os
PYOPENGL_PLATFORM = os.environ.get('PYOPENGL_PLATFORM')
if not PYOPENGL_PLATFORM:
os.environ['PYOPENGL_PLATFORM'] = 'osmesa'
elif PYOPENGL_PLATFORM.lower() != 'osmesa':
raise ImportError(
'Cannot use OSMesa rendering platform. '
'The PYOPENGL_PLATFORM environment variable is set to {!r} '
'(should be either unset or \'osmesa\').'
.format(PYOPENGL_PLATFORM))
# pylint: disable=g-import-not-at-top
from OpenGL import GL
from OpenGL import osmesa
from OpenGL.GL import arrays
_DEPTH_BITS = 24
_STENCIL_BITS = 8
_ACCUM_BITS = 0
class GLContext:
"""An OSMesa context for software-based OpenGL rendering."""
def __init__(self, max_width, max_height):
"""Initializes this OSMesa context."""
self._context = osmesa.OSMesaCreateContextExt(
osmesa.OSMESA_RGBA,
_DEPTH_BITS,
_STENCIL_BITS,
_ACCUM_BITS,
None, # sharelist
)
if not self._context:
raise RuntimeError('Failed to create OSMesa GL context.')
self._height = max_height
self._width = max_width
# Allocate a buffer to render into.
self._buffer = arrays.GLfloatArray.zeros((max_height, max_width, 4))
def make_current(self):
if self._context:
success = osmesa.OSMesaMakeCurrent(
self._context,
self._buffer,
GL.GL_FLOAT,
self._width,
self._height)
if not success:
raise RuntimeError('Failed to make OSMesa context current.')
def free(self):
"""Frees resources associated with this context."""
if self._context and self._context == osmesa.OSMesaGetCurrentContext():
osmesa.OSMesaMakeCurrent(None, None, GL.GL_FLOAT, 0, 0)
osmesa.OSMesaDestroyContext(self._context)
self._buffer = None
self._context = None
def __del__(self):
self.free()
python/mujoco/raw.h
+60
View File
@@ -0,0 +1,60 @@
// 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.
#ifndef MUJOCO_PYTHON_RAW_H_
#define MUJOCO_PYTHON_RAW_H_
#include <mjdata.h>
#include <mjmodel.h>
#include <mjrender.h>
#include <mjvisualize.h>
// Type aliases for MuJoCo C structs to allow us refer to consistently refer
// to them under the "raw" namespace.
namespace mujoco::raw {
using MjContact = ::mjContact;
using MjData = ::mjData;
using MjLROpt = ::mjLROpt;
using MjModel = ::mjModel;
using MjOption = ::mjOption;
using MjSolverStat = ::mjSolverStat;
using MjStatistic = ::mjStatistic;
using MjTimerStat = ::mjTimerStat;
using MjVisual = ::mjVisual;
using MjVisualGlobal = decltype(::mjVisual::global);
using MjVisualQuality = decltype(::mjVisual::quality);
using MjVisualHeadlight = decltype(::mjVisual::headlight);
using MjVisualMap = decltype(::mjVisual::map);
using MjVisualScale = decltype(::mjVisual::scale);
using MjVisualRgba = decltype(::mjVisual::rgba);
using MjWarningStat = ::mjWarningStat;
// From mjrender.h
using MjrRect = ::mjrRect;
using MjrContext = ::mjrContext;
// From mjvisualize.h
using MjvPerturb = ::mjvPerturb;
using MjvCamera = ::mjvCamera;
using MjvGLCamera = ::mjvGLCamera;
using MjvGeom = ::mjvGeom;
using MjvLight = ::mjvLight;
using MjvOption = ::mjvOption;
using MjvScene = ::mjvScene;
using MjvFigure = ::mjvFigure;
} // namespace mujoco::raw
#endif // MUJOCO_PYTHON_RAW_H_
python/mujoco/render.cc
+276
View File
@@ -0,0 +1,276 @@
// 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.
#include <array>
#include <Eigen/Core>
#include <mjrender.h>
#include <mujoco.h>
#include "errors.h"
#include "function_traits.h"
#include "functions.h"
#include "raw.h"
#include "structs.h"
#include <pybind11/eigen.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
namespace mujoco::python {
namespace _impl {
template <>
class MjWrapper<raw::MjrContext> : public WrapperBase<raw::MjrContext> {
public:
MjWrapper();
MjWrapper(const MjModelWrapper& model, int fontscale);
MjWrapper(const MjWrapper&) = delete;
MjWrapper(MjWrapper&&) = default;
~MjWrapper() = default;
#define X(var) py_array_or_tuple_t<mjtNum> var
X(fogRGBA);
X(auxWidth);
X(auxHeight);
X(auxSamples);
X(auxFBO);
X(auxFBO_r);
X(auxColor);
X(auxColor_r);
X(textureType);
X(texture);
X(skinvertVBO);
X(skinnormalVBO);
X(skintexcoordVBO);
X(skinfaceVBO);
X(charWidth);
X(charWidthBig);
#undef X
};
using MjrContextWrapper = MjWrapper<raw::MjrContext>;
template <>
struct enable_if_mj_struct<raw::MjrContext> {
using type = void;
};
static void MjrContextCapsuleDestructor(PyObject* pyobj) {
auto* ptr =
static_cast<raw::MjrContext*>(PyCapsule_GetPointer(pyobj, nullptr));
mjr_freeContext(ptr);
delete ptr;
}
#define X(var) var(InitPyArray(ptr_->var, owner_))
#define X_SKIN(var) var(InitPyArray(std::array{ptr_->nskin}, ptr_->var, owner_))
MjrContextWrapper::MjWrapper()
: WrapperBase([]() {
raw::MjrContext *const ctx = new raw::MjrContext;
mjr_defaultContext(ctx);
return ctx;
}()),
X(fogRGBA),
X(auxWidth),
X(auxHeight),
X(auxSamples),
X(auxFBO),
X(auxFBO_r),
X(auxColor),
X(auxColor_r),
X(textureType),
X(texture),
X_SKIN(skinvertVBO),
X_SKIN(skinnormalVBO),
X_SKIN(skintexcoordVBO),
X_SKIN(skinfaceVBO),
X(charWidth),
X(charWidthBig) {}
MjrContextWrapper::MjWrapper(const MjModelWrapper& model, int fontscale)
: WrapperBase([fontscale](const raw::MjModel* m) {
raw::MjrContext *const ctx = new raw::MjrContext;
mjr_defaultContext(ctx);
InterceptMjErrors(mjr_makeContext)(m, ctx, fontscale);
return ctx;
}(model.get()), &MjrContextCapsuleDestructor),
X(fogRGBA),
X(auxWidth),
X(auxHeight),
X(auxSamples),
X(auxFBO),
X(auxFBO_r),
X(auxColor),
X(auxColor_r),
X(textureType),
X(texture),
X_SKIN(skinvertVBO),
X_SKIN(skinnormalVBO),
X_SKIN(skintexcoordVBO),
X_SKIN(skinfaceVBO),
X(charWidth),
X(charWidthBig) {}
#undef X_SKIN
#undef X
} // namespace _impl
namespace {
PYBIND11_MODULE(_render, pymodule) {
namespace py = ::pybind11;
namespace traits = python_traits;
using _impl::MjModelWrapper;
using _impl::MjrContextWrapper;
// Import the _structs module so that pybind11 knows about Python bindings
// for MjWrapper types and therefore generates prettier docstrings.
py::module::import("mujoco._structs");
py::class_<raw::MjrRect> mjrRect(pymodule, "MjrRect");
mjrRect.def(py::init([](int left, int bottom, int width, int height) {
return raw::MjrRect{left, bottom, width, height};
}),
py::arg("left"), py::arg("bottom"), py::arg("width"),
py::arg("height"));
mjrRect.def("__copy__",
[](const raw::MjrRect& other) { return raw::MjrRect(other); });
mjrRect.def("__deepcopy__", [](const raw::MjrRect& other, py::dict) {
return raw::MjrRect(other);
});
#define X(var) mjrRect.def_readwrite(#var, &raw::MjrRect::var)
X(left);
X(bottom);
X(width);
X(height);
#undef X
py::class_<MjrContextWrapper> mjrContext(pymodule, "MjrContext");
mjrContext.def(py::init<>());
mjrContext.def(py::init<const MjModelWrapper&, int>());
#define X(var) \
mjrContext.def_property( \
#var, [](const MjrContextWrapper& c) { return c.get()->var; }, \
[](MjrContextWrapper& c, mjtNum rhs) { c.get()->var = rhs; })
X(lineWidth);
X(shadowClip);
X(shadowScale);
X(fogStart);
X(fogEnd);
X(shadowSize);
X(offWidth);
X(offHeight);
X(offSamples);
X(fontScale);
X(offFBO);
X(offFBO_r);
X(offColor);
X(offColor_r);
X(offDepthStencil);
X(offDepthStencil_r);
X(shadowFBO);
X(shadowTex);
X(ntexture);
X(basePlane);
X(baseMesh);
X(baseHField);
X(baseBuiltin);
X(baseFontNormal);
X(baseFontShadow);
X(baseFontBig);
X(rangePlane);
X(rangeMesh);
X(rangeHField);
X(rangeBuiltin);
X(rangeFont);
X(nskin);
X(charHeight);
X(charHeightBig);
X(glewInitialized);
X(windowAvailable);
X(windowSamples);
X(windowStereo);
X(windowDoublebuffer);
X(currentBuffer);
#undef X
#define X(var) \
mjrContext.def_property_readonly( \
#var, [](const MjrContextWrapper& c) { return c.get()->var; })
X(nskin);
#undef X
#define X(var) DefinePyArray(mjrContext, #var, &MjrContextWrapper::var)
X(fogRGBA);
X(auxWidth);
X(auxHeight);
X(auxSamples);
X(auxFBO);
X(auxFBO_r);
X(auxColor);
X(auxColor_r);
X(textureType);
X(texture);
X(skinvertVBO);
X(skinnormalVBO);
X(skintexcoordVBO);
X(skinfaceVBO);
X(charWidth);
X(charWidthBig);
#undef X
using EigenUnsignedCharVectorX = Eigen::Vector<unsigned char, Eigen::Dynamic>;
using EigenFloatVectorX = Eigen::Vector<float, Eigen::Dynamic>;
// Skipped: mjr_defaultContext (have MjrContext.__init__)
// Skipped: mjr_makeContext (have MjrContext.__init__)
Def<traits::mjr_changeFont>(pymodule);
Def<traits::mjr_addAux>(pymodule);
// Skipped: mjr_freeContext (have MjrContext.__del__)
Def<traits::mjr_uploadTexture>(pymodule);
Def<traits::mjr_uploadMesh>(pymodule);
Def<traits::mjr_uploadHField>(pymodule);
Def<traits::mjr_restoreBuffer>(pymodule);
Def<traits::mjr_setBuffer>(pymodule);
Def<traits::mjr_readPixels>(
pymodule, [](std::optional<py::array_t<uint8_t>> rgb,
std::optional<py::array_t<float>> depth,
const raw::MjrRect* viewport, const raw::MjrContext* con) {
return InterceptMjErrors(::mjr_readPixels)(
rgb.has_value() ? rgb->mutable_data() : nullptr,
depth.has_value() ? depth->mutable_data() : nullptr, *viewport,
con);
});
Def<traits::mjr_drawPixels>(
pymodule,
[](std::optional<Eigen::Ref<const EigenUnsignedCharVectorX>> rgb,
std::optional<Eigen::Ref<const EigenFloatVectorX>> depth,
const raw::MjrRect* viewport, const raw::MjrContext* con) {
return InterceptMjErrors(::mjr_drawPixels)(
rgb.has_value() ? rgb->data() : nullptr,
depth.has_value() ? depth->data() : nullptr, *viewport, con);
});
Def<traits::mjr_blitBuffer>(pymodule);
Def<traits::mjr_setAux>(pymodule);
Def<traits::mjr_blitAux>(pymodule);
Def<traits::mjr_text>(pymodule);
Def<traits::mjr_overlay>(pymodule);
Def<traits::mjr_maxViewport>(pymodule);
Def<traits::mjr_rectangle>(pymodule);
Def<traits::mjr_label>(pymodule);
Def<traits::mjr_figure>(pymodule);
Def<traits::mjr_render>(pymodule);
Def<traits::mjr_finish>(pymodule);
Def<traits::mjr_getError>(pymodule);
Def<traits::mjr_findRect>(pymodule);
} // PYBIND11_MODULE
} // namespace
} // namespace mujoco::python
python/mujoco/render_test.py
+81
View File
@@ -0,0 +1,81 @@
# 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.
# ==============================================================================
"""Tests for MuJoCo Python rendering."""
from absl.testing import absltest
import mujoco
import numpy as np
@absltest.skipUnless(hasattr(mujoco, 'mjr_render'),
'MuJoCo rendering is disabled')
class MuJoCoRenderTest(absltest.TestCase):
def setUp(self):
super().setUp()
self.gl = mujoco.GLContext(640, 480)
self.gl.make_current()
def tearDown(self):
super().tearDown()
del self.gl
def test_can_render(self):
"""Test that the bindings can successfully render a simple image.
This test sets up a basic MuJoCo rendering context similar to the example in
https://mujoco.readthedocs.io/en/latest/programming.html#visualization
It calls `mjr_rectangle` rather than `mjr_render` so that we can assert an
exact rendered image without needing golden data. The purpose of this test
is to ensure that the bindings can correctly return pixels in Python, rather
than to test MuJoCo's rendering pipeline itself.
"""
self.model = mujoco.MjModel.from_xml_string('<mujoco><worldbody/></mujoco>')
self.data = mujoco.MjData(self.model)
scene = mujoco.MjvScene(self.model, maxgeom=0)
mujoco.mjv_updateScene(
self.model, self.data, mujoco.MjvOption(), mujoco.MjvPerturb(),
mujoco.MjvCamera(), mujoco.mjtCatBit.mjCAT_ALL.value, scene)
context = mujoco.MjrContext(
self.model,
mujoco.mjtFontScale.mjFONTSCALE_150.value)
mujoco.mjr_setBuffer(
mujoco.mjtFramebuffer.mjFB_OFFSCREEN.value, context)
# MuJoCo's default render buffer size is 640x480.
full_rect = mujoco.MjrRect(0, 0, 640, 480)
mujoco.mjr_rectangle(full_rect, 0, 0, 0, 1)
blue_rect = mujoco.MjrRect(56, 67, 234, 123)
mujoco.mjr_rectangle(blue_rect, 0, 0, 1, 1)
expected_upside_down_image = np.zeros((480, 640, 3), dtype=np.uint8)
expected_upside_down_image[67:67+123, 56:56+234, 2] = 255
upside_down_image = np.empty((480, 640, 3), dtype=np.uint8)
mujoco.mjr_readPixels(upside_down_image, None, full_rect, context)
np.testing.assert_array_equal(upside_down_image, expected_upside_down_image)
# Check that mjr_readPixels can accept a flattened array.
upside_down_image[:] = 0
mujoco.mjr_readPixels(
np.reshape(upside_down_image, -1), None, full_rect, context)
np.testing.assert_array_equal(upside_down_image, expected_upside_down_image)
if __name__ == '__main__':
absltest.main()
python/mujoco/rollout.cc
+263
View File
@@ -0,0 +1,263 @@
// 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.
#include <array>
#include <cstdio>
#include <iostream>
#include <optional>
#include <sstream>
#include <string>
#include "functions.h"
#include "raw.h"
#include <pybind11/buffer_info.h>
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
namespace mujoco::python {
namespace {
namespace py = ::pybind11;
const auto rollout_doc = R"(
Roll out open-loop trajectories from initial states, get subsequent states and sensor values.
input arguments (required):
model an instance of MjModel
data an associated instance of MjData
nstate an integer, number of initial states from which to roll out trajectories
nstep an integer, number of steps to be taken for each trajectory
input arguments (optional):
initial_state (nstate x nqva) nstate initial state vectors, nqva=nq+nv+na
initial_time (nstate x 1) nstate initial times
initial_warmstart (nstate x nv) nstate qacc_warmstart vectors
ctrl (nstate x nstep x nu) nstate length-nstep controls
qfrc_applied (nstate x nstep x nv) nstate length-nstep generalized forces
xfrc_applied (nstate x nstep x nbody*6) nstate length-nstep Cartesian wrenches
mocap (nstate x nstep x nmocap*7) nstate length-nstep mocap body poses
output arguments (optional):
state (nstate x nstep x nqva) nstate length-nstep states
sensordata (nstate x nstep x nsendordata) nstate length-nstep sensordatas
)";
// C-style rollout function, assumes all arguments are valid
// all input fields of d are initialised, contents at call time do not matter
// after returning, d will contain the last step of the last rollout
void _unsafe_rollout(const mjModel* m, mjData* d, int nstate, int nstep,
const mjtNum* state0, const mjtNum* ctrl,
const mjtNum* qfrc, const mjtNum* xfrc,
const mjtNum* mocap, const mjtNum* time0,
const mjtNum* warmstart0,
mjtNum* state, mjtNum* sensordata) {
// model sizes
int nq = m->nq;
int nv = m->nv;
int na = m->na;
int nqva = nq + nv + na;
int nu = m->nu;
int nbody = m->nbody;
int nmocap = m->nmocap;
int nsensordata = m->nsensordata;
// loop over initial states
for (int s=0; s < nstate; s++) {
// set initial state
if (state0) {
mju_copy(d->qpos, state0 + s*nqva, nq);
mju_copy(d->qvel, state0 + s*nqva + nq, nv);
mju_copy(d->act, state0 + s*nqva + nq + nv, na);
} else {
mju_copy(d->qpos, m->qpos0, nq);
mju_zero(d->qvel, nv);
mju_zero(d->act, na);
}
// set initial time
d->time = time0 ? time0[s] : 0;
// set warmstart accelerations
if (warmstart0) {
mju_copy(d->qacc_warmstart, warmstart0 + s*nv, nv);
} else {
mju_zero(d->qacc_warmstart, nv);
}
// clear control inputs if unspecified
if (s == 0) {
if (!ctrl) {
mju_zero(d->ctrl, nu);
}
if (!qfrc) {
mju_zero(d->qfrc_applied, nv);
}
if (!xfrc) {
mju_zero(d->xfrc_applied, 6*nbody);
}
if (!mocap) {
for (int j=0; j<nbody; j++) {
int id = m->body_mocapid[j];
if (id>=0) {
mju_copy3(d->mocap_pos+3*id, m->body_pos+3*j);
mju_copy4(d->mocap_quat+4*id, m->body_quat+4*j);
}
}
}
}
// roll out trajectories
for (int t = 0; t < nstep; t++) {
// controls
if (ctrl) {
mju_copy(d->ctrl, ctrl + s*nstep*nu + t*nu, nu);
}
// generalized forces
if (qfrc) {
mju_copy(d->qfrc_applied, qfrc + s*nstep*nv + t*nv, nv);
}
// Cartesian wrenches
if (xfrc) {
mju_copy(d->xfrc_applied, xfrc + s*nstep*6*nbody + t*6*nbody, 6*nbody);
}
// mocap bodies
if (mocap) {
mju_copy(d->mocap_pos,
mocap + s*nstep*7*nmocap + t*7*nmocap, 3*nmocap);
mju_copy(d->mocap_quat,
mocap + s*nstep*7*nmocap + t*7*nmocap + 3*nmocap, 4*nmocap);
}
// step
mj_step(m, d);
// copy out new state
if (state) {
mju_copy(state + s*nstep*nqva + t*nqva, d->qpos, nq);
mju_copy(state + s*nstep*nqva + t*nqva + nq, d->qvel, nv);
mju_copy(state + s*nstep*nqva + t*nqva + nq + nv, d->act, na);
}
// copy out sensor values
if (sensordata) {
mju_copy(sensordata + s*nstep*nsensordata + t*nsensordata,
d->sensordata, nsensordata);
}
}
}
}
// check size of optional argument to rollout(), return raw pointer
mjtNum* get_array_ptr(std::optional<const py::array_t<mjtNum>> arg,
const char* name, int nstate, int nstep, int dim) {
// if empty return nullptr
if (!arg.has_value()) {
return nullptr;
}
// get info
py::buffer_info info = arg->request();
// check size
int expected_size = nstate * nstep * dim;
if (info.size != expected_size) {
std::ostringstream msg;
msg << name << ".size should be " << expected_size << ", got " << info.size;
throw py::value_error(msg.str());
}
return static_cast<mjtNum*>(info.ptr);
}
PYBIND11_MODULE(_rollout, pymodule) {
namespace py = ::pybind11;
using PyCArray = py::array_t<mjtNum, py::array::c_style>;
// roll out open loop trajectories from multiple initial states
// get subsequent states and corresponding sensor values
pymodule.def(
"rollout",
[](const MjModelWrapper& m, MjDataWrapper& d, int nstate, int nstep,
std::optional<const PyCArray> init_state,
std::optional<const PyCArray> init_time,
std::optional<const PyCArray> init_warmstart,
std::optional<const PyCArray> ctrl,
std::optional<const PyCArray> qfrc,
std::optional<const PyCArray> xfrc,
std::optional<const PyCArray> mocap,
std::optional<const PyCArray> state,
std::optional<const PyCArray> sensordata
) {
const raw::MjModel* model = m.get();
raw::MjData* data = d.get();
// check that some steps need to be taken, return if not
if (nstate < 1 || nstep < 1) {
return;
}
// get raw pointers
int nqva = model->nq + model->nv + model->na;
mjtNum* init_state_ptr =
get_array_ptr(init_state, "initial_state", nstate, 1, nqva);
mjtNum* ctrl_ptr = get_array_ptr(ctrl, "ctrl", nstate, nstep, model->nu);
mjtNum* qfrc_ptr =
get_array_ptr(qfrc, "qfrc_applied", nstate, nstep, model->nv);
mjtNum* xfrc_ptr =
get_array_ptr(xfrc, "xfrc_applied", nstate, nstep, 6*model->nbody);
mjtNum* mocap_ptr =
get_array_ptr(mocap, "mocap", nstate, nstep, 7*model->nmocap);
mjtNum* init_time_ptr =
get_array_ptr(init_time, "init_time", nstate, 1, 1);
mjtNum* init_warmstart_ptr =
get_array_ptr(init_warmstart, "init_warmstart", nstate, 1, model->nv);
mjtNum* state_ptr = get_array_ptr(state, "state", nstate, nstep, nqva);
mjtNum* sensordata_ptr =
get_array_ptr(sensordata, "sensordata", nstate, nstep, model->nsensordata);
// perform rollouts
{
// release the GIL
py::gil_scoped_release no_gil;
// call unsafe rollout function
InterceptMjErrors(_unsafe_rollout)(
model, data, nstate, nstep, init_state_ptr, ctrl_ptr, qfrc_ptr,
xfrc_ptr, mocap_ptr, init_time_ptr, init_warmstart_ptr, state_ptr,
sensordata_ptr);
}
},
py::arg("model"),
py::arg("data"),
py::arg("nstate"),
py::arg("nstep"),
py::arg("initial_state") = py::none(),
py::arg("initial_time") = py::none(),
py::arg("initial_warmstart") = py::none(),
py::arg("ctrl") = py::none(),
py::arg("qfrc_applied") = py::none(),
py::arg("xfrc_applied") = py::none(),
py::arg("mocap") = py::none(),
py::arg("state") = py::none(),
py::arg("sensordata") = py::none(),
py::doc(rollout_doc)
);
} // namespace
}
}
python/mujoco/rollout.py
+205
View File
@@ -0,0 +1,205 @@
# 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.
# ==============================================================================
"""Roll out open-loop trajectories from initial states, get subsequent states and sensor values."""
from mujoco import _rollout
import numpy as np
def rollout(model, data, initial_state=None, ctrl=None,
*, # require following arguments to be named
skip_checks=False,
nstate=None,
nstep=None,
initial_time=None,
initial_warmstart=None,
qfrc_applied=None,
xfrc_applied=None,
mocap=None,
state=None,
sensordata=None):
"""Roll out open-loop trajectories from initial states, get subsequent states and sensor values.
This function serves as a Python wrapper for the C++ functionality in
`rollout.cc`, please see documentation therein. This python funtion will
infer `nstate` and `nstep`, tile input arguments with singleton dimensions,
and allocate output arguments if none are given.
"""
# don't infer nstate/nstep, don't support singleton expansion, don't allocate
# output arrays, just call rollout
if skip_checks:
_rollout.rollout(model, data, nstate, nstep, initial_state, initial_time,
initial_warmstart, ctrl, qfrc_applied, xfrc_applied, mocap,
state, sensordata)
return state, sensordata
# check types
if nstate and not isinstance(nstate, int):
raise ValueError('nstate must be an integer')
if nstep and not isinstance(nstep, int):
raise ValueError('nstep must be an integer')
_check_must_be_numeric(
initial_state=initial_state,
initial_time=initial_time,
initial_warmstart=initial_warmstart,
ctrl=ctrl,
qfrc_applied=qfrc_applied,
xfrc_applied=xfrc_applied,
mocap=mocap,
state=state,
sensordata=sensordata)
# check number of dimensions
_check_number_of_dimensions(2,
initial_state=initial_state,
initial_time=initial_time,
initial_warmstart=initial_warmstart)
_check_number_of_dimensions(3,
ctrl=ctrl,
qfrc_applied=qfrc_applied,
xfrc_applied=xfrc_applied,
mocap=mocap,
state=state,
sensordata=sensordata)
# ensure 2D, make contiguous, row-major (C ordering)
initial_state = _ensure_2d(initial_state)
initial_time = _ensure_2d(initial_time)
initial_warmstart = _ensure_2d(initial_warmstart)
# ensure 3D, make contiguous, row-major (C ordering)
ctrl = _ensure_3d(ctrl)
qfrc_applied = _ensure_3d(qfrc_applied)
xfrc_applied = _ensure_3d(xfrc_applied)
mocap = _ensure_3d(mocap)
state = _ensure_3d(state)
sensordata = _ensure_3d(sensordata)
# check trailing dimensions
_check_trailing_dimension(model.nq + model.nv + model.na,
initial_state=initial_state, state=state)
_check_trailing_dimension(1, initial_time=initial_time)
_check_trailing_dimension(model.nu, ctrl=ctrl)
_check_trailing_dimension(model.nv, qfrc_applied=qfrc_applied)
_check_trailing_dimension(model.nbody*6, xfrc_applied=xfrc_applied)
_check_trailing_dimension(model.nmocap*7, mocap=mocap)
_check_trailing_dimension(model.nsensordata, sensordata=sensordata)
# infer nstate, check for incompatibilities
nstate = _infer_dimension(0, nstate or 1,
initial_state=initial_state,
initial_time=initial_time,
initial_warmstart=initial_warmstart,
ctrl=ctrl,
qfrc_applied=qfrc_applied,
xfrc_applied=xfrc_applied,
mocap=mocap,
state=state,
sensordata=sensordata)
# infer nstep, check for incompatibilities
nstep = _infer_dimension(1, nstep or 1,
ctrl=ctrl,
qfrc_applied=qfrc_applied,
xfrc_applied=xfrc_applied,
mocap=mocap,
state=state,
sensordata=sensordata)
# tile input arrays if required (singleton expansion)
initial_state = _tile_if_required(initial_state, nstate)
initial_time = _tile_if_required(initial_time, nstate)
initial_warmstart = _tile_if_required(initial_warmstart, nstate)
ctrl = _tile_if_required(ctrl, nstate, nstep)
qfrc_applied = _tile_if_required(qfrc_applied, nstate, nstep)
xfrc_applied = _tile_if_required(xfrc_applied, nstate, nstep)
mocap = _tile_if_required(mocap, nstate, nstep)
# allocate output if not provided
if state is None:
state = np.empty((nstate, nstep, model.nq + model.nv + model.na))
if sensordata is None:
sensordata = np.empty((nstate, nstep, model.nsensordata))
# call rollout
_rollout.rollout(model, data, nstate, nstep, initial_state, initial_time,
initial_warmstart, ctrl, qfrc_applied, xfrc_applied, mocap,
state, sensordata)
# return squeezed outputs
return state.squeeze(), sensordata.squeeze()
def _check_must_be_numeric(**kwargs):
for key, value in kwargs.items():
if value is None:
continue
if not isinstance(value, np.ndarray) and not isinstance(value, float):
raise ValueError(f'{key} must be a numpy array or float')
def _check_number_of_dimensions(ndim, **kwargs):
for key, value in kwargs.items():
if value is None:
continue
if value.ndim > ndim:
raise ValueError(f'{key} can have at most {ndim} dimensions')
def _check_trailing_dimension(dim, **kwargs):
for key, value in kwargs.items():
if value is None:
continue
if value.shape[-1] != dim:
raise ValueError(f'trailing dimension of {key} must be {dim}, got {value.shape[-1]}')
def _ensure_2d(arg):
if arg is None:
return None
else:
return np.ascontiguousarray(np.atleast_2d(arg), dtype=np.float64)
def _ensure_3d(arg):
if arg is None:
return None
else:
# np.atleast_3d adds both leading and trailing dims, we want only leading
if arg.ndim == 0:
arg = arg[np.newaxis, np.newaxis, np.newaxis, ...]
elif arg.ndim == 1:
arg = arg[np.newaxis, np.newaxis, ...]
elif arg.ndim == 2:
arg = arg[np.newaxis, ...]
return np.ascontiguousarray(arg, dtype=np.float64)
def _infer_dimension(dim, value, **kwargs):
for name, array in kwargs.items():
if array is None:
continue
if array.shape[dim] != value:
if value == 1:
value = array.shape[dim]
elif array.shape[dim] != 1:
raise ValueError(
f'dimension {dim} inferred as {value} but {name} has {array.shape[dim]}'
)
return value
def _tile_if_required(array, dim0, dim1=None):
if array is None:
return
reps = np.ones(array.ndim, dtype=int)
if array.shape[0] == 1:
reps[0] = dim0
if dim1 is not None and array.shape[1] == 1:
reps[1] = dim1
return np.tile(array, reps)
python/mujoco/rollout_test.py
+488
View File
@@ -0,0 +1,488 @@
# 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.
# ==============================================================================
"""tests for rollout function."""
from absl.testing import absltest
from absl.testing import parameterized
import mujoco
import numpy as np
import concurrent.futures
import threading
from mujoco import rollout
#--------------------------- models used for testing ---------------------------
TEST_XML = r"""
<mujoco>
<worldbody>
<light pos="0 0 2"/>
<geom type="plane" size="5 5 .1"/>
<body pos="0 0 .1">
<joint name="yaw" axis="0 0 1"/>
<joint name="pitch" axis="0 1 0"/>
<geom type="capsule" size=".02" fromto="0 0 0 1 0 0"/>
<geom type="box" pos="1 0 0" size=".1 .1 .1"/>
<site name="site" pos="1 0 0"/>
</body>
</worldbody>
<actuator>
<general joint="pitch" gainprm="100"/>
<general joint="yaw" dyntype="filter" dynprm="1" gainprm="100"/>
</actuator>
<sensor>
<accelerometer site="site"/>
</sensor>
</mujoco>
"""
TEST_XML_NO_SENSORS = r"""
<mujoco>
<worldbody>
<light pos="0 0 2"/>
<geom type="plane" size="5 5 .1"/>
<body pos="0 0 .1">
<joint name="yaw" axis="0 0 1"/>
<joint name="pitch" axis="0 1 0"/>
<geom type="capsule" size=".02" fromto="0 0 0 1 0 0"/>
<geom type="box" pos="1 0 0" size=".1 .1 .1"/>
<site name="site" pos="1 0 0"/>
</body>
</worldbody>
<actuator>
<general joint="pitch" gainprm="100"/>
<general joint="yaw" dyntype="filter" dynprm="1" gainprm="100"/>
</actuator>
</mujoco>
"""
TEST_XML_NO_ACTUATORS = r"""
<mujoco>
<worldbody>
<light pos="0 0 2"/>
<geom type="plane" size="5 5 .1"/>
<body pos="0 0 .1">
<joint name="yaw" axis="0 0 1"/>
<joint name="pitch" axis="0 1 0"/>
<geom type="capsule" size=".02" fromto="0 0 0 1 0 0"/>
<geom type="box" pos="1 0 0" size=".1 .1 .1"/>
<site name="site" pos="1 0 0"/>
</body>
</worldbody>
<sensor>
<accelerometer site="site"/>
</sensor>
</mujoco>
"""
TEST_XML_MOCAP = r"""
<mujoco>
<worldbody>
<body name="1" mocap="true">
</body>
<body name="2" mocap="true">
</body>
</worldbody>
<sensor>
<framepos objtype="xbody" objname="1"/>
<framequat objtype="xbody" objname="2"/>
</sensor>
</mujoco>
"""
TEST_XML_EMPTY = r"""
<mujoco>
</mujoco>
"""
ALL_MODELS = {'TEST_XML': TEST_XML,
'TEST_XML_NO_SENSORS': TEST_XML_NO_SENSORS,
'TEST_XML_NO_ACTUATORS': TEST_XML_NO_ACTUATORS,
'TEST_XML_EMPTY': TEST_XML_EMPTY}
#------------------------------- tests -----------------------------------------
class MuJoCoRolloutTest(parameterized.TestCase):
def setUp(self):
super().setUp()
np.random.seed(42)
#----------------------------- test basic operation
@parameterized.parameters(ALL_MODELS.keys())
def test_single_step(self, model_name):
model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
data = mujoco.MjData(model)
initial_state = np.random.randn(model.nq + model.nv + model.na)
ctrl = np.random.randn(model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
mujoco.mj_resetData(model, data)
py_state, py_sensordata = step(model, data, initial_state, ctrl=ctrl)
np.testing.assert_array_equal(state, py_state)
np.testing.assert_array_equal(sensordata, py_sensordata)
@parameterized.parameters(ALL_MODELS.keys())
def test_single_rollout(self, model_name):
nstep = 3 # number of timesteps
model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
data = mujoco.MjData(model)
initial_state = np.random.randn(model.nq + model.nv + model.na)
ctrl = np.random.randn(nstep, model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
py_state, py_sensordata = single_rollout(model, data, initial_state,
ctrl=ctrl)
np.testing.assert_array_equal(state, np.asarray(py_state))
np.testing.assert_array_equal(sensordata, np.asarray(py_sensordata))
@parameterized.parameters(ALL_MODELS.keys())
def test_multi_step(self, model_name):
model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
data = mujoco.MjData(model)
nstate = 5 # number of initial states
initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
ctrl = np.random.randn(nstate, 1, model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
mujoco.mj_resetData(model, data)
py_state, py_sensordata = multi_rollout(model, data, initial_state,
ctrl=ctrl)
np.testing.assert_array_equal(state, py_state)
np.testing.assert_array_equal(sensordata, py_sensordata)
@parameterized.parameters(ALL_MODELS.keys())
def test_single_rollout_fixed_ctrl(self, model_name):
nstep = 3
model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
data = mujoco.MjData(model)
initial_state = np.random.randn(model.nq + model.nv + model.na)
ctrl = np.random.randn(model.nu)
state = np.empty((nstep, model.nq + model.nv + model.na))
sensordata = np.empty((nstep, model.nsensordata))
rollout.rollout(model, data, initial_state, ctrl,
state=state, sensordata=sensordata)
ctrl = np.tile(ctrl, (nstep, 1)) # repeat??
py_state, py_sensordata = single_rollout(model, data, initial_state,
ctrl=ctrl)
np.testing.assert_array_equal(state, py_state)
np.testing.assert_array_equal(sensordata, py_sensordata)
@parameterized.parameters(ALL_MODELS.keys())
def test_multi_rollout(self, model_name):
model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
data = mujoco.MjData(model)
nstate = 2 # number of initial states
nstep = 3 # number of timesteps
initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
ctrl = np.random.randn(nstate, nstep, model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
py_state, py_sensordata = multi_rollout(model, data, initial_state,
ctrl=ctrl)
np.testing.assert_array_equal(py_state, py_state)
np.testing.assert_array_equal(py_sensordata, py_sensordata)
@parameterized.parameters(ALL_MODELS.keys())
def test_multi_rollout_fixed_ctrl_infer_from_output(self, model_name):
model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
data = mujoco.MjData(model)
nstate = 2 # number of initial states
nstep = 3 # number of timesteps
initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
ctrl = np.random.randn(nstate, 1, model.nu) # 1 control in the time dimension
state = np.empty((nstate, nstep, model.nq + model.nv + model.na))
state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
state=state)
ctrl = np.repeat(ctrl, nstep, axis=1)
py_state, py_sensordata = multi_rollout(model, data, initial_state,
ctrl=ctrl)
np.testing.assert_array_equal(state, py_state)
np.testing.assert_array_equal(sensordata, py_sensordata)
@parameterized.product(arg_nstep=[[3, 1, 1], [3, 3, 1], [3, 1, 3]],
model_name=list(ALL_MODELS.keys()))
def test_multi_rollout_multiple_inputs(self, arg_nstep, model_name):
model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
data = mujoco.MjData(model)
nstate = 4 # number of initial states
initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
# arg_nstep is the horizon for {ctrl, qfrc_applied, xfrc_applied}, respectively
ctrl = np.random.randn(nstate, arg_nstep[0], model.nu)
qfrc_applied = np.random.randn(nstate, arg_nstep[1], model.nv)
xfrc_applied = np.random.randn(nstate, arg_nstep[2], model.nbody*6)
state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
qfrc_applied=qfrc_applied,
xfrc_applied=xfrc_applied)
# tile singleton arguments
nstep = max(arg_nstep)
if arg_nstep[0] == 1:
ctrl = np.repeat(ctrl, nstep, axis=1)
if arg_nstep[1] == 1:
qfrc_applied = np.repeat(qfrc_applied, nstep, axis=1)
if arg_nstep[2] == 1:
xfrc_applied = np.repeat(xfrc_applied, nstep, axis=1)
py_state, py_sensordata = multi_rollout(model, data, initial_state,
ctrl=ctrl,
qfrc_applied=qfrc_applied,
xfrc_applied=xfrc_applied)
np.testing.assert_array_equal(state, py_state)
np.testing.assert_array_equal(sensordata, py_sensordata)
#----------------------------- test threaded operation
def test_threading(self):
model = mujoco.MjModel.from_xml_string(TEST_XML)
num_workers = 32
nstate = 10000
nstep = 5
initial_state = np.random.randn(nstate, model.nq+model.nv+model.na)
state = np.zeros((nstate, nstep, model.nq+model.nv+model.na))
sensordata = np.zeros((nstate, nstep, model.nsensordata))
ctrl = np.random.randn(nstate, nstep, model.nu)
thread_local = threading.local()
def thread_initializer():
thread_local.data = mujoco.MjData(model)
def call_rollout(initial_state, ctrl, state):
rollout.rollout(model, thread_local.data, skip_checks=True,
nstate=initial_state.shape[0], nstep=nstep,
initial_state=initial_state, ctrl=ctrl, state=state)
n = initial_state.shape[0] // num_workers # integer division
chunks = [] # a list of tuples, one per worker
for i in range(num_workers-1):
chunks.append(
(initial_state[i*n:(i+1)*n], ctrl[i*n:(i+1)*n], state[i*n:(i+1)*n]))
# last chunk, absorbing the remainder:
chunks.append(
(initial_state[(num_workers-1)*n:], ctrl[(num_workers-1)*n:],
state[(num_workers-1)*n:]))
with concurrent.futures.ThreadPoolExecutor(
max_workers=num_workers, initializer=thread_initializer) as executor:
futures = []
for chunk in chunks:
futures.append(executor.submit(call_rollout, *chunk))
for future in concurrent.futures.as_completed(futures):
future.result()
data = mujoco.MjData(model)
py_state, py_sensordata = multi_rollout(model, data, initial_state,
ctrl=ctrl)
np.testing.assert_array_equal(state, py_state)
#----------------------------- test advanced operation
def test_time(self):
model = mujoco.MjModel.from_xml_string(TEST_XML)
data = mujoco.MjData(model)
nstate = 1
nstep = 3
initial_time = np.array([[2.]])
initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
ctrl = np.random.randn(nstate, nstep, model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
initial_time=initial_time)
self.assertAlmostEqual(data.time, 2 + nstep*model.opt.timestep)
def test_warmstart(self):
model = mujoco.MjModel.from_xml_string(TEST_XML)
data = mujoco.MjData(model)
state0 = np.zeros(model.nq + model.nv + model.na)
ctrl = np.zeros(model.nu)
state1, _ = step(model, data, state0, ctrl=ctrl)
initial_warmstart = data.qacc_warmstart.copy()
state2, _ = step(model, data, state1, ctrl=ctrl)
state, _ = rollout.rollout(model, data, state1, ctrl)
assert np.linalg.norm(state-state2) > 0
state, _ = rollout.rollout(model, data, state1, ctrl,
initial_warmstart=initial_warmstart)
np.testing.assert_array_equal(state, state2)
def test_mocap(self):
model = mujoco.MjModel.from_xml_string(TEST_XML_MOCAP)
data = mujoco.MjData(model)
initial_state = np.zeros(model.nq + model.nv + model.na)
pos1 = np.array((1., 2., 3.))
quat1 = np.array((1., 2., 3., 4.))
quat1 /= np.linalg.norm(quat1)
pos2 = np.array((2., 3., 4.))
quat2 = np.array((2., 3., 4., 5.))
quat2 /= np.linalg.norm(quat2)
mocap = np.hstack((pos1, quat1, pos2, quat2))
state, sensordata = rollout.rollout(model, data, initial_state, mocap=mocap)
np.testing.assert_array_almost_equal(sensordata[:3], pos1)
np.testing.assert_array_almost_equal(sensordata[3:], quat2)
#----------------------------- test correctness
def test_intercept_mj_errors(self):
model = mujoco.MjModel.from_xml_string(TEST_XML)
data = mujoco.MjData(model)
initial_state = np.zeros(model.nq + model.nv + model.na)
ctrl = np.zeros((3, model.nu))
model.opt.solver = 10 # invalid solver type
with self.assertRaisesWithLiteralMatch(mujoco.FatalError,
'Unknown solver type 10'):
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
def test_invalid(self):
model = mujoco.MjModel.from_xml_string(TEST_XML)
data = mujoco.MjData(model)
initial_state = np.zeros(model.nq + model.nv + model.na)
ctrl = 'string'
with self.assertRaisesWithLiteralMatch(
ValueError, 'ctrl must be a numpy array or float'):
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
qfrc_applied = np.zeros((2, 3, 4, 5))
with self.assertRaisesWithLiteralMatch(
ValueError, 'qfrc_applied can have at most 3 dimensions'):
state, sensordata = rollout.rollout(model, data, initial_state,
qfrc_applied=qfrc_applied)
def test_bad_sizes(self):
model = mujoco.MjModel.from_xml_string(TEST_XML)
data = mujoco.MjData(model)
initial_state = np.random.randn(model.nq + model.nv + model.na+1)
with self.assertRaisesWithLiteralMatch(
ValueError, 'trailing dimension of initial_state must be 5, got 6'):
state, sensordata = rollout.rollout(model, data, initial_state)
initial_state = np.random.randn(model.nq + model.nv + model.na)
ctrl = np.random.randn(model.nu+1)
with self.assertRaisesWithLiteralMatch(
ValueError, 'trailing dimension of ctrl must be 2, got 3'):
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
ctrl = np.random.randn(2, model.nu)
qfrc_applied = np.random.randn(3, model.nv) # incompatible horizon
with self.assertRaisesWithLiteralMatch(
ValueError, 'dimension 1 inferred as 2 but qfrc_applied has 3'):
state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
qfrc_applied=qfrc_applied)
def test_stateless(self):
model = mujoco.MjModel.from_xml_string(TEST_XML)
model.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_WARMSTART.value
data = mujoco.MjData(model)
# call step with a clean mjData
initial_state = np.random.randn(model.nq + model.nv + model.na)
ctrl = np.random.randn(model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
# fill mjData with some debug value, see that we still get the same outputs
mujoco.mj_resetDataDebug(model, data, 255)
debug_state, debug_sensordata = rollout.rollout(model, data, initial_state,
ctrl)
np.testing.assert_array_equal(state, debug_state)
np.testing.assert_array_equal(sensordata, debug_sensordata)
#--------------- Python implementation of rollout functionality ----------------
def get_state(data):
return np.hstack((data.qpos, data.qvel, data.act))
def set_state(model, data, state):
data.qpos = state[:model.nq]
data.qvel = state[model.nq:model.nq+model.nv]
data.act = state[model.nq+model.nv:model.nq+model.nv+model.na]
def step(model, data, state, **kwargs):
if state is not None:
set_state(model, data, state)
for key, value in kwargs.items():
if value is not None:
setattr(data, key, np.reshape(value, getattr(data, key).shape))
mujoco.mj_step(model, data)
return (get_state(data), data.sensordata)
def single_rollout(model, data, initial_state, **kwargs):
arg_nstep = set([a.shape[0] for a in kwargs.values()])
assert len(arg_nstep) == 1 # nstep dimensions must match
nstep = arg_nstep.pop()
state = np.empty((nstep, model.nq + model.nv + model.na))
sensordata = np.empty((nstep, model.nsensordata))
mujoco.mj_resetData(model, data)
for t in range(nstep):
kwargs_t = {}
for key, value in kwargs.items():
kwargs_t[key] = value[0 if value.ndim == 1 else t]
state[t], sensordata[t] = step(model, data,
initial_state if t==0 else None,
**kwargs_t)
return state, sensordata
def multi_rollout(model, data, initial_state, **kwargs):
nstate = initial_state.shape[0]
arg_nstep = set([a.shape[1] for a in kwargs.values()])
assert len(arg_nstep) == 1 # nstep dimensions must match
nstep = arg_nstep.pop()
state = np.empty((nstate, nstep, model.nq + model.nv + model.na))
sensordata = np.empty((nstate, nstep, model.nsensordata))
for s in range(nstate):
kwargs_s = {key : value[s] for key, value in kwargs.items()}
state_s, sensordata_s = single_rollout(model, data, initial_state[s],
**kwargs_s)
state[s] = state_s
sensordata[s] = sensordata_s
return state.squeeze(), sensordata.squeeze()
if __name__ == '__main__':
absltest.main()
python/mujoco/serialization.h
+72
View File
@@ -0,0 +1,72 @@
// 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.
#ifndef MUJOCO_PYTHON_SERIALIZATION_H_
#define MUJOCO_PYTHON_SERIALIZATION_H_
#include <iostream>
#include <mjtnum.h>
namespace mujoco::python::_impl {
// For now, the serialization format is architecture- and version-dependent:
// It assumes that the writer and the reader have the same endianess, same
// numeric value sizes and same struct definitions.
// It is not safe to load serialized data from another version of MuJoCo or from
// a machine with a different architecture.
static_assert(sizeof(char) == 1);
static_assert(sizeof(int) == 4);
static_assert(sizeof(mjtNum) == 8);
inline void WriteChar(std::ostream& output, char c) {
output.write(&c, 1);
}
inline char ReadChar(std::istream& input) {
char c = '\0';
input.read(&c, 1);
return c;
}
inline void WriteInt(std::ostream& output, int i) {
output.write(reinterpret_cast<char*>(&i), sizeof(int));
}
inline int ReadInt(std::istream& input) {
int i = 0;
input.read(reinterpret_cast<char*>(&i), sizeof(int));
return i;
}
inline void WriteBytes(std::ostream& output, const void* src, size_t nbytes) {
// Start by writing nbytes itself, so it can be validated at the time of
// reading.
WriteInt(output, nbytes);
output.write(reinterpret_cast<const char*>(src), nbytes);
}
inline void ReadBytes(std::istream& input, void* dest, size_t nbytes) {
size_t actual_nbytes = ReadInt(input);
if (actual_nbytes != nbytes) {
input.setstate(input.rdstate() | std::ios_base::failbit);
return;
}
input.read(reinterpret_cast<char*>(dest), nbytes);
}
} // namespace mujoco::python::_impl
#endif // MUJOCO_PYTHON_SERIALIZATION_H_
python/mujoco/structs.cc
+2221
View File
File diff suppressed because it is too large Load Diff
python/mujoco/structs.h
+978
View File
@@ -0,0 +1,978 @@
// 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.
#ifndef MUJOCO_PYTHON_STRUCTS_H_
#define MUJOCO_PYTHON_STRUCTS_H_
#include <array>
#include <istream>
#include <memory>
#include <optional>
#include <ostream>
#include <unordered_map>
#include <string>
#include <vector>
#include <absl/types/span.h>
#include <mujoco.h>
#include <mjxmacro.h>
#include "indexers.h"
#include "mjdata_meta.h"
#include "raw.h"
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
namespace mujoco::python {
namespace _impl {
template <typename T>
class WrapperBase {
public:
WrapperBase(WrapperBase&& other) = default;
T* get() { return ptr_; }
const T* get() const { return ptr_; }
const pybind11::handle owner() const { return owner_; }
protected:
static void DefaultCapsuleDestructor(PyObject* pyobj) {
T* ptr = static_cast<T*>(PyCapsule_GetPointer(pyobj, nullptr));
if (ptr) {
delete ptr;
}
}
// Takes ownership of ptr.
explicit WrapperBase(T* ptr,
void (*destructor)(PyObject*) = DefaultCapsuleDestructor)
: ptr_(ptr),
owner_(pybind11::capsule(ptr_, /* name = */ nullptr, destructor)) {}
// `ptr` is owned by `owner`.
WrapperBase(T* ptr, pybind11::handle owner)
: ptr_(ptr),
owner_(pybind11::reinterpret_borrow<pybind11::object>(owner)) {}
T* ptr_;
pybind11::object owner_;
};
template <typename T, typename = void>
struct py_array_or_tuple {
using type = pybind11::tuple;
};
template <typename T>
struct py_array_or_tuple<T, std::enable_if_t<std::is_arithmetic_v<T>>> {
using type = pybind11::array_t<T>;
};
// A type that resolves to a NumPy array if the dtype is numeric, and
// a Python tuple otherwise.
template <typename T> using py_array_or_tuple_t =
typename py_array_or_tuple<T>::type;
template <typename T>
struct enable_if_mj_struct {};
template <typename T>
class MjWrapper {};
template <typename T>
class StructListBase {
public:
StructListBase(T* ptr, int num, pybind11::handle owner) : ptr_(ptr) {
for (int i = 0; i < num; ++i) {
wrappers_.push_back(std::make_shared<MjWrapper<T>>(&ptr[i], owner));
}
}
StructListBase(const StructListBase& other) = delete;
StructListBase(StructListBase&& other) = default;
MjWrapper<T>& operator[](int i) {
if (i < 0 || i >= wrappers_.size()) {
throw pybind11::index_error();
}
return *wrappers_[i];
}
int size() const {
return wrappers_.size();
}
protected:
// Slicing
StructListBase(StructListBase& other, pybind11::slice slice) {
pybind11::size_t start, stop, step, slicelength;
slice.compute(other.size(), &start, &stop, &step, &slicelength);
ptr_ = &other.ptr_[start];
for (int i = start; i < stop; i += step) {
wrappers_.push_back(other.wrappers_[i]);
}
}
T* ptr_;
// Using shared_ptr here so that we get identical Python objects when slicing.
std::vector<std::shared_ptr<MjWrapper<T>>> wrappers_;
};
template <typename T>
struct is_mj_struct_list { static constexpr bool value = false; };
template <typename T>
class MjStructList {};
// ==================== MJOPTION ===============================================
template <>
class MjWrapper<raw::MjOption> : public WrapperBase<raw::MjOption> {
public:
MjWrapper();
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&) = default;
MjWrapper(raw::MjOption* ptr, pybind11::handle owner);
~MjWrapper() = default;
#define X(var, dim) py_array_or_tuple_t<mjtNum> var;
MJOPTION_VECTORS
#undef X
};
using MjOptionWrapper = MjWrapper<raw::MjOption>;
template <>
struct enable_if_mj_struct<raw::MjOption> { using type = void; };
// ==================== MJVISUAL ===============================================
template <>
class MjWrapper<raw::MjVisualHeadlight>
: public WrapperBase<raw::MjVisualHeadlight> {
public:
MjWrapper();
MjWrapper(const MjWrapper& other);
MjWrapper(MjWrapper&&) = default;
MjWrapper(raw::MjVisualHeadlight* ptr, pybind11::handle owner);
~MjWrapper() = default;
#define X(var) py_array_or_tuple_t<float> var
X(ambient);
X(diffuse);
X(specular);
#undef X
};
using MjVisualHeadlightWrapper = MjWrapper<raw::MjVisualHeadlight>;
template <>
struct enable_if_mj_struct<raw::MjVisualHeadlight> { using type = void; };
template <>
class MjWrapper<raw::MjVisualRgba> : public WrapperBase<raw::MjVisualRgba> {
public:
MjWrapper();
MjWrapper(const MjWrapper& other);
MjWrapper(MjWrapper&&) = default;
MjWrapper(raw::MjVisualRgba* ptr, pybind11::handle owner);
~MjWrapper() = default;
#define X(var) py_array_or_tuple_t<float> var
X(fog);
X(haze);
X(force);
X(inertia);
X(joint);
X(actuator);
X(actuatornegative);
X(actuatorpositive);
X(com);
X(camera);
X(light);
X(selectpoint);
X(connect);
X(contactpoint);
X(contactforce);
X(contactfriction);
X(contacttorque);
X(contactgap);
X(rangefinder);
X(constraint);
X(slidercrank);
X(crankbroken);
#undef X
};
using MjVisualRgbaWrapper = MjWrapper<raw::MjVisualRgba>;
template <>
struct enable_if_mj_struct<raw::MjVisualRgba> { using type = void; };
template <>
class MjWrapper<raw::MjVisual> : public WrapperBase<raw::MjVisual> {
public:
MjWrapper();
MjWrapper(const MjWrapper& other);
MjWrapper(MjWrapper&&) = default;
MjWrapper(raw::MjVisual* ptr, pybind11::handle owner);
~MjWrapper() = default;
MjVisualHeadlightWrapper headlight;
MjVisualRgbaWrapper rgba;
};
using MjVisualWrapper = MjWrapper<raw::MjVisual>;
template <>
struct enable_if_mj_struct<raw::MjVisual> { using type = void; };
// ==================== MJSTATISTIC ============================================
template <>
class MjWrapper<raw::MjStatistic> : public WrapperBase<raw::MjStatistic> {
public:
MjWrapper();
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&) = default;
MjWrapper(raw::MjStatistic* ptr, pybind11::handle owner);
~MjWrapper() = default;
#define X(var) py_array_or_tuple_t<mjtNum> var
X(center);
#undef X
};
using MjStatisticWrapper = MjWrapper<raw::MjStatistic>;
template <>
struct enable_if_mj_struct<raw::MjStatistic> { using type = void; };
// ==================== MJWARNINGSTAT ==========================================
template <>
class MjWrapper<raw::MjWarningStat> : public WrapperBase<raw::MjWarningStat> {
public:
MjWrapper();
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&) = default;
MjWrapper(raw::MjWarningStat* ptr, pybind11::handle owner);
~MjWrapper() = default;
};
using MjWarningStatWrapper = MjWrapper<raw::MjWarningStat>;
template <>
struct enable_if_mj_struct<raw::MjWarningStat> { using type = void; };
template <>
class MjStructList<raw::MjWarningStat>
: public StructListBase<raw::MjWarningStat> {
public:
MjStructList(raw::MjWarningStat* ptr, int num, pybind11::handle owner);
using StructListBase::operator[];
using StructListBase::size;
MjStructList Slice(pybind11::slice slice) {
return MjStructList(*this, slice);
}
#define X(type, var) pybind11::array_t<type> var
X(int, lastinfo);
X(int, number);
#undef X
protected:
MjStructList(MjStructList& other, pybind11::slice slice);
};
using MjWarningStatList = MjStructList<raw::MjWarningStat>;
template <>
struct py_array_or_tuple<raw::MjWarningStat> {
using type = MjWarningStatList;
};
template <>
struct is_mj_struct_list<raw::MjWarningStat> {
static constexpr bool value = true;
};
// ==================== MJTIMERSTAT ============================================
template <>
class MjWrapper<raw::MjTimerStat> : public WrapperBase<raw::MjTimerStat> {
public:
MjWrapper();
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&) = default;
MjWrapper(raw::MjTimerStat* ptr, pybind11::handle owner);
~MjWrapper() = default;
};
using MjTimerStatWrapper = MjWrapper<raw::MjTimerStat>;
template <>
struct enable_if_mj_struct<raw::MjTimerStat> { using type = void; };
template <>
class MjStructList<raw::MjTimerStat> : public StructListBase<raw::MjTimerStat> {
public:
MjStructList(raw::MjTimerStat* ptr, int num, pybind11::handle owner);
using StructListBase::operator[];
using StructListBase::size;
MjStructList Slice(pybind11::slice slice) {
return MjStructList(*this, slice);
}
#define X(type, var) pybind11::array_t<type> var
X(mjtNum, duration);
X(int, number);
#undef X
protected:
MjStructList(MjStructList& other, pybind11::slice slice);
};
using MjTimerStatList = MjStructList<raw::MjTimerStat>;
template <>
struct py_array_or_tuple<raw::MjTimerStat> {
using type = MjTimerStatList;
};
template <>
struct is_mj_struct_list<raw::MjTimerStat> {
static constexpr bool value = true;
};
// ==================== MJSOLVERSTAT ===========================================
template <>
class MjWrapper<raw::MjSolverStat> : public WrapperBase<raw::MjSolverStat> {
public:
MjWrapper();
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&) = default;
MjWrapper(raw::MjSolverStat* ptr, pybind11::handle owner);
~MjWrapper() = default;
};
using MjSolverStatWrapper = MjWrapper<raw::MjSolverStat>;
template <>
struct enable_if_mj_struct<raw::MjSolverStat> { using type = void; };
template <>
class MjStructList<raw::MjSolverStat>
: public StructListBase<raw::MjSolverStat> {
public:
MjStructList(raw::MjSolverStat* ptr, int num, pybind11::handle owner);
using StructListBase::operator[];
using StructListBase::size;
MjStructList Slice(pybind11::slice slice) {
return MjStructList(*this, slice);
}
#define X(type, var) pybind11::array_t<type> var
X(mjtNum, improvement);
X(mjtNum, gradient);
X(mjtNum, lineslope);
X(int, nactive);
X(int, nchange);
X(int, neval);
X(int, nupdate);
#undef X
protected:
MjStructList(MjStructList& other, pybind11::slice slice);
};
using MjSolverStatList = MjStructList<raw::MjSolverStat>;
template <>
struct py_array_or_tuple<raw::MjSolverStat> {
using type = MjSolverStatList;
};
template <>
struct is_mj_struct_list<raw::MjSolverStat> {
static constexpr bool value = true;
};
// ==================== MJMODEL ================================================
template <>
class MjWrapper<raw::MjModel> : public WrapperBase<raw::MjModel> {
public:
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&);
~MjWrapper();
MjModelIndexer& indexer() { return indexer_; }
void Serialize(std::ostream& output) const;
static MjWrapper<raw::MjModel> Deserialize(std::istream& input);
static MjWrapper LoadXMLFile(
const std::string& filename,
const std::optional<
std::unordered_map<std::string, pybind11::bytes>>& assets);
static MjWrapper LoadBinaryFile(
const std::string& filename,
const std::optional<
std::unordered_map<std::string, pybind11::bytes>>& assets);
static MjWrapper LoadXML(
const std::string& xml,
const std::optional<
std::unordered_map<std::string, pybind11::bytes>>& assets);
static constexpr char kFromRawPointer[] =
"__MUJOCO_STRUCTS_MJMODELWRAPPER_LOOKUP";
static MjWrapper* FromRawPointer(raw::MjModel* m) noexcept;
MjOptionWrapper opt;
MjVisualWrapper vis;
MjStatisticWrapper stat;
#define X(dtype, var, dim0, dim1) py_array_or_tuple_t<dtype> var;
MJMODEL_POINTERS
#undef X
// TODO(nimrod): Exclude text_data and names from the MJMODEL_POINTERS macro.
pybind11::bytes text_data_bytes;
pybind11::bytes names_bytes;
private:
explicit MjWrapper(raw::MjModel* ptr);
MjModelIndexer indexer_;
};
using MjModelWrapper = MjWrapper<raw::MjModel>;
template <>
struct enable_if_mj_struct<raw::MjModel> { using type = void; };
// ==================== MJCONTACT ==============================================
template <>
class MjWrapper<raw::MjContact> : public WrapperBase<raw::MjContact> {
public:
MjWrapper();
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&) = default;
MjWrapper(raw::MjContact* ptr, pybind11::handle owner);
~MjWrapper() = default;
#define X(var) py_array_or_tuple_t<mjtNum> var
X(pos);
X(frame);
X(friction);
X(solref);
X(solimp);
X(H);
#undef X
};
using MjContactWrapper = MjWrapper<raw::MjContact>;
template <>
struct enable_if_mj_struct<raw::MjContact> { using type = void; };
template <>
class MjStructList<raw::MjContact> : public StructListBase<raw::MjContact> {
public:
MjStructList(raw::MjContact* ptr, int num, pybind11::handle owner);
using StructListBase::operator[];
using StructListBase::size;
MjStructList Slice(pybind11::slice slice) {
return MjStructList(*this, slice);
}
#define X(type, var) pybind11::array_t<type> var
X(mjtNum, dist);
X(mjtNum, pos);
X(mjtNum, frame);
X(mjtNum, includemargin);
X(mjtNum, friction);
X(mjtNum, solref);
X(mjtNum, solimp);
X(mjtNum, mu);
X(mjtNum, H);
X(int, dim);
X(int, geom1);
X(int, geom2);
X(int, exclude);
X(int, efc_address);
#undef X
protected:
MjStructList(MjStructList& other, pybind11::slice slice);
};
using MjContactList = MjStructList<raw::MjContact>;
template <>
struct py_array_or_tuple<raw::MjContact> {
using type = MjContactList;
};
template <>
struct is_mj_struct_list<raw::MjContact> {
static constexpr bool value = true;
};
// ==================== MJDATA =================================================
template <>
class MjWrapper<raw::MjData>: public WrapperBase<raw::MjData> {
public:
explicit MjWrapper(const MjModelWrapper& model);
MjWrapper(const MjWrapper& other);
MjWrapper(MjWrapper&&);
~MjWrapper();
MjDataIndexer& indexer() { return indexer_; }
void Serialize(std::ostream& output) const;
static MjWrapper<raw::MjData> Deserialize(std::istream& input);
static constexpr char kFromRawPointer[] =
"__MUJOCO_STRUCTS_MJDATAWRAPPER_LOOKUP";
static MjWrapper* FromRawPointer(raw::MjData* m) noexcept;
#define X(dtype, var, dim0, dim1) py_array_or_tuple_t<dtype> var;
MJDATA_POINTERS
#undef X
py_array_or_tuple_t<raw::MjWarningStat> warning;
py_array_or_tuple_t<raw::MjTimerStat> timer;
py_array_or_tuple_t<raw::MjSolverStat> solver;
py_array_or_tuple_t<mjtNum> solver_fwdinv;
py_array_or_tuple_t<mjtNum> energy;
private:
// Internal constructor which takes ownership of given mjData pointer.
// Used for deserialization.
explicit MjWrapper(MjDataMetadata&& metadata, raw::MjData* d);
raw::MjData* Copy() const;
MjDataMetadata metadata_;
MjDataIndexer indexer_;
};
using MjDataWrapper = MjWrapper<raw::MjData>;
template <>
struct enable_if_mj_struct<raw::MjData> { using type = void; };
// ==================== MJVPERTURB =============================================
template <>
class MjWrapper<raw::MjvPerturb> : public WrapperBase<raw::MjvPerturb> {
public:
MjWrapper();
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&) = default;
~MjWrapper() = default;
#define X(var) py_array_or_tuple_t<mjtNum> var
X(refpos);
X(refquat);
X(localpos);
#undef X
};
using MjvPerturbWrapper = MjWrapper<raw::MjvPerturb>;
template <>
struct enable_if_mj_struct<raw::MjvPerturb> { using type = void; };
// ==================== MJVCAMERA ==============================================
template <>
class MjWrapper<raw::MjvCamera> : public WrapperBase<raw::MjvCamera> {
public:
MjWrapper();
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&) = default;
~MjWrapper() = default;
#define X(var) py_array_or_tuple_t<mjtNum> var
X(lookat);
#undef X
};
using MjvCameraWrapper = MjWrapper<raw::MjvCamera>;
template <>
struct enable_if_mj_struct<raw::MjvCamera> { using type = void; };
// ==================== MJVGLCAMERA ============================================
template <>
class MjWrapper<raw::MjvGLCamera> : public WrapperBase<raw::MjvGLCamera> {
public:
MjWrapper();
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&) = default;
MjWrapper(raw::MjvGLCamera* ptr, pybind11::handle owner);
explicit MjWrapper(raw::MjvGLCamera&& other);
~MjWrapper() = default;
#define X(var) py_array_or_tuple_t<float> var
X(pos);
X(forward);
X(up);
#undef X
};
using MjvGLCameraWrapper = MjWrapper<raw::MjvGLCamera>;
template <>
struct enable_if_mj_struct<raw::MjvGLCamera> { using type = void; };
// ==================== MJVGEOM ================================================
template <>
class MjWrapper<raw::MjvGeom> : public WrapperBase<raw::MjvGeom> {
public:
MjWrapper();
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&) = default;
MjWrapper(raw::MjvGeom* ptr, pybind11::handle owner);
~MjWrapper() = default;
#define X(var) py_array_or_tuple_t<float> var
X(texrepeat);
X(size);
X(pos);
X(mat);
X(rgba);
#undef X
};
using MjvGeomWrapper = MjWrapper<raw::MjvGeom>;
template <>
struct enable_if_mj_struct<raw::MjvGeom> { using type = void; };
// ==================== MJVLIGHT ===============================================
template <>
class MjWrapper<raw::MjvLight> : public WrapperBase<raw::MjvLight> {
public:
MjWrapper();
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&) = default;
MjWrapper(raw::MjvLight* ptr, pybind11::handle owner);
~MjWrapper() = default;
#define X(var) py_array_or_tuple_t<float> var
X(pos);
X(dir);
X(attenuation);
X(ambient);
X(diffuse);
X(specular);
#undef X
};
using MjvLightWrapper = MjWrapper<raw::MjvLight>;
template <>
struct enable_if_mj_struct<raw::MjvLight> { using type = void; };
// ==================== MJVOPTION ==============================================
template <>
class MjWrapper<raw::MjvOption> : public WrapperBase<raw::MjvOption> {
public:
MjWrapper();
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&) = default;
~MjWrapper() = default;
#define X(var) py_array_or_tuple_t<mjtByte> var
X(geomgroup);
X(sitegroup);
X(jointgroup);
X(tendongroup);
X(actuatorgroup);
X(flags);
#undef X
};
using MjvOptionWrapper = MjWrapper<raw::MjvOption>;
template <>
struct enable_if_mj_struct<raw::MjvOption> { using type = void; };
// ==================== MJVSCENE ===============================================
template <>
class MjWrapper<raw::MjvScene> : public WrapperBase<raw::MjvScene> {
public:
MjWrapper();
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&) = default;
MjWrapper(const MjModelWrapper& model, int maxgeom);
~MjWrapper() = default;
int nskinvert;
#define X(dtype, var) py_array_or_tuple_t<dtype> var
X(mjvGeom, geoms);
X(int, geomorder);
X(int, skinfacenum);
X(int, skinvertadr);
X(int, skinvertnum);
X(float, skinvert);
X(float, skinnormal);
#undef X
#define X(dtype, var) py_array_or_tuple_t<dtype> var
X(raw::MjvLight, lights);
X(raw::MjvGLCamera, camera);
X(float, translate);
X(float, rotate);
X(mjtByte, flags);
X(float, framergb);
#undef X
};
using MjvSceneWrapper = MjWrapper<raw::MjvScene>;
template <>
struct enable_if_mj_struct<raw::MjvScene> { using type = void; };
// ==================== MJVFIGURE ==============================================
template <>
class MjWrapper<raw::MjvFigure> : public WrapperBase<raw::MjvFigure> {
public:
MjWrapper();
MjWrapper(const MjWrapper&);
MjWrapper(MjWrapper&&) = default;
~MjWrapper() = default;
#define X(dtype, var) py_array_or_tuple_t<dtype> var
X(int, flg_ticklabel);
X(int, gridsize);
X(float, gridrgb);
X(float, figurergba);
X(float, panergba);
X(float, legendrgba);
X(float, textrgb);
X(float, linergb);
X(float, range);
X(int, highlight);
X(int, linepnt);
X(float, linedata);
X(int, xaxispixel);
X(int, yaxispixel);
X(float, xaxisdata);
X(float, yaxisdata);
#undef X
pybind11::array linename;
};
using MjvFigureWrapper = MjWrapper<raw::MjvFigure>;
template <>
struct enable_if_mj_struct<raw::MjvFigure> { using type = void; };
} // namespace _impl
template <typename T>
using MjWrapper = typename std::conditional_t<
std::is_const_v<T>, const _impl::MjWrapper<std::remove_const_t<T>>,
_impl::MjWrapper<T>>;
template <typename T>
using enable_if_mj_struct_t =
typename _impl::enable_if_mj_struct<std::remove_const_t<T>>::type;
using _impl::MjOptionWrapper;
using _impl::MjVisualHeadlightWrapper;
using _impl::MjVisualRgbaWrapper;
using _impl::MjVisualWrapper;
using _impl::MjStatisticWrapper;
using _impl::MjWarningStatWrapper;
using _impl::MjTimerStatWrapper;
using _impl::MjSolverStatWrapper;
using _impl::MjModelWrapper;
using _impl::MjDataWrapper;
using _impl::MjContactWrapper;
using _impl::MjvPerturbWrapper;
using _impl::MjvCameraWrapper;
using _impl::MjvGLCameraWrapper;
using _impl::MjvGeomWrapper;
using _impl::MjvLightWrapper;
using _impl::MjvOptionWrapper;
using _impl::MjvSceneWrapper;
using _impl::MjvFigureWrapper;
template <typename T>
using MjStructList = typename std::conditional_t<
std::is_const_v<T>, const _impl::MjStructList<std::remove_const_t<T>>,
_impl::MjStructList<T>>;
template <typename T>
static constexpr bool is_mj_struct_list_v =
_impl::is_mj_struct_list<std::remove_const_t<T>>::value;
using _impl::MjContactList;
using _impl::MjWarningStatList;
using _impl::MjTimerStatList;
using _impl::MjSolverStatList;
// ==================== HELPER FUNCTIONS FOR BINDING ARRAYS ====================
// Python array initialization.
// If T is an arithmetic (i.e. numeric) type, returns a NumPy array that wraps
// around a preexisting data buffer.
template <typename T, typename Shape>
std::enable_if_t<std::is_arithmetic_v<T>, pybind11::array_t<T>>
static InitPyArray(Shape&& shape, T* buf, pybind11::handle owner) {
int size = 1;
for (const auto& i : shape) {
size *= i;
}
if (shape.empty() || size == 0) {
return pybind11::array_t<T>(shape);
} else {
return pybind11::array_t<T>(shape, buf, owner);
}
}
// Same as above, but where we can determine array dimensions through the
// C array type directly.
template <typename T, typename Int, Int N>
std::enable_if_t<(N > 0) && std::is_arithmetic_v<T>, pybind11::array_t<T>>
static InitPyArray(T (&buf)[N], pybind11::handle owner) {
return pybind11::array_t<T>({N}, &buf[0], owner);
}
template <typename T, typename Int, Int N1, Int N2>
std::enable_if_t<(N1*N2 > 0) && std::is_arithmetic_v<T>, pybind11::array_t<T>>
static InitPyArray(T (&buf)[N1][N2], pybind11::handle owner) {
return pybind11::array_t<T>(
{N1, N2}, &buf[0][0], owner);
}
template <typename T, typename Shape>
std::enable_if_t<is_mj_struct_list_v<T>, MjStructList<T>>
static InitPyArray(Shape&& shape, T* buf, pybind11::handle owner) {
return MjStructList<T>(buf, shape[0], owner);
}
// For arrays of non-arithmetic type, we create tuple of tuples of MjWrapper<T>.
template <typename T, typename Shape>
std::enable_if_t<!std::is_arithmetic_v<T> && !is_mj_struct_list_v<T>,
pybind11::tuple>
static InitPyArray(Shape&& shape, T* buf, pybind11::handle owner) {
int size = 1;
for (const auto& i : shape) {
size *= i;
}
if (shape.empty() || !size) {
return pybind11::tuple();
}
pybind11::list out;
const auto n = shape[0];
if (shape.size() == 1) {
for (int i = 0; i < n; ++i) {
out.append(MjWrapper<T>(&buf[i], owner));
}
} else {
auto block_shape = absl::MakeConstSpan(shape).subspan(1);
auto block_size = std::accumulate(
block_shape.begin(), block_shape.end(),
1, std::multiplies<decltype(n)>());
for (int i = 0; i < n; ++i) {
out.append(InitPyArray(block_shape, &buf[i * block_size], owner));
}
}
return out;
}
// Same as above, but where we can determine array dimensions through the
// C array type directly.
template <typename T, typename Int, Int N>
std::enable_if_t<(N > 0) && !std::is_arithmetic_v<T> &&
!std::is_array_v<T> && !is_mj_struct_list_v<T>,
pybind11::tuple>
static InitPyArray(T (&buf)[N], pybind11::handle owner) {
return InitPyArray(std::array{N}, buf, owner);
}
template <typename T, typename Int, Int N1, Int N2>
std::enable_if_t<(N1*N2 > 0) && !std::is_arithmetic_v<T> &&
!std::is_array_v<T> && !is_mj_struct_list_v<T>,
pybind11::tuple>
static InitPyArray(T (&buf)[N1][N2], pybind11::handle owner) {
return InitPyArray(std::array{N1, N2}, buf, owner);
}
template <typename T, typename Int, Int N>
std::enable_if_t<is_mj_struct_list_v<T>, MjStructList<T>>
static InitPyArray(T (&buf)[N], pybind11::handle owner) {
return MjStructList<T>(buf, N, owner);
}
// Helpers for defining array/tuple properties in pybind11 classes.
//
// Defines a NumPy array property of a Python class that supports assignments.
// Specifically, we implement the setter such that `obj.arr = val` is the same
// as `obj.arr[:] = val`.
//
// Use `DefinePyArray(c, "somearray", &MjStructHolder<C>::somearray)`
// as a drop-in replacement for
// `c.def_readonly("somearray", &MjStructHolder<C>::somearray)`.
template <typename T, typename C, typename... O>
static void DefinePyArray(pybind11::class_<C, O...> c, const char* name,
pybind11::array_t<T> C::* arr) {
namespace py = pybind11;
c.def_property(
name,
[arr](const C& wrapper) { return wrapper.*arr; },
[arr](const C& wrapper, py::handle rhs) -> void {
(wrapper.*arr)[py::slice(py::none(), py::none(), py::none())] = rhs;
}
);
}
// For array of non-arithmetic type, we bind to tuples rather than NumPy array.
// These can't be assigned to directly so we just use def_readonly.
template <typename T, typename C, typename... O>
static void DefinePyArray(pybind11::class_<C, O...> c, const char* name,
T C::* arr) {
c.def_property_readonly(
name,
[arr](const C& wrapper) -> auto& { return wrapper.*arr; });
}
template <typename Raw, int N, typename C, typename... O>
static void DefinePyStr(pybind11::class_<C, O...> c, const char* name,
char (Raw::* arr)[N]) {
c.def_property(
name,
[arr](const C& c) { return pybind11::str(c.get()->*arr); },
[name = std::string(name), arr](C& c, std::string_view rhs) {
constexpr int kMaxLen = sizeof(c.get()->*arr);
const int actual_len = rhs.size();
if (actual_len >= kMaxLen) {
std::ostringstream msg;
msg << "len(" << name << ") cannot exceed " << kMaxLen - 1
<< ": got length " << actual_len;
throw pybind11::value_error(msg.str());
}
rhs.copy(c.get()->*arr, actual_len);
(c.get()->*arr)[actual_len] = '\0';
});
}
} // namespace mujoco::python
#endif // MUJOCO_PYTHON_STRUCTS_H_
python/mujoco/util/CMakeLists.txt
+96
View File
@@ -0,0 +1,96 @@
# 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
#
# https://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.
if(MSVC AND MSVC_VERSION GREATER_EQUAL 1927)
set(CMAKE_CXX_STANDARD 20) # For forceinline lambdas.
else()
set(CMAKE_CXX_STANDARD 17)
endif()
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)
# INTERPROCEDURAL_OPTIMIZATION is enforced when enabled.
set(CMAKE_POLICY_DEFAULT_CMP0069 NEW)
if(APPLE)
add_compile_options(-Werror=partial-availability -Werror=unguarded-availability)
add_link_options(-Wl,-no_weak_imports)
endif()
add_library(crossplatform INTERFACE crossplatform.h)
set_target_properties(crossplatform PROPERTIES PUBLIC_HEADER crossplatform.h)
target_include_directories(crossplatform INTERFACE ${mujoco_SOURCE_DIR}/mujoco)
add_library(array_traits INTERFACE array_traits.h)
set_target_properties(array_traits PROPERTIES PUBLIC_HEADER array_traits.h)
target_include_directories(array_traits INTERFACE ${mujoco_SOURCE_DIR}/mujoco)
target_link_libraries(array_traits INTERFACE crossplatform Eigen3::Eigen)
add_library(func_traits INTERFACE func_traits.h)
set_target_properties(func_traits PROPERTIES PUBLIC_HEADER func_traits.h)
target_include_directories(func_traits INTERFACE ${mujoco_SOURCE_DIR}/mujoco)
add_library(tuple_tools INTERFACE tuple_tools.h)
set_target_properties(tuple_tools PROPERTIES PUBLIC_HEADER tuple_tools.h)
target_include_directories(tuple_tools INTERFACE ${mujoco_SOURCE_DIR}/mujoco)
target_link_libraries(tuple_tools INTERFACE crossplatform)
add_library(func_wrap INTERFACE func_wrap.h)
set_target_properties(func_wrap PROPERTIES PUBLIC_HEADER func_wrap.h)
target_include_directories(func_wrap INTERFACE ${mujoco_SOURCE_DIR}/mujoco)
target_link_libraries(
func_wrap
INTERFACE crossplatform
Eigen3::Eigen
array_traits
func_traits
)
if(MUJOCO_TEST_PYTHON_UTIL)
add_executable(array_traits_test array_traits_test.cc)
target_link_libraries(
array_traits_test
array_traits
gmock
gtest_main
)
gtest_add_tests(TARGET array_traits_test SOURCES array_traits_test.cc)
add_executable(func_traits_test func_traits_test.cc)
target_link_libraries(
func_traits_test
func_traits
gmock
gtest_main
)
gtest_add_tests(TARGET func_traits_test SOURCES func_traits_test.cc)
add_executable(func_wrap_test func_wrap_test.cc)
target_link_libraries(
func_wrap_test
func_wrap
gmock
gtest_main
)
gtest_add_tests(TARGET func_wrap_test SOURCES func_wrap_test.cc)
add_executable(tuple_tools_test tuple_tools_test.cc)
target_link_libraries(
tuple_tools_test
func_wrap
gmock
gtest_main
)
gtest_add_tests(TARGET tuple_tools_test SOURCES tuple_tools_test.cc)
endif()
python/mujoco/util/array_traits.h
+140
View File
@@ -0,0 +1,140 @@
// 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.
#ifndef MUJOCO_PYTHON_UTIL_ARRAY_TRAITS_H_
#define MUJOCO_PYTHON_UTIL_ARRAY_TRAITS_H_
#include <type_traits>
#include "crossplatform.h"
#include <Eigen/Eigen>
#include <unsupported/Eigen/CXX11/Tensor>
namespace mujoco::util {
// Forward declaration so that the public interface appears at the top of file.
namespace _impl {
template <typename T, int... N> struct c_array;
template <typename T> struct c_array_traits;
} // namespace _impl
// Array type from scalar type and extents. This is intended to be used to
// deduce array extents as template integer parameters.
// For example c_array_t<double, 9, 4, 7> is the same as double[9][4][7].
template <typename T, int... N>
using c_array_t = typename _impl::c_array<T, N...>::type;
// Scalar type from an array, reference-to-array, or pointer-to-array type.
template <typename T>
using array_scalar_t = typename _impl::c_array_traits<T>::scalar_type;
// The number of dimensions of an array type. If the array type is regarded as
// a tensor then this corresponds to the tensor rank.
template <typename T>
static constexpr int array_ndim_v = _impl::c_array_traits<T>::ndim;
// Makes an Eigen::Matrix or Eigen::Tensor object with the same data type and
// shape as the given array type. The Eigen object returned is always row-major
// (i.e. C ordering) and dense. If the array type is one- or two-dimensional
// then an Eigen::Matrix with compile-time constant shape is returned.
// Otherwise, an Eigen::Tensor is returned.
template <typename ArrType>
constexpr auto MUJOCO_ALWAYS_INLINE MakeEigen() {
return _impl::c_array_traits<ArrType>::template MakeEigen<>();
}
template <typename ArrType>
using array_eigen_t = std::conditional_t<
std::is_const_v<ArrType>,
const decltype(MakeEigen<std::remove_const_t<ArrType>>()),
// Still need remove_const here since the type substitution in the false
// branch always occurs regardless of the condition, and
// MakeEigen<const T>() is invalid.
decltype(MakeEigen<std::remove_const_t<ArrType>>())>;
// =====================================================================
// IMPLEMENTATION DETAIL. FOR INTERNAL USE WITHIN THIS HEADER FILE ONLY.
// =====================================================================
namespace _impl {
template <typename T, int... N>
struct c_array {};
template <typename T>
struct c_array<T> {
using type = T;
static constexpr int ndim = 0;
};
template <typename T, int M, int... N>
struct c_array<T, M, N...> {
using type = typename c_array<T, N...>::type[M];
static constexpr int ndim = c_array<T, N...>::ndim + 1;
};
template <typename T>
struct c_array_traits {
static constexpr int ndim = 0;
using scalar_type = std::remove_reference_t<T>;
template <int... N>
static constexpr auto MUJOCO_ALWAYS_INLINE MakeEigen() {
if constexpr (c_array<T, N...>::ndim <= 2) {
return Eigen::Matrix<T, N..., Eigen::RowMajor>();
} else {
return Eigen::Tensor<
T, c_array<T, N...>::ndim, Eigen::RowMajor, Eigen::DenseIndex>(N...);
}
}
};
template <typename T, int N>
struct c_array_traits<T[N]> {
// Recursively peel off the innermost extent.
static constexpr int ndim = c_array_traits<T>::ndim + 1;
using scalar_type = typename c_array_traits<T>::scalar_type;
template <int... M>
static constexpr auto MUJOCO_ALWAYS_INLINE MakeEigen() {
return c_array_traits<T>::template MakeEigen<M..., N>();
}
};
template <typename T, int N>
struct c_array_traits<T(&)[N]> {
// Delegate everything to the T[N] case.
static constexpr int ndim = c_array_traits<T[N]>::ndim;
using scalar_type = typename c_array_traits<T[N]>::scalar_type;
template <int... M>
static constexpr auto MUJOCO_ALWAYS_INLINE MakeEigen() {
return c_array_traits<T[N]>::template MakeEigen<M...>();
}
};
template <typename T, int N>
struct c_array_traits<T(*)[N]> {
// Delegate everything to the T[N] case.
static constexpr int ndim = c_array_traits<T[N]>::ndim;
using scalar_type = typename c_array_traits<T[N]>::scalar_type;
template <int... M>
static constexpr auto MUJOCO_ALWAYS_INLINE MakeEigen() {
return c_array_traits<T[N]>::template MakeEigen<M...>();
}
};
} // namespace _impl
} // namespace mujoco::util
#endif // MUJOCO_PYTHON_UTIL_ARRAY_TRAITS_H_
python/mujoco/util/array_traits_test.cc
+95
View File
@@ -0,0 +1,95 @@
// 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.
#include "array_traits.h"
#include <type_traits>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
namespace mujoco::util {
namespace {
using ::testing::ElementsAre;
TEST(ArrayTraitsTest, CArrayType) {
struct Foo {};
static_assert(std::is_same_v<c_array_t<int, 3, 4, 5>,
int[3][4][5]>);
static_assert(std::is_same_v<c_array_t<double>,
double>);
static_assert(std::is_same_v<c_array_t<Foo, 7>,
Foo[7]>);
static_assert(std::is_same_v<c_array_t<Foo*, 3, 4, 5, 6>,
Foo*[3][4][5][6]>);
}
TEST(ArrayTraitsTest, ArrayNdim) {
struct Foo {};
EXPECT_EQ(array_ndim_v<int[3][4][5]>, 3);
EXPECT_EQ(array_ndim_v<double(*)[3][4]>, 2);
EXPECT_EQ(array_ndim_v<Foo*[3][4]>, 2);
EXPECT_EQ(array_ndim_v<Foo(&)[3][4][5][6]>, 4);
}
TEST(ArrayTraitsTest, ArrayScalarType) {
struct Foo {};
static_assert(std::is_same_v<
array_scalar_t<int[3][4][5]>,
int
>);
static_assert(std::is_same_v<
array_scalar_t<double(*)[3][4]>,
double
>);
static_assert(std::is_same_v<
array_scalar_t<Foo*[3][4]>,
Foo*
>);
static_assert(std::is_same_v<
array_scalar_t<Foo(&)[3][4][5][6]>,
Foo
>);
}
TEST(ArrayTraitsTest, MakeEigen) {
{
auto eigen = MakeEigen<float[3]>();
static_assert(std::is_same_v<
decltype(eigen),
Eigen::Vector3f
>);
}
{
auto eigen = MakeEigen<int[2][3]>();
static_assert(std::is_same_v<
decltype(eigen),
Eigen::Matrix<int, 2, 3, Eigen::RowMajor>
>);
}
{
auto eigen = MakeEigen<double[2][3][4]>();
static_assert(std::is_same_v<
decltype(eigen),
Eigen::Tensor<double, 3, Eigen::RowMajor, Eigen::DenseIndex>
>);
EXPECT_THAT(eigen.dimensions(), ElementsAre(2, 3, 4));
}
}
} // namespace
} // namespace mujoco::util
python/mujoco/util/crossplatform.h
+47
View File
@@ -0,0 +1,47 @@
// 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.
#ifndef MUJOCO_PYTHON_UTIL_CROSSPLATFORM_H_
#define MUJOCO_PYTHON_UTIL_CROSSPLATFORM_H_
#ifdef __has_attribute
#define MUJOCO_HAS_ATTRIBUTE(x) __has_attribute(x)
#else
#define MUJOCO_HAS_ATTRIBUTE(x) 0
#endif
#if MUJOCO_HAS_ATTRIBUTE(always_inline) || \
(defined(__GNUC__) && !defined(__clang__))
#define MUJOCO_ALWAYS_INLINE __attribute__((always_inline))
#define MUJOCO_ALWAYS_INLINE_LAMBDA MUJOCO_ALWAYS_INLINE
#define MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE MUJOCO_ALWAYS_INLINE_LAMBDA mutable
#elif defined(_MSC_VER)
#define MUJOCO_ALWAYS_INLINE __forceinline
#if _MSC_VER >= 1927 && _MSVC_LANG >= 202002L
#define MUJOCO_ALWAYS_INLINE_LAMBDA [[msvc::forceinline]]
#endif
#define MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE mutable MUJOCO_ALWAYS_INLINE_LAMBDA
#else
#define MUJOCO_ALWAYS_INLINE
#endif
#ifndef MUJOCO_ALWAYS_INLINE_LAMBDA
#define MUJOCO_ALWAYS_INLINE_LAMBDA
#endif
#ifndef MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE
#define MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE
#endif
#endif // MUJOCO_PYTHON_UTIL_CROSSPLATFORM_H_
python/mujoco/util/func_traits.h
+114
View File
@@ -0,0 +1,114 @@
// 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.
#ifndef MUJOCO_PYTHON_FUNC_TRAITS_H_
#define MUJOCO_PYTHON_FUNC_TRAITS_H_
#include <tuple>
#include <type_traits>
namespace mujoco::util {
// Forward declaration so that the public interface appears at the top of file.
namespace _impl {
template <typename, typename = void> struct is_callable;
template <typename, int, bool = false> struct func_arg;
} // namespace _impl
// True if T is callable, i.e. if T is either a function pointer/reference or
// is an instance of a type with operator().
template <typename T>
static constexpr bool is_callable_v =
_impl::is_callable<std::remove_reference_t<T>>::value;
// Type of the Nth argument of a function or functor, where N=0 refers to the
// first argument. If N exceeds the number arguments for F then
// func_arg_t<F, N> is void.
template <typename F, int N = 0>
using func_arg_t = typename _impl::func_arg<F, N, (N > 0)>::type;
template <typename F>
static constexpr int func_arg_count_v = _impl::func_arg<F, 0>::count;
// =====================================================================
// IMPLEMENTATION DETAIL. FOR INTERNAL USE WITHIN THIS HEADER FILE ONLY.
// =====================================================================
namespace _impl {
template <typename T, typename>
struct is_callable {
static constexpr bool value = false;
};
template <typename T>
struct is_callable<T, std::void_t<decltype(&T::operator())>> {
static constexpr bool value = true;
};
template <typename Return, typename... Args>
struct is_callable<Return(Args...)> {
static constexpr bool value = true;
};
template <typename Return, typename... Args>
struct is_callable<Return (*)(Args...)> {
static constexpr bool value = true;
};
// Support functors by looking at its member function Func::operator().
template <typename Func, int N, bool Recursing>
struct func_arg {
using call = decltype(
&std::remove_const_t<std::remove_reference_t<Func>>::operator());
using type = typename func_arg<call, N>::type;
static constexpr int count = func_arg<call, N>::count;
};
// Base case (N == 0) for function: resolve to Arg0.
template <typename Ret, typename Arg0, typename... Args>
struct func_arg<Ret(Arg0, Args...), 0> {
using type = Arg0;
static constexpr int count = 1 + std::tuple_size_v<std::tuple<Args...>>;
};
// Recursive case (N > 0) for function: discard Arg0 it and resolve to N-1.
template <typename Ret, int N, typename Arg0, typename... Args>
struct func_arg<Ret(Arg0, Args...), N, true> {
using type = typename func_arg<Ret(Args...), N - 1, (N > 1)>::type;
static constexpr int count = 1 + std::tuple_size_v<std::tuple<Args...>>;
};
// Specialization for non-const member functions.
template <typename C, typename Ret, int N, typename... Args>
struct func_arg<Ret (C::*)(Args...), N> {
using type = typename func_arg<Ret(Args...), N, (N > 0)>::type;
static constexpr int count = std::tuple_size_v<std::tuple<Args...>>;
};
// Specialization for const member functions (matches lambda::operator()).
template <typename C, typename Ret, int N, typename... Args>
struct func_arg<Ret (C::*)(Args...) const, N> {
using type = typename func_arg<Ret(Args...), N, (N > 0)>::type;
static constexpr int count = std::tuple_size_v<std::tuple<Args...>>;
};
// Functions with no argument: always resolve to void.
template <typename Ret, int N, bool Recursing>
struct func_arg<Ret(), N, Recursing> {
using type = void;
static constexpr int count = 0;
};
} // namespace _impl
} // namespace mujoco::util
#endif // MUJOCO_PYTHON_FUNC_TRAITS_H_
python/mujoco/util/func_traits_test.cc
+156
View File
@@ -0,0 +1,156 @@
// 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.
#include "func_traits.h"
#include <functional>
#include <type_traits>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
namespace mujoco::util {
namespace {
TEST(FuncTraitsTest, IsCallable) {
EXPECT_FALSE(is_callable_v<int>);
EXPECT_TRUE(is_callable_v<double(double)>);
EXPECT_TRUE(is_callable_v<void(*)(void)>);
EXPECT_TRUE((is_callable_v<int(&)(int, float)>));
EXPECT_TRUE(is_callable_v<std::function<void(void)>>);
{
auto lambda = [](){};
EXPECT_TRUE(is_callable_v<decltype(lambda)>);
}
{
auto mutable_lambda = []() mutable {};
EXPECT_TRUE(is_callable_v<decltype(mutable_lambda)>);
}
{
struct Functor { void operator()() {} };
EXPECT_TRUE(is_callable_v<Functor>);
}
{
struct ConstFunctor { void operator()() const {} };
EXPECT_TRUE(is_callable_v<ConstFunctor>);
}
{
struct NotCallable { void Foo() {} };
EXPECT_FALSE(is_callable_v<NotCallable>);
EXPECT_FALSE(is_callable_v<decltype(&NotCallable::Foo)>);
}
}
TEST(FuncTraitsTest, FuncArgType) {
static_assert(std::is_same_v<
func_arg_t<void()>,
void
>);
static_assert(std::is_same_v<
func_arg_t<bool(int)>,
int
>);
static_assert(std::is_same_v<
func_arg_t<bool(int&&, char&, const float&)>,
int&&
>);
static_assert(std::is_same_v<
func_arg_t<bool(int&&, char&, const float&), 0>,
int&&
>);
static_assert(std::is_same_v<
func_arg_t<bool(int&&, char&, const float&), 1>,
char&
>);
static_assert(std::is_same_v<
func_arg_t<bool(int&&, char&, const float&), 2>,
const float&
>);
static_assert(std::is_same_v<
func_arg_t<bool(int&&, char&, const float&), 3>,
void
>);
static_assert(std::is_same_v<
func_arg_t<bool(int&&, char&, const float&), 7>,
void
>);
{
auto lambda = [](bool, double&, float&&){};
static_assert(std::is_same_v<
func_arg_t<decltype(lambda)>,
bool
>);
static_assert(std::is_same_v<
func_arg_t<decltype(lambda), 0>,
bool
>);
static_assert(std::is_same_v<
func_arg_t<decltype(lambda), 1>,
double&
>);
static_assert(std::is_same_v<
func_arg_t<decltype(lambda), 2>,
float&&
>);
static_assert(std::is_same_v<
func_arg_t<decltype(lambda), 3>,
void
>);
static_assert(std::is_same_v<
func_arg_t<decltype(lambda), 10>,
void
>);
}
{
struct Functor { void operator()(char, void*) {} };
static_assert(std::is_same_v<
func_arg_t<Functor>,
char
>);
static_assert(std::is_same_v<
func_arg_t<Functor, 0>,
char
>);
static_assert(std::is_same_v<
func_arg_t<Functor, 1>,
void*
>);
static_assert(std::is_same_v<
func_arg_t<Functor, 2>,
void
>);
static_assert(std::is_same_v<
func_arg_t<Functor, 5>,
void
>);
}
}
TEST(FuncTraitsTest, FuncArgCount) {
EXPECT_EQ(func_arg_count_v<void()>, 0);
EXPECT_EQ(func_arg_count_v<bool(int)>, 1);
EXPECT_EQ(func_arg_count_v<bool(int&&, char&, const float&)>, 3);
{
auto lambda = [](bool, double&, float&&){};
EXPECT_EQ(func_arg_count_v<decltype(lambda)>, 3);
}
{
struct Functor { void operator()(char, void*) {} };
EXPECT_EQ(func_arg_count_v<Functor>, 2);
}
}
} // namespace
} // namespace mujoco::util
python/mujoco/util/func_wrap.h
+193
View File
@@ -0,0 +1,193 @@
// 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.
#ifndef MUJOCO_PYTHON_UTIL_FUNC_WRAP_H_
#define MUJOCO_PYTHON_UTIL_FUNC_WRAP_H_
#include <type_traits>
#include <utility>
#include <Eigen/Eigen>
#include "array_traits.h"
#include "crossplatform.h"
#include "func_traits.h"
namespace mujoco::util {
// Represents an argument type T of a C++ function that is callable from Python
// via pybind11. Template specializations of this struct defines how to unwrap
// arguments from pybind11 before passing them the underlying C++ function.
//
// This is used to help bind functions whose argument types are not related
// to the types that are registered with pybind11, and where it is not
// desirable/appropriate for the function's argument types to be registered.
//
// Usage:
// In the compilation unit that binds a function, specialize the template for
// each argument type that needs to be unwrapped, e.g. if a function expects
// an argument of type SomeArgType, but the type that is known to pybind11
// is SomeWrapperType, then the specialization looks like:
//
// template <> wrapped py_arg<SomeArgType*> {
// static constexpr SomeArgType* unwrap(SomeWrapperType* wrapped_arg) {
// return wrapped_arg->get_the_underlying_thing();
// }
// };
template <typename T, typename = void>
struct wrapped {
MUJOCO_ALWAYS_INLINE
static constexpr T unwrap(T arg) {
return arg;
}
};
// The wrapper type for T that can be unwrapped via wrapped<T>::unwrap.
template <typename T> using wrapper_t =
typename util::func_arg_t<decltype(wrapped<T>::unwrap)>;
namespace _impl {
template <typename T, typename = void>
struct arg_type_deducer {
static_assert(util::is_callable_v<T>, "not a Callable type");
template <typename WrapOp>
static constexpr auto WrapFunc(T&& callable) {
using Call = decltype(&std::remove_reference_t<T>::operator());
return arg_type_deducer<T, Call>::template WrapFunc<WrapOp>(
std::forward<T>(callable));
}
};
template <typename Return, typename... Args>
using func_t = Return(Args...);
// Specializations to deduce argument types for vanilla function references.
template <typename Return, typename... Args>
struct arg_type_deducer<func_t<Return, Args...>&> {
template <typename WrapOp>
static constexpr auto WrapFunc(Return (&func)(Args...)) {
return WrapOp::template WrapFunc<Return, Args...>(func);
}
};
// Specializations to deduce argument types for vanilla function pointers.
template <typename Return, typename... Args>
struct arg_type_deducer<Return (*)(Args...)> {
template <typename WrapOp>
static constexpr auto WrapFunc(Return (*func)(Args...)) {
return WrapOp::template WrapFunc<Return, Args...>(*func);
}
};
// Specialization to deduce argument types for non-const operator().
template <typename Callable, typename Return, typename... Args>
struct arg_type_deducer<
Callable, Return (std::remove_reference_t<Callable>::*)(Args...)> {
template <typename WrapOp>
static constexpr auto WrapFunc(Callable&& callable) {
return WrapOp::template WrapFunc<Return, Args...>(
std::forward<Callable>(callable));
}
};
// Specialization to deduce argument types for const operator().
template <typename Callable, typename Return, typename... Args>
struct arg_type_deducer<
Callable, Return (std::remove_reference_t<Callable>::*)(Args...) const> {
template <typename WrapOp>
static constexpr auto WrapFunc(Callable&& callable) {
return WrapOp::template WrapFunc<Return, Args...>(
std::forward<Callable>(callable));
}
};
template <typename WrapOp, typename Callable>
constexpr auto WrapFunc(Callable&& callable) {
return arg_type_deducer<Callable>::template WrapFunc<WrapOp>(
std::forward<Callable>(callable));
}
struct UnwrapArgs {
template <typename Return, typename... Args, typename Callable>
static constexpr auto WrapFunc(Callable&& callable) {
return [callable](wrapper_t<Args>... wrapped_args)
MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE {
return callable(wrapped<Args>::unwrap(wrapped_args)...);
};
}
};
template <bool OutArgProvided>
struct ReturnArrayArg0 {
template <typename Return, typename OutArg, typename... InArgs,
typename Callable>
static constexpr auto WrapFunc(Callable&& callable) {
using OutArray = std::remove_reference_t<std::remove_pointer_t<OutArg>>;
using OutScalar = util::array_scalar_t<OutArray>;
static_assert(
std::is_array_v<OutArray> && std::is_arithmetic_v<OutScalar>,
"output is not an array of arithmetic type");
static_assert(
std::is_void_v<Return>,
"callable under ReturnArrayArg0 cannot return a value");
// MSVC has a bug with `if constexpr`, as a workaround we precompute the
// condition into a constexpr variable first.
// https://developercommunity.visualstudio.com/t/1509806
constexpr bool OutArgIsRef = std::is_same_v<OutArg, OutArray&>;
if constexpr (OutArgProvided) {
using EigenOutType = Eigen::Ref<decltype(util::MakeEigen<OutArray>())>;
return [callable](InArgs... args, EigenOutType eigen_out)
MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE {
if constexpr (OutArgIsRef) {
callable(*reinterpret_cast<OutArray*>(eigen_out.data()), args...);
} else {
callable(reinterpret_cast<OutArray*>(eigen_out.data()), args...);
}
};
} else {
return [callable](InArgs... args) MUJOCO_ALWAYS_INLINE_LAMBDA_MUTABLE {
auto eigen_out = util::MakeEigen<OutArray>();
if constexpr (OutArgIsRef) {
callable(*reinterpret_cast<OutArray*>(eigen_out.data()), args...);
} else {
callable(reinterpret_cast<OutArray*>(eigen_out.data()), args...);
}
return eigen_out;
};
}
}
};
} // namespace _impl
// Makes a callable that unwraps each argument before passing it to the
// given callable. Specifically, given f(T1 x1, T2 x2, ...) this function
// returns a callable
// g(wrapper_t<T1> w1, wrapper_t<T2> w2, ...) = f(unwrap(w1), unwrap(w2), ...).
template <typename Callable>
constexpr auto UnwrapArgs(Callable&& callable) {
return _impl::WrapFunc<_impl::UnwrapArgs>(std::forward<Callable>(callable));
}
template <bool OutArgProvided = false, typename Callable>
constexpr auto ReturnArrayArg0(Callable&& callable) {
return _impl::WrapFunc<_impl::ReturnArrayArg0<OutArgProvided>>(
std::forward<Callable>(callable));
}
} // namespace mujoco::util
#endif // MUJOCO_PYTHON_UTIL_FUNC_WRAP_H_
python/mujoco/util/func_wrap_test.cc
+102
View File
@@ -0,0 +1,102 @@
// 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.
#include "func_wrap.h"
#include <string>
#include <type_traits>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <Eigen/Eigen>
#include "func_traits.h"
namespace {
struct BoxedDouble {
double value;
};
struct BoxedInt {
int value;
};
} // namespace
namespace mujoco::util {
template <> struct wrapped<BoxedDouble> {
static BoxedDouble unwrap(const std::string& wrapped) {
return BoxedDouble{std::stod(wrapped)};
}
};
template <> struct wrapped<BoxedInt> {
static BoxedInt unwrap(int wrapped) {
return BoxedInt{wrapped};
}
};
template <typename T, int N> struct wrapped<T(*)[N]> {
using Array = T[N];
static Array* unwrap(Eigen::Ref<Eigen::Vector<T, N>> wrapped) {
return reinterpret_cast<Array*>(wrapped.data());
}
};
template <typename T, int N> struct wrapped<const T(*)[N]> {
using Array = const T[N];
static Array* unwrap(const Eigen::Vector<T, N>& wrapped) {
return reinterpret_cast<Array*>(wrapped.data());
}
};
} // namespace mujoco::util
namespace {
using ::mujoco::util::func_arg_t;
using ::mujoco::util::UnwrapArgs;
using ::mujoco::util::ReturnArrayArg0;
double add(BoxedDouble x, float y, BoxedInt z) {
return x.value + y + z.value;
}
void add_array4(double (*out)[4], const double (*x)[4], const double (*y)[4]) {
for (int i = 0; i < 4; ++i) {
(*out)[i] = (*x)[i] + (*y)[i];
}
}
TEST(FuncWrapTest, UnwrapArgs) {
{
auto wrapped_add = UnwrapArgs(add);
static_assert(std::is_same_v<
func_arg_t<decltype(wrapped_add), 0>, const std::string&
>);
static_assert(std::is_same_v<
func_arg_t<decltype(wrapped_add), 1>, float
>);
static_assert(std::is_same_v<
func_arg_t<decltype(wrapped_add), 2>, int
>);
// Use binary powers so that we can do exact floating point comparison.
EXPECT_EQ(wrapped_add("1.6e+1", 5e-1, 2), 18.5);
}
{
Eigen::Vector4d out;
UnwrapArgs(add_array4)(out, {1, 3, 5, 7}, {2, 6, 9, 11});
EXPECT_THAT(out, ::testing::ElementsAre(3, 9, 14, 18));
}
}
TEST(FuncWrapTest, ReturnArrayArg0) {
auto out = UnwrapArgs(ReturnArrayArg0(add_array4))({1, 3, 5, 7},
{2, 6, 9, 11});
static_assert(std::is_same_v<decltype(out), Eigen::Vector4d>);
EXPECT_THAT(out, ::testing::ElementsAre(3, 9, 14, 18));
}
} // namespace
python/mujoco/util/tuple_tools.h
+159
View File
@@ -0,0 +1,159 @@
// 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.
#ifndef MUJOCO_PYTHON_UTIL_TUPLE_TOOLS_H_
#define MUJOCO_PYTHON_UTIL_TUPLE_TOOLS_H_
#include <optional>
#include <string_view>
#include <tuple>
#include <type_traits>
#include <utility>
#include "crossplatform.h"
namespace mujoco::util {
// Forward declaration so that the public interface appears at the top of file.
namespace _impl {
template <int N> struct head_slicer;
template <int N> struct tail_slicer;
} // namespace _impl
// Removes the first `Begin` elements from the from a tuple.
// If `Begin` is negative, the resulting tuple is obtained by removing the first
// `N - Begin` elements, where N is the size of the input tuple.
//
// For consistency with the <Begin, End> form (see below), the second optional
// template argument can be explicitly spelled out as void,
// e.g. tuple_slice<2, void>(tuple), to indicate that the slice takes every
// element up to the end of the tuple.
//
// This function has the same semantic as Python's list slicing x[Begin:].
template <int Begin, typename End = void, typename Tuple>
MUJOCO_ALWAYS_INLINE
constexpr auto tuple_slice(Tuple&& tuple) {
static_assert(
std::is_void_v<End>,
"argument End must be either omitted or explicitly specified as void");
constexpr int Size = std::tuple_size_v<std::remove_reference_t<Tuple>>;
static_assert(Begin >= -Size && Begin <= Size);
constexpr int NHead = (Begin >= 0) ? Begin : (Size + Begin);
return std::apply(_impl::head_slicer<NHead>(), std::forward<Tuple>(tuple));
}
// Extracts a contiguous slice of elements from a tuple so that the resulting
// tuple begins with the element at index `Begin` and ends with the element
// immediately before the one at index `End`. A negative value of `Begin` or
// `End` is interpreted as indexing an element from the end of the input tuple,
// where -1 refers to the last element.
//
// This function has the same semantic as Python's list slicing x[Begin:End].
template <int Begin, int End, typename Tuple>
MUJOCO_ALWAYS_INLINE
constexpr auto tuple_slice(Tuple&& tuple) {
constexpr int Size = std::tuple_size_v<std::remove_reference_t<Tuple>>;
static_assert(End >= -Size && End <= Size);
constexpr int NTail = (End >= 0) ? (Size - End) : (-End);
static_assert(
(Begin >= -Size && Begin <= -NTail) ||
(Begin >= 0 && Begin <= Size - NTail),
"Begin should refer to an element that comes before End");
constexpr int NHead = (Begin >= 0) ? Begin : (Size + Begin);
return tuple_slice<NHead>(
std::apply(_impl::tail_slicer<NTail>(), std::forward<Tuple>(tuple)));
}
// Compile-time function to check whether a string occurs in a tuple.
// Should ideally be declared consteval if we switch to C++20.
template <typename Str, typename Tuple>
static constexpr bool string_is_in_tuple(Str str, Tuple&& tuple) {
if constexpr (std::tuple_size_v<std::remove_reference_t<Tuple>> == 0) {
return false;
} else if (std::string_view(str) == std::string_view(std::get<0>(tuple))) {
return true;
} else {
return string_is_in_tuple(str, util::tuple_slice<1, void>(tuple));
}
}
// Compile-time function to check whether the elements of one tuple is a subset
// another. Should ideally be declared consteval if we switch to C++20.
template <typename Tuple1, typename Tuple2>
constexpr bool is_subset_strings(Tuple1 tuple1, Tuple2 tuple2) {
if constexpr (std::tuple_size_v<Tuple1> == 0) {
return true;
} else if (string_is_in_tuple(std::get<0>(tuple1), tuple2)) {
return is_subset_strings(util::tuple_slice<1, void>(tuple1), tuple2);
} else {
return false;
}
}
// =====================================================================
// IMPLEMENTATION DETAIL. FOR INTERNAL USE WITHIN THIS HEADER FILE ONLY.
// =====================================================================
namespace _impl{
template <int N>
struct head_slicer {
template <typename T, typename... U>
MUJOCO_ALWAYS_INLINE
constexpr auto operator()(T&& t, U&&... u) const {
return head_slicer<N-1>()(std::forward<U>(u)...);
}
};
template <>
struct head_slicer<0> {
template <typename... T>
MUJOCO_ALWAYS_INLINE
constexpr auto operator()(T&&... t) const {
return std::forward_as_tuple(t...);
}
};
template <int N>
struct move_head_to_tail {
template <typename T, typename... U>
MUJOCO_ALWAYS_INLINE
static constexpr auto move(T&& t, U&&... u) {
return move_head_to_tail<N-1>::move(
std::forward<U>(u)..., std::forward<T>(t));
}
};
template <>
struct move_head_to_tail<0> {
template <typename... T>
MUJOCO_ALWAYS_INLINE
static constexpr auto move(T&&... t) {
return std::forward_as_tuple(t...);
}
};
template <int N>
struct tail_slicer {
template <typename... T>
MUJOCO_ALWAYS_INLINE
constexpr auto operator()(T&&... t) const {
constexpr int Size = std::tuple_size_v<std::tuple<T...>>;
constexpr int NHead = Size - N;
return std::apply(head_slicer<N>(),
move_head_to_tail<NHead>::move(std::forward<T>(t)...));
}
};
} // namespace _impl
} // namespace mujoco::util
#endif // MUJOCO_PYTHON_UTIL_TUPLE_TOOLS_H_
python/mujoco/util/tuple_tools_test.cc
+54
View File
@@ -0,0 +1,54 @@
// 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.
#include "tuple_tools.h"
#include <tuple>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
namespace mujoco::util {
namespace {
TEST(TupleToolsTest, Slice) {
auto tuple = std::make_tuple(2, 3, 5, 7, 11, 13);
EXPECT_EQ((tuple_slice<1, 5>(tuple)), std::make_tuple(3, 5, 7, 11));
EXPECT_EQ((tuple_slice<2, 3>(tuple)), std::make_tuple(5));
// negative indices
EXPECT_EQ((tuple_slice<1, -2>(tuple)), std::make_tuple(3, 5, 7));
EXPECT_EQ((tuple_slice<-4, 6>(tuple)), std::make_tuple(5, 7, 11, 13));
EXPECT_EQ((tuple_slice<-3, -1>(tuple)), std::make_tuple(7, 11));
// empty slices
EXPECT_EQ((tuple_slice<0, 0>(tuple)), std::make_tuple());
EXPECT_EQ((tuple_slice<3, 3>(tuple)), std::make_tuple());
EXPECT_EQ((tuple_slice<-2, -2>(tuple)), std::make_tuple());
// specify void as the End argument
EXPECT_EQ((tuple_slice<2, void>(tuple)), std::make_tuple(5, 7, 11, 13));
EXPECT_EQ((tuple_slice<-2, void>(tuple)), std::make_tuple(11, 13));
// omit the End argument
EXPECT_EQ((tuple_slice<2>(tuple)), std::make_tuple(5, 7, 11, 13));
EXPECT_EQ((tuple_slice<-2>(tuple)), std::make_tuple(11, 13));
// empty input tuples
EXPECT_EQ((tuple_slice<0>(std::make_tuple())), std::make_tuple());
EXPECT_EQ((tuple_slice<0, 0>(std::make_tuple())), std::make_tuple());
}
} // namespace
} // namespace mujoco::util
python/setup.py
+301
View File
@@ -0,0 +1,301 @@
# 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.
# ==============================================================================
"""Install script for MuJoCo."""
import fnmatch
import os
import platform
import random
import re
import shutil
import string
import subprocess
import sys
import sysconfig
import setuptools
from setuptools import find_packages
from setuptools import setup
from setuptools.command import build_ext
__version__ = '2.1.2'
MUJOCO_CMAKE = 'MUJOCO_CMAKE'
MUJOCO_CMAKE_ARGS = 'MUJOCO_CMAKE_ARGS'
MUJOCO_PATH = 'MUJOCO_PATH'
EXT_PREFIX = 'mujoco.'
def get_long_description():
"""Creates a long description for the package from bundled markdown files."""
current_dir = os.path.dirname('__file__')
with open(os.path.join(current_dir, 'README.md')) as f:
description = f.read()
try:
with open(os.path.join(current_dir, 'LICENSES_THIRD_PARTY.md')) as f:
description = f'{description}\n{f.read()}'
except FileNotFoundError:
pass
return description
def get_mujoco_lib_pattern():
if platform.system() == 'Windows':
return 'mujoco*.lib'
elif platform.system() == 'Darwin':
return 'libmujoco*.dylib'
else:
return 'libmujoco*.so*'
def get_external_lib_patterns():
if platform.system() == 'Windows':
return ['mujoco*.dll']
elif platform.system() == 'Darwin':
return ['libmujoco*.dylib']
else:
return ['libmujoco*.so*', 'libglew*.so']
def start_and_end(iterable):
it = iter(iterable)
while True:
try:
first = next(it)
second = next(it)
yield first, second
except StopIteration:
return
def tokenize_quoted_substr(input_string, quote_char, placeholders=None):
"""Replace quoted substrings with random text placeholders with no spaces."""
# Matches quote characters not proceded with a backslash.
pattern = re.compile(r'(?<!\\)' + quote_char)
quote_positions = [m.start() for m in pattern.finditer(input_string)]
if len(quote_positions) % 2:
raise ValueError(f'unbalanced quotes {quote_char}...{quote_char}')
output_string = ''
placeholders = placeholders if placeholders is not None else dict()
prev_end = -1
for start, end in start_and_end(quote_positions):
output_string += input_string[prev_end+1:start]
while True:
placeholder = ''.join(random.choices(string.ascii_lowercase, k=5))
if placeholder not in input_string and placeholder not in output_string:
break
output_string += placeholder
placeholders[placeholder] = input_string[start+1:end]
prev_end = end
output_string += input_string[prev_end+1:]
return output_string, placeholders
def parse_cmake_args_from_environ(env_var_name=MUJOCO_CMAKE_ARGS):
"""Parses CMake arguments from an environment variable."""
raw_args = os.environ.get(env_var_name, '').strip()
unquoted, placeholders = tokenize_quoted_substr(raw_args, '"')
unquoted, placeholders = tokenize_quoted_substr(unquoted, "'", placeholders)
parts = re.split(r'\s+', unquoted.strip())
out = []
for part in parts:
for k, v in placeholders.items():
part = part.replace(k, v)
part = part.replace('\\"', '"').replace("\\'", "'")
if part:
out.append(part)
return out
class CMakeExtension(setuptools.Extension):
"""A Python extension that has been prebuilt by CMake.
We do not want distutils to handle the build process for our extensions, so
so we pass an empty list to the super constructor.
"""
def __init__(self, name):
super().__init__(name, sources=[])
class BuildCMakeExtension(build_ext.build_ext):
"""Uses CMake to build extensions."""
def run(self):
self._mujoco_library_path, self._mujoco_include_path = self._find_mujoco()
self._configure_cmake()
for ext in self.extensions:
assert ext.name.startswith(EXT_PREFIX)
assert '.' not in ext.name[len(EXT_PREFIX):]
self.build_extension(ext)
self._copy_external_libraries()
self._copy_mujoco_headers()
def _find_mujoco(self):
if MUJOCO_PATH not in os.environ:
raise RuntimeError(f'{MUJOCO_PATH} environment variable is not set')
library_path = None
include_path = None
for directory, _, filenames in os.walk(os.environ['MUJOCO_PATH']):
if fnmatch.filter(filenames, get_mujoco_lib_pattern()):
library_path = directory
if fnmatch.filter(filenames, 'mujoco.h'):
include_path = directory
if library_path and include_path:
return library_path, include_path
raise RuntimeError('Cannot find MuJoCo library and/or include paths')
def _copy_external_libraries(self):
dst = os.path.dirname(self.get_ext_fullpath(self.extensions[0].name))
for directory, _, filenames in os.walk(os.environ['MUJOCO_PATH']):
for pattern in get_external_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)),
'include')
os.mkdir(dst)
for directory, _, filenames in os.walk(self._mujoco_include_path):
for filename in fnmatch.filter(filenames, '*.h'):
shutil.copyfile(os.path.join(directory, filename),
os.path.join(dst, filename))
def _configure_cmake(self):
"""Check for CMake."""
cmake = os.environ.get(MUJOCO_CMAKE, 'cmake')
build_cfg = 'Debug' if self.debug else 'Release'
cmake_module_path = os.path.join(os.path.dirname(__file__), 'cmake')
cmake_args = [
f'-DPython3_ROOT_DIR={sys.prefix}',
f'-DPython3_EXECUTABLE={sys.executable}',
f'-DMUJOCO_LIBRARY_DIR={self._mujoco_library_path}',
f'-DMUJOCO_INCLUDE_DIR={self._mujoco_include_path}',
f'-DCMAKE_MODULE_PATH={cmake_module_path}',
f'-DCMAKE_BUILD_TYPE={build_cfg}',
f'-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={self.build_temp}',
f'-DCMAKE_INTERPROCEDURAL_OPTIMIZATION={"OFF" if self.debug else "ON"}',
'-DCMAKE_Fortran_COMPILER:STRING=',
'-DCMAKE_VERBOSE_MAKEFILE=ON',
'-DBUILD_TESTING=OFF',
]
if platform.system() != 'Windows':
cmake_args.extend([
f'-DPython3_LIBRARY={sysconfig.get_paths()["stdlib"]}',
f'-DPython3_INCLUDE_DIR={sysconfig.get_paths()["include"]}',
])
if platform.system() == 'Darwin' and os.environ.get('ARCHFLAGS'):
osx_archs = []
if '-arch x86_64' in os.environ['ARCHFLAGS']:
osx_archs.append('x86_64')
if '-arch arm64' in os.environ['ARCHFLAGS']:
osx_archs.append('arm64')
cmake_args.append(f'-DCMAKE_OSX_ARCHITECTURES={";".join(osx_archs)}')
cmake_args.extend(parse_cmake_args_from_environ())
os.makedirs(self.build_temp, exist_ok=True)
if platform.system() == 'Windows':
cmake_args = [arg.replace('\\', '/') for arg in cmake_args]
print('Configuring CMake with the following arguments:')
for arg in cmake_args:
print(f' {arg}')
subprocess.check_call(
[cmake] + cmake_args +
[os.path.join(os.path.dirname(__file__), 'mujoco')],
cwd=self.build_temp)
print('Building all extensions with CMake')
subprocess.check_call(
[cmake, '--build', '.', f'-j{os.cpu_count()}', '--config', build_cfg],
cwd=self.build_temp)
def build_extension(self, ext):
dest_path = self.get_ext_fullpath(ext.name)
build_path = os.path.join(self.build_temp, os.path.basename(dest_path))
subprocess.check_call(['cp', build_path, dest_path])
def find_data_files(package_dir, patterns):
"""Recursively finds files whose names match the given shell patterns."""
paths = set()
for directory, _, filenames in os.walk(package_dir):
for pattern in patterns:
for filename in fnmatch.filter(filenames, pattern):
# NB: paths must be relative to the package directory.
relative_dirpath = os.path.relpath(directory, package_dir)
paths.add(os.path.join(relative_dirpath, filename))
return list(paths)
setup(
name='mujoco',
version=__version__,
author='DeepMind',
author_email='mujoco@deepmind.com',
description='MuJoCo Physics Simulator',
long_description=get_long_description(),
long_description_content_type='text/markdown',
url='https://github.com/deepmind/mujoco',
license='Apache License 2.0',
classifiers=[
'License :: OSI Approved :: Apache Software License',
],
cmdclass=dict(build_ext=BuildCMakeExtension),
ext_modules=[
CMakeExtension('mujoco._callbacks'),
CMakeExtension('mujoco._constants'),
CMakeExtension('mujoco._enums'),
CMakeExtension('mujoco._errors'),
CMakeExtension('mujoco._functions'),
CMakeExtension('mujoco._render'),
CMakeExtension('mujoco._rollout'),
CMakeExtension('mujoco._structs'),
],
python_requires='>=3.7',
install_requires=[
'absl-py',
'glfw',
'numpy',
'pyopengl',
],
tests_require=[
'absl-py',
'glfw',
'numpy',
'pyopengl',
],
test_suite='mujoco',
packages=find_packages(),
package_data={
'mujoco':
find_data_files(
package_dir='mujoco',
patterns=[
'libmujoco*.dylib',
'libmujoco*.so*',
'mujoco*.dll',
'libglew*.so*',
'mujoco.h',
'mj*.h',
]),
},
)