Source code for brom_drake.systems.end_effector_wrench_calculator

import numpy as np
from pydrake.multibody.tree import JacobianWrtVariable
from pydrake.systems.framework import LeafSystem, BasicVector




class EndEffectorWrenchCalculator(LeafSystem):
    """
    **Description**

    A simple system which takes as input joint torques and outputs the corresponding
    wrench applied to the end-effector.

    **Diagram**

    The system can be illustrated with the following block: ::

                            ---------------------------------
                            |                               |
                            |                               |
                            |                               |
        joint_positions --> |  EndEffectorWrenchCalculator  | ---> end_effector_wrench
        joint_angles -----> |                               |
        joint_torques ----> |                               |
                            |                               |
                            |                               |
                            ---------------------------------
    """

    def __init__(self, plant, end_effector_frame):
        LeafSystem.__init__(self)

        self.plant = plant
        self.context = self.plant.CreateDefaultContext()
        self.ee_frame = end_effector_frame

        # Inputs are joint positions, angles and torques
        self.q_port = self.DeclareVectorInputPort(
            "joint_positions", BasicVector(plant.num_positions())
        )
        self.v_port = self.DeclareVectorInputPort(
            "joint_velocities", BasicVector(plant.num_velocities())
        )
        self.tau_port = self.DeclareVectorInputPort(
            "joint_torques", BasicVector(plant.num_actuators())
        )

        # Output is applied wrench at the end-effector
        self.DeclareVectorOutputPort(
            "end_effector_wrench", BasicVector(6), self.CalcEndEffectorWrench
        )

    def CalcEndEffectorWrench(self, context, output):
        # Gather inputs
        q = self.q_port.Eval(context)
        v = self.v_port.Eval(context)
        tau = self.tau_port.Eval(context)

        # Set internal model state
        self.plant.SetPositions(self.context, q)
        self.plant.SetVelocities(self.context, v)

        # Some dynamics computations
        M = self.plant.CalcMassMatrixViaInverseDynamics(self.context)
        tau_g = -self.plant.CalcGravityGeneralizedForces(self.context)

        # Compute end-effector jacobian
        J = self.plant.CalcJacobianSpatialVelocity(
            self.context,
            JacobianWrtVariable.kV,
            self.ee_frame,
            np.zeros(3),
            self.plant.world_frame(),
            self.plant.world_frame(),
        )

        # Compute jacobian pseudoinverse
        Minv = np.linalg.inv(M)
        Lambda = np.linalg.pinv(J @ Minv @ J.T)
        Jbar = Lambda @ J @ Minv

        # Compute wrench (spatial force) applied at end-effector
        w = Jbar @ (tau - tau_g)

        output.SetFromVector(w)