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

PiperOrigin-RevId: 434731612
Change-Id: I0cfda3e7a3d1c72036764986efc252ffa1b8c6b0

python/mujoco/util/CMakeLists.txt

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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