ami-iit / adam

adam implements a collection of algorithms for calculating rigid-body dynamics in Jax, CasADi, PyTorch, and Numpy.
https://adam-docs.readthedocs.io/en/latest/
BSD 3-Clause "New" or "Revised" License
131 stars 20 forks source link

Forward kinematics is so slow #90

Closed richardrl closed 4 months ago

richardrl commented 4 months ago

Hi, Is there any way to get a fast forward kinematics from this library?

Here's the speed for inverse kinematics and the speed for forward kinematics: IK time 0.0021796226501464844 FK time 0.22966623306274414

This is our IK function using scipy optimization:

def retarget(self, target_se3_in_world_frame, return_cost=False, close_joints_coefficient=1e-4):
        """
        r_h: vector of fingertip targets
        """

        # print(r_h)

        res = minimize(self.fast_c,
                       self.start_angles, method="SLSQP", args=(target_se3_in_world_frame, self.prev_angles,
                                                                close_joints_coefficient),
                       jac=self.grad_c, tol=1e-10, options={'maxiter': 6000}, bounds=self.bounds,
                       )
        self.prev_angles = self.start_angles
        self.start_angles = res.x

        print("arm retargeter")
        print("optimal value")
        print(res.fun)
        print("\n")

        if return_cost:
            return res.x, res.fun  # ,  [fk1[:3,-1], fk2[:3,-1], fk3[:3,-1], fk4[:3,-1]]
        else:
            return res.x

forward kinematics:

        self.fk = kinDyn.forward_kinematics_fun('tool0')

    def compute_fk(self, angles):
        return self.fk(jnp.eye(4), angles)

We are using jax.

The FK is 100x slower than IK. I understand there is probably some symbolic computation happening for this FK, but is there a faster version somewhere in this library?

traversaro commented 4 months ago

What is your use case? If you are not interested in batching, I guess numpy or casadi probably should give you a faster fk (for casadi especially if you are using common subexpression elimination, see https://github.com/ami-iit/adam/pull/86).

traversaro commented 4 months ago

By the way, how are you computing self.fast_c in your snippet? To get some a fast optimization-based IK, I guess you have some sort of fast FK to use inside the cost function?

richardrl commented 4 months ago

@traversaro Thanks for your comment - the problem was not using the jitted version of forward kinematics. It is close to 0 compute time with the jitted version.