Explicitly specify actdim on PID plugin actuators, to allow them to be used with keyframes.

Support actdim attribute on plugin actuators. Remove the actuator_actdim callback. This callback leads to crashes when loading keyframes for models that have a stateful plugin. PiperOrigin-RevId: 650649505 Change-Id: I40d1994d02ae4e20ce631b2702a984766cf2628e
2024-07-09 09:07:44 -07:00
parent e0d33eb1d0
commit 393895bb29
16 changed files with 201 additions and 101 deletions
@@ -5666,6 +5666,8 @@ Associate this actuator with an :ref:`engine plugin<exPlugin>`. Either :at:`plug

 .. _actuator-plugin-user:

+.. _actuator-plugin-actdim:
+
 .. _actuator-plugin-dynprm:

 .. _actuator-plugin-actearly:
@@ -5673,7 +5675,7 @@ Associate this actuator with an :ref:`engine plugin<exPlugin>`. Either :at:`plug
 .. |actuator/plugin attrib list| replace:: :at:`name`, :at:`class`, :at:`group`, :at:`actlimited`, :at:`ctrllimited`,
   :at:`forcelimited`, :at:`ctrlrange`, :at:`forcerange`, :at:`lengthrange`, :at:`gear`, :at:`cranklength`,
   :at:`joint`, :at:`jointinparent`, :at:`site`, :at:`tendon`, :at:`cranksite`, :at:`slidersite`, :at:`user`,
-   :at:`dynprm`, :at:`actearly`
+   :at:`actdim`, :at:`dynprm`, :at:`actearly`

 |actuator/plugin attrib list|
   Same as in actuator/ :ref:`general <actuator-general>`.
@@ -816,9 +816,11 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`gear<actuator-plugin-gear>`                               | :ref:`cranklength<actuator-plugin-cranklength>`                 | :ref:`joint<actuator-plugin-joint>`                             | :ref:`jointinparent<actuator-plugin-jointinparent>`             | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`site<actuator-plugin-site>`                               | :ref:`dyntype<actuator-plugin-dyntype>`                         | :ref:`dynprm<actuator-plugin-dynprm>`                           | :ref:`tendon<actuator-plugin-tendon>`                           | |
+|                                    |    |    | :ref:`site<actuator-plugin-site>`                               | :ref:`actdim<actuator-plugin-actdim>`                           | :ref:`dyntype<actuator-plugin-dyntype>`                         | :ref:`dynprm<actuator-plugin-dynprm>`                           | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`cranksite<actuator-plugin-cranksite>`                     | :ref:`slidersite<actuator-plugin-slidersite>`                   | :ref:`user<actuator-plugin-user>`                               | :ref:`actearly<actuator-plugin-actearly>`                       | |
+|                                    |    |    | :ref:`tendon<actuator-plugin-tendon>`                           | :ref:`cranksite<actuator-plugin-cranksite>`                     | :ref:`slidersite<actuator-plugin-slidersite>`                   | :ref:`user<actuator-plugin-user>`                               | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`actearly<actuator-plugin-actearly>`                       |                                                                 |                                                                 |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_2| plugin |br| |_2| |L|          |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -22,10 +22,15 @@ General
      :ref:`mju_mulMatTVec3`. These function names and argument order are more consistent with the rest of the API.
      The older functions have been removed from the Python bindings and will be removed from the C API in the next
      release.
+   4. Removed the ``actuator_actdim`` callback from actuator plugins. They now have the ``actdim`` attribute, which
+      must be used with actuators that write state to the ``act`` array. This fixed a crash which happend when
+      keyframes were used in a model with stateful actuator plugins. The PID plugin will give an error when the wrong
+      value of actdim is provided.

-4. The :ref:`VFS<Virtualfilesystem>` implementation has been rewritten in C++ and is now considerably more efficient in
+5. The :ref:`VFS<Virtualfilesystem>` implementation has been rewritten in C++ and is now considerably more efficient in
   speed and memory footprint.
-5. Added a new API for :doc:`procedural model manipulation<programming/modeledit>`. Fixes :github:issue:`364`.
+
+6. Added a new API for :doc:`procedural model manipulation<programming/modeledit>`. Fixes :github:issue:`364`.
   Still missing:

   - Detailed documentation.
@@ -34,27 +39,27 @@ General
   :align: right
   :width: 240px

-6. Added support for orthographic cameras. This is available for both fixed cameras and the free camera, using the
+7. Added support for orthographic cameras. This is available for both fixed cameras and the free camera, using the
   :ref:`camera/orthographic<body-camera-orthographic>` and :ref:`global/orthographic<visual-global-orthographic>`
   attributes, respectively.
-7. Added :ref:`maxhullvert<asset-mesh-maxhullvert>`, the maximum number of vertices in a mesh's convex hull.
-8. Added :ref:`mj_setKeyframe` for saving the current state into a model keyframe.
-9. Added support for ``ball`` joints in the URDF parser ("spherical" in URDF).
-10. Replaced ``mjUSEDOUBLE`` which was previously hard-coded in
+8. Added :ref:`maxhullvert<asset-mesh-maxhullvert>`, the maximum number of vertices in a mesh's convex hull.
+9. Added :ref:`mj_setKeyframe` for saving the current state into a model keyframe.
+10. Added support for ``ball`` joints in the URDF parser ("spherical" in URDF).
+11. Replaced ``mjUSEDOUBLE`` which was previously hard-coded in
    `mjtnum.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtnum.h>`__
    with the build-time flag ``mjUSESINGLE``. If this symbol is not defined, MuJoCo will use double-precision floating
    point, as usual. If ``mjUSESINGLE`` is defined, MuJoCo will use single-precision floating point. See :ref:`mjtNum`.

    Relatedly, fixed various type errors that prevented building with single-precision.
-11. Quaternions in ``mjData->qpos`` and ``mjData->mocap_quat`` are no longer normalized in-place by
+12. Quaternions in ``mjData->qpos`` and ``mjData->mocap_quat`` are no longer normalized in-place by
    :ref:`mj_kinematics`. Instead they are normalized when they are used. After the first step, quaternions in
    ``mjData->qpos`` will be normalized.

 MJX
 ~~~
-12. Added support for :ref:`elliptic friction cones<option-cone>`.
-13. Fixed a bug that resulted in less-optimal linesearch solutions for some difficult constraint settings.
-14. Fixed a bug in the Newton solver that sometimes resulted in less-optimal gradients.
+13. Added support for :ref:`elliptic friction cones<option-cone>`.
+14. Fixed a bug that resulted in less-optimal linesearch solutions for some difficult constraint settings.
+15. Fixed a bug in the Newton solver that sometimes resulted in less-optimal gradients.


 .. youtube:: P83tKA1iz2Y
@@ -63,14 +68,14 @@ MJX

 Simulate
 ^^^^^^^^
-15. Added improved tutorial video.
-16. Improved the Brownian noise generator.
+16. Added improved tutorial video.
+17. Improved the Brownian noise generator.

 |br| |br| |br| |br|

 Python bindings
 ^^^^^^^^^^^^^^^
-17. Fixed a memory leak when using ``copy.deepcopy()`` on a ``mujoco.MjData`` instance (:github:issue:`1572`).
+18. Fixed a memory leak when using ``copy.deepcopy()`` on a ``mujoco.MjData`` instance (:github:issue:`1572`).

 Version 3.1.6 (Jun 3, 2024)
 ---------------------------
@@ -1445,9 +1445,6 @@ struct mjpPlugin_ {

  // methods specific to actuators (optional)

-  // dimension of the actuator state for the plugin (excluding state from actuator's dyntype)
-  int (*actuator_actdim)(const mjModel*m, int instance, int actuator_id);
-
  // updates the actuator plugin's entries in act_dot
  // called after native act_dot is computed and before the compute callback
  void (*actuator_act_dot)(const mjModel* m, mjData* d, int instance);
@@ -2117,7 +2114,6 @@ typedef struct mjsActuator_ {      // actuator specification
  mjtDyn dyntype;                  // dynamics type
  double dynprm[mjNDYN];           // dynamics parameters
  int actdim;                      // number of activation variables
-  int plugin_actdim;               // actuator state size for plugins
  mjtByte actearly;                // apply next activations to qfrc

  // transmission
@@ -112,9 +112,6 @@ struct mjpPlugin_ {

  // methods specific to actuators (optional)

-  // dimension of the actuator state for the plugin (excluding state from actuator's dyntype)
-  int (*actuator_actdim)(const mjModel*m, int instance, int actuator_id);
-
  // updates the actuator plugin's entries in act_dot
  // called after native act_dot is computed and before the compute callback
  void (*actuator_act_dot)(const mjModel* m, mjData* d, int instance);
@@ -628,7 +628,6 @@ typedef struct mjsActuator_ {      // actuator specification
  mjtDyn dyntype;                  // dynamics type
  double dynprm[mjNDYN];           // dynamics parameters
  int actdim;                      // number of activation variables
-  int plugin_actdim;               // actuator state size for plugins
  mjtByte actearly;                // apply next activations to qfrc

  // transmission
@@ -10528,11 +10528,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                 type=ValueType(name='int'),
                 doc='number of activation variables',
             ),
-             StructFieldDecl(
-                 name='plugin_actdim',
-                 type=ValueType(name='int'),
-                 doc='actuator state size for plugins',
-             ),
             StructFieldDecl(
                 name='actearly',
                 type=ValueType(name='mjtByte'),
@@ -62,7 +62,7 @@
  <actuator>
    <plugin joint="j1" plugin="mujoco.pid" instance="pid1" ctrlrange="0 0.15" />
    <plugin joint="j2" plugin="mujoco.pid" instance="pid2" ctrlrange="0 0.15" />
-    <plugin joint="j3" plugin="mujoco.pid" instance="pid3" ctrlrange="0 0.15" />
-    <plugin joint="j4" plugin="mujoco.pid" instance="pid4" ctrlrange="0 0.15" />
+    <plugin joint="j3" plugin="mujoco.pid" instance="pid3" ctrlrange="0 0.15" actdim="1" />
+    <plugin joint="j4" plugin="mujoco.pid" instance="pid4" ctrlrange="0 0.15" actdim="1" />
  </actuator>
 </mujoco>
@@ -117,6 +117,23 @@ std::unique_ptr<Pid> Pid::Create(const mjModel* m, int instance) {
    mju_warning("actuator not found for plugin instance %d", instance);
    return nullptr;
  }
+  // Validate actnum values for all actuators:
+  for (int actuator_id : actuators) {
+    int actnum = m->actuator_actnum[actuator_id];
+    int expected_actnum = Pid::ActDim(m, instance, actuator_id);
+    int dyntype = m->actuator_dyntype[actuator_id];
+    if (dyntype == mjDYN_FILTER || dyntype == mjDYN_FILTEREXACT ||
+        dyntype == mjDYN_INTEGRATOR) {
+      expected_actnum++;
+    }
+    if (actnum != expected_actnum) {
+      mju_warning(
+          "actuator %d has actdim %d, expected %d. Add actdim=\"%d\" to the "
+          "actuator plugin element.",
+          actuator_id, actnum, expected_actnum, expected_actnum);
+      return nullptr;
+    }
+  }
  return std::unique_ptr<Pid>(new Pid(config, std::move(actuators)));
 }

@@ -240,8 +257,6 @@ void Pid::RegisterPlugin() {
                                         kAttrIMax, kAttrSlewMax};
  plugin.nattribute = attributes.size();
  plugin.attributes = attributes.data();
-
-  plugin.actuator_actdim = Pid::ActDim;
  plugin.nstate = Pid::StateSize;

  plugin.init = +[](const mjModel* m, mjData* d, int instance) {
@@ -52,9 +52,6 @@ class Pid {
  // Returns the number of state variables for the plugin instance
  static int StateSize(const mjModel* m, int instance);

-  // Returns the number of activation variables for the plugin instance
-  static int ActDim(const mjModel* m, int instance, int actuator_id);
-
  // Resets the C++ Pid instance's state.
  // plugin_state is a C array pointer into mjData->plugin_state, with a size
  // equal to the value returned from StateSize.
@@ -76,6 +73,9 @@ class Pid {
 private:
  Pid(PidConfig config, std::vector<int> actuators);

+  // Returns the expected number of activation variables for the instance.
+  static int ActDim(const mjModel* m, int instance, int actuator_id);
+
  struct State {
    mjtNum previous_ctrl = 0;
    // if using slew rate limits, mjData.act will contain an activation variable
@@ -1343,7 +1343,7 @@ void mjCModel::SetSizes() {
  // nu, na
  for (int i=0; i<actuators_.size(); i++) {
    nu++;
-    na += actuators_[i]->actdim + actuators_[i]->plugin_actdim;
+    na += actuators_[i]->actdim;
  }

  // nbvh, nbvhstatic, nbvhdynamic
@@ -2639,7 +2639,7 @@ void mjCModel::CopyObjects(mjModel* m) {
    m->actuator_biastype[i] = pac->biastype;
    m->actuator_trnid[2*i] = pac->trnid[0];
    m->actuator_trnid[2*i+1] = pac->trnid[1];
-    m->actuator_actnum[i] = pac->actdim + pac->plugin_actdim;
+    m->actuator_actnum[i] = pac->actdim;
    m->actuator_actadr[i] = m->actuator_actnum[i] ? adr : -1;
    pac->actadr_ = m->actuator_actadr[i];
    pac->actnum_ = m->actuator_actnum[i];
@@ -3535,13 +3535,6 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
          this->nsensordata += nsensordata;
        }
      }
-      if ((plugin->capabilityflags & mjPLUGIN_ACTUATOR) && plugin->actuator_actdim) {
-        for (int actuator_id : plugin_to_actuators[i]) {
-          int plugin_actdim = plugin->actuator_actdim(m, i, actuator_id);
-          actuators_[actuator_id]->plugin_actdim = plugin_actdim;
-          this->na += plugin_actdim;
-        }
-      }
    }
    m->npluginstate = stateadr;
  }
@@ -5208,14 +5208,16 @@ void mjCActuator::Compile(void) {
  }

  // check and set actdim
-  if (actdim > 1 && dyntype != mjDYN_USER) {
-    throw mjCError(this, "actdim > 1 is only allowed for dyntype 'user' in actuator");
-  }
-  if (actdim == 1 && dyntype == mjDYN_NONE) {
-    throw mjCError(this, "invalid actdim 1 in stateless actuator");
-  }
-  if (actdim == 0 && dyntype != mjDYN_NONE) {
-    throw mjCError(this, "invalid actdim 0 in stateful actuator");
+  if (!plugin.active) {
+    if (actdim > 1 && dyntype != mjDYN_USER) {
+      throw mjCError(this, "actdim > 1 is only allowed for dyntype 'user' in actuator");
+    }
+    if (actdim == 1 && dyntype == mjDYN_NONE) {
+      throw mjCError(this, "invalid actdim 1 in stateless actuator");
+    }
+    if (actdim == 0 && dyntype != mjDYN_NONE) {
+      throw mjCError(this, "invalid actdim 0 in stateful actuator");
+    }
  }

  // set actdim
@@ -423,10 +423,11 @@ const char* MJCF[nMJCF][mjXATTRNUM] = {
            "lmin", "lmax", "vmax", "fpmax", "fvmax"},
        {"adhesion", "*", "9", "name", "class", "group",
            "forcelimited", "ctrlrange", "forcerange", "user", "body", "gain"},
-        {"plugin", "*", "24", "name", "class",  "plugin", "instance", "group",
+        {"plugin", "*", "25", "name", "class",  "plugin", "instance", "group",
            "ctrllimited", "forcelimited", "actlimited", "ctrlrange", "forcerange", "actrange",
            "lengthrange", "gear", "cranklength", "joint", "jointinparent",
-            "site", "dyntype", "dynprm", "tendon", "cranksite", "slidersite", "user", "actearly"},
+            "site", "actdim", "dyntype", "dynprm", "tendon", "cranksite", "slidersite", "user",
+            "actearly"},
        {"<"},
          {"config", "*", "2", "key", "value"},
        {">"},
@@ -2234,6 +2235,7 @@ void mjXReader::OneActuator(XMLElement* elem, mjsActuator* pact) {
      pact->actearly = (n==1);
    }
    ReadAttr(elem, "dynprm", mjNDYN, pact->dynprm, text, false, false);
+    ReadAttrInt(elem, "actdim", &pact->actdim);
  }

  else {          // SHOULD NOT OCCUR
@@ -757,6 +757,13 @@ void mjXWriter::OneActuator(XMLElement* elem, const mjCActuator* pact, mjCDef* d
  WriteAttr(elem, "cranklength", 1, &pact->cranklength, &def->Actuator().cranklength);
  WriteAttrKey(elem, "actearly", bool_map, 2, pact->actearly,
               def->Actuator().actearly);
+  // special handling of actdim which has default value of -1
+  if (writingdefaults) {
+    WriteAttrInt(elem, "actdim", pact->actdim, def->Actuator().actdim);
+  } else {
+    int default_actdim = pact->dyntype == mjDYN_NONE ? 0 : 1;
+    WriteAttrInt(elem, "actdim", pact->actdim, default_actdim);
+  }
  WriteAttrKey(elem, "dyntype", dyn_map, dyn_sz, pact->dyntype, def->Actuator().dyntype);
  WriteAttr(elem, "dynprm", mjNDYN, pact->dynprm, def->Actuator().dynprm);

@@ -767,13 +774,6 @@ void mjXWriter::OneActuator(XMLElement* elem, const mjCActuator* pact, mjCDef* d

  // non-plugins: write actuator parameters
  else {
-    // special handling of actdim which has default value of -1
-    if (writingdefaults) {
-      WriteAttrInt(elem, "actdim", pact->actdim, def->Actuator().actdim);
-    } else {
-      int default_actdim = pact->dyntype == mjDYN_NONE ? 0 : 1;
-      WriteAttrInt(elem, "actdim", pact->actdim, default_actdim);
-    }
    WriteAttrKey(elem, "gaintype", gain_map, gain_sz, pact->gaintype, def->Actuator().gaintype);
    WriteAttrKey(elem, "biastype", bias_map, bias_sz, pact->biastype, def->Actuator().biastype);
    WriteAttr(elem, "gainprm", mjNGAIN, pact->gainprm, def->Actuator().gainprm, true);
@@ -149,12 +149,6 @@ class TestActuator : public BaseTestPlugin {
    } else {
      multiplier = kDefaultMultiplier;
    }
-    s = mj_getPluginConfig(m, instance, "actdim");
-    if (*s) {
-      std::stringstream(s) >> actdim_;
-    } else {
-      actdim_ = 0;
-    }
    for (int i = 0; i < m->nu; ++i) {
      if (m->actuator_plugin[i] == instance) {
        actuators.push_back(&d->actuator_force[i]);
@@ -187,7 +181,13 @@ class TestActuator : public BaseTestPlugin {
      if (m->actuator_plugin[i] != instance_) {
        continue;
      }
-      mju_fill(d->act_dot + m->actuator_actadr[i], kActDotValue, actdim_);
+      // If the user specified an actdim, fill the act_dot values that belong
+      // to the plugin.
+      int actnum = m->actuator_actnum[i];
+      if (m->actuator_dyntype[i] != mjDYN_NONE) {
+        actnum--;
+      }
+      mju_fill(d->act_dot + m->actuator_actadr[i], kActDotValue, actnum);
    }
  }

@@ -195,7 +195,6 @@ class TestActuator : public BaseTestPlugin {
  mjtNum multiplier;
  std::vector<mjtNum*> actuators;
  int instance_;
-  int actdim_;

  void WriteActuatorForce() {
    for (mjtNum* actuator_force : actuators) {
@@ -263,7 +262,7 @@ int RegisterActuatorPlugin() {

  plugin.name = "mujoco.test.actuator";

-  const char* attributes[] = {"stride", "multiplier", "actdim"};
+  const char* attributes[] = {"stride", "multiplier"};
  plugin.nattribute = sizeof(attributes) / sizeof(*attributes);
  plugin.attributes = attributes;

@@ -296,16 +295,6 @@ int RegisterActuatorPlugin() {
    auto actuator = reinterpret_cast<TestActuator*>(d->plugin_data[instance]);
    actuator->Advance();
  };
-  plugin.actuator_actdim =
-      +[](const mjModel* m, int instance, int actuator_id) {
-    const char* actdim_str = mj_getPluginConfig(m, instance, "actdim");
-    if (actdim_str) {
-      int actdim = 0;
-      std::stringstream(actdim_str) >> actdim;
-      return actdim;
-    }
-    return 0;
-  };
  plugin.actuator_act_dot = +[](const mjModel* m, mjData* d, int instance) {
    auto actuator = reinterpret_cast<TestActuator*>(d->plugin_data[instance]);
    actuator->ActDot(m, d);
@@ -383,7 +372,6 @@ constexpr char xml[] = R"(
      <instance name="actuator2">
        <config key="stride" value="2"/>
        <config key="multiplier" value="0.125"/>
-        <config key="actdim" value="3"/>
      </instance>
    </plugin>
    <plugin plugin="mujoco.test.passive"/>
@@ -416,8 +404,8 @@ constexpr char xml[] = R"(
      <config key="stride" value="4"/>
      <config key="multiplier" value="0.03125"/>
    </plugin>
-    <plugin joint="h2" instance="actuator2"/>
-    <plugin joint="h2" dyntype="filter" dynprm="0.9" instance="actuator2"/>
+    <plugin joint="h2" actdim="3" instance="actuator2"/>
+    <plugin joint="h2" actdim="4" dyntype="filter" dynprm="0.9" instance="actuator2"/>
  </actuator>
 </mujoco>
 )";
@@ -640,7 +628,7 @@ TEST_F(EnginePluginTest, FilteredActuatorPlugin) {

  // Expecting 7 actuator state variables: 3x2 from actuator2 instances, and 1
  // from setting dyntype="filter" on one of the plugin actuators
-  EXPECT_EQ(m->na, 7);
+  ASSERT_EQ(m->na, 7);
  EXPECT_EQ(m->actuator_actnum[0], 0);
  EXPECT_EQ(m->actuator_actnum[1], 3);
  EXPECT_EQ(m->actuator_actnum[2], 4);
@@ -14,12 +14,14 @@

 // Tests for the PID controller plugin

+#include <string>
 #include <string_view>
 #include <vector>

 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <absl/cleanup/cleanup.h>
+#include <absl/strings/str_replace.h>
 #include <absl/strings/string_view.h>
 #include <mujoco/mujoco.h>
 #include "test/fixture.h"
@@ -109,7 +111,7 @@ TEST_F(PidTest, PGainWithFilterExact) {
  </worldbody>

  <actuator>
-    <plugin joint="j1" plugin="mujoco.pid" instance="pid1"
+    <plugin joint="j1" plugin="mujoco.pid" instance="pid1" actdim="1"
        dyntype="filterexact" dynprm="0.1" actearly="true"/>
    <general joint="j2" gainprm="4.0 0 0" biastype="affine" biasprm="0 -4.0 0"
        dyntype="filterexact" dynprm="0.1" actearly="true"/>
@@ -169,7 +171,7 @@ TEST_F(PidTest, SlewMaxRate) {
  </worldbody>

  <actuator>
-    <plugin joint="j1" plugin="mujoco.pid" instance="pid1" />
+    <plugin joint="j1" plugin="mujoco.pid" instance="pid1" actdim="1" />
    <position joint="j2" kp="4.0" />
  </actuator>
  </mujoco>
@@ -235,10 +237,10 @@ TEST_F(PidTest, IntegratedVelocitySlewMaxRate) {
  </worldbody>

  <actuator>
-    <plugin joint="j1" plugin="mujoco.pid" instance="pid1" />
+    <plugin joint="j1" plugin="mujoco.pid" instance="pid1" actdim="1"/>
    <!-- make an integrated velocity controller using the PID plugin -->
    <plugin joint="j2" plugin="mujoco.pid" instance="pid2"
-        dyntype="integrator" dynprm="1 0 0" actearly="true"/>
+        dyntype="integrator" dynprm="1 0 0" actearly="true" actdim="1"/>
  </actuator>
  </mujoco>
  )";
@@ -250,9 +252,9 @@ TEST_F(PidTest, IntegratedVelocitySlewMaxRate) {

  // having a slew rate means that there should be one extra state variable
  // for the plugin.
-  EXPECT_EQ(m->actuator_actnum[0], 1);
+  ASSERT_EQ(m->actuator_actnum[0], 1);
  // The integrated-velocity controller should have one activation variable too.
-  EXPECT_EQ(m->actuator_actnum[1], 1);
+  ASSERT_EQ(m->actuator_actnum[1], 1);

  mjData* d = mj_makeData(m);
  absl::Cleanup d_deleter = [d] { mj_deleteData(d); };
@@ -299,7 +301,9 @@ TEST_F(PidTest, SlewMaxRateUsesFirstCtrl) {
    </body>
  </worldbody>

-  <actuator><plugin joint="j1" plugin="mujoco.pid" instance="pid1" /></actuator>
+  <actuator>
+    <plugin joint="j1" plugin="mujoco.pid" instance="pid1" actdim="1"/>
+  </actuator>
  </mujoco>
  )";

@@ -363,8 +367,8 @@ TEST_F(PidTest, ITerm) {

  <actuator>
    <plugin joint="j1" plugin="mujoco.pid" instance="pid1" />
-    <plugin joint="j2" plugin="mujoco.pid" instance="pid2" />
-    <plugin joint="j3" plugin="mujoco.pid" instance="pid3" />
+    <plugin joint="j2" plugin="mujoco.pid" instance="pid2" actdim="1" />
+    <plugin joint="j3" plugin="mujoco.pid" instance="pid3" actdim="1" />
  </actuator>
  </mujoco>
  )";
@@ -413,7 +417,7 @@ TEST_F(PidTest, FiniteDifferencing) {
  </worldbody>

  <actuator>
-    <plugin joint="j" plugin="mujoco.pid" instance="pid" />
+    <plugin joint="j" plugin="mujoco.pid" instance="pid" actdim="2" />
  </actuator>
  </mujoco>
  )";
@@ -425,7 +429,7 @@ TEST_F(PidTest, FiniteDifferencing) {

  // actuators with an I term and max slew rate should have 2 activation
  // variables.
-  EXPECT_EQ(m->actuator_actnum[0], 2);
+  ASSERT_EQ(m->actuator_actnum[0], 2);

  mjData* d = mj_makeData(m);
  absl::Cleanup d_deleter = [d] { mj_deleteData(d); };
@@ -471,7 +475,7 @@ TEST_F(PidTest, CtrlClamp) {

  <actuator>
    <plugin joint="j1" plugin="mujoco.pid" instance="pid"
-        ctrlrange="0.25 0.75" />
+        ctrlrange="0.25 0.75" actdim="1" />
  </actuator>
  </mujoco>
  )";
@@ -534,7 +538,7 @@ TEST_F(PidTest, CopyData) {

  <actuator>
    <plugin joint="j1" plugin="mujoco.pid" instance="pid1" />
-    <plugin joint="j2" plugin="mujoco.pid" instance="pid2" />
+    <plugin joint="j2" plugin="mujoco.pid" instance="pid2" actdim="2"/>
  </actuator>
  </mujoco>
  )";
@@ -588,8 +592,8 @@ TEST_F(PidTest, MultipleActuatorsSamePlugin) {
  </worldbody>

  <actuator>
-    <plugin joint="j1" plugin="mujoco.pid" instance="pid" />
-    <plugin joint="j2" plugin="mujoco.pid" instance="pid" />
+    <plugin joint="j1" plugin="mujoco.pid" instance="pid" actdim="1" />
+    <plugin joint="j2" plugin="mujoco.pid" instance="pid" actdim="1" />
  </actuator>
  </mujoco>
  )";
@@ -682,5 +686,105 @@ TEST_F(PidTest, InvalidSlew) {
  EXPECT_THAT(std::string_view(error), HasSubstr("plugin"));
  EXPECT_THAT(std::string_view(error), HasSubstr("slewmax"));
 }
+
+TEST_F(PidTest, WrongActdim) {
+  // XML where PLACEHOLDER is going to be replaced with various things
+  constexpr absl::string_view kBaseXml = R"(
+  <mujoco>
+  <extension>
+    <plugin plugin="mujoco.pid">
+      <instance name="pid1">
+        <config key="kp" value="4.0"/>
+        <config key="slewmax" value="0.75"/>
+      </instance>
+    </plugin>
+  </extension>
+
+  <worldbody>
+    <body>
+      <joint name="j1" type="slide"/>
+      <geom size="0.01"/>
+    </body>
+  </worldbody>
+
+  <actuator>
+    <plugin joint="j1" plugin="mujoco.pid" instance="pid1" PLACEHOLDER />
+  </actuator>
+  </mujoco>
+  )";
+
+  char error[1024] = {0};
+  {
+    std::string no_actdim =
+        absl::StrReplaceAll(kBaseXml, {{"PLACEHOLDER", ""}});
+    mjModel* m = LoadModelFromString(no_actdim, error, sizeof(error));
+    EXPECT_THAT(m, IsNull());
+    EXPECT_THAT(std::string_view(error), HasSubstr("actdim=\"1\""));
+  }
+
+  {
+    std::string big_actdim =
+        absl::StrReplaceAll(kBaseXml, {{"PLACEHOLDER", "actdim=\"2\""}});
+    mjModel* m = LoadModelFromString(big_actdim, error, sizeof(error));
+    EXPECT_THAT(m, IsNull());
+    EXPECT_THAT(std::string_view(error), HasSubstr("actdim=\"1\""));
+  }
+
+  {
+    std::string dyntype_integrator = absl::StrReplaceAll(
+        kBaseXml, {{"PLACEHOLDER", "dyntype=\"integrator\" actdim=\"1\""}});
+    mjModel* m = LoadModelFromString(dyntype_integrator, error, sizeof(error));
+    EXPECT_THAT(m, IsNull());
+    EXPECT_THAT(std::string_view(error), HasSubstr("actdim=\"2\""));
+  }
+}
+
+// Regression test: loading models with PID plugin and keyframes used to crash.
+TEST_F(PidTest, Keyframe) {
+  constexpr absl::string_view kModelXml = R"(
+  <mujoco>
+  <extension>
+    <plugin plugin="mujoco.pid">
+      <instance name="pid1">
+        <config key="kp" value="4.0"/>
+        <config key="slewmax" value="0.75"/>
+      </instance>
+    </plugin>
+  </extension>
+
+  <worldbody>
+    <body>
+      <joint name="j1" type="slide"/>
+      <geom size="0.01"/>
+    </body>
+  </worldbody>
+
+  <actuator>
+    <plugin joint="j1" plugin="mujoco.pid" instance="pid1" actdim="1" />
+  </actuator>
+  <keyframe>
+    <key name="home" qpos="0" act="1" />
+  </keyframe>
+  </mujoco>
+  )";
+
+  char error[1024] = {0};
+  mjModel* m = LoadModelFromString(kModelXml, error, sizeof(error));
+  ASSERT_THAT(m, NotNull()) << error;
+  absl::Cleanup m_deleter = [m] { mj_deleteModel(m); };
+
+  // having a slew rate means that there should be one extra state variable
+  // for the plugin.
+  EXPECT_EQ(m->actuator_actnum[0], 1);
+  EXPECT_EQ(m->na, 1);
+  ASSERT_EQ(m->nkey, 1);
+  EXPECT_EQ(m->key_act[0], 1.0);
+
+  mjData* d = mj_makeData(m);
+  absl::Cleanup d_deleter = [d] { mj_deleteData(d); };
+
+  mj_resetDataKeyframe(m, d, 0);
+  EXPECT_EQ(d->act[0], 1.0);
+}
 }  // namespace
 }  // namespace mujoco