From ea98c57921ec213aed0a61995f65034f7b8888d9 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Wed, 25 Dec 2024 07:43:26 -0800
Subject: [PATCH] Clean up engine_core_smooth_benchmark_test.

PiperOrigin-RevId: 709573203
Change-Id: I461f259c33d74b3eaf0b6fe58124fe09418c5303
---
 .../engine_core_smooth_benchmark_test.cc      | 29 +++++++------------
 1 file changed, 10 insertions(+), 19 deletions(-)

diff --git a/test/benchmark/engine_core_smooth_benchmark_test.cc b/test/benchmark/engine_core_smooth_benchmark_test.cc
index 9a80ae8e..3b728169 100644
--- a/test/benchmark/engine_core_smooth_benchmark_test.cc
+++ b/test/benchmark/engine_core_smooth_benchmark_test.cc
@@ -14,7 +14,6 @@
 
 // A benchmark for comparing different implementations of mj_solveLD.
 
-#include <vector>
 #include <benchmark/benchmark.h>
 #include <absl/base/attributes.h>
 #include <mujoco/mjdata.h>
@@ -25,9 +24,6 @@
 namespace mujoco {
 namespace {
 
-// number of steps to roll out before benchmarking
-static const int kNumWarmupSteps = 200;
-
 // number of steps to benchmark
 static const int kNumBenchmarkSteps = 50;
 
@@ -42,37 +38,31 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
   }
 
   mjData* d = mj_makeData(m);
+  mj_forward(m, d);
 
-  // warm-up rollout to get a typical state
-  for (int i=0; i < kNumWarmupSteps; i++) {
-    mj_step(m, d);
-  }
-
-  // allocate gradient
+  // allocate input and output vectors
   mj_markStack(d);
-  mjtNum *grad = mj_stackAllocNum(d, m->nv);
-  mjtNum *Ma = mj_stackAllocNum(d, m->nv);
+  mjtNum *vec = mj_stackAllocNum(d, m->nv);
   mjtNum *res = mj_stackAllocNum(d, m->nv);
 
-  // compute gradient
-  mj_mulM(m, d, Ma, d->qacc);
+  // arbitrary input vector
   for (int i=0; i < m->nv; i++) {
-    grad[i] = Ma[i] - d->qfrc_smooth[i] - d->qfrc_constraint[i];
+    vec[i] = 0.2 + 0.3*i;
   }
 
-  // CSR matrix
+  // make CSR matrix
   mjtNum* LDs = mj_stackAllocNum(d, m->nC);
   for (int i=0; i < m->nC; i++) {
     LDs[i] = d->qLD[d->mapM2C[i]];
   }
 
-  // reset state, benchmark subsequent kNumBenchmarkSteps steps
+  // benchmark
   while (state.KeepRunningBatch(kNumBenchmarkSteps)) {
     for (int i=0; i < kNumBenchmarkSteps; i++) {
       if (featherstone) {
-        mj_solveM(m, d, res, grad, 1);
+        mj_solveM(m, d, res, vec, 1);
       } else {
-        mju_copy(res, grad, m->nv);
+        mju_copy(res, vec, m->nv);
         mj_solveLDs(res, LDs, d->qLDiagInv, m->nv,
                     d->C_rownnz, d->C_rowadr, d->C_diag, d->C_colind);
       }
@@ -82,6 +72,7 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
   // finalize
   mj_freeStack(d);
   mj_deleteData(d);
+  mj_deleteModel(m);
   state.SetItemsProcessed(state.iterations());
 }