From 7e16ecaa254e69894e41e010c0a4c8227ad8324c Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Thu, 2 Oct 2025 05:36:44 -0700
Subject: [PATCH] Updates for testspeed.cc ctrl noise.

PiperOrigin-RevId: 814192687
Change-Id: I0498c8efc8e71cd578e6b9d40ffa1411447d3e43
---
 sample/testspeed.cc | 59 +++++++++++++++++++++++++++++++++------------
 1 file changed, 44 insertions(+), 15 deletions(-)

diff --git a/sample/testspeed.cc b/sample/testspeed.cc
index c845ffd9..e7b6aeda 100644
--- a/sample/testspeed.cc
+++ b/sample/testspeed.cc
@@ -60,19 +60,40 @@ int finish(const char* msg = NULL, mjModel* m = NULL) {
   return EXIT_SUCCESS;
 }
 
-std::vector<mjtNum> CtrlNoise(const mjModel* m, int nsteps, mjtNum ctrlnoise) {
+std::vector<mjtNum> CtrlNoise(const mjModel* m, int nsteps, mjtNum ctrl_noise_std, mjtNum ctrl_noise_rate, int key) {
   std::vector<mjtNum> ctrl;
+
+  // convert rate and scale to discrete time (Ornstein–Uhlenbeck)
+  mjtNum rate = mju_exp(-m->opt.timestep / ctrl_noise_rate);
+  mjtNum scale = ctrl_noise_std * mju_sqrt(1-rate*rate);
+
   for (int step=0; step < nsteps; step++) {
     for (int i = 0; i < m->nu; i++) {
-      mjtNum center = 0.0;
-      mjtNum radius = 1.0;
+      mjtNum midpoint = 0.0;
+      mjtNum halfrange = 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;
+        midpoint =  0.5 * (range[1] + range[0]);
+        halfrange = 0.5 * (range[1] - range[0]);
       }
-      radius *= ctrlnoise;
-      ctrl.push_back(center + radius * (2 * mju_Halton(step, i+2) - 1));
+
+      // overwrite midpoint with keyframe, if given
+      if (key >= 0) {
+        midpoint = m->key_ctrl[key*m->nu+i];
+      }
+
+      // exponential convergence to midpoint at ctrl_noise_rate
+      mjtNum ctrl_ = step > 0 ? rate * ctrl[(step-1)*m->nu+i] + (1-rate) * midpoint : midpoint;
+
+      // add noise
+      ctrl_ += scale * halfrange * (2 * mju_Halton(step, i+2) - 1);
+
+      // clip to range if limited
+      if (m->actuator_ctrllimited[i]) {
+        ctrl_ = mju_clip(ctrl_, range[0], range[1]);
+      }
+
+      ctrl.push_back(ctrl_);
     }
   }
   return ctrl;
@@ -127,7 +148,8 @@ int main(int argc, char** argv) {
       "  modelfile                 path to model (required)\n"
       "  nstep         10000       number of steps per rollout\n"
       "  nthread       1           number of threads running parallel rollouts\n"
-      "  ctrlnoise     0.01        scale of pseudo-random noise injected into actuators\n"
+      "  ctrlnoisestd  0.01        scale of pseudo-random noise injected into actuators\n"
+      "  ctrlnoiserate 0.1         rate of convergence to ctrl keyframe/midpoint\n"
       "  npoolthread   0           number of threads in engine-internal threadpool\n"
       "\n"
       "Note: If the model has a keyframe named \"test\", it will be loaded prior to simulation\n");
@@ -136,22 +158,29 @@ int main(int argc, char** argv) {
   // read arguments
   int nstep = 10000, nthread = 0, npoolthread = 0;
   // inject small noise by default, to avoid fixed contact state
-  double ctrlnoise = 0.01;
+  double ctrlnoisestd = 0.01;
+  double ctrlnoiserate = 0.1;
   if (argc > 2 && (std::sscanf(argv[2], "%d", &nstep) != 1 || nstep <= 0)) {
     return finish("Invalid nstep argument");
   }
   if (argc > 3 && std::sscanf(argv[3], "%d", &nthread) != 1) {
     return finish("Invalid nthread argument");
   }
-  if (argc > 4 && std::sscanf(argv[4], "%lf", &ctrlnoise) != 1) {
-    return finish("Invalid ctrlnoise argument");
+  if (argc > 4 && std::sscanf(argv[4], "%lf", &ctrlnoisestd) != 1) {
+    return finish("Invalid ctrlnoisestd argument");
   }
-  if (argc > 5 && std::sscanf(argv[5], "%d", &npoolthread) != 1) {
+  if (argc > 5 && std::sscanf(argv[5], "%lf", &ctrlnoiserate) != 1) {
+    return finish("Invalid ctrlnoiserate argument");
+  }
+  if (argc > 6 && std::sscanf(argv[6], "%d", &npoolthread) != 1) {
     return finish("Invalid npoolthread argument");
   }
 
-  // clamp ctrlnoise to [0.0, 1.0]
-  ctrlnoise = mjMAX(0.0, mjMIN(ctrlnoise, 1.0));
+  // clamp ctrlnoisestd to [0.0, 1.0]
+  ctrlnoisestd = mju_clip(ctrlnoisestd, 0.0, 1.0);
+
+  // clamp ctrlnoiserate to [0.0, 1.0]
+  ctrlnoiserate = mju_clip(ctrlnoiserate, 0.0, 1.0);
 
   // clamp nthread to [1, maxthread]
   nthread = mjMAX(1, mjMIN(maxthread, nthread));
@@ -216,7 +245,7 @@ int main(int argc, char** argv) {
   std::printf("...\n\n");
 
   // create pseudo-random control sequence
-  std::vector<mjtNum> ctrl = CtrlNoise(m, nstep, ctrlnoise);
+  std::vector<mjtNum> ctrl = CtrlNoise(m, nstep, ctrlnoisestd, ctrlnoiserate, testkey);
 
   // run simulation, record total time
   std::thread th[maxthread];