[QUESTION] <How to get Jacobian wrt action>

[lyuyiqi]

Hi Warp group, thank you for providing this platform. We wanted to use Warp in my project, but I encountered some issues. We modified the cartpole example and replaced the model with the ant model. We tried to compute the Jacobian matrix of joint_q wrt action, but the computation took a long time. When we printed the Jacobian matrix, we got a zero matrix. Could you give us some guidance on this?

Here is the code snippet:

    def simulate(self, frame_num):
        for _ in range(self.sim_substeps):
            self.state_0.clear_forces()
            tape = wp.Tape()
            with tape:
                wp.sim.collide(self.model, self.state_0)
                self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt, control=self.control)
                self.compute_ee_position()
                jacobians = np.empty((self.num_envs, 15, 8), dtype=np.float32)
                for output_index in range(15):
                    select_index = np.zeros(15)
                    select_index[output_index] = 1.0
                    e = wp.array(np.tile(select_index, self.num_envs), dtype=wp.vec3)
                    joint_q = wp.array(self.state_1.joint_q.numpy().reshape(self.num_envs, -1).astype(np.float32), requires_grad=True)
                    tape.backward(grads={joint_q: e})
                    q_grad_i = tape.gradients[self.control.joint_act]
                    jacobians[:, output_index, :] = q_grad_i.numpy().reshape(self.num_envs, 8)
                    tape.zero()

                self.control.joint_act = wp.array(np.random.uniform(-200, 200, size=8).astype(np.float32))
                self.control.joint_act.requires_grad = True