[Feature] Swap mplib for something else integrated with sapien/maniskill

[StoneT2000]

Currently mplib is essentially no longer maintained. I am the owner of the mplib repo now but do not have time nor the expertise to maintain it properly.

I've been considering switching to one of mink or pyroki but open to any other options. I am leaning towards mink since it looks less jax dependent (and jax does not always play nice with torch dependencies). Mink also has plenty of pose reaching/collision avoidance examples.

Some concerns over those options:

• Mink: Heavily tied to mujoco. Code for handling collisions like https://github.com/kevinzakka/mink/blob/ace0d723676898a20cbeda517a9a7a47863314d1/src/mink/limits/collision_avoidance_limit.py#L118 would have to be refactored out.
• Pyroki: No sampling based planner (RRT or screw motion planning is often used in maniskill). That being said perhaps it is not too had to add support for RRT to this? Another concern is the heavy jax dependencies.

[jstmn]

Yo! I few thoughts:

1. I think avoiding mujoco would be wise. It's always a pain to work with. The xml configuration scheme is especially terrible imo. Why would the collision checking code you linked need to be refactored out?
2. Pyroki seems like a cool project. The JIT overhead (see here) might be annoying? And I wasn't aware of the jax/torch issue but good to be aware of. Adding a RRT / a screw planner shouldn't be hard at all though. This seems like a good idea to me.
3. A third option is to spin up our own thing. I have some code already that goes from a robot_descriptions.py robot to a collection of trimesh collision geometries. I've had no need to expand the projects functionality yet (i'm only using it as a collision checker), but again, adding a planner wouldn't be hard. For maniskill, we could skip the robot_descriptions.py, and just go straight from urdf to trimesh objects. Once we have the scene in trimesh, we can easily implement whatever planner we want. The benefits here are reduced dependencies and more flexibility. The downsides are 1. we wouldn't benefit from any features added in a different project, and 2. it could be more work.
4. What exactly are the specifications here btw? Do we just need a screw planner and RRT? Or also IK?
5. Do you have any idea how much time / how difficult it would be to just solve the mplib bug? It's working great except for that issue AFAICT

[StoneT2000]

I last recall torch and jax having some incompatible versions. I could be wrong though, maybe its been fixed.

I'd like to avoid spinning up our own thing as much as possible, try and stick with some existing APIs and behaviors people are used to (something similar to mplib, or if using pyroki/mink then their APIs as both are already popular). This project is not entirely meant to be a big mono-repo robotics project, thus I think it is best to avoid redoing this part.

For requirements, all we need for a motion planning library is the following for ManiSkill

• Be able to reliably and accurately control a robot at its current joint config to some one of
  1. target joint config
  2. target link pose
• Make plans with collision avoidance (allow user to input primitive shapes at minimum like boxes or a point cloud / spheres, convex meshes would also be great)
• Make plans with collision avoidance including grasped objects
• Be instantiated when given a URDF or some kinematic structure/definition (for maniskill I'm fine with adding any code to extract kinematic structure from the agent class object in env.unwrapped.agent if it means we don't need to add the rest of the MP library to maniskill's codebase). Also permit a SRDF style definition to disable some impossible/unlikely self collisions.

[jstmn]

Cool, makes sense. Then I think Pyroki is the most promising option. Looking through the documentation, they have IK, collision checking, and planning functionality which are the three main components we'll need. It reads from urdf too.

I'm happy to look into this in february once the RSS push is over, i've been meaning to try out pyroki for a while now.

[StoneT2000]

Will also have to wait until after RSS haha, we can sync here after that to discuss where we could collaborate on this perhaps and identify key things we need to add.

I also noticed that pyroki works with jax CPU version (not sure if that is much slower or anything), so in case someone can't install jax cuda version with torch cuda together, CPU can be a good fallback.