[copybara] Import robotpy examples (#8608)

GitOrigin-RevId: 9ba4bc3040fa7e772f5a594039e78fc6c43d114e

robotpyExamples/MecanumDrivePoseEstimator/drivetrain.py

@@ -0,0 +1,163 @@
+#
+# Copyright (c) FIRST and other WPILib contributors.
+# Open Source Software; you can modify and/or share it under the terms of
+# the WPILib BSD license file in the root directory of this project.
+#
+
+import math
+
+import wpilib
+import wpimath
+import wpimath.units
+
+from exampleglobalmeasurementsensor import ExampleGlobalMeasurementSensor
+
+
+class Drivetrain:
+    """Represents a mecanum drive style drivetrain."""
+
+    kMaxSpeed = 3.0  # 3 meters per second
+    kMaxAngularSpeed = math.pi  # 1/2 rotation per second
+
+    def __init__(self) -> None:
+        self.frontLeftMotor = wpilib.PWMSparkMax(1)
+        self.frontRightMotor = wpilib.PWMSparkMax(2)
+        self.backLeftMotor = wpilib.PWMSparkMax(3)
+        self.backRightMotor = wpilib.PWMSparkMax(4)
+
+        self.frontLeftEncoder = wpilib.Encoder(0, 1)
+        self.frontRightEncoder = wpilib.Encoder(2, 3)
+        self.backLeftEncoder = wpilib.Encoder(4, 5)
+        self.backRightEncoder = wpilib.Encoder(6, 7)
+
+        self.frontLeftLocation = wpimath.Translation2d(0.381, 0.381)
+        self.frontRightLocation = wpimath.Translation2d(0.381, -0.381)
+        self.backLeftLocation = wpimath.Translation2d(-0.381, 0.381)
+        self.backRightLocation = wpimath.Translation2d(-0.381, -0.381)
+
+        self.frontLeftPIDController = wpimath.PIDController(1, 0, 0)
+        self.frontRightPIDController = wpimath.PIDController(1, 0, 0)
+        self.backLeftPIDController = wpimath.PIDController(1, 0, 0)
+        self.backRightPIDController = wpimath.PIDController(1, 0, 0)
+
+        self.imu = wpilib.OnboardIMU(wpilib.OnboardIMU.MountOrientation.kFlat)
+
+        self.kinematics = wpimath.MecanumDriveKinematics(
+            self.frontLeftLocation,
+            self.frontRightLocation,
+            self.backLeftLocation,
+            self.backRightLocation,
+        )
+
+        # Here we use MecanumDrivePoseEstimator so that we can fuse odometry readings. The numbers
+        # used below are robot specific, and should be tuned.
+        self.poseEstimator = wpimath.MecanumDrivePoseEstimator(
+            self.kinematics,
+            self.imu.getRotation2d(),
+            self.getCurrentDistances(),
+            wpimath.Pose2d(),
+            (0.05, 0.05, wpimath.units.degreesToRadians(5)),
+            (0.5, 0.5, wpimath.units.degreesToRadians(30)),
+        )
+
+        # Gains are for example purposes only - must be determined for your own robot!
+        self.feedforward = wpimath.SimpleMotorFeedforwardMeters(1, 3)
+
+        self.imu.resetYaw()
+
+        # We need to invert one side of the drivetrain so that positive voltages
+        # result in both sides moving forward. Depending on how your robot's
+        # gearbox is constructed, you might have to invert the left side instead.
+        self.frontRightMotor.setInverted(True)
+        self.backRightMotor.setInverted(True)
+
+    def getCurrentState(self) -> wpimath.MecanumDriveWheelSpeeds:
+        """Returns the current state of the drivetrain.
+
+        :returns: The current state of the drivetrain.
+        """
+        return wpimath.MecanumDriveWheelSpeeds(
+            self.frontLeftEncoder.getRate(),
+            self.frontRightEncoder.getRate(),
+            self.backLeftEncoder.getRate(),
+            self.backRightEncoder.getRate(),
+        )
+
+    def getCurrentDistances(self) -> wpimath.MecanumDriveWheelPositions:
+        """Returns the current distances measured by the drivetrain.
+
+        :returns: The current distances measured by the drivetrain.
+        """
+        positions = wpimath.MecanumDriveWheelPositions()
+        positions.frontLeft = self.frontLeftEncoder.getDistance()
+        positions.frontRight = self.frontRightEncoder.getDistance()
+        positions.rearLeft = self.backLeftEncoder.getDistance()
+        positions.rearRight = self.backRightEncoder.getDistance()
+        return positions
+
+    def setSpeeds(self, speeds: wpimath.MecanumDriveWheelSpeeds) -> None:
+        """Set the desired speeds for each wheel.
+
+        :param speeds: The desired wheel speeds.
+        """
+        frontLeftFeedforward = self.feedforward.calculate(speeds.frontLeft)
+        frontRightFeedforward = self.feedforward.calculate(speeds.frontRight)
+        backLeftFeedforward = self.feedforward.calculate(speeds.rearLeft)
+        backRightFeedforward = self.feedforward.calculate(speeds.rearRight)
+
+        frontLeftOutput = self.frontLeftPIDController.calculate(
+            self.frontLeftEncoder.getRate(), speeds.frontLeft
+        )
+        frontRightOutput = self.frontRightPIDController.calculate(
+            self.frontRightEncoder.getRate(), speeds.frontRight
+        )
+        backLeftOutput = self.backLeftPIDController.calculate(
+            self.backLeftEncoder.getRate(), speeds.rearLeft
+        )
+        backRightOutput = self.backRightPIDController.calculate(
+            self.backRightEncoder.getRate(), speeds.rearRight
+        )
+
+        self.frontLeftMotor.setVoltage(frontLeftOutput + frontLeftFeedforward)
+        self.frontRightMotor.setVoltage(frontRightOutput + frontRightFeedforward)
+        self.backLeftMotor.setVoltage(backLeftOutput + backLeftFeedforward)
+        self.backRightMotor.setVoltage(backRightOutput + backRightFeedforward)
+
+    def drive(
+        self,
+        xSpeed: float,
+        ySpeed: float,
+        rot: float,
+        fieldRelative: bool,
+        period: float,
+    ) -> None:
+        """Method to drive the robot using joystick info.
+
+        :param xSpeed: Speed of the robot in the x direction (forward).
+        :param ySpeed: Speed of the robot in the y direction (sideways).
+        :param rot: Angular rate of the robot.
+        :param fieldRelative: Whether the provided x and y speeds are relative to the field.
+        """
+        chassisSpeeds = wpimath.ChassisSpeeds(xSpeed, ySpeed, rot)
+        if fieldRelative:
+            chassisSpeeds = chassisSpeeds.toRobotRelative(
+                self.poseEstimator.getEstimatedPosition().rotation()
+            )
+        self.setSpeeds(
+            self.kinematics.toWheelSpeeds(chassisSpeeds.discretize(period)).desaturate(
+                self.kMaxSpeed
+            )
+        )
+
+    def updateOdometry(self) -> None:
+        """Updates the field relative position of the robot."""
+        self.poseEstimator.update(self.imu.getRotation2d(), self.getCurrentDistances())
+
+        # Also apply vision measurements. We use 0.3 seconds in the past as an example -- on
+        # a real robot, this must be calculated based either on latency or timestamps.
+        self.poseEstimator.addVisionMeasurement(
+            ExampleGlobalMeasurementSensor.getEstimatedGlobalPose(
+                self.poseEstimator.getEstimatedPosition()
+            ),
+            wpilib.Timer.getTimestamp() - 0.3,
+        )

robotpyExamples/MecanumDrivePoseEstimator/exampleglobalmeasurementsensor.py

@@ -0,0 +1,34 @@
+#
+# Copyright (c) FIRST and other WPILib contributors.
+# Open Source Software; you can modify and/or share it under the terms of
+# the WPILib BSD license file in the root directory of this project.
+#
+
+import random
+
+import wpimath
+import wpimath.units
+
+
+class ExampleGlobalMeasurementSensor:
+    """This dummy class represents a global measurement sensor, such as a computer vision
+    solution.
+    """
+
+    def __init__(self) -> None:
+        raise RuntimeError("Utility class")
+
+    @staticmethod
+    def getEstimatedGlobalPose(estimatedRobotPose: wpimath.Pose2d) -> wpimath.Pose2d:
+        """Get a "noisy" fake global pose reading.
+
+        :param estimatedRobotPose: The robot pose.
+        """
+        rand_x = random.gauss(0.0, 0.5)
+        rand_y = random.gauss(0.0, 0.5)
+        rand_rot = random.gauss(0.0, wpimath.units.degreesToRadians(30))
+        return wpimath.Pose2d(
+            estimatedRobotPose.x + rand_x,
+            estimatedRobotPose.y + rand_y,
+            estimatedRobotPose.rotation().rotateBy(wpimath.Rotation2d(rand_rot)),
+        )

robotpyExamples/MecanumDrivePoseEstimator/robot.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) FIRST and other WPILib contributors.
+# Open Source Software; you can modify and/or share it under the terms of
+# the WPILib BSD license file in the root directory of this project.
+#
+
+import wpilib
+import wpimath
+
+from drivetrain import Drivetrain
+
+
+class MyRobot(wpilib.TimedRobot):
+    def __init__(self) -> None:
+        super().__init__()
+
+        self.controller = wpilib.NiDsXboxController(0)
+        self.mecanum = Drivetrain()
+
+        # Slew rate limiters to make joystick inputs more gentle; 1/3 sec from 0 to 1.
+        self.xspeedLimiter = wpimath.SlewRateLimiter(3)
+        self.yspeedLimiter = wpimath.SlewRateLimiter(3)
+        self.rotLimiter = wpimath.SlewRateLimiter(3)
+
+    def autonomousPeriodic(self) -> None:
+        self.driveWithJoystick(False)
+        self.mecanum.updateOdometry()
+
+    def teleopPeriodic(self) -> None:
+        self.driveWithJoystick(True)
+
+    def driveWithJoystick(self, fieldRelative: bool) -> None:
+        # Get the x speed. We are inverting this because Xbox controllers return
+        # negative values when we push forward.
+        xSpeed = -self.xspeedLimiter.calculate(self.controller.getLeftY())
+        xSpeed *= Drivetrain.kMaxSpeed
+
+        # Get the y speed or sideways/strafe speed. We are inverting this because
+        # we want a positive value when we pull to the left. Xbox controllers
+        # return positive values when you pull to the right by default.
+        ySpeed = -self.yspeedLimiter.calculate(self.controller.getLeftX())
+        ySpeed *= Drivetrain.kMaxSpeed
+
+        # Get the rate of angular rotation. We are inverting this because we want a
+        # positive value when we pull to the left (remember, CCW is positive in
+        # mathematics). Xbox controllers return positive values when you pull to
+        # the right by default.
+        rot = -self.rotLimiter.calculate(self.controller.getRightX())
+        rot *= Drivetrain.kMaxAngularSpeed
+
+        self.mecanum.drive(xSpeed, ySpeed, rot, fieldRelative, self.getPeriod())