From c17375fa32a6e799e512050c46672fb009ce8fdf Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Tue, 22 Nov 2022 07:40:06 -0800
Subject: [PATCH] Stop profiling mju_halton in testspeed and bechmarks.

PiperOrigin-RevId: 490239351
Change-Id: Id73c101672e5b974f16b6b4da3aae757a657de15
---
 sample/testspeed.cc                   | 39 +++++++++++++++++----------
 test/benchmark/step_benchmark_test.cc | 10 ++++---
 test/fixture.cc                       | 25 ++++++++++-------
 test/fixture.h                        |  3 ++-
 4 files changed, 49 insertions(+), 28 deletions(-)

diff --git a/sample/testspeed.cc b/sample/testspeed.cc
index 867de278..5f51cae1 100644
--- a/sample/testspeed.cc
+++ b/sample/testspeed.cc
@@ -17,6 +17,7 @@
 #include <cstring>
 #include <string>
 #include <thread>
+#include <vector>
 
 #include <mujoco/mujoco.h>
 
@@ -58,9 +59,27 @@ int finish(const char* msg = NULL, mjModel* m = NULL) {
   return 0;
 }
 
+std::vector<mjtNum> CtrlNoise(const mjModel* m, int nsteps, mjtNum ctrlnoise) {
+  std::vector<mjtNum> ctrl;
+  for (int step=0; step < nsteps; 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 *= ctrlnoise;
+      ctrl.push_back(center + radius * (2 * mju_Halton(step, i+2) - 1));
+    }
+  }
+  return ctrl;
+}
+
 
 // thread function
-void simulate(int id, int nstep, mjtNum ctrlnoise) {
+void simulate(int id, int nstep, mjtNum* ctrl) {
   // clear statistics
   contacts[id] = 0;
   constraints[id] = 0;
@@ -69,18 +88,7 @@ void simulate(int id, int nstep, mjtNum ctrlnoise) {
   double start = gettm();
   for (int i=0; i<nstep; i++) {
     // inject pseuso-random control noise
-    if (ctrlnoise)
-      for (int j=0; j<m->nu; j++) {
-        mjtNum center = 0.0;
-        mjtNum radius = 1.0;
-        mjtNum* range = m->actuator_ctrlrange + 2*j;
-        if (m->actuator_ctrllimited[j]) {
-          center = (range[1] + range[0]) / 2;
-          radius = (range[1] - range[0]) / 2;
-        }
-        radius *= ctrlnoise;
-        d[id]->ctrl[j] = center + radius * (2*mju_Halton(i, j+2) - 1);
-      }
+    mju_copy(d[id]->ctrl, ctrl + i*m->nu, m->nu);
 
     // advance simulation
     mj_step(m, d[id]);
@@ -172,11 +180,14 @@ int main(int argc, char** argv) {
     std::printf("\nRunning %d steps at dt = %g ...\n\n", nstep, m->opt.timestep);
   }
 
+  // create pseudo-random control sequence
+  std::vector<mjtNum> ctrl = CtrlNoise(m, nstep, ctrlnoise);
+
   // run simulation, record total time
   std::thread th[maxthread];
   double starttime = gettm();
   for (int id=0; id<nthread; id++) {
-    th[id] = std::thread(simulate, id, nstep, ctrlnoise);
+    th[id] = std::thread(simulate, id, nstep, ctrl.data());
   }
   for (int id=0; id<nthread; id++) {
     th[id].join();
diff --git a/test/benchmark/step_benchmark_test.cc b/test/benchmark/step_benchmark_test.cc
index 860fabf1..ef022ed1 100644
--- a/test/benchmark/step_benchmark_test.cc
+++ b/test/benchmark/step_benchmark_test.cc
@@ -42,9 +42,13 @@ std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
 static void run_step_benchmark(const mjModel* model, benchmark::State& state) {
   mjData* data = mj_makeData(model);
 
+  // compute noise
+  int nsteps = kNumWarmupSteps+kNumBenchmarkSteps;
+  std::vector<mjtNum> ctrl = GetCtrlNoise(model, nsteps);
+
   // warm-up rollout to get a typcal state
   for (int i=0; i < kNumWarmupSteps; i++) {
-    AddCtrlNoise(model, data, i);
+    mju_copy(data->ctrl, ctrl.data()+model->nu*i, model->nu);
     mj_step(model, data);
   }
   // save state
@@ -60,8 +64,8 @@ static void run_step_benchmark(const mjModel* model, benchmark::State& state) {
     mju_copy(data->act, act.data(), model->na);
     mju_copy(data->qacc_warmstart, warmstart.data(), model->nv);
 
-    for (int i=0; i < kNumBenchmarkSteps; i++) {
-      AddCtrlNoise(model, data, i+kNumWarmupSteps);
+    for (int i=kNumWarmupSteps; i < nsteps; i++) {
+      mju_copy(data->ctrl, ctrl.data()+model->nu*i, model->nu);
       mj_step(model, data);
     }
   }
diff --git a/test/fixture.cc b/test/fixture.cc
index 937adc10..d95a5289 100644
--- a/test/fixture.cc
+++ b/test/fixture.cc
@@ -134,18 +134,23 @@ const std::string SaveAndReadXml(const mjModel* model) {
   return contents;
 }
 
-void AddCtrlNoise(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;
+std::vector<mjtNum> GetCtrlNoise(const mjModel* m, int nsteps,
+                                 mjtNum ctrlnoise) {
+  std::vector<mjtNum> ctrl;
+  for (int step=0; step < nsteps; 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 *= ctrlnoise;
+      ctrl.push_back(center + radius * (2 * mju_Halton(step, i+2) - 1));
     }
-    radius *= 0.01;
-    d->ctrl[i] = center + radius * (2 * mju_Halton(step, i + 2) - 1);
   }
+  return ctrl;
 }
 
 }  // namespace mujoco
diff --git a/test/fixture.h b/test/fixture.h
index c32c200c..755c3baf 100644
--- a/test/fixture.h
+++ b/test/fixture.h
@@ -62,7 +62,8 @@ mjModel* LoadModelFromPath(const char* model_path);
 const std::string SaveAndReadXml(const mjModel* model);
 
 // Adds control noise.
-void AddCtrlNoise(const mjModel* m, mjData* d, int step);
+std::vector<mjtNum> GetCtrlNoise(const mjModel* m, int nsteps,
+                                 mjtNum ctrlnoise = 0.01);
 
 }  // namespace mujoco
 #endif  // MUJOCO_TEST_FIXTURE_H_