chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Clean up engine_core_smooth_benchmark_test.

Browse Source 
PiperOrigin-RevId: 709573203
Change-Id: I461f259c33d74b3eaf0b6fe58124fe09418c5303
This commit is contained in:
Yuval Tassa
2024-12-25 07:43:26 -08:00
committed by Copybara-Service
parent 273ced428a
commit ea98c57921
1 changed files with 10 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/benchmark/engine_core_smooth_benchmark_test.cc
+10 -19
View File
@@ -14,7 +14,6 @@
// A benchmark for comparing different implementations of mj_solveLD. // A benchmark for comparing different implementations of mj_solveLD.
#include <vector>
#include <benchmark/benchmark.h> #include <benchmark/benchmark.h>
#include <absl/base/attributes.h> #include <absl/base/attributes.h>
#include <mujoco/mjdata.h> #include <mujoco/mjdata.h>
@@ -25,9 +24,6 @@
namespace mujoco { namespace mujoco {
namespace { namespace {
// number of steps to roll out before benchmarking
static const int kNumWarmupSteps = 200;
// number of steps to benchmark // number of steps to benchmark
static const int kNumBenchmarkSteps = 50; 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); mjData* d = mj_makeData(m);
mj_forward(m, d);
// warm-up rollout to get a typical state // allocate input and output vectors
for (int i=0; i < kNumWarmupSteps; i++) {
mj_step(m, d);
}
// allocate gradient
mj_markStack(d); mj_markStack(d);
mjtNum *grad = mj_stackAllocNum(d, m->nv); mjtNum *vec = mj_stackAllocNum(d, m->nv);
mjtNum *Ma = mj_stackAllocNum(d, m->nv);
mjtNum *res = mj_stackAllocNum(d, m->nv); mjtNum *res = mj_stackAllocNum(d, m->nv);
// compute gradient // arbitrary input vector
mj_mulM(m, d, Ma, d->qacc);
for (int i=0; i < m->nv; i++) { 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); mjtNum* LDs = mj_stackAllocNum(d, m->nC);
for (int i=0; i < m->nC; i++) { for (int i=0; i < m->nC; i++) {
LDs[i] = d->qLD[d->mapM2C[i]]; LDs[i] = d->qLD[d->mapM2C[i]];
} }
// reset state, benchmark subsequent kNumBenchmarkSteps steps // benchmark
while (state.KeepRunningBatch(kNumBenchmarkSteps)) { while (state.KeepRunningBatch(kNumBenchmarkSteps)) {
for (int i=0; i < kNumBenchmarkSteps; i++) { for (int i=0; i < kNumBenchmarkSteps; i++) {
if (featherstone) { if (featherstone) {
mj_solveM(m, d, res, grad, 1); mj_solveM(m, d, res, vec, 1);
} else { } else {
mju_copy(res, grad, m->nv); mju_copy(res, vec, m->nv);
mj_solveLDs(res, LDs, d->qLDiagInv, m->nv, mj_solveLDs(res, LDs, d->qLDiagInv, m->nv,
d->C_rownnz, d->C_rowadr, d->C_diag, d->C_colind); 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 // finalize
mj_freeStack(d); mj_freeStack(d);
mj_deleteData(d); mj_deleteData(d);
mj_deleteModel(m);
state.SetItemsProcessed(state.iterations()); state.SetItemsProcessed(state.iterations());
} }