Small refactoring

Author: rdesarz

biped_walking_controller/plot.py

@@ -3,6 +3,8 @@
     compute_single_support_polygon,
 )
 
+from biped_walking_controller.preview_control import WalkingState
+
 
 def plot_steps(axes, steps_pose, step_shape):
     # Plot double support polygon
@@ -187,11 +189,11 @@ def traj2d(arr):  # drop last 2 samples as in your code
 
 
 def plot_contact_forces(
-    t,
-    force_rf,
-    force_lf,
-    states,
-    title="Contact force Fx",
+    t: float,
+    force_rf: float,
+    force_lf: float,
+    states: ,
+    title: str ="Contact force Fx",
 ):
     t = np.asarray(t).ravel()
     force_rf = np.asarray(force_rf).ravel()

biped_walking_controller/preview_control.py

@@ -257,7 +257,7 @@ def update_control(ctrl_mat: PreviewControllerMatrices, current_zmp, zmp_ref, x,
     return u, x_next, y_next
 
 
-class State(Enum):
+class WalkingState(Enum):
     INIT = 0
     DS = 1
     SS_LEFT = 2
@@ -275,10 +275,10 @@ class WalkingStateMachineParams:
 
 
 class WalkingStateMachine:
-    def __init__(self, params: WalkingStateMachineParams, initial_state=State.INIT):
+    def __init__(self, params: WalkingStateMachineParams, initial_state=WalkingState.INIT):
         self.params = params
         self.state = initial_state
-        self.next_ss_state = State.SS_RIGHT
+        self.next_ss_state = WalkingState.SS_RIGHT
         self.t_start = 0.0
         self.steps = None
         self.step_idx = None
@@ -298,29 +298,31 @@ def update(
 
         delta_t = t - self.t_start
 
-        if self.state == State.INIT:
+        if self.state == WalkingState.INIT:
             self._try_transition_to_single_support(t, self.params.t_init)
 
-        elif self.state == State.DS:
+        elif self.state == WalkingState.DS:
             self._try_transition_to_single_support(t, self.params.t_ds)
 
-        elif self.state == State.SS_RIGHT:
+        elif self.state == WalkingState.SS_RIGHT:
             self._try_transition_to_ds_or_end(t, lf_contact_force)
 
-        elif self.state == State.SS_LEFT:
+        elif self.state == WalkingState.SS_LEFT:
             self._try_transition_to_ds_or_end(t, rf_contact_force)
 
-        elif self.state == State.END:
+        elif self.state == WalkingState.END:
             if delta_t > self.params.t_end:
                 self.steps = None
 
-    def get_current_state(self) -> State:
+    def get_current_state(self) -> WalkingState:
         return self.state
 
     def _switch_single_support_leg(self, t: float):
         self.t_start = t
         self.state = self.next_ss_state
-        self.next_ss_state = State.SS_LEFT if self.state == State.SS_RIGHT else State.SS_RIGHT
+        self.next_ss_state = (
+            WalkingState.SS_LEFT if self.state == WalkingState.SS_RIGHT else WalkingState.SS_RIGHT
+        )
 
     def _try_transition_to_single_support(self, t: float, duration: float):
         if t - self.t_start > duration:
@@ -335,5 +337,5 @@ def _try_transition_to_ds_or_end(self, t: float, contact_force: float):
             and contact_force > self.params.force_threshold
         ):
             self.t_start = t
-            self.state = State.END if self._is_last_step() else State.DS
+            self.state = WalkingState.END if self._is_last_step() else WalkingState.DS
             self.step_idx += 1

tests/test_preview_control.py

@@ -6,7 +6,7 @@
     compute_preview_control_matrices,
     WalkingStateMachineParams,
     WalkingStateMachine,
-    State,
+    WalkingState,
 )
 
 from numpy.linalg import eigvals
@@ -104,48 +104,48 @@ def setUp(self):
         self.steps = np.ones([5, 3])  # Five steps
 
     def test_begin_init(self):
-        wsm = WalkingStateMachine(self.params, initial_state=State.INIT)
+        wsm = WalkingStateMachine(self.params, initial_state=WalkingState.INIT)
         wsm.update_steps(self.steps)
 
         wsm.update(t=0.0, rf_contact_force=100.0, lf_contact_force=100.0)
 
-        self.assertEqual(wsm.get_current_state(), State.INIT)
+        self.assertEqual(wsm.get_current_state(), WalkingState.INIT)
 
     def test_init_to_ds(self):
-        wsm = WalkingStateMachine(self.params, initial_state=State.INIT)
+        wsm = WalkingStateMachine(self.params, initial_state=WalkingState.INIT)
         wsm.update_steps(self.steps)
         wsm.update(t=0.0, rf_contact_force=100.0, lf_contact_force=100.0)
 
         wsm.update(t=self.params.t_init + 0.1, rf_contact_force=0.0, lf_contact_force=100.0)
 
-        self.assertEqual(wsm.get_current_state(), State.SS_RIGHT)
+        self.assertEqual(wsm.get_current_state(), WalkingState.SS_RIGHT)
 
     def test_begin_double_support(self):
-        wsm = WalkingStateMachine(self.params, initial_state=State.DS)
+        wsm = WalkingStateMachine(self.params, initial_state=WalkingState.DS)
         wsm.update_steps(self.steps)
 
         wsm.update(t=0.0, rf_contact_force=100.0, lf_contact_force=100.0)
 
-        self.assertEqual(wsm.get_current_state(), State.DS)
+        self.assertEqual(wsm.get_current_state(), WalkingState.DS)
 
     def test_switch_to_single_support(self):
-        wsm = WalkingStateMachine(self.params, initial_state=State.DS)
+        wsm = WalkingStateMachine(self.params, initial_state=WalkingState.DS)
         wsm.update_steps(self.steps)
 
         wsm.update(t=self.params.t_ds + 0.1, rf_contact_force=0.0, lf_contact_force=0.0)
 
-        self.assertEqual(wsm.get_current_state(), State.SS_RIGHT)
+        self.assertEqual(wsm.get_current_state(), WalkingState.SS_RIGHT)
 
     def test_do_not_switch_to_ds_if_beginning_of_phase(self):
-        wsm = WalkingStateMachine(self.params, initial_state=State.SS_RIGHT)
+        wsm = WalkingStateMachine(self.params, initial_state=WalkingState.SS_RIGHT)
         wsm.update_steps(self.steps)
 
         wsm.update(t=0.0, rf_contact_force=0.0, lf_contact_force=self.params.force_threshold + 10)
 
-        self.assertEqual(wsm.get_current_state(), State.SS_RIGHT)
+        self.assertEqual(wsm.get_current_state(), WalkingState.SS_RIGHT)
 
     def test_do_not_switch_to_ds_if_force_too_low(self):
-        wsm = WalkingStateMachine(self.params, initial_state=State.SS_RIGHT)
+        wsm = WalkingStateMachine(self.params, initial_state=WalkingState.SS_RIGHT)
         wsm.update_steps(self.steps)
 
         wsm.update(
@@ -154,10 +154,10 @@ def test_do_not_switch_to_ds_if_force_too_low(self):
             lf_contact_force=self.params.force_threshold - 10,
         )
 
-        self.assertEqual(wsm.get_current_state(), State.SS_RIGHT)
+        self.assertEqual(wsm.get_current_state(), WalkingState.SS_RIGHT)
 
     def test_switch_to_ds_if_force_too_low_and_phase_close_to_end(self):
-        wsm = WalkingStateMachine(self.params, initial_state=State.SS_LEFT)
+        wsm = WalkingStateMachine(self.params, initial_state=WalkingState.SS_LEFT)
         wsm.update_steps(self.steps)
 
         wsm.update(
@@ -166,7 +166,7 @@ def test_switch_to_ds_if_force_too_low_and_phase_close_to_end(self):
             lf_contact_force=0.0,
         )
 
-        self.assertEqual(wsm.get_current_state(), State.DS)
+        self.assertEqual(wsm.get_current_state(), WalkingState.DS)
 
 
 if __name__ == "__main__":