diff --git a/leo_docking/leo_docking/states/dock.py b/leo_docking/leo_docking/states/dock.py
index 0d12102..55e1388 100755
--- a/leo_docking/leo_docking/states/dock.py
+++ b/leo_docking/leo_docking/states/dock.py
@@ -53,6 +53,8 @@ class Dock(BaseDockingState):
     - the effort on the wheel motors is high enough - the rover is pushing against the base
     """
 
+    params: DockParams
+
     def __init__(
         self,
         global_params: GlobalParams,
@@ -78,12 +80,14 @@ def __init__(
         self.end_time = 0.0
 
         self.charging = False
-        self.battery_reference = None
+        self.battery_reference: Optional[float] = None
         self.acc_data = 0.0
         self.counter = 0
 
         self.effort_stop = False
-        self.effort_buf = Queue(maxsize=self.global_params.effort_buffer_size)
+        self.effort_buf: Queue[float] = Queue(
+            maxsize=self.global_params.effort_buffer_size
+        )
         super().reset_state()
 
     def reset_state(self):
diff --git a/leo_docking/leo_docking/states/docking_state_machine.py b/leo_docking/leo_docking/states/docking_state_machine.py
index 3a4dab3..1ed9493 100644
--- a/leo_docking/leo_docking/states/docking_state_machine.py
+++ b/leo_docking/leo_docking/states/docking_state_machine.py
@@ -18,7 +18,7 @@
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
-from typing import Callable
+from typing import Callable, Any
 
 from smach import StateMachine, Sequence
 
@@ -49,7 +49,7 @@ def __init__(
         debug_visualizations_cb: Callable,
     ):
         self.params = state_machine_params
-        self.states = {
+        self.states: dict[str, dict[str, Any]] = {
             "Start": {
                 "state": StartState(self.params.start_params, logger),
                 "transitions": {
diff --git a/leo_docking/leo_docking/states/reach_docking_area.py b/leo_docking/leo_docking/states/reach_docking_area.py
index 680beec..f2fe94a 100755
--- a/leo_docking/leo_docking/states/reach_docking_area.py
+++ b/leo_docking/leo_docking/states/reach_docking_area.py
@@ -245,6 +245,8 @@ class RotateToDockArea(BaseDockAreaState):
     responsible for rotating the rover towards target point in the area (default: 2m in straight
     line from docking base)."""
 
+    params: RotateToDockAreaParams
+
     def __init__(
         self,
         local_params: RotateToDockAreaParams,
@@ -270,6 +272,8 @@ class RideToDockArea(BaseDockAreaState):
     responsible for driving the rover to the target point in the area (default: 2m in straight line
     from docking base)"""
 
+    params: RideToDockAreaParams
+
     def __init__(
         self,
         local_params: RideToDockAreaParams,
@@ -293,6 +297,8 @@ class RotateToBoard(BaseDockAreaState):
     """The third state of the sequence state machine getting rover to docking area;
     responsible for rotating the rover toward board on the docking base"""
 
+    params: RotateToBoardParams
+
     def __init__(
         self,
         local_params: RotateToBoardParams,
diff --git a/leo_docking/leo_docking/states/reach_docking_pose.py b/leo_docking/leo_docking/states/reach_docking_pose.py
index 09ad3f2..9624e7e 100755
--- a/leo_docking/leo_docking/states/reach_docking_pose.py
+++ b/leo_docking/leo_docking/states/reach_docking_pose.py
@@ -208,9 +208,10 @@ def aruco_detection_cb(self, data: ArucoDetection) -> None:
                         board, distance=self.global_params.docking_point_dist
                     )
                     board_normalized = normalize_board(board)
-                    self.debug_visualizations_cb(
-                        docking_point, docking_orientation, board_normalized
-                    )
+                    if self.debug_visualizations_cb:
+                        self.debug_visualizations_cb(
+                            docking_point, docking_orientation, board_normalized
+                        )
 
                 with self.route_lock:
                     self.calculate_route_left(board)
@@ -308,6 +309,8 @@ class RotateToDockingPoint(BaseDockingState):
     (default: 0.6m in straight line from docking base).
     """
 
+    params: RotateToDockingPointParams
+
     def __init__(
         self,
         global_params: GlobalParams,
@@ -355,6 +358,8 @@ class ReachDockingPoint(BaseDockingState):
     to be always looking on the docking point.
     """
 
+    params: ReachDockingPointParams
+
     def __init__(
         self,
         global_params: GlobalParams,
@@ -444,6 +449,8 @@ class ReachDockingOrientation(BaseDockingState):
     rover needs just to drive forward to reach the base.
     """
 
+    params: ReachDockingOrientationParams
+
     def __init__(
         self,
         global_params: GlobalParams,