external_force_aggregator_system.py

from pydrake.all import (
    ExternallyAppliedSpatialForce,
    LeafSystem,
    List,
    Value,
)


def aggregate_force(force_a, force_b):
    """
    Adds two forces on the same body.

    @throws if force_a and force_b do not share common origin, because that
    would require plant kinematics to shift using orientation (p_BoBq_B and
    F_Bq_W).
    """
    assert force_a.body_index == force_b.body_index
    assert (force_a.p_BoBq_B == force_b.p_BoBq_B).all()
    force_out = ExternallyAppliedSpatialForce()
    force_out.body_index = force_a
    force_out.p_BoBq_B = force_a.p_BoBq_B
    force_out.F_Bq_W = force_a.F_Bq_W + force_b.F_Bq_W
    return force_out


def aggregate_forces(forces_list):
    """Naively aggregates forces."""
    force_map = {}
    for forces in forces_list:
        for force in forces:
            existing_force = force_map.get(force.body_index)
            if existing_force is not None:
                force = aggregate_force(existing_force, force)
            force_map[force.body_index] = force
    return list(force_map.values())


class ForceAggregator(LeafSystem):
    def __init__(self, num_inputs):
        super().__init__()
        forces_cls = Value[List[ExternallyAppliedSpatialForce]]
        self._num_inputs = num_inputs
        for i in range(self._num_inputs):
            self.DeclareAbstractInputPort(
                f"forces_in_{i}",
                model_value=forces_cls(),
            )
        self.DeclareAbstractOutputPort(
            "forces_out",
            alloc=forces_cls,
            calc=self._calc_forces,
        )

    def _calc_force_input(self, context, i):
        return self.get_input_port(i).Eval(context)

    def _calc_forces(self, context, output):
        forces_list = []
        for i in range(self._num_inputs):
            forces = self._calc_force_input(context, i)
            forces_list.append(forces)
        forces_agg = aggregate_forces(forces_list)
        output.set_value(forces_agg)

    @staticmethod
    def AddToBuilder(builder, force_ports):
        num_inputs = len(force_ports)
        system = builder.AddSystem(ForceAggregator(num_inputs))
        for i, port in enumerate(force_ports):
            builder.Connect(
                port,
                system.get_input_port(i),
            )
        return system


def example_usage():
    ...
    force_ports = [...]
    force_agg = ForceAggregator.AddToBuilder(builder, force_ports)
    builder.Connect(
        force_agg.get_output_port(),
        plant.get_applied_spatial_force_input_port(),
    )

simple_floating_body_controller.py

import numpy as np

from pydrake.common.cpp_param import List
from pydrake.common.value import Value
from pydrake.math import RigidTransform
from pydrake.multibody.math import SpatialForce, SpatialVelocity
from pydrake.multibody.plant import ExternallyAppliedSpatialForce
from pydrake.multibody.tree import JacobianWrtVariable
from pydrake.systems.framework import LeafSystem

from anzu.perception.pose_util import rotation_matrix_to_axang3


def get_frame_spatial_velocity(plant, context, frame_T, frame_F, frame_E=None):
    """
    Returns:
        SpatialVelocity of frame F's origin w.r.t. frame T, expressed in E
        (which is frame T if unspecified).
    """
    if frame_E is None:
        frame_E = frame_T
    Jv_TF_E = plant.CalcJacobianSpatialVelocity(
        context,
        with_respect_to=JacobianWrtVariable.kV,
        frame_B=frame_F,
        p_BP=[0, 0, 0],
        frame_A=frame_T,
        frame_E=frame_E,
    )
    v = plant.GetVelocities(context)
    V_TF_E = SpatialVelocity(Jv_TF_E @ v)
    return V_TF_E


def reexpress_to_matrix(R_AE, I_BP_E):
    I_BP_A = I_BP_E.ReExpress(R_AE)
    return I_BP_A.CopyToFullMatrix3()


class FloatingBodyPoseController(LeafSystem):
    """
    Controls for the pose of a single-body floating model using frame P w.r.t.
    inertial frame T.

    Inputs:
        X_TPdes: Desired pose.
        V_TPdes: Desired velocity.
    Outputs:
        forces:
            Spatial forces to apply to body to track reference
            trajectories.
    """

    def __init__(self, plant, model_instance, frame_T, frame_P):
        super().__init__()
        nx = plant.num_positions(model_instance) + plant.num_velocities(
            model_instance
        )
        # N.B. This will be in the `controller` closure, and thus kept alive.
        context = plant.CreateDefaultContext()

        M_PPo_P = plant.CalcSpatialInertia(
            context, frame_P, plant.GetBodyIndices(model_instance)
        )

        M = M_PPo_P.get_mass() * np.eye(3)
        I_PPo_P = M_PPo_P.CalcRotationalInertia()
        # TODO(eric.cousineau): Hoist these parameters somewhere.
        Kp_xyz = 500.0 * M
        Kd_xyz = 40.0 * M
        Kp_rot = lambda R_WP: 500.0 * reexpress_to_matrix(R_WP, I_PPo_P)
        Kd_rot = lambda R_WP: 40.0 * reexpress_to_matrix(R_WP, I_PPo_P)

        def control_math(x, X_TPdes, V_TPdes):
            # Cribbed from:
            # https://github.com/gizatt/drake_hydra_interact/tree/cce3ecbb
            frame_W = plant.world_frame()
            plant.SetPositionsAndVelocities(context, model_instance, x)

            X_WT = plant.CalcRelativeTransform(context, frame_W, frame_T)
            V_WT = me.get_frame_spatial_velocity(
                plant, context, frame_W, frame_T
            )
            X_WPdes = X_WT @ X_TPdes
            V_WPdes = V_WT.ComposeWithMovingFrameVelocity(
                X_WT.translation(), V_TPdes.Rotate(X_WT.rotation())
            )

            X_WP = plant.CalcRelativeTransform(context, frame_W, frame_P)
            R_WP = X_WP.rotation()
            V_WP = me.get_frame_spatial_velocity(
                plant, context, frame_W, frame_P
            )
            # Transform to "negative error": desired w.r.t. actual,
            # expressed in world frame (for applying the force).
            p_PPdes_W = X_WPdes.translation() - X_WP.translation()
            # N.B. We don't project away symmetry here because we're expecting
            # smooth trajectories (for now).
            R_PPdes = R_WP.inverse() @ X_WPdes.rotation()
            axang3_PPdes = rotation_matrix_to_axang3(R_PPdes)
            axang3_PPdes_W = R_WP @ axang3_PPdes
            V_PPdes_W = V_WPdes - V_WP
            # Compute wrench components.
            f_P_W = Kp_xyz @ p_PPdes_W + Kd_xyz @ V_PPdes_W.translational()
            tau_P_W = (
                Kp_rot(R_WP) @ axang3_PPdes_W
                + Kd_rot(R_WP) @ V_PPdes_W.rotational()
            )
            F_P_W_feedback = SpatialForce(tau=tau_P_W, f=f_P_W)
            # Add gravity-compensation term.
            g_W = plant.gravity_field().gravity_vector()
            F_Pcm_W = SpatialForce(tau=[0, 0, 0], f=-g_W * M_PPo_P.get_mass())
            p_PoPcm_W = R_WP @ M_PPo_P.get_com()
            p_PcmP_W = -p_PoPcm_W
            F_P_W_feedforward = F_Pcm_W.Shift(p_PcmP_W)
            # Package it up.
            F_P_W = F_P_W_feedback + F_P_W_feedforward
            external_force = ExternallyAppliedSpatialForce()
            external_force.body_index = frame_P.body().index()
            external_force.F_Bq_W = F_P_W
            external_force.p_BoBq_B = (
                frame_P.GetFixedPoseInBodyFrame().translation()
            )
            return external_force

        self.plant_state_input = self.DeclareVectorInputPort("plant_state", nx)
        self.X_TPdes_input = self.DeclareAbstractInputPort(
            "X_TPdes", Value[RigidTransform]()
        )
        self.V_TPdes_input = self.DeclareAbstractInputPort(
            "V_TPdes", Value[SpatialVelocity]()
        )

        def control_calc(sys_context, output):
            plant_state = self.plant_state_input.Eval(sys_context)
            X_TPdes = self.X_TPdes_input.Eval(sys_context)
            V_TPdes = self.V_TPdes_input.Eval(sys_context)
            external_force = control_math(plant_state, X_TPdes, V_TPdes)
            output.set_value([external_force])

        forces_cls = Value[DrakeList[ExternallyAppliedSpatialForce]]
        self.forces_output = self.DeclareAbstractOutputPort(
            "forces_out", alloc=forces_cls, calc=control_calc,
        )

    @staticmethod
    def AddToBuilder(
        builder,
        plant,
        model_instance,
        frame_T,
        frame_P,
        *,
        connect_to_plant=True,
        name="controller",
    ):
        controller = FloatingBodyPoseController(
            plant, model_instance, frame_T, frame_P
        )
        controller.set_name(name)
        builder.AddSystem(controller)
        builder.Connect(
            plant.get_state_output_port(model_instance),
            controller.plant_state_input,
        )
        if connect_to_plant:
            builder.Connect(
                controller.forces_output,
                plant.get_applied_spatial_force_input_port(),
            )
        return controller