allow picking one of multiple end-effectors to control at a time

abr_control/controllers/osc.py

@@ -146,7 +146,7 @@ def _Mx(self, M, J, threshold=1e-3):
 
         return Mx, M_inv
 
-    def _calc_orientation_forces(self, target_abg, q):
+    def _calc_orientation_forces(self, target_abg, q, ref_frame="EE"):
         """Calculate the desired Euler angle forces to apply to the arm to
         move the end-effector to the target orientation
 
@@ -167,7 +167,7 @@ def _calc_orientation_forces(self, target_abg, q):
                 )
             )
             # get the quaternion for the end effector
-            q_e = self.robot_config.quaternion("EE", q=q)
+            q_e = self.robot_config.quaternion(ref_frame, q=q)
             q_r = transformations.quaternion_multiply(
                 q_d, transformations.quaternion_conjugate(q_e)
             )
@@ -176,7 +176,7 @@ def _calc_orientation_forces(self, target_abg, q):
         elif self.orientation_algorithm == 1:
             # From (Caccavale et al, 1997) Section IV Quaternion feedback
             # get rotation matrix for the end effector orientation
-            R_e = self.robot_config.R("EE", q)
+            R_e = self.robot_config.R(ref_frame, q)
             # get rotation matrix for the target orientation
             R_d = transformations.euler_matrix(
                 target_abg[0], target_abg[1], target_abg[2], axes="rxyz"
@@ -256,7 +256,7 @@ def generate(
 
         # if orientation is being controlled
         if np.sum(self.ctrlr_dof[3:]) > 0:
-            u_task[3:] = self._calc_orientation_forces(target[3:], q)
+            u_task[3:] = self._calc_orientation_forces(target[3:], q, ref_frame=ref_frame)
 
         # task space integrated error term
         if self.ki != 0:

abr_control/controllers/path_planners/inverse_kinematics.py

@@ -34,6 +34,7 @@ def generate_path(
         plot=False,
         method=3,
         axes="rxyz",
+        EE="EE",
     ):
         """
 
@@ -84,13 +85,13 @@ def generate_path(
 
         q = np.copy(position)
         for ii in range(n_timesteps):
-            J = self.robot_config.J("EE", q=q)
-            Tx = self.robot_config.Tx("EE", q=q)
+            J = self.robot_config.J(EE, q=q)
+            Tx = self.robot_config.Tx(EE, q=q)
             ee_track.append(Tx)
 
             dx = target_position[:3] - Tx
 
-            Qe = self.robot_config.quaternion("EE", q=q)
+            Qe = self.robot_config.quaternion(EE, q=q)
             # Method 4
             dr = Qe[0] * Qd[1:] - Qd[0] * Qe[1:] - np.cross(Qd[1:], Qe[1:])
 

abr_control/interfaces/mujoco.py

@@ -44,7 +44,7 @@ def __init__(
         # if we want the offscreen render context
         self.create_offscreen_rendercontext = create_offscreen_rendercontext
 
-    def connect(self, joint_names=None, camera_id=-1):
+    def connect(self, joint_names=None, camera_id=-1, EE="EE"):
         """
         joint_names: list, optional (Default: None)
             list of joint names to send control signal to and get feedback from
@@ -69,7 +69,7 @@ def connect(self, joint_names=None, camera_id=-1):
             bodyid = mujoco.mj_name2id(
                 self.model,
                 mujoco.mjtObj.mjOBJ_BODY,
-                "EE"
+                EE
             )
             # and working back to the world body
             while self.model.body_parentid[bodyid] != 0:
@@ -138,7 +138,7 @@ def get_joint_dyn_addrs(self, jntadr):
         joint_dyn_addr = list(range(first_dyn, first_dyn + dynvec_length))[::-1]
         return joint_dyn_addr
 
-    def send_forces(self, u, update_display=True, use_joint_dyn_addrs=True):
+    def send_forces(self, u, update_display=True, use_joint_dyn_addrs=True, EE="EE"):
         """Apply the specified torque to the robot joints, move the simulation
         one time step forward, and update the position of the hand object.
 
@@ -161,11 +161,11 @@ def send_forces(self, u, update_display=True, use_joint_dyn_addrs=True):
 
         # Update position of hand object
         feedback = self.get_feedback()
-        hand_xyz = self.robot_config.Tx(name="EE", q=feedback["q"])
+        hand_xyz = self.robot_config.Tx(name=EE, q=feedback["q"])
         self.set_mocap_xyz("hand", hand_xyz)
 
         # Update orientation of hand object
-        hand_quat = self.robot_config.quaternion(name="EE", q=feedback["q"])
+        hand_quat = self.robot_config.quaternion(name=EE, q=feedback["q"])
         self.set_mocap_orientation("hand", hand_quat)
 
         if self.visualize and update_display: