Add a few jacobian tests

examples/pybullet/unittests/unittests.py

@@ -40,7 +40,79 @@ class TestPybulletMethods(unittest.TestCase):
         self.assertLess(vel[1][1],1e-10)
         self.assertLess(vel[1][2],1e-10)
         p.disconnect()
-       
+
+class TestPybulletJacobian(unittest.TestCase):
+
+    def getMotorJointStates(self, robot):
+        import pybullet as p
+        joint_states = p.getJointStates(robot, range(p.getNumJoints(robot)))
+        joint_infos = [p.getJointInfo(robot, i) for i in range(p.getNumJoints(robot))]
+        joint_states = [j for j, i in zip(joint_states, joint_infos) if i[3] > -1]
+        joint_positions = [state[0] for state in joint_states]
+        joint_velocities = [state[1] for state in joint_states]
+        joint_torques = [state[3] for state in joint_states]
+        return joint_positions, joint_velocities, joint_torques
+
+    def setJointPosition(self, robot, position, kp=1.0, kv=0.3):
+        import pybullet as p
+        num_joints = p.getNumJoints(robot)
+        zero_vec = [0.0] * num_joints
+        if len(position) == num_joints:
+            p.setJointMotorControlArray(robot, range(num_joints), p.POSITION_CONTROL,
+                targetPositions=position, targetVelocities=zero_vec,
+                positionGains=[kp] * num_joints, velocityGains=[kv] * num_joints)
+
+    def testJacobian(self):
+        try:
+            import numpy as np
+        except ImportError:
+            return
+        import pybullet as p
+
+        clid = p.connect(p.SHARED_MEMORY)
+        if (clid<0):
+            p.connect(p.DIRECT)
+
+        time_step = 0.001
+        gravity_constant = -9.81
+
+        urdfs = ["TwoJointRobot_w_fixedJoints.urdf",
+                 "TwoJointRobot_w_fixedJoints.urdf",
+                 "kuka_iiwa/model.urdf",
+                 "kuka_lwr/kuka.urdf"]
+        for urdf in urdfs:
+            p.resetSimulation()
+            p.setTimeStep(time_step)
+            p.setGravity(0.0, 0.0, gravity_constant)
+
+            robotId = p.loadURDF(urdf, useFixedBase=True)
+            p.resetBasePositionAndOrientation(robotId,[0,0,0],[0,0,0,1])
+            numJoints = p.getNumJoints(robotId)
+            endEffectorIndex = numJoints - 1
+
+            # Set a joint target for the position control and step the sim.
+            self.setJointPosition(robotId, [0.1 * (i % 3)
+                                            for i in range(numJoints)])
+            p.stepSimulation()
+
+            # Get the joint and link state directly from Bullet.
+            mpos, mvel, mtorq = self.getMotorJointStates(robotId)
+
+            result = p.getLinkState(robotId, endEffectorIndex,
+                    computeLinkVelocity=1, computeForwardKinematics=1)
+            link_trn, link_rot, com_trn, com_rot, frame_pos, frame_rot, link_vt, link_vr = result
+            # Get the Jacobians for the CoM of the end-effector link.
+            # Note that in this example com_rot = identity, and we would need to use com_rot.T * com_trn.
+            # The localPosition is always defined in terms of the link frame coordinates.
+
+            zero_vec = [0.0] * len(mpos)
+            jac_t, jac_r = p.calculateJacobian(robotId, endEffectorIndex,
+                    com_trn, mpos, zero_vec, zero_vec)
+
+            assert(np.allclose(np.array(jac_t).dot(mvel), link_vt))
+            assert(np.allclose(np.array(jac_r).dot(mvel), link_vr))
+        p.disconnect()
+
 if __name__ == '__main__':
     unittest.main()