clean up, and updating to be compatible with jrl==0.1.0

.gitignore

@@ -7,6 +7,7 @@ thirdparty/mdn/mdn/__pycache__/models.cpython-38.pyc
 # Virtual Environments
 venv/
 venvTemp/
+venv-test/
 
 # IDE Config
 .vscode/

ikflow/evaluation_utils.py

@@ -70,37 +70,30 @@ def solution_pose_errors(
 
     Args:
         robot (Robot): The Robot which contains the FK function we will use
-        solutions (Union[torch.Tensor, np.ndarray]): [n x 7] IK solutions for the given target pose
+        solutions torch.Tensor: [n x 7] IK solutions for the given target pose
         target_pose (np.ndarray): [7] the target pose the IK solutions were generated for
 
     Returns:
         Tuple[np.ndarray, np.ndarray]: The L2, and angular (rad) errors of IK solutions for the given target_pose
     """
-    assert isinstance(
-        target_poses, (np.ndarray, torch.Tensor)
-    ), f"target_poses must be a torch.Tensor or np.ndarray (got {type(target_poses)})"
+    assert isinstance(target_poses, torch.Tensor), f"target_poses must be a torch.Tensor (got {type(target_poses)})"
     assert isinstance(solutions, torch.Tensor), f"solutions must be a torch.Tensor (got {type(solutions)})"
     n_solutions = solutions.shape[0]
     if n_solutions >= 1000:
         print("Heads up: It may be faster to run solution_pose_errors() with pytorch directly on the cpu/gpu")
 
-    if isinstance(target_poses, torch.Tensor):
-        target_poses = target_poses.detach().cpu().numpy()
     target_poses = _get_target_pose_batch(target_poses, solutions.shape[0])
-
-    ee_pose_ikflow = robot.forward_kinematics(solutions[:, 0 : robot.ndof].detach().cpu().numpy())
+    ee_pose_ikflow = robot.forward_kinematics(solutions[:, 0 : robot.ndof])
     rot_output = ee_pose_ikflow[:, 3:]
 
     # Positional Error
-    l2_errors = np.linalg.norm(ee_pose_ikflow[:, 0:3] - target_poses[:, 0:3], axis=1)
+    l2_errors = torch.norm(ee_pose_ikflow[:, 0:3] - target_poses[:, 0:3], dim=1)
     rot_target = target_poses[:, 3:]
     assert rot_target.shape == rot_output.shape
 
     # Surprisingly, this is almost always faster to calculate on the gpu than on the cpu. I would expect the opposite
     # for low number of solutions (< 200).
-    q_target_pt = torch.tensor(rot_target, device=DEVICE, dtype=torch.float32)
-    q_current_pt = torch.tensor(rot_output, device=DEVICE, dtype=torch.float32)
-    ang_errors = geodesic_distance_between_quaternions(q_target_pt, q_current_pt).detach().cpu().numpy()
+    ang_errors = geodesic_distance_between_quaternions(rot_target, rot_output)
     return l2_errors, ang_errors
 
 
@@ -145,9 +138,7 @@ def evaluate_solutions(
             robot, ee_pose_target, samples
         )
     """
-    assert isinstance(
-        target_poses, (np.ndarray, torch.Tensor)
-    ), f"target_poses must be a torch.Tensor or np.ndarray (got {type(target_poses)})"
+    assert isinstance(target_poses, torch.Tensor), f"target_poses must be a torch.Tensor (got {type(target_poses)})"
     assert isinstance(solutions, torch.Tensor), f"solutions must be a torch.Tensor (got {type(solutions)})"
     target_poses = _get_target_pose_batch(target_poses, solutions.shape[0])
     l2_errors, angular_errors = solution_pose_errors(robot, solutions, target_poses)

ikflow/ikflow_solver.py

@@ -71,15 +71,6 @@ def conditional_size(self) -> int:
         """
         return self.dim_cond
 
-        """Refine a batch of IK solutions using the klampt IK solver
-        Args:
-            ikflow_solutions (torch.Tensor): A batch of IK solutions of the form [batch x n_dofs]
-            target_pose (Union[List[float], np.ndarray]): The target endpose(s). Must either be of the form
-                                                            [x, y, z, q0, q1, q2, q3] or be a [batch x 7] numpy array
-        Returns:
-            torch.Tensor: A batch of IK refined solutions [batch x n_dofs]
-        """
-
     def _run_inference(
         self,
         latent: torch.Tensor,
@@ -109,7 +100,7 @@ def _run_inference(
 
     def _calculate_pose_error(self, qs: torch.Tensor, target_poses: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         # check error
-        pose_realized = self.robot.forward_kinematics_batch(qs)
+        pose_realized = self.robot.forward_kinematics(qs)
         pos_errors = torch.norm(pose_realized[:, 0:3] - target_poses[:, 0:3], dim=1)
         rot_errors = geodesic_distance_between_quaternions(target_poses[:, 3:], pose_realized[:, 3:])
         return pos_errors, rot_errors
@@ -201,10 +192,10 @@ def _generate_exact_ik_solutions(
                 assert len(q) == n_invalid * repeat_count
                 t02 = time()
                 if run_lma_on_cpu and not do_run_entirely_on_cpu:
-                    q = self.robot.inverse_kinematics_single_step_levenburg_marquardt(target_poses_tiled.cpu(), q.cpu())
+                    q = self.robot.inverse_kinematics_step_levenburg_marquardt(target_poses_tiled.cpu(), q.cpu())
                     q = q.to(device)
                 else:
-                    q = self.robot.inverse_kinematics_single_step_levenburg_marquardt(target_poses_tiled, q)
+                    q = self.robot.inverse_kinematics_step_levenburg_marquardt(target_poses_tiled, q)
                 t_lma += time() - t02
                 pos_errors, rot_errors = self._calculate_pose_error(q, target_poses_tiled)
                 valids_i_tiled = torch.logical_and(pos_errors < pos_error_threshold, rot_errors < rot_error_threshold)
@@ -253,7 +244,7 @@ def _generate_exact_ik_solutions(
     def generate_ik_solutions(
         self,
         y: torch.Tensor,
-        n: Optional[int],
+        n: Optional[int] = None,
         latent: Optional[torch.Tensor] = None,
         latent_distribution: str = "gaussian",
         latent_scale: float = 1.0,

ikflow/model_descriptions.yaml

@@ -66,7 +66,7 @@ fetch_arm__large__mh186_9.25m:
 
 
 
-# TODO: add weights back once Iiwa7 is added to Jrl (currently ommited because missing capsules)
+# TODO: add weights back once Iiwa7 is added to jrl (currently ommited because missing capsules)
 
 # ===========================================================
 # === Iiwa7

ikflow/training/lt_data.py

@@ -27,8 +27,8 @@ def __init__(
 
         dataset_directory = get_dataset_directory(self._robot_name)
         assert os.path.isdir(dataset_directory), (
-            f"Directory '{dataset_directory}' doesn't exist - have you created the dataset for this robot yet? (try"
-            f" `python scripts/build_dataset.py --robot_name={robot_name} --training_set_size=10000000`)"
+            f"Directory '{dataset_directory}' doesn't exist - have you created the dataset for this robot yet? Your probably want to run"
+            f" `python scripts/build_dataset.py --robot_name={robot_name} --training_set_size=10000000 --only_non_self_colliding`)"
         )
         samples_tr_file_path, poses_tr_file_path, samples_te_file_path, poses_te_file_path, _ = get_dataset_filepaths(
             dataset_directory, dataset_tags

ikflow/training/lt_model.py

@@ -228,7 +228,9 @@ def on_after_backward(self):
     def validation_step(self, batch, batch_idx) -> Dict[str, torch.Tensor]:
         del batch_idx
         _, y = batch
-        ee_pose_target = y.cpu().detach().numpy()[0]
+        ee_pose_target = y[0]
+
+        # Generates m=self.hparams.samples_per_pose solutions for a single pose
         solutions, model_runtime = self.generate_ik_solutions(
             ee_pose_target, self.hparams.samples_per_pose, return_runtime=True
         )
@@ -339,8 +341,8 @@ def generate_ik_solutions(
 
         # Note: No code change required here to handle using/not using softflow.
         conditional = torch.zeros(m, self.ik_solver.dim_cond)
-        conditional[:, 0:3] = torch.FloatTensor(y[:3])
-        conditional[:, 3 : 3 + 4] = torch.FloatTensor(np.array([y[3:]]))
+        conditional[:, 0:3] = y[:3]
+        conditional[:, 3 : 3 + 4] = y[3:]
         conditional = conditional.to(config.device)
 
         shape = (m, self.dim_tot)

ikflow/visualizations.py

@@ -129,7 +129,7 @@ def loop_fn(worlds, _demo_state):
         else:
             # (B x 3*(n+1) )
             x_torch = torch.from_numpy(x_random).float().to(device)
-            fk = robot.forward_kinematics_batch(x_torch)
+            fk = robot.forward_kinematics(x_torch)
             ee_pose = fk[0, 0:3]
             vis.add("ee", (so3.identity(), ee_pose[0:3]), length=0.15, width=2)
 
@@ -260,8 +260,8 @@ def loop_fn(worlds, _demo_state):
         ik_solutions = ik_solver.generate_ik_solutions(_demo_state.target_pose, nb_sols, latent=latent)
         l2_errors, ang_errors = solution_pose_errors(ik_solver.robot, ik_solutions, _demo_state.target_pose)
 
-        _demo_state.ave_l2_error = np.mean(l2_errors) * 1000
-        _demo_state.ave_ang_error = np.rad2deg(np.mean(ang_errors))
+        _demo_state.ave_l2_error = l2_errors.mean().item() * 1000
+        _demo_state.ave_ang_error = np.rad2deg(ang_errors.mean().item())
 
         # Update viz with solutions
         qs = robot._x_to_qs(ik_solutions.detach().cpu().numpy())
@@ -331,7 +331,6 @@ def __init__(self):
 
     def setup_fn(worlds):
         vis.add("robot", worlds[0].robot(0))
-        vis.setColor("robot", 1, 0.1, 0.1, 1)
         assert len(worlds) == 1
 
     def loop_fn(worlds, _demo_state):
@@ -369,6 +368,5 @@ def viz_update_fn(worlds, _demo_state):
         viz_update_fn,
         time_p_loop=time_p_loop,
         title=title,
-        load_terrain=True,
         demo_state=demo_state,
     )