Initial open sourcing of MuJoCo.

PiperOrigin-RevId: 450374687
Change-Id: Ie3225a46ce095fc28ae8e63c326a640261f562bb

test/benchmark/CMakeLists.txt

@@ -0,0 +1,57 @@
+# 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
+#
+#     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.
+
+# Macro for benchmarks that don't use GL (same as mujoco_api_test, but uses
+# benchmark::benchmark_main).
+
+macro(mujoco_benchmark_test name)
+  add_executable(${name} ${name}.cc)
+  target_link_libraries(
+    ${name}
+    benchmark::benchmark_main
+    mujoco
+    absl::core_headers
+  )
+  target_include_directories(${name} PRIVATE ${MUJOCO_TEST_INCLUDE})
+  # TODO(fraromano) Check RPATH settings
+  set_target_properties(${name} PROPERTIES BUILD_RPATH ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
+  # gtest_discover_tests is recommended over gtest_add_tests, but has some issues in Windows.
+  gtest_add_tests(
+    TARGET ${name}
+    SOURCES ${name}.cc
+    WORKING_DIRECTORY ${MUJOCO_TEST_WORKING_DIR}
+    TEST_LIST testList
+  )
+  if(WIN32)
+    set_tests_properties(
+      ${testList} PROPERTIES ENVIRONMENT "PATH=$<TARGET_FILE_DIR:mujoco>;$ENV{PATH}"
+    )
+  endif()
+endmacro()
+
+mujoco_benchmark_test(step_benchmark_test)
+target_link_libraries(
+  step_benchmark_test
+  fixture
+  gmock
+  benchmark::benchmark
+)
+
+mujoco_benchmark_test(parse_benchmark_test)
+target_link_libraries(
+  parse_benchmark_test
+  fixture
+  gmock
+  benchmark::benchmark
+)

test/benchmark/parse_benchmark_test.cc

@@ -0,0 +1,68 @@
+// 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.
+
+// A benchmark for parsing and compiling models from XML.
+
+#include <array>
+
+#include <benchmark/benchmark.h>
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <absl/base/attributes.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mujoco.h>
+#include "test/fixture.h"
+
+namespace mujoco {
+namespace {
+
+using ::testing::NotNull;
+
+static void run_parse_benchmark(const char *model_path, benchmark::State& state) {
+  MujocoErrorTestGuard guard;  // Fail test if there are any mujoco errors
+
+  const std::string xml_path = GetModelPath(model_path);
+  std::array<char, 1024> error;
+  for (auto s : state) {
+    // TODO(nimrod): Load the models from VFS rather than from disk, to
+    // limit the benchmark to the parsing and model compilation speed.
+    mjModel* model =
+      mj_loadXML(xml_path.data(), nullptr, error.data(), error.size());
+    ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data();
+    mj_deleteModel(model);
+  }
+  state.SetLabel(model_path);
+}
+
+// Use ABSL_ATTRIBUTE_NO_TAIL_CALL to make sure the benchmark functions appear
+// separately in CPU profiles (and don't get replaced with raw calls to
+// run_parse_benchmark).
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_ParseCloth(benchmark::State& state) {
+  run_parse_benchmark("composite/cloth.xml", state);
+}
+BENCHMARK(BM_ParseCloth);
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_ParseHumanoid(benchmark::State& state) {
+  run_parse_benchmark("humanoid/humanoid.xml", state);
+}
+BENCHMARK(BM_ParseHumanoid);
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_ParseHumanoid100(benchmark::State& state) {
+  run_parse_benchmark("humanoid100/humanoid100.xml", state);
+}
+BENCHMARK(BM_ParseHumanoid100);
+
+}  // namespace
+}  // namespace mujoco

test/benchmark/run_ablation.py

@@ -0,0 +1,223 @@
+#!/usr/bin/env python3
+# 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.
+
+"""A script for running step_benchmark_test for a variety of build settings.
+
+This is not meant to be run in google3, but in the git repository's root.
+
+This will be used for release 2.1.1, and adapted for future use when more
+decisions about the build need to be made.
+"""
+
+import argparse
+import os
+import re
+import shutil
+import subprocess
+import sys
+from typing import NamedTuple, Tuple
+
+
+class _Setting(NamedTuple):
+  name: str
+  cmake_args: Tuple[str]
+
+
+def run_benchmark_variants(output_directory: str, build_directory: str,
+                           benchmark_repetitions: int, delete_build: bool):
+  """Builds different variants of the benchmark test, and runs them."""
+  # For each of the settings, the first option is what we use as a baseline.
+  # We chose to pick the settings that were likely best (so this is an ablation
+  # study), so interactions between settings are tested.
+  settings_ranges = {
+      "avx": [
+          _Setting("on", ("-DMUJOCO_ENABLE_AVX=ON",)),
+          _Setting("off", ("-DMUJOCO_ENABLE_AVX=OFF",))
+      ],
+      "avx_intrinsics": [
+          _Setting("on", ("-DMUJOCO_ENABLE_AVX_INTRINSICS=ON",)),
+          _Setting("off", ("-DMUJOCO_ENABLE_AVX_INTRINSICS=OFF",))
+      ],
+      "lto": [
+          _Setting("on", ("-DCMAKE_INTERPROCEDURAL_OPTIMIZATION=ON",)),
+          _Setting("off", ("-DCMAKE_INTERPROCEDURAL_OPTIMIZATION=OFF",)),
+      ],
+      "compiler": compiler_cmake_flags_variants(),
+  }
+
+  baseline_settings = {
+      key: value[0] for key, value in settings_ranges.items()
+  }
+  run_benchmark(
+      f"{sys.platform}_baseline",
+      baseline_settings,
+      output_directory=output_directory,
+      build_directory=f"{build_directory}_baseline",
+      benchmark_repetitions=benchmark_repetitions,
+      delete_build=delete_build,
+  )
+
+  for key in sorted(settings_ranges):
+    settings = baseline_settings.copy()
+    for setting in settings_ranges[key][1:]:
+      settings[key] = setting
+      run_benchmark(
+          f"{sys.platform}_{key}_{setting.name}",
+          settings,
+          output_directory=output_directory,
+          build_directory=f"{build_directory}_{key}_{setting.name}",
+          benchmark_repetitions=benchmark_repetitions,
+          delete_build=delete_build,
+      )
+
+
+def run_benchmark(name, settings, output_directory: str, build_directory: str,
+                  benchmark_repetitions: int, delete_build: bool):
+  """Builds and runs a single benchmark."""
+  current_dir = os.getcwd()
+  try:
+    if delete_build:
+      shutil.rmtree(build_directory, ignore_errors=True)
+    os.makedirs(build_directory, exist_ok=True)
+    os.makedirs(output_directory, exist_ok=True)
+    os.chdir(build_directory)
+    cmake_args = [
+        "cmake", "..", "-DMUJOCO_BUILD_TESTS=ON", "-DCMAKE_BUILD_TYPE=Release"
+    ]
+    cmake_args.extend(os_specific_cmake_flags())
+    for _, setting in sorted(settings.items()):
+      cmake_args.extend(setting.cmake_args)
+
+    print(f"Running ({name}):", " ".join(cmake_args))
+    subprocess.check_call(cmake_args)
+    subprocess.check_call([
+        "cmake", "--build", ".", "-j8", "--config=Release", "-t",
+        "step_benchmark_test"
+    ])
+    output_path = os.path.join(output_directory, f"{name}.json")
+
+    os.chdir("../mujoco/test")
+    if sys.platform == "win32":
+      cygwin_run_benchmark(build_directory, output_path, benchmark_repetitions)
+    else:
+      binary_path = os.path.join(build_directory, "bin", "step_benchmark_test")
+      subprocess.check_call([
+          binary_path,
+          "--benchmark_filter=all",
+          "--benchmark_enable_random_interleaving=true",
+          "--benchmark_min_time=0.5",
+          f"--benchmark_repetitions={benchmark_repetitions}",
+          "--benchmark_format=json",
+          f"--benchmark_out={output_path}",
+      ])
+  finally:
+    os.chdir(current_dir)
+
+
+def os_specific_cmake_flags():
+  """CMake args that should be passed to all benchmarks on current OS."""
+  if sys.platform == "win32":
+    return ("-A", "x64", "-Thost=x86",
+            "-DCMAKE_MSVC_RUNTIME_LIBRARY=MultiThreaded",
+            "-DCMAKE_SYSTEM_VERSION=10.0.19041.0")
+  elif sys.platform.startswith("linux"):
+    return ()
+  elif sys.platform == "darwin":
+    # macOS
+    return ()
+  else:
+    raise ValueError(f"Unknown OS: {sys.platform}")
+
+
+def compiler_cmake_flags_variants():
+  """Returns a list of benchmark settings for different compilers."""
+  if sys.platform == "win32":
+    return [
+        _Setting("VS2022", ("-G", "Visual Studio 17 2022")),
+        # MSVC 2017 doesn't seem to support C17 standard.
+        # _Setting("VS2017", ("-G", "Visual Studio 15 2017")),
+    ]
+  elif sys.platform.startswith("linux"):
+    return [
+        _Setting("clang-11",
+                 ("-DCMAKE_C_COMPILER=clang-11",
+                  "-DCMAKE_CXX_COMPILER=clang++-11", "-DCMAKE_LINKER=lld-11")),
+        _Setting("clang-8",
+                 ("-DCMAKE_C_COMPILER=clang-8",
+                  "-DCMAKE_CXX_COMPILER=clang++-8", "-DCMAKE_LINKER=lld-8")),
+    ]
+  elif sys.platform == "darwin":
+    # macOS
+    return [_Setting("clang", ())]
+  else:
+    raise ValueError(f"Unknown OS: {sys.platform}")
+
+
+def cygwin_run_benchmark(build_directory, output_path, benchmark_repetitions):
+  """Runs the benchmark command under the cygwin environment."""
+  # The subprocess module interacts badly with cygwin.
+  # Rather than trying to run the binary directly through Popen, use bash on
+  # cygwin.
+  lib_path = cygwin_path(os.path.join(build_directory, "lib", "Release"))
+  binary_path = cygwin_path(
+      os.path.join(build_directory, "bin", "Release",
+                   "step_benchmark_test"))
+  cygwin = subprocess.Popen(["bash"], stdin=subprocess.PIPE)
+  command = (
+      f'PATH="$PATH:{lib_path}" {binary_path} --benchmark_filter=all '
+      "--benchmark_enable_random_interleaving=true ",
+      "--benchmark_min_time=0.5 "
+      f"--benchmark_repetitions={benchmark_repetitions} "
+      "--benchmark_format=json "
+      f"--benchmark_out='{output_path}'")
+  cygwin.communicate(input=bytes(command, "utf-8"))
+  if cygwin.returncode:
+    raise ValueError(f"Benchmark returned error code: {cygwin.returncode}")
+
+
+def cygwin_path(windows_path):
+  path = windows_path.replace("\\", "/")
+  path = re.sub(r".*cygwin64", "", path)
+  return path
+
+
+def main(argv):
+  parser = argparse.ArgumentParser()
+  parser.add_argument("-o", "--output_directory", default="benchmark_results")
+  parser.add_argument(
+      "-b",
+      "--build_directory",
+      default="build",
+      help="Base path for build directories. Benchmark names will be appended."
+  )
+  parser.add_argument("-n", "--repetitions", type=int, default=30)
+  parser.add_argument(
+      "-d",
+      "--delete_build",
+      nargs="?",
+      type=bool,
+      const=True,
+      default=False,
+      help="Delete the build directory before building benchmark.")
+  args = parser.parse_args(argv[1:])
+  run_benchmark_variants(
+      output_directory=os.path.abspath(args.output_directory),
+      build_directory=os.path.abspath(args.build_directory),
+      delete_build=args.delete_build,
+      benchmark_repetitions=args.repetitions)
+
+
+if __name__ == "__main__":
+  main(sys.argv)

test/benchmark/step_benchmark_test.cc

@@ -0,0 +1,98 @@
+// 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.
+
+// A benchmark which steps various models without rendering, and measures speed.
+
+#include <array>
+
+#include <benchmark/benchmark.h>
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <absl/base/attributes.h>
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mujoco.h>
+#include "test/fixture.h"
+
+namespace mujoco {
+namespace {
+
+using ::testing::NotNull;
+
+static void add_ctrl_noise(const mjModel* m, mjData* d, int step) {
+  for (int i = 0; i < m->nu; i++) {
+    mjtNum center = 0.0;
+    mjtNum radius = 1.0;
+    mjtNum* range = m->actuator_ctrlrange + 2 * i;
+    if (m->actuator_ctrllimited[i]) {
+      center = (range[1] + range[0]) / 2;
+      radius = (range[1] - range[0]) / 2;
+    }
+    radius *= 0.01;
+    d->ctrl[i] = center + radius * (2 * mju_Halton(step, i + 2) - 1);
+  }
+}
+
+static void assert_model_not_null(mjModel* model,
+                                  const std::array<char, 1024>& error) {
+  ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data();
+}
+
+static mjModel* load_model(const char* model_path) {
+  const std::string xml_path = GetModelPath(model_path);
+  std::array<char, 1024> error;
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error.data(), error.size());
+  assert_model_not_null(model, error);
+  return model;
+}
+
+static void run_step_benchmark(const mjModel* model, benchmark::State& state) {
+  mjData* data = mj_makeData(model);
+  int i = 0;
+  for (auto s : state) {
+    add_ctrl_noise(model, data, i++);
+    mj_step(model, data);
+  }
+  mj_deleteData(data);
+  state.SetItemsProcessed(state.iterations());
+}
+
+// Use ABSL_ATTRIBUTE_NO_TAIL_CALL to make sure the benchmark functions appear
+// separately in CPU profiles (and don't get replaced with raw calls to
+// run_step_benchmark).
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_StepCloth(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  static mjModel* model = load_model("composite/cloth.xml");
+  run_step_benchmark(model, state);
+}
+BENCHMARK(BM_StepCloth);
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_StepHumanoid(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  static mjModel* model = load_model("humanoid/humanoid.xml");
+  run_step_benchmark(model, state);
+}
+BENCHMARK(BM_StepHumanoid);
+BENCHMARK(BM_StepHumanoid)->ThreadRange(2, 16);
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_StepHumanoid100(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  static mjModel* model = load_model("humanoid100/humanoid100.xml");
+  run_step_benchmark(model, state);
+}
+BENCHMARK(BM_StepHumanoid100);
+
+}  // namespace
+}  // namespace mujoco