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,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_