mink + JAX or PyTorch?

[jonzamora]

Hi @kevinzakka

I'm wondering if there are any plans to eventually port mink over to a JAX or PyTorch-based implementation?

Another great work, https://github.com/UM-ARM-Lab/pytorch_kinematics, implements IK solving with PyTorch, and I would imagine this is faster than a CPU-based IK solver like what mink currently uses.

I bring this up because i'd like to use a mink controller for some Reinforcement Learning tasks i'm building, but I noticed that the simulation FPS slows down by ~6.5x when using it for controlling my robot. On my 4090 Desktop with an Intel i9-13900K CPU, sim FPS goes from ~4000 FPS simulating a Robot MJCF (no mink controller, just RGB rendering at 30 FPS) to ~600 FPS (with mink controller + RGB rendering at 30 FPS).

• This includes some low-level kinks too, like changing max_iters=1 for the IK solver, and changing sim timestep to model.opt.timestep = 1 / 100 instead of MuJoCo's default 1 / 500. These changes didn't seem to affect the End-Effector Displacement for my tasks, and it resulted in faster sim FPS with mink.

I understand that IK will inevitably slow down FPS for Simulation, but if solving IK with PyTorch or JAX speeds things up by a reasonable margin, that'd be great! I'm wondering your thoughts on this, i'm happy to share more info and continue this discussion : )