Add wx250s robot for testing

Author: LemonPi

tests/test_inverse_kinematics.py

@@ -10,8 +10,9 @@
 import pybullet as p
 import pybullet_data
 
-visualize = False
+visualize = True
 
+TEST_DIR = os.path.dirname(__file__)
 
 def _make_robot_translucent(robot_id, alpha=0.4):
     def make_transparent(link):
@@ -206,6 +207,73 @@ def test_ik_in_place_no_err():
     assert torch.allclose(sol.err_rot[0], torch.zeros(1, device=device), atol=1e-6)
 
 
+def test_extract_serial_chain_from_tree():
+    pytorch_seed.seed(2)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # device = "cpu"
+    urdf = "widowx/wx250s.urdf"
+    full_urdf = os.path.join(TEST_DIR, urdf)
+    chain = pk.build_serial_chain_from_urdf(open(full_urdf, mode="rb").read(), "ee_gripper_link")
+    # chain = pk.SerialChain(chain, "ee_gripper_link", "base_link")
+    chain = chain.to(device=device)
+
+    # full chain should have DOF = 8, however since we are creating just a serial chain to ee_gripper_link, should be 6
+    # TODO pretty tree print
+    """
+    /
+    └── gripper_bar_link
+        └── fingers_link
+            ├── left_finger_link
+            ├── right_finger_link
+            └── ee_gripper_link
+    """
+
+    # robot frame
+    pos = torch.tensor([0.0, 0.0, 0.0], device=device)
+    rot = torch.tensor([0.0, 0.0, 0.0], device=device)
+    rob_tf = pk.Transform3d(pos=pos, rot=rot, device=device)
+
+    # world frame goal
+    M = 1000
+    # 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]
+
+    # get ee pose (in robot frame)
+    goal_in_rob_frame_tf = chain.forward_kinematics(goal_q)
+
+    # transform to world frame for visualization
+    goal_tf = rob_tf.compose(goal_in_rob_frame_tf)
+    goal = goal_tf.get_matrix()
+    goal_pos = goal[..., :3, 3]
+    goal_rot = pk.matrix_to_euler_angles(goal[..., :3, :3], "XYZ")
+
+    num_retries = 10
+    ik = pk.PseudoInverseIK(chain, max_iterations=30, num_retries=num_retries,
+                            joint_limits=lim.T,
+                            early_stopping_any_converged=True,
+                            early_stopping_no_improvement="all",
+                            # line_search=pk.BacktrackingLineSearch(max_lr=0.2),
+                            debug=False,
+                            lr=0.2)
+
+    # do IK
+    timer_start = timer()
+    sol = ik.solve(goal_in_rob_frame_tf)
+    timer_end = timer()
+    print("IK took %f seconds" % (timer_end - timer_start))
+    print("IK converged number: %d / %d" % (sol.converged.sum(), sol.converged.numel()))
+    print("IK took %d iterations" % sol.iterations)
+    print("IK solved %d / %d goals" % (sol.converged_any.sum(), M))
+
+    # check that solving again produces the same solutions
+    sol_again = ik.solve(goal_in_rob_frame_tf)
+    assert torch.allclose(sol.solutions, sol_again.solutions)
+    assert torch.allclose(sol.converged, sol_again.converged)
+
+
 if __name__ == "__main__":
-    test_jacobian_follower()
-    test_ik_in_place_no_err()
+    # test_jacobian_follower()
+    # test_ik_in_place_no_err()
+    test_extract_serial_chain_from_tree()

tests/widowx/README.md

@@ -0,0 +1,5 @@
+# WidowX 250 S Robot Description (URDF)
+
+The robot model here is based on the real2sim project: https://github.com/simpler-env/SimplerEnv
+
+

tests/widowx/interbotix_black.png

@@ -0,0 +1,13 @@
+# Blender MTL File: 'None'
+# Material Count: 1
+
+newmtl Material.018
+Ns 250.000000
+Ka 1.000000 1.000000 1.000000
+Kd 0.800000 0.800000 0.800000
+Ks 0.500000 0.500000 0.500000
+Ke 0.000000 0.000000 0.000000
+Ni 1.450000
+d 1.000000
+illum 2
+map_Kd ../interbotix_black.png