Add Forward Dynamics computation via ABA

src/adam/casadi/casadi_like.py

@@ -93,6 +93,12 @@ def __getitem__(self, idx) -> "CasadiLike":
         """Overrides get item operator"""
         return CasadiLike(self.array[idx])
 
+    def __eq__(self, other: Union["CasadiLike", npt.ArrayLike]) -> bool:
+        """Overrides == operator"""
+        if type(self) is not type(other):
+            return self.array == other
+        return self.array == other.array
+
     @property
     def T(self) -> "CasadiLike":
         """
@@ -149,6 +155,39 @@ def skew(x: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike":
         else:
             return CasadiLike(cs.skew(x))
 
+    @staticmethod
+    def vee(x: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike":
+        """
+        Args:
+            x (Union[CasadiLike, npt.ArrayLike]): 3x3 matrix
+
+        Returns:
+            CasadiLike: the vector from the skew symmetric matrix x
+        """
+        vee = lambda x: cs.vertcat(x[2, 1], x[0, 2], x[1, 0])
+        if isinstance(x, CasadiLike):
+            return CasadiLike(vee(x.array))
+        else:
+            return CasadiLike(vee(x))
+
+    @staticmethod
+    def inv(x: Union["CasadiLike", npt.ArrayLike]) -> "CasadiLike":
+        """
+        Args:
+            x (Union[CasadiLike, npt.ArrayLike]): matrix
+
+        Returns:
+            CasadiLike: inverse of x
+        """
+        if isinstance(x, CasadiLike):
+            return CasadiLike(cs.inv(x.array))
+        else:
+            return (
+                CasadiLike(cs.inv(x))
+                if x.size1() <= 5
+                else CasadiLike(self.solve(x, self.eye(x.size1())))
+            )
+
     @staticmethod
     def sin(x: npt.ArrayLike) -> "CasadiLike":
         """
@@ -206,6 +245,18 @@ def horzcat(*x) -> "CasadiLike":
         y = [xi.array if isinstance(xi, CasadiLike) else xi for xi in x]
         return CasadiLike(cs.horzcat(*y))
 
+    @staticmethod
+    def solve(A: "CasadiLike", b: "CasadiLike") -> "CasadiLike":
+        """
+        Args:
+            A (CasadiLike): matrix
+            b (CasadiLike): vector
+
+        Returns:
+            CasadiLike: solution of A*x=b
+        """
+        return CasadiLike(cs.solve(A.array, b.array))
+
 
 if __name__ == "__main__":
     math = SpatialMath()

src/adam/casadi/computations.py

@@ -411,3 +411,33 @@ def CoM_position(
             CoM (Union[cs.SX, cs.DM]): The CoM position
         """
         return self.rbdalgos.CoM_position(base_transform, joint_positions).array
+
+    def forward_dynamics(
+        self,
+        base_transform: Union[cs.SX, cs.DM],
+        base_velocity: Union[cs.SX, cs.DM],
+        joint_positions: Union[cs.SX, cs.DM],
+        joint_velocities: Union[cs.SX, cs.DM],
+        joint_torques: Union[cs.SX, cs.DM],
+    ) -> Union[cs.SX, cs.DM]:
+        """Returns base and joints accelerations of the floating-base dynamics equation
+
+        Args:
+            base_transform (Union[cs.SX, cs.DM]): The homogenous transform from base to world frame
+            base_velocity (Union[cs.SX, cs.DM]): The base velocity in mixed representation
+            joint_positions (Union[cs.SX, cs.DM]): The joints position
+            joint_velocities (Union[cs.SX, cs.DM]): The joints velocity
+            joint_torques (Union[cs.SX, cs.DM]): The joints torque
+
+        Returns:
+            base_acceleration (Union[cs.SX, cs.DM]): The base acceleration in mixed representation
+            joint_accelerations (Union[cs.SX, cs.DM]): The joints acceleration
+        """
+        return self.rbdalgos.aba(
+            base_transform,
+            base_velocity,
+            joint_positions,
+            joint_velocities,
+            joint_torques,
+            self.g,
+        )

src/adam/core/rbd_algorithms.py

@@ -5,7 +5,7 @@
 
 from adam.core.constants import Representations
 from adam.core.spatial_math import SpatialMath
-from adam.model import Model, Node
+from adam.model import Model, Node, Link
 
 
 class RBDAlgorithms:
@@ -457,5 +457,170 @@ def rnea(
         tau[:6] = B_X_BI.T @ tau[:6]
         return tau
 
-    def aba(self):
-        raise NotImplementedError
+    def aba(
+        self,
+        base_transform: npt.ArrayLike,
+        base_velocity: npt.ArrayLike,
+        joint_positions: npt.ArrayLike,
+        joint_velocities: npt.ArrayLike,
+        tau: npt.ArrayLike,
+        g: npt.ArrayLike,
+    ) -> npt.ArrayLike:
+        """Implementation of Articulated Body Algorithm
+
+        Args:
+            base_transform (T): The homogenous transform from base to world frame
+            base_velocity (T): The base velocity in mixed representation
+            joint_positions (T): The joints position
+            joint_velocities (T): The joints velocity
+            tau (T): The generalized force variables
+            g (T): The 6D gravity acceleration
+
+        Returns:
+            base_acceleration (T): The base acceleration in mixed representation
+            joint_accelerations (T): The joints acceleration
+        """
+        model = self.model.reduce(self.model.actuated_joints)
+        joints = list(model.joints.values())
+
+        NB = model.N
+
+        i_X_pi = self.math.factory.zeros(NB, 6, 6)
+        v = self.math.factory.zeros(NB, 6, 1)
+        c = self.math.factory.zeros(NB, 6, 1)
+        pA = self.math.factory.zeros(NB, 6, 1)
+        IA = self.math.factory.zeros(NB, 6, 6)
+        U = self.math.factory.zeros(NB, 6, 1)
+        D = self.math.factory.zeros(NB, 1, 1)
+        u = self.math.factory.zeros(NB, 1, 1)
+
+        a = self.math.factory.zeros(NB, 6, 1)
+        sdd = self.math.factory.zeros(NB, 1, 1)
+        B_X_W = self.math.adjoint_mixed(base_transform)
+
+        if model.floating_base:
+            IA[0] = model.tree.get_node_from_name(self.root_link).link.spatial_inertia()
+            v[0] = B_X_W @ base_velocity
+            pA[0] = (
+                self.math.spatial_skew_star(v[0]) @ IA[0] @ v[0]
+            )  # - self.math.adjoint_inverse(B_X_W).T @ f_ext[0]
+
+        def get_tree_transform(self, joints) -> "Array":
+            """returns the tree transform
+
+            Returns:
+                Array: the tree transform
+            """
+            relative_transform = lambda j: self.math.inv(
+                model.tree.graph[j.child].parent_arc.spatial_transform(0)
+            ) @ j.spatial_transform(0)
+
+            return self.math.vertcat(
+                [self.math.factory.eye(6).array]
+                + list(
+                    map(
+                        lambda j: (
+                            relative_transform(j).array
+                            if j.parent != self.root_link
+                            else self.math.factory.eye(6).array
+                        ),
+                        joints,
+                    )
+                )
+            )
+
+        tree_transform = get_tree_transform(self, joints)
+        # p = lambda i: list(model.tree.graph).index(joints[i].parent)
+
+        # def p(i):
+        #     jp = joints[i].parent
+
+        #     try:
+        #         if model.tree.graph[jp].parent_arc.name in model.actuated_joints:
+        #             return model.tree.graph[jp].parent_arc.idx
+        #         else:
+        #             return p(model.tree.graph[jp].parent_arc.parent.idx)
+        #     except AttributeError:
+        #         if model.tree.graph[jp].parent_arc is None:
+        #             return -1  # root link
+        #         else:
+        #             return p(model.tree.graph[jp].parent_arc.idx)
+
+        # Pass 1
+        for i, node in enumerate(model.tree):
+            if node.name == self.root_link:
+                continue
+            # link, parent_arc, parent
+            link_i, joint_i, link_pi = node.get_elements()
+            pi = model.tree.get_idx_from_name(link_pi.name)
+
+            q = joint_positions[i - 1]
+            q_dot = joint_velocities[i - 1]
+
+            # Parent-child transform
+            i_X_pi[i] = tree_transform[i] @ joint_i.spatial_transform(q)
+            v_J = joint_i.motion_subspace() * q_dot
+
+            v[i] = i_X_pi[i] @ v[pi] + v_J
+            c[i] = i_X_pi[i] @ c[pi] + self.math.spatial_skew(v[i]) @ v_J
+
+            IA[i] = link_i.spatial_inertia()
+
+            pA[i] = IA[i] @ c[i] + self.math.spatial_skew_star(v[i]) @ IA[i] @ v[i]
+
+        LUMPED = False
+
+        # Pass 2
+        for i, node in enumerate(reversed(model.tree)):
+            if node.name == self.root_link:
+                continue
+
+            link_i, joint_i, link_pi = node.get_elements()
+            pi = model.tree.get_idx_from_name(link_pi.name)
+
+            U[i] = IA[i] @ joint_i.motion_subspace()
+            D[i] = joint_i.motion_subspace().T @ U[i]
+            u[i] = self.math.vertcat(tau[i - 1]) - joint_i.motion_subspace().T @ pA[i]
+
+            Ia = IA[i] - U[i] / D[i] @ U[i].T
+            pa = pA[i] + Ia @ c[i] + U[i] * u[i] / D[i]
+
+            if joint_i.parent != self.root_link or not model.floating_base:
+                IA[pi] += i_X_pi[i].T @ Ia @ i_X_pi[i]
+                pA[pi] += i_X_pi[i].T @ pa
+                continue
+
+        a[0] = B_X_W @ g if model.floating_base else self.math.solve(-IA[0], pA[0])
+
+        # Pass 3
+        for i, node in enumerate(model.tree):
+            if node.name == self.root_link:
+                continue
+
+            link_i, joint_i, link_pi = node.get_elements()
+            pi = model.tree.get_idx_from_name(link_pi.name)
+
+            sdd[i - 1] = (u[i] - U[i].T @ a[i]) / D[i]
+
+            a[i] += i_X_pi[i].T @ a[pi] + joint_i.motion_subspace() * sdd[i - 1] + c[i]
+
+        # Squeeze sdd
+        s_ddot = self.math.vertcat(*[sdd[i] for i in range(sdd.shape[0])])
+
+        if (
+            self.frame_velocity_representation
+            == Representations.BODY_FIXED_REPRESENTATION
+        ):
+            return (a[0], s_ddot)
+
+        elif self.frame_velocity_representation == Representations.MIXED_REPRESENTATION:
+            return (
+                self.math.horzcat(
+                    (
+                        (self.math.solve(B_X_W, a[0]) + g)[:, None]
+                        if model.floating_base
+                        else self.math.zeros(6, 1)
+                    ),
+                ),
+                s_ddot,
+            )

src/adam/core/spatial_math.py

@@ -158,6 +158,26 @@ def cos(x: npt.ArrayLike) -> npt.ArrayLike:
     def skew(x):
         pass
 
+    @abc.abstractmethod
+    def vee(x):
+        pass
+
+    @abc.abstractmethod
+    def inv(x):
+        pass
+
+    @abc.abstractmethod
+    def solve(A: npt.ArrayLike, b: npt.ArrayLike) -> npt.ArrayLike:
+        """
+        Args:
+            A (npt.ArrayLike): matrix
+            b (npt.ArrayLike): vector
+
+        Returns:
+            npt.ArrayLike: solution of the linear system Ax=b
+        """
+        pass
+
     def R_from_axis_angle(self, axis: npt.ArrayLike, q: npt.ArrayLike) -> npt.ArrayLike:
         """
         Args:

src/adam/jax/computations.py

@@ -239,3 +239,35 @@ def get_total_mass(self) -> float:
             mass: The total mass
         """
         return self.rbdalgos.get_total_mass()
+
+    def forward_dynamics(
+        self,
+        base_transform: jnp.array,
+        base_velocity: jnp.array,
+        joint_positions: jnp.array,
+        joint_velocities: jnp.array,
+        joint_torques: jnp.array,
+    ) -> jnp.array:
+        """Returns base and joints accelerations of the floating-base dynamics equation
+
+        Args:
+            base_transform (jnp.array): The homogenous transform from base to world frame
+            base_velocity (jnp.array): The base velocity in mixed representation
+            joint_positions (jnp.array): The joints position
+            joint_velocities (jnp.array): The joints velocity
+            joint_torques (jnp.array): The joints torques
+
+        Returns:
+            base_acceleration (jnp.array): The base acceleration in mixed representation
+            joint_accelerations (jnp.array): The joints acceleration
+        """
+        base_acceleration, joint_accelerations = self.rbdalgos.aba(
+            base_transform,
+            base_velocity,
+            joint_positions,
+            joint_velocities,
+            joint_torques,
+            self.g,
+        )
+
+        return base_acceleration.array.squeeze(), joint_accelerations.array.squeeze()