Merge pull request #39 from thedropbears/add_shooter_state_machine

Add shooter state machine

components/chassis.py

@@ -27,6 +27,7 @@
 from magicbot import feedback
 
 from utilities.functions import constrain_angle, rate_limit_module
+from utilities.game import is_red
 from utilities.ctre import FALCON_FREE_RPS
 from ids import CancoderIds, TalonIds
 
@@ -206,6 +207,9 @@ class Chassis:
     LENGTH = 1.0105
     WIDTH = 0.8705
     DRIVE_CURRENT_THRESHOLD = 35
+
+    HEADING_TOLERANCE = math.radians(5)
+
     # maxiumum speed for any wheel
     max_wheel_speed = FALCON_FREE_RPS * SwerveModule.DRIVE_MOTOR_REV_TO_METRES
 
@@ -220,13 +224,19 @@ class Chassis:
     do_smooth = magicbot.tunable(True)
     swerve_lock = magicbot.tunable(False)
 
+    RED_TEST_POSE = Pose2d(13.6, 5.5, 0)
+    BLUE_TEST_POSE = Pose2d(2.6, 5.5, math.pi)
+
     def setup(self) -> None:
         self.imu = navx.AHRS.create_spi()
         self.heading_controller = ProfiledPIDControllerRadians(
             1, 0, 0, TrapezoidProfileRadians.Constraints(2, 2)
         )
         self.heading_controller.enableContinuousInput(-math.pi, math.pi)
         self.snapping_to_heading = False
+        self.heading_controller.setTolerance(self.HEADING_TOLERANCE)
+
+        self.on_red_alliance = False
 
         self.modules = [
             # Front Left
@@ -272,19 +282,21 @@ def setup(self) -> None:
         self.sync_all()
         self.imu.zeroYaw()
         self.imu.resetDisplacement()
+
+        initial_pose = Chassis.RED_TEST_POSE if is_red() else Chassis.BLUE_TEST_POSE
         self.estimator = SwerveDrive4PoseEstimator(
             self.kinematics,
             self.imu.getRotation2d(),
             self.get_module_positions(),
-            Pose2d(0, 0, 0),
+            initial_pose,
             stateStdDevs=(0.05, 0.05, 0.01),
             visionMeasurementStdDevs=(0.4, 0.4, math.inf),
         )
         self.field_obj = self.field.getObject("fused_pose")
         self.module_objs: list[wpilib.FieldObject2d] = []
         for idx, _module in enumerate(self.modules):
             self.module_objs.append(self.field.getObject("s_module_" + str(idx)))
-        self.set_pose(Pose2d(0, 0, 0))
+        self.set_pose(initial_pose)
 
     def drive_field(self, vx: float, vy: float, omega: float) -> None:
         """Field oriented drive commands"""
@@ -381,6 +393,15 @@ def get_velocity(self) -> ChassisSpeeds:
         )
 
     def update_odometry(self) -> None:
+        # Check whether our alliance has "changed"
+        # If so, it means we have an update from the FMS and need to re-init the odom
+        if is_red() != self.on_red_alliance:
+            self.on_red_alliance = is_red()
+            if self.on_red_alliance:
+                self.set_pose(Chassis.RED_TEST_POSE)
+            else:
+                self.set_pose(Chassis.BLUE_TEST_POSE)
+
         self.estimator.update(self.imu.getRotation2d(), self.get_module_positions())
         self.field_obj.setPose(self.get_pose())
         if self.send_modules:
@@ -441,3 +462,7 @@ def get_rotation(self) -> Rotation2d:
     @feedback
     def get_drive_current(self) -> float:
         return sum(abs(x.get_drive_current()) for x in self.modules)
+
+    @feedback
+    def at_desired_heading(self) -> bool:
+        return self.heading_controller.atGoal()

components/shooter.py

@@ -10,7 +10,7 @@
 import phoenix6.hardware
 
 
-class Shooter:
+class ShooterComponent:
     flywheel_speed = tunable(0.0)
     inject_speed = tunable(0.0)
 

controllers/shooter.py

@@ -0,0 +1,36 @@
+from magicbot import StateMachine, state
+from components.chassis import Chassis
+
+from utilities.game import get_goal_speaker_position
+
+from math import atan2
+
+
+class Shooter(StateMachine):
+    chassis: Chassis
+
+    def setup(self) -> None:
+        pass
+
+    @state(first=True)
+    def acquiring(self) -> None:
+        # determine heading required for goal
+        translation_to_goal = (
+            get_goal_speaker_position().toTranslation2d()
+            - self.chassis.get_pose().translation()
+        )
+        bearing_to_speaker = atan2(translation_to_goal.y, translation_to_goal.x)
+        # set to appropriate heading
+        self.chassis.snap_to_heading(bearing_to_speaker)
+
+        # progress state machine if within tolerance
+        if self.chassis.at_desired_heading():
+            self.next_state("shooting")
+
+    @state(must_finish=True)
+    def shooting(self) -> None:
+        # commence shooting action
+        self.done()
+
+    def shoot(self) -> None:
+        self.engage()

robot.py

@@ -5,18 +5,22 @@
 import magicbot
 
 from components.chassis import Chassis
-from components.shooter import Shooter
+from components.shooter import ShooterComponent
+
+from controllers.shooter import Shooter
+
 import math
 
 from utilities.scalers import rescale_js
 
 
 class MyRobot(magicbot.MagicRobot):
     # Controllers
+    shooter: Shooter
 
     # Components
     chassis: Chassis
-    shooter: Shooter
+    shooter_component: ShooterComponent
 
     max_speed = magicbot.tunable(Chassis.max_wheel_speed * 0.95)
 
@@ -52,7 +56,7 @@ def teleopPeriodic(self) -> None:
         drive_x = -rescale_js(self.gamepad.getLeftY(), 0.1) * self.max_speed
         drive_y = -rescale_js(self.gamepad.getLeftX(), 0.1) * self.max_speed
         drive_z = -rescale_js(self.gamepad.getRightX(), 0.1, exponential=2) * spin_rate
-        local_driving = self.gamepad.getBButton()
+        local_driving = self.gamepad.getYButton()
         driver_inputs = (drive_x, drive_y, drive_z)
         if local_driving:
             self.chassis.drive_local(*driver_inputs)