Cal zim/treadward

README.md

@@ -10,8 +10,6 @@ It is expected that you have a `hebi-py` version of at least `2.2.0` for this re
 
 **Note that this revision does not support Python 2!**
 
-Additionally, some of the kits examples require the [PySDL2 Package](https://pypi.org/project/PySDL2/). You will need the `SDL2` shared library for this to work. For Windows, this library is included for you in this repository.
-
 ### Nightly API
 
 If you want access to the newest API, you can download one of our nightly builds by doing:

kits/arms/ar_control_sm.py

@@ -1,8 +1,7 @@
 
 from enum import Enum, auto
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from time import time, sleep
-import os
 import numpy as np
 from scipy.spatial.transform import Rotation as R
 
@@ -12,23 +11,24 @@
 import typing
 if typing.TYPE_CHECKING:
     from typing import Callable
-    import numpy.typing as npt
+    from numpy.typing import NDArray
     from hebi._internal.mobile_io import MobileIO
     from hebi.arm import Arm, Gripper
 
 
 class ArmControlState(Enum):
     STARTUP = auto()
     HOMING = auto()
+    IDLE = auto()
     TELEOP = auto()
     DISCONNECTED = auto()
     EXIT = auto()
 
 
 @dataclass(kw_only=True)
 class ArmMobileIOInputs:
-    phone_pos: 'npt.NDArray[np.float32]' = field(default_factory=lambda: np.array([0., 0., 0.]))
-    phone_rot: 'npt.NDArray[np.float64]' = field(default_factory=lambda: np.eye(3))
+    phone_pos: 'NDArray[np.float32]'
+    phone_rot: 'NDArray[np.float64]'
     ar_scaling: float = 1.0
     lock_toggle: bool = False
     locked: bool = True
@@ -37,38 +37,31 @@ class ArmMobileIOInputs:
 
 
 class ArmMobileIOControl:
-    def __init__(self, arm: 'Arm', gripper: 'Gripper | None' = None, homing_time=5.0, traj_duration=1.0, xyz_scale: 'npt.NDArray[np.float64]' = np.ones(3)):
+    def __init__(self,
+                 arm: 'Arm',
+                 gripper: 'Gripper | None' = None,
+                 homing_time=5.0,
+                 traj_duration=1.0,
+                 xyz_scale: 'NDArray[np.float64]' = np.ones(3),
+                 ):
+
         self.namespace = ''
 
         self.state = ArmControlState.STARTUP
         self.arm = arm
         self.gripper = gripper
 
-        # 5 DoF home
-        # self.arm_xyz_home = [0.34, 0.0, 0.23]
-        # 6 DoF home
-        self.arm_xyz_home = [0.5, 0.0, 0.0]
-        self.arm_rot_home: 'npt.NDArray[np.float64]' = (R.from_euler(
-            'z', np.pi / 2) * R.from_euler('x', np.pi)).as_matrix()
         self.homing_time = homing_time
         self.traj_duration = traj_duration
-
-        # 5 DoF seed
-        # self.arm_seed_ik = np.array([0.25, -1.0, 0, -0.75, 0])
-        # 6 DoF seed
-        self.arm_seed_ik = np.array([0.3, 1.2, 2.2, 2.9, -1.57, 0])
-        self.arm_home = self.arm.ik_target_xyz_so3(
-            self.arm_seed_ik,
-            self.arm_xyz_home,
-            self.arm_rot_home)
-
         self.xyz_scale = xyz_scale
 
-        self.last_locked_xyz = self.arm_xyz_home.copy()
-        self.last_locked_rot = self.arm_rot_home.copy()
-        self.last_locked_seed = self.arm_home.copy()
+        self.arm_ik_seed = np.array([0.3, 1.2, 2.2, 2.9, -1.57, 0])
+
+        home_orientation = R.from_euler('z', np.pi / 2) * R.from_euler('x', np.pi)
+        self.set_arm_home(
+            np.array([0.5, 0.0, 0.0]),
+            home_orientation.as_matrix())
 
-        self.locked = True
         self._transition_handlers: 'list[Callable[[ArmMobileIOControl, ArmControlState], None]]' = [
         ]
         self.last_cmd_t = time()
@@ -77,6 +70,13 @@ def __init__(self, arm: 'Arm', gripper: 'Gripper | None' = None, homing_time=5.0
     def running(self):
         return self.state is not self.state.EXIT
 
+    def set_arm_home(self, xyz, rot):
+        self.arm_xyz_home = xyz
+        self.arm_rot_home = rot
+        self.arm_home = self.arm.ik_target_xyz_so3(
+            self.arm_ik_seed,
+            xyz, rot)
+
     def send(self):
         self.arm.send()
         if self.gripper:
@@ -97,13 +97,13 @@ def update(self, t_now: float, arm_input: 'ArmMobileIOInputs | None'):
         # Reset the timeout
         self.last_cmd_t = t_now
         # Transition to teleop if mobileIO is reconnected
-        last_pos = self.arm.last_feedback.position
+        last_pos = self.arm.last_feedback.position_command
+        if np.any(np.isnan(last_pos)):
+            last_pos = self.arm.last_feedback.position
         if self.state is self.state.DISCONNECTED:
             self.last_cmd_t = t_now
             print(self.namespace + 'Controller reconnected, demo continued.')
-            # Lock arm when reconnected
-            self.locked = True
-            self.transition_to(t_now, self.state.TELEOP)
+            self.transition_to(t_now, self.state.IDLE)
 
         # After startup, transition to homing
         elif self.state is self.state.STARTUP:
@@ -112,56 +112,44 @@ def update(self, t_now: float, arm_input: 'ArmMobileIOInputs | None'):
         # If homing is complete, transition to teleop
         elif self.state is self.state.HOMING:
             if self.arm.at_goal:
-                self.phone_xyz_home = arm_input.phone_pos
-                self.phone_rot_home = arm_input.phone_rot
-
-                self.last_locked_seed = last_pos.copy()
-                xyz = np.empty(3)
-                orientation = np.empty((3, 3))
-                self.arm.FK(last_pos,
-                            xyz_out=xyz,
-                            orientation_out=orientation)
+                self.transition_to(t_now, self.state.IDLE)
 
-                self.last_locked_xyz = xyz
-                self.last_locked_rot = orientation
+        # Idle mode
+        elif self.state is self.state.IDLE:
+            if arm_input.home:
+                self.transition_to(t_now, self.state.HOMING)
+                return
 
+            if arm_input.lock_toggle and not arm_input.locked:
+                self.phone_xyz_home = arm_input.phone_pos
+                self.phone_rot_home = arm_input.phone_rot
+                self.set_locked_pose(last_pos)
                 self.transition_to(t_now, self.state.TELEOP)
+                return
 
         # Teleop mode
         elif self.state is self.state.TELEOP:
-            if arm_input.home:
-                self.transition_to(t_now, self.state.HOMING)
+            if arm_input.lock_toggle and arm_input.locked:
+                g = hebi.arm.Goal(self.arm.size)
+                g.add_waypoint(t=0.3,
+                               position=self.arm.last_feedback.position_command,
+                               velocity=np.zeros(g.dof_count),
+                               acceleration=np.zeros(g.dof_count))
+                self.arm.set_goal(g)
+                self.transition_to(t_now, self.state.IDLE)
                 return
 
-            if arm_input.lock_toggle:
-                self.locked = arm_input.locked
+            if (arm_goal := self.compute_arm_goal(arm_input)) is not None:
+                self.arm.set_goal(arm_goal)
 
-            if not self.locked:
-                arm_goal = self.compute_arm_goal(arm_input)
-                if arm_goal is not None:
-                    self.arm.set_goal(arm_goal)
-            else:
-                self.phone_xyz_home = arm_input.phone_pos
-                self.phone_rot_home = arm_input.phone_rot
-
-                self.last_locked_seed = last_pos.copy()
-                xyz = np.zeros(3)
-                orientation = np.zeros((3, 3))
-                self.arm.FK(last_pos,
-                            xyz_out=xyz,
-                            orientation_out=orientation)
-
-                self.last_locked_xyz = xyz
-                self.last_locked_rot = orientation
-
-            if self.gripper is not None:
-                gripper_closed = self.gripper.state == 1.0
-                if arm_input.gripper_closed and not gripper_closed:
-                    print('Gripper Close')
-                    self.gripper.close()
-                elif not arm_input.gripper_closed and gripper_closed:
-                    print('Gripper Open')
-                    self.gripper.open()
+        if self.gripper is not None:
+            gripper_closed = self.gripper.state == 1.0
+            if arm_input.gripper_closed and not gripper_closed:
+                print('Gripper Close')
+                self.gripper.close()
+            elif not arm_input.gripper_closed and gripper_closed:
+                print('Gripper Open')
+                self.gripper.open()
 
     def transition_to(self, t_now: float, state: ArmControlState):
         # self transitions are noop
@@ -172,6 +160,9 @@ def transition_to(self, t_now: float, state: ArmControlState):
             print(self.namespace + "TRANSITIONING TO HOMING")
             self.home(self.homing_time)
 
+        elif state is self.state.IDLE:
+            print(self.namespace + "TRANSITIONING TO IDLE")
+
         elif state is self.state.TELEOP:
             print(self.namespace + "TRANSITIONING TO TELEOP")
 
@@ -185,7 +176,22 @@ def transition_to(self, t_now: float, state: ArmControlState):
             handler(self, state)
         self.state = state
 
+    def set_locked_pose(self, position):
+        xyz = np.empty(3)
+        orientation = np.empty((3, 3))
+        self.arm.FK(position,
+                    xyz_out=xyz,
+                    orientation_out=orientation)
+
+        self.last_locked_seed = position.copy()
+        self.last_locked_xyz = xyz
+        self.last_locked_rot = orientation
+
     def compute_arm_goal(self, arm_input: ArmMobileIOInputs):
+        if np.any(np.isnan(arm_input.phone_pos)) or np.any(np.isnan(arm_input.phone_rot)):
+            print(f'NaN detected in pose from AR device:\n{arm_input.phone_pos}\n{arm_input.phone_rot}')
+            return None
+
         phone_offset = arm_input.phone_pos - self.phone_xyz_home
         rot_mat = self.phone_rot_home
         arm_xyz_target = self.last_locked_xyz + arm_input.ar_scaling * \
@@ -198,10 +204,20 @@ def compute_arm_goal(self, arm_input: ArmMobileIOInputs):
         if arm_input.ar_scaling == 0.0:
             self.phone_xyz_home = arm_input.phone_pos
 
-        joint_target = self.arm.ik_target_xyz_so3(
-            self.last_locked_seed,
-            arm_xyz_target,
-            arm_rot_target)
+        try:
+            joint_target = self.arm.ik_target_xyz_so3(
+                self.last_locked_seed,
+                arm_xyz_target,
+                arm_rot_target)
+        except Exception as e:
+            print('-----------------')
+            print(arm_xyz_target)
+            print('-----------------')
+            print(self.last_locked_seed)
+            print('-----------------')
+            print(arm_rot_target)
+            print('-----------------')
+            raise e
 
         arm_goal = hebi.arm.Goal(self.arm.size)
         arm_goal.add_waypoint(t=self.traj_duration, position=joint_target)
@@ -232,8 +248,7 @@ def parse_mobile_feedback(m: 'MobileIO'):
     if m.get_button_state(8):
         return True, None
 
-    if m.get_button_state(1):
-        return False, ArmMobileIOInputs(home=True)
+    user_home = m.get_button_state(1)
 
     try:
         wxyz = m.orientation
@@ -244,6 +259,7 @@ def parse_mobile_feedback(m: 'MobileIO'):
         rotation = np.eye(3)
 
     arm_input = ArmMobileIOInputs(
+        home=user_home,
         phone_pos=np.copy(m.position),
         phone_rot=rotation,
         lock_toggle=m.get_button_state(5),