From a8db22f0d077aee86f771808fd24ea60a148c93f Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Thu, 28 Mar 2024 08:56:26 -0700
Subject: [PATCH] Add local vector clamping function in engine_forward.c

PiperOrigin-RevId: 619948315
Change-Id: Ie4b92f5c8facaa2c169f01982940e84a0248b9c9
---
 src/engine/engine_forward.c | 49 ++++++++++++++++---------------------
 1 file changed, 21 insertions(+), 28 deletions(-)

diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c
index 5f51c5a9..bcaa55ff 100644
--- a/src/engine/engine_forward.c
+++ b/src/engine/engine_forward.c
@@ -243,6 +243,19 @@ static mjtNum nextActivation(const mjModel* m, const mjData* d,
 
 
 
+// clamp vector to range
+static void mju_clamp(mjtNum* vec, const mjtNum* range, const mjtByte* limited, int n,
+                      const int* index) {
+  for (int i=0; i < n; i++) {
+    int j = index ? index[i] : i;
+    if (limited[i]) {
+      vec[j] = mju_clip(vec[j], range[2*i], range[2*i + 1]);
+    }
+  }
+}
+
+
+
 // (qpos, qvel, ctrl, act) => (qfrc_actuator, actuator_force, act_dot)
 void mj_fwdActuation(const mjModel* m, mjData* d) {
   TM_START;
@@ -262,18 +275,9 @@ void mj_fwdActuation(const mjModel* m, mjData* d) {
   // local, clamped copy of ctrl
   mj_markStack(d);
   mjtNum *ctrl = mj_stackAllocNum(d, nu);
-  if (mjDISABLED(mjDSBL_CLAMPCTRL)) {
-    mju_copy(ctrl, d->ctrl, nu);
-  } else {
-    for (int i=0; i < nu; i++) {
-      // clamp ctrl
-      if (m->actuator_ctrllimited[i]) {
-        mjtNum *ctrlrange = m->actuator_ctrlrange + 2*i;
-        ctrl[i] = mju_clip(d->ctrl[i], ctrlrange[0], ctrlrange[1]);
-      } else {
-        ctrl[i] = d->ctrl[i];
-      }
-    }
+  mju_copy(ctrl, d->ctrl, nu);
+  if (!mjDISABLED(mjDSBL_CLAMPCTRL)) {
+    mju_clamp(ctrl, m->actuator_ctrlrange, m->actuator_ctrllimited, nu, NULL);
   }
 
   // check controls, set all to 0 if any are bad
@@ -459,25 +463,13 @@ void mj_fwdActuation(const mjModel* m, mjData* d) {
   }
 
   // clamp actuator_force
-  for (int i=0; i < nu; i++) {
-    if (m->actuator_forcelimited[i]) {
-      mjtNum *forcerange = m->actuator_forcerange + 2*i;
-      force[i] = mju_clip(force[i], forcerange[0], forcerange[1]);
-    }
-  }
+  mju_clamp(force, m->actuator_forcerange, m->actuator_forcelimited, nu, NULL);
 
   // qfrc_actuator = moment' * force
   mju_mulMatTVec(d->qfrc_actuator, moment, force, nu, nv);
 
   // clamp qfrc_actuator
-  int njnt = m->njnt;
-  for (int i=0; i < njnt; i++) {
-    if (m->jnt_actfrclimited[i]) {
-      mjtNum *forcerange = m->jnt_actfrcrange + 2*i;
-      mjtNum *qfrc = d->qfrc_actuator + m->jnt_dofadr[i];
-      qfrc[0] = mju_clip(qfrc[0], forcerange[0], forcerange[1]);
-    }
-  }
+  mju_clamp(d->qfrc_actuator, m->jnt_actfrcrange, m->jnt_actfrclimited, m->njnt, m->jnt_dofadr);
 
   mj_freeStack(d);
   TM_END(mjTIMER_ACTUATION);
@@ -706,9 +698,10 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
 // advance state and time given activation derivatives, acceleration, and optional velocity
 static void mj_advance(const mjModel* m, mjData* d,
                        const mjtNum* act_dot, const mjtNum* qacc, const mjtNum* qvel) {
-  // advance activations and clamp
+  // advance activations
   if (m->na && !mjDISABLED(mjDSBL_ACTUATION)) {
-    for (int i=0; i < m->nu; i++) {
+    int nu = m->nu;
+    for (int i=0; i < nu; i++) {
       int actadr = m->actuator_actadr[i];
       int actadr_end = actadr + m->actuator_actnum[i];
       for (int j=actadr; j < actadr_end; j++) {