Make `eval_fk` independent of joint ordering

[eric-heiden]

The kinematics functions newton.eval_fk and newton.eval_ik currently assume the joints in an articulation are ordered so that joints parent to joints deeper in an articulation tree appear earlier.
This causes problems when joints are in arbitrary order, which is possible by programmatically declaring such articulations or loading certain USD/URDF files with joint_order=None.

Several solutions are possible:

• We can enforce a certain joint ordering either when finalizing a ModelBuilder or closing an articulation by raising an error or reshuffling those joints (which could be bad for users expecting stable joint indices)
• We can introduce another Model attribute that encodes the joint ordering, so inside the kinematics kernels we can query the joints in topological order, e.g. via joint_attribute[joint_order[i]]

See the following example that shows the incorrect kinematics from Newton, while MuJoCo uses the correctly ordered joints and therefore doesn't exhibit this problem in its own renderer based on its own kinematics function:

import warp as wp

import newton
import newton.examples


class Example:
    def __init__(self, viewer):
        # setup simulation parameters first
        self.fps = 100
        self.frame_dt = 1.0 / self.fps
        self.sim_time = 0.0
        self.sim_substeps = 10
        self.sim_dt = self.frame_dt / self.sim_substeps

        self.viewer = viewer

        builder = newton.ModelBuilder()

        # common geometry settings
        cuboid_hx = 0.1
        cuboid_hy = 0.1
        cuboid_hz = 0.75
        upper_hz = 0.25 * cuboid_hz

        # layout positions (y-rows)
        rows = [-3.0, 0.0, 3.0]
        drop_z = 2.0

        # -----------------------------
        # REVOLUTE (hinge) joint demo
        # -----------------------------
        y = rows[0]

        a_rev = builder.add_body()
        b_rev = builder.add_body()
        c_rev = builder.add_body()
        builder.add_shape_box(a_rev, hx=cuboid_hx, hy=cuboid_hy, hz=upper_hz)
        builder.add_shape_box(b_rev, hx=cuboid_hx, hy=cuboid_hy, hz=cuboid_hz)
        builder.add_shape_box(c_rev, hx=cuboid_hx, hy=cuboid_hy, hz=cuboid_hz)

        builder.add_joint_revolute(
            parent=b_rev,
            child=c_rev,
            parent_xform=wp.transform(p=wp.vec3(0.0, 0.0, -cuboid_hz)),
            child_xform=wp.transform(p=wp.vec3(0.0, 0.0, -cuboid_hz)),
            key="revolute_b_c",
        )
        builder.joint_q[-1] = -0.4
        builder.add_joint_fixed(
            parent=-1,
            child=a_rev,
            parent_xform=wp.transform(p=wp.vec3(0.0, y, drop_z + upper_hz)),
            key="fixed_revolute_anchor",
        )
        builder.add_joint_revolute(
            parent=a_rev,
            child=b_rev,
            parent_xform=wp.transform(p=wp.vec3(0.0, 0.0, -upper_hz)),
            child_xform=wp.transform(p=wp.vec3(0.0, 0.0, +cuboid_hz)),
            key="revolute_a_b",
        )
        builder.joint_q[-1] = wp.pi * 0.5


        # finalize model
        self.model = builder.finalize()

        self.solver = newton.solvers.SolverMuJoCo(self.model, njmax=100, disable_contacts=True)

        self.state_0 = self.model.state()
        self.state_1 = self.model.state()
        self.control = self.model.control()
        self.contacts = self.model.collide(self.state_0)

        self.viewer.set_model(self.model)

        # not required for MuJoCo, but required for other solvers
        newton.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, self.state_0)

        self.capture()

    def capture(self):
        if wp.get_device().is_cuda:
            with wp.ScopedCapture() as capture:
                self.simulate()
            self.graph = capture.graph
        else:
            self.graph = None

    def simulate(self):
        for _ in range(self.sim_substeps):
            self.state_0.clear_forces()

            # apply forces to the model
            self.viewer.apply_forces(self.state_0)

            self.contacts = self.model.collide(self.state_0)
            self.solver.step(self.state_0, self.state_1, self.control, self.contacts, self.sim_dt)

            # swap states
            self.state_0, self.state_1 = self.state_1, self.state_0

    def step(self):
        if self.graph:
            wp.capture_launch(self.graph)
        else:
            self.simulate()

        self.sim_time += self.frame_dt

    def render(self):
        self.viewer.begin_frame(self.sim_time)
        self.viewer.log_state(self.state_0)
        self.viewer.log_contacts(self.contacts, self.state_0)
        self.viewer.end_frame()

        self.solver.render_mujoco_viewer()


if __name__ == "__main__":
    # Parse arguments and initialize viewer
    viewer, args = newton.examples.init()

    # Create viewer and run
    example = Example(viewer)

    newton.examples.run(example, args)

[adenzler-nvidia]

we have some warning now but it's turned off by default, maybe if we turn it on we can bump this to after the release @eric-heiden ?