Update test_inverse_kinematics.py

Author: StoneT2000

tests/test_inverse_kinematics.py

@@ -24,16 +24,27 @@ def make_transparent(link):
     for link in visual_data:
         make_transparent(link)
 
-
-def test_jacobian_follower():
+def create_test_chain(robot="kuka_iiwa", device="cpu"):
+    if robot == "kuka_iiwa":
+        urdf = "kuka_iiwa/model.urdf"
+        search_path = pybullet_data.getDataPath()
+        full_urdf = os.path.join(search_path, urdf)
+        chain = pk.build_serial_chain_from_urdf(open(full_urdf).read(), "lbr_iiwa_link_7")
+        chain = chain.to(device=device)
+    elif robot == "widowx":
+        urdf = "widowx/wx250s.urdf"
+        full_urdf = urdf
+        chain = pk.build_serial_chain_from_urdf(open(full_urdf, "rb").read(), "ee_gripper_link")
+        chain = chain.to(device=device)
+    else:
+        raise NotImplementedError(f"Robot {robot} not implemented")
+    return chain, urdf
+
+def test_jacobian_follower(robot="kuka_iiwa"):
     pytorch_seed.seed(2)
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    # device = "cpu"
-    urdf = "kuka_iiwa/model.urdf"
     search_path = pybullet_data.getDataPath()
-    full_urdf = os.path.join(search_path, urdf)
-    chain = pk.build_serial_chain_from_urdf(open(full_urdf).read(), "lbr_iiwa_link_7")
-    chain = chain.to(device=device)
+    chain, urdf = create_test_chain(robot=robot, device=device)
 
     # robot frame
     pos = torch.tensor([0.0, 0.0, 0.0], device=device)
@@ -45,7 +56,7 @@ def test_jacobian_follower():
     # generate random goal joint angles (so these are all achievable)
     # use the joint limits to generate random joint angles
     lim = torch.tensor(chain.get_joint_limits(), device=device)
-    goal_q = torch.rand(M, 7, device=device) * (lim[1] - lim[0]) + lim[0]
+    goal_q = torch.rand(M, lim.shape[1], device=device) * (lim[1] - lim[0]) + lim[0]
 
     # get ee pose (in robot frame)
     goal_in_rob_frame_tf = chain.forward_kinematics(goal_q)
@@ -209,5 +220,6 @@ def test_ik_in_place_no_err():
 
 
 if __name__ == "__main__":
-    test_jacobian_follower()
+    # test_jacobian_follower(robot="kuka_iiwa")
+    test_jacobian_follower(robot="widowx")
     test_ik_in_place_no_err()