Implementation of Stable PD controller: RNE, CRB and code migration

[rohit-kumar-j]

I want to implement stable PD Stable Proportional-Derivative Controllers proposed here, which require the mass matrix and the forces needed to reach the given joint accelerations.

Reason to do this:Using larger time steps while keeping large gains stable throughout the simulation

Would It be right to assume that mjData.qfrc_inverse is equivalent to pybullet.calculateInverseDynamics()
and
mujoco.mj_crb(model,data) to pybullet.calculateMassMatrix()
(and accessed through mjData.crb) in terms of implementing the Stable PD algorithm? Hence replacing them respectively in the following below would do the trick?

Simplified Stable PD from the code here and Leftmost in this deom video:

Simplified Implementation

def computePD(self, bodyUniqueId, jointIndices, desiredPositions, desiredVelocities, kps, kds,
                maxForces, timeStep):
    numJoints = self._pb.getNumJoints(bodyUniqueId)
    jointStates = self._pb.getJointStates(bodyUniqueId, jointIndices)
    q1 = []
    qdot1 = []
    zeroAccelerations = []
    for i in range(numJoints):
      q1.append(jointStates[i][0])
      qdot1.append(jointStates[i][1])
      zeroAccelerations.append(0)

    q = np.array(q1)
    qdot = np.array(qdot1)
    qdes = np.array(desiredPositions)
    qdotdes = np.array(desiredVelocities)

    qError = qdes - q
    qdotError = qdotdes - qdot

    Kp = np.diagflat(kps)
    Kd = np.diagflat(kds)

    # Compute -Kp(q + qdot - qdes)
    p_term = Kp.dot(qError - qdot*timeStep)
    # Compute -Kd(qdot - qdotdes)
    d_term = Kd.dot(qdotError)

    # Compute Inertia matrix M(q)
    M = self._pb.calculateMassMatrix(bodyUniqueId, q1)
    M = np.array(M)
    # Given: M(q) * qddot + C(q, qdot) = T_ext + T_int
    # Compute Coriolis and External (Gravitational) terms G = C - T_ext
    G = self._pb.calculateInverseDynamics(bodyUniqueId, q1, qdot1, zeroAccelerations)
    G = np.array(G)
    # Obtain estimated generalized accelerations, considering Coriolis and Gravitational forces, and stable PD actions
    qddot = np.linalg.solve(a=(M + Kd * timeStep),
                            b=(-G + p_term + d_term))
    # Compute control generalized forces (T_int)
    tau = p_term + d_term - (Kd.dot(qddot) * timeStep)
    # Clip generalized forces to actuator limits
    maxF = np.array(maxForces)
    generalized_forces = np.clip(tau, -maxF, maxF)
    return generalized_forces

pybullet.calculateInverseDynamics(): < Source: Pg-61,62 >

calculateInverseDynamics will compute the forces needed to reach the given joint accelerations, starting from specified joint positions and velocities. The inverse dynamics is computed using the Recursive Newton Euler algorithm (RNEA). calculateInverseDynamics returns a list of joint forces for each degree of freedom.

pybullet.calculateMassMatrix(): < Source: Pg-63 >

calculateMassMatrix will compute the system inertia for an articulated body given its joint positions. The composite rigid body algorithm (CBRA) is used to compute the mass matrix. The result is the square mass matrix with dimensions dofCount * dofCount, stored as a list of dofCount rows, each row is a list of dofCount mass matrix elements.

Centrifugal force, gravity, etc: (C(q,q˙) - τext, equation 9 from the paper) computed in the solver(mujoco doc reference below)< Source >

The bias force cc includes Coriolis, centrifugal and gravitational forces. Their sum is computed using the Recursive Newton-Euler (RNE) algorithm with acceleration set to 0.

The joint-space inertia matrix MM is computed using the Composite Rigid-Body (CRB) algorithm. This matrix is usually quite sparse, and we represent it as such, in a custom format tailored to kinematic trees.

[yuvaltassa]

If I understand correctly, you are looking for the bias forces and the mass matrix. These are computer during mj_forward() and then you can just read them out from mjData. You can also call sub-component functions if you want this to be faster. The bias forces live in mjData.qfrc_bias (computed by mj_rne()). The mass matrix lives in mjData.qM (computed by mj_crb()) but it is in sparse form. Call mj_fullM() to get the dense representation.

[saran-t]

@Balint-H @rohit-kumar-j If you're using the mujoco Python bindings, you should prefer the named access API over mj_name2id.

See e.g. https://colab.research.google.com/drive/1V5OdtiN1e2GlVd4RCp_KwSO6mZ0m5AGM#scrollTo=vXTUIbJy8kTr

[Balint-H]

@rohit-kumar-j

I'm planning to put together the neccesary files for sharing but for that I first need to isolate the scripts so they are standalone and dont depend on other things in my project.

Until then, I found one bug in my suggestion now that I had a chance to look at it closer, MM should be indexed a different way to get the submatrix: M[dof_indices, :][:, dof_indices] (as numpy zips together multiple iterable indices in a single indexing operation, which we avoid with this)

Furthermore, I noticed you are using actuators to apply forces. If the joints are underactuated, then you should only store the dof indices of actuated dofs. Lets say there is nu<nv of them. So the dof_indices array should be (nu,) shaped. You should ensure the order of dof_indices matches the order of your collection of actuators.

Then indexing the mass matrix the above way should give you a nuxnu matrix.

Kd should be the same shape as well.

[yuvaltassa]

I think @saran-t meant to link to the relevant docs section. Also note that you get infomative messages if you do (in the above colab) m.joint('abdomen_x') or m.joint()

[saran-t]

@yuvaltassa is right, I'll correct the link in my previous comment for future readers :)

[rohit-kumar-j]

@Balint-H , thanks for the explanation!
I was able to get a stable simulation at timestep=0.01s: (Will upload the code once I clean it up)

mujoco_final_spd.mp4

@yuvaltassa , @saran-t , thanks for the support!