From 236dcb4ffd57a534adc1869a120472fb88fe26bc Mon Sep 17 00:00:00 2001
From: Nimrod Gileadi <nimrod@google.com>
Date: Mon, 4 Dec 2023 08:30:41 -0800
Subject: [PATCH] Fix a bug with PID actuators, when plugin definition order
 differs from actuator order.

PiperOrigin-RevId: 587729879
Change-Id: If71c7e374f36ef08c5147e7316066d6f5365925f
---
 plugin/actuator/pid.cc           |  8 ++++----
 plugin/actuator/pid.h            |  2 +-
 test/plugin/actuator/pid_test.cc | 12 ++++++++----
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/plugin/actuator/pid.cc b/plugin/actuator/pid.cc
index 54260fe01d..62fba64733 100644
--- a/plugin/actuator/pid.cc
+++ b/plugin/actuator/pid.cc
@@ -160,7 +160,7 @@ mjtNum Pid::GetCtrl(const mjModel* m, const mjData* d, const State& state,
 }
 
 void Pid::ActDot(const mjModel* m, mjData* d, int instance) const {
-  State state = GetState(m, d, instance);
+  State state = GetState(m, d, actuator_idx_);
   mjtNum ctrl = GetCtrl(m, d, state, /*actearly=*/false);
   mjtNum error = ctrl - d->actuator_length[actuator_idx_];
 
@@ -180,7 +180,7 @@ void Pid::ActDot(const mjModel* m, mjData* d, int instance) const {
 }
 
 void Pid::Compute(const mjModel* m, mjData* d, int instance) {
-  State state = GetState(m, d, instance);
+  State state = GetState(m, d, actuator_idx_);
   mjtNum ctrl = GetCtrl(m, d, state, m->actuator_actearly[actuator_idx_]);
 
   mjtNum error = ctrl - d->actuator_length[actuator_idx_];
@@ -217,9 +217,9 @@ int Pid::ActDim(const mjModel* m, int instance, int actuator_id) {
   return (i_gain ? 1 : 0) + (HasSlew(m, instance) ? 1 : 0);
 }
 
-Pid::State Pid::GetState(const mjModel* m, mjData* d, int instance) const {
+Pid::State Pid::GetState(const mjModel* m, mjData* d, int actuator_idx) const {
   State state;
-  int state_idx = m->actuator_actadr[instance];
+  int state_idx = m->actuator_actadr[actuator_idx];
   if (config_.i_gain) {
     state.integral = d->act[state_idx++];
   }
diff --git a/plugin/actuator/pid.h b/plugin/actuator/pid.h
index 1e2c5f24db..185e07f3b1 100644
--- a/plugin/actuator/pid.h
+++ b/plugin/actuator/pid.h
@@ -84,7 +84,7 @@ class Pid {
     mjtNum integral = 0;
   };
   // Reads data from d->act and returns it as a State struct.
-  State GetState(const mjModel* m, mjData* d, int instance) const;
+  State GetState(const mjModel* m, mjData* d, int actuator_idx) const;
 
   // Returns the PID setpoint, which is normally d->ctrl, but can be d->act for
   // actuators with dyntype != none.
diff --git a/test/plugin/actuator/pid_test.cc b/test/plugin/actuator/pid_test.cc
index 5e13c43976..06938d2f08 100644
--- a/test/plugin/actuator/pid_test.cc
+++ b/test/plugin/actuator/pid_test.cc
@@ -322,10 +322,6 @@ TEST_F(PidTest, ITerm) {
   <mujoco>
   <extension>
     <plugin plugin="mujoco.pid">
-      <instance name="pid1">
-        <config key="kp" value="40.0"/>
-        <config key="kd" value="4"/>
-      </instance>
       <instance name="pid2">
         <config key="kp" value="40.0"/>
         <config key="ki" value="40"/>
@@ -337,6 +333,14 @@ TEST_F(PidTest, ITerm) {
         <config key="kd" value="4"/>
         <config key="imax" value="5"/>
       </instance>
+      <!-- intentionally put the plugin instances out of order, to test
+           any accidentaly dependency on the plugin definition order matching
+           the actuator order.
+      -->
+      <instance name="pid1">
+        <config key="kp" value="40.0"/>
+        <config key="kd" value="4"/>
+      </instance>
     </plugin>
   </extension>