From 62ec802ec172dc3fff8c53e1b9dd8cde64e258bc Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Tue, 1 Oct 2024 04:07:33 -0700
Subject: [PATCH] Tweaks to benchmarks.

- Modify parse_benchmark_test to not time the first instance of file loading.
- Improve readability of step_benchmark_test.

PiperOrigin-RevId: 680953562
Change-Id: I4569d1b946ecb49b6590ea313f0b598ffc897888
---
 test/benchmark/parse_benchmark_test.cc | 9 ++++++---
 test/benchmark/step_benchmark_test.cc  | 7 +++----
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/test/benchmark/parse_benchmark_test.cc b/test/benchmark/parse_benchmark_test.cc
index c5855865..059f0c58 100644
--- a/test/benchmark/parse_benchmark_test.cc
+++ b/test/benchmark/parse_benchmark_test.cc
@@ -54,11 +54,14 @@ static void run_parse_benchmark(const std::string xml_path,
 
   ASSERT_THAT(vfs_errno, Eq(0)) << "Failed to add file to VFS: " << vfs_errmsg;
 
+  // load once to warm up filesystem and compiler cache
   std::array<char, 1024> error;
-  for (auto s : state) {
-    mjModel* model =
+  mjModel* model =
       mj_loadXML(xml_path.data(), vfs.get(), error.data(), error.size());
-    ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data();
+  ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data();
+
+  for (auto s : state) {
+    mjModel* model = mj_loadXML(xml_path.data(), vfs.get(), 0, 0);
     mj_deleteModel(model);
   }
   state.SetLabel(xml_path);
diff --git a/test/benchmark/step_benchmark_test.cc b/test/benchmark/step_benchmark_test.cc
index 0c4c7bca..a2039d80 100644
--- a/test/benchmark/step_benchmark_test.cc
+++ b/test/benchmark/step_benchmark_test.cc
@@ -17,8 +17,6 @@
 #include <vector>
 
 #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>
@@ -57,8 +55,9 @@ static void run_step_benchmark(const mjModel* model, benchmark::State& state) {
   while (state.KeepRunningBatch(kNumBenchmarkSteps)) {
     mj_setState(model, data, initial_state.data(), spec);
 
-    for (int i=kNumWarmupSteps; i < nsteps; i++) {
-      mju_copy(data->ctrl, ctrl.data()+model->nu*i, model->nu);
+    for (int i=0; i < kNumBenchmarkSteps; i++) {
+      mjtNum* ctrl_data = ctrl.data()+model->nu*(i+kNumWarmupSteps);
+      mju_copy(data->ctrl, ctrl_data, model->nu);
       mj_step(model, data);
     }
   }