justinberi
commented
2 years ago Some other references in case they are useful.
TDS implementation - https://github.com/google-research/tiny-differentiable-simulator/blob/master/src/rb_constraint_solver.hpp
which references Michael Cline's thesis - https://open.library.ubc.ca/soa/cIRcle/collections/ubctheses/831/items/1.0051676
Use something like the Featherstone algorithm to calculate joint state in reduced coordinates - one coordinate for each rotational/translational degree of freedom for each joint.
This would be a major rewrite and has a number of subtasks - the first of which is to determine whether the underlying algorithms (mass matrix calculation, forward kinematics, etc) are all XLA/TPU-friendly.