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

PiperOrigin-RevId: 434731612 Change-Id: I0cfda3e7a3d1c72036764986efc252ffa1b8c6b0
2022-03-15 14:04:44 +00:00
parent 175d25cd9f
commit 3577e2cf8b
304 changed files with 23906 additions and 1013 deletions
@@ -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_