Update arm simulation example

- Fixes #79

Author: virtuald

arm-simulation/constants.py

@@ -0,0 +1,31 @@
+# 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
+
+from wpimath import units
+
+
+class Constants:
+    kMotorPort = 0
+    kEncoderAChannel = 0
+    kEncoderBChannel = 1
+    kJoystickPort = 0
+
+    kArmPositionKey = "ArmPosition"
+    kArmPKey = "ArmP"
+
+    # The P gain for the PID controller that drives this arm.
+    kDefaultArmKp = 50.0
+    kDefaultArmSetpointDegrees = 75.0
+
+    # distance per pulse = (angle per revolution) / (pulses per revolution)
+    #  = (2 * PI rads) / (4096 pulses)
+    kArmEncoderDistPerPulse = 2.0 * math.pi / 4096
+
+    kArmReduction = 200
+    kArmMass = 8.0  # Kilograms
+    kArmLength = units.inchesToMeters(30)
+    kMinAngleRads = units.degreesToRadians(-75)
+    kMaxAngleRads = units.degreesToRadians(255)

arm-simulation/physics.py

@@ -24,6 +24,8 @@
 if typing.TYPE_CHECKING:
     from robot import MyRobot
 
+from constants import Constants
+
 
 class PhysicsEngine:
     """
@@ -43,34 +45,39 @@ def __init__(self, physics_controller: PhysicsInterface, robot: "MyRobot"):
         self.armGearbox = wpimath.system.plant.DCMotor.vex775Pro(2)
 
         # Simulation classes help us simulate what's going on, including gravity.
-        # This sim represents an arm with 2 775s, a 600:1 reduction, a mass of 5kg,
-        # 30in overall arm length, range of motion in [-75, 255] degrees, and noise
-        # with a standard deviation of 1 encoder tick.
+        # This arm sim represents an arm that can travel from -75 degrees (rotated down front)
+        # to 255 degrees (rotated down in the back).
         self.armSim = wpilib.simulation.SingleJointedArmSim(
             self.armGearbox,
-            600.0,
-            wpilib.simulation.SingleJointedArmSim.estimateMOI(0.762, 5),
-            0.762,
-            math.radians(-75),
-            math.radians(255),
+            Constants.kArmReduction,
+            wpilib.simulation.SingleJointedArmSim.estimateMOI(
+                Constants.kArmLength, Constants.kArmMass
+            ),
+            Constants.kArmLength,
+            Constants.kMinAngleRads,
+            Constants.kMaxAngleRads,
             True,
-            0,
+            # Add noise with a std-dev of 1 tick
+            Constants.kArmEncoderDistPerPulse,
         )
-        self.encoderSim = wpilib.simulation.EncoderSim(robot.encoder)
-        self.motorSim = wpilib.simulation.PWMSim(robot.motor.getChannel())
+        self.encoderSim = wpilib.simulation.EncoderSim(robot.arm.encoder)
+        self.motorSim = wpilib.simulation.PWMSim(robot.arm.motor.getChannel())
 
-        # Create a Mechanism2d display of an Arm
+        # Create a Mechanism2d display of an Arm with a fixed ArmTower and moving Arm.
         self.mech2d = wpilib.Mechanism2d(60, 60)
-        self.armBase = self.mech2d.getRoot("ArmBase", 30, 30)
-        self.armTower = self.armBase.appendLigament(
-            "Arm Tower", 30, -90, 6, wpilib.Color8Bit(wpilib.Color.kBlue)
-        )
-        self.arm = self.armBase.appendLigament(
-            "Arm", 30, self.armSim.getAngle(), 6, wpilib.Color8Bit(wpilib.Color.kYellow)
+        self.armPivot = self.mech2d.getRoot("ArmPivot", 30, 30)
+        self.armTower = self.armPivot.appendLigament("Arm Tower", 30, -90)
+        self.arm = self.armPivot.appendLigament(
+            "Arm",
+            30,
+            math.degrees(self.armSim.getAngle()),
+            6,
+            wpilib.Color8Bit(wpilib.Color.kYellow),
         )
 
         # Put Mechanism to SmartDashboard
         wpilib.SmartDashboard.putData("Arm Sim", self.mech2d)
+        self.armTower.setColor(wpilib.Color8Bit(wpilib.Color.kBlue))
 
     def update_sim(self, now: float, tm_diff: float) -> None:
         """
@@ -94,9 +101,9 @@ def update_sim(self, now: float, tm_diff: float) -> None:
         # voltage
         self.encoderSim.setDistance(self.armSim.getAngle())
         # SimBattery estimates loaded battery voltage
-        # wpilib.simulation.RoboRioSim.setVInVoltage(
-        #     wpilib.simulation.BatterySim
-        # )
+        wpilib.simulation.RoboRioSim.setVInVoltage(
+            wpilib.simulation.BatterySim.calculate([self.armSim.getCurrentDraw()])
+        )
 
         # Update the mechanism arm angle based on the simulated arm angle
         # -> setAngle takes degrees, getAngle returns radians... >_>

arm-simulation/robot.py

@@ -4,69 +4,31 @@
 # 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.controller
 
+from subsytems.arm import Arm
+from constants import Constants
 
-class MyRobot(wpilib.TimedRobot):
-    kMotorPort = 0
-    kEncoderAChannel = 0
-    kEncoderBChannel = 1
-    kJoystickPort = 0
-
-    kArmPositionKey = "ArmPosition"
-    kArmPKey = "ArmP"
-
-    # distance per pulse = (angle per revolution) / (pulses per revolution)
-    #  = (2 * PI rads) / (4096 pulses)
-    kArmEncoderDistPerPulse = 2.0 * math.pi / 4096.0
-
-    def robotInit(self) -> None:
-        # The P gain for the PID controller that drives this arm.
-        self.kArmKp = 50.0
 
-        self.armPosition = 75.0
-
-        # standard classes for controlling our arm
-        self.controller = wpimath.controller.PIDController(self.kArmKp, 0, 0)
-        self.encoder = wpilib.Encoder(self.kEncoderAChannel, self.kEncoderBChannel)
-        self.motor = wpilib.PWMSparkMax(self.kMotorPort)
-        self.joystick = wpilib.Joystick(self.kJoystickPort)
-
-        self.encoder.setDistancePerPulse(self.kArmEncoderDistPerPulse)
-
-        # Set the Arm position setpoint and P constant to Preferences if the keys
-        # don't already exist
-        if not wpilib.Preferences.containsKey(self.kArmPositionKey):
-            wpilib.Preferences.setDouble(self.kArmPositionKey, self.armPosition)
-        if not wpilib.Preferences.containsKey(self.kArmPKey):
-            wpilib.Preferences.setDouble(self.kArmPKey, self.kArmKp)
+class MyRobot(wpilib.TimedRobot):
+    def robotInit(self):
+        self.arm = Arm()
+        self.joystick = wpilib.Joystick(Constants.kJoystickPort)
 
-    def teleopInit(self) -> None:
-        # Read Preferences for Arm setpoint and kP on entering Teleop
-        self.armPosition = wpilib.Preferences.getDouble(
-            self.kArmPositionKey, self.armPosition
-        )
-        if self.kArmKp != wpilib.Preferences.getDouble(self.kArmPKey, self.kArmKp):
-            self.kArmKp = wpilib.Preferences.getDouble(self.kArmPKey, self.kArmKp)
-            self.controller.setP(self.kArmKp)
+    def teleopInit(self):
+        self.arm.loadPreferences()
 
-    def teleopPeriodic(self) -> None:
+    def teleopPeriodic(self):
         if self.joystick.getTrigger():
-            # Here we run PID control like normal, with a setpoint read from
-            # preferences in degrees
-            pidOutput = self.controller.calculate(
-                self.encoder.getDistance(), math.radians(self.armPosition)
-            )
-            self.motor.setVoltage(pidOutput)
+            # Here, we run PID control like normal.
+            self.arm.reachSetpoint()
         else:
-            # Otherwise we disable the motor
-            self.motor.set(0.0)
+            # Otherwise, we disable the motor.
+            self.arm.stop()
 
-    def disabledInit(self) -> None:
-        # This just makes sure that our simulation code knows that the motor is off
-        self.motor.set(0)
+    def disabledInit(self):
+        # This just makes sure that our simulation code knows that the motor's off.
+        self.arm.stop()
 
 
 if __name__ == "__main__":

arm-simulation/subsytems/arm.py

@@ -0,0 +1,58 @@
+# 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
+
+from wpimath.controller import PIDController
+from wpimath.system.plant import DCMotor
+from wpimath import units
+
+from wpilib.simulation import BatterySim, EncoderSim, RoboRioSim, SingleJointedArmSim
+
+from constants import Constants
+
+
+class Arm:
+    def __init__(self):
+        # The P gain for the PID controller that drives this arm.
+        self.armKp = Constants.kDefaultArmKp
+        self.armSetpointDegrees = Constants.kDefaultArmSetpointDegrees
+
+        # The arm gearbox represents a gearbox containing two Vex 775pro motors.
+        self.armGearbox = DCMotor.vex775Pro(2)
+
+        # Standard classes for controlling our arm
+        self.controller = PIDController(self.armKp, 0, 0)
+        self.encoder = wpilib.Encoder(
+            Constants.kEncoderAChannel, Constants.kEncoderBChannel
+        )
+        self.motor = wpilib.PWMSparkMax(Constants.kMotorPort)
+
+        # Subsystem constructor.
+        self.encoder.setDistancePerPulse(Constants.kArmEncoderDistPerPulse)
+
+        # Set the Arm position setpoint and P constant to Preferences if the keys don't already exist
+        wpilib.Preferences.initDouble(
+            Constants.kArmPositionKey, self.armSetpointDegrees
+        )
+        wpilib.Preferences.initDouble(Constants.kArmPKey, self.armKp)
+
+    def loadPreferences(self):
+        # Read Preferences for Arm setpoint and kP on entering Teleop
+        self.armSetpointDegrees = wpilib.Preferences.getDouble(
+            Constants.kArmPositionKey, self.armSetpointDegrees
+        )
+        if self.armKp != wpilib.Preferences.getDouble(Constants.kArmPKey, self.armKp):
+            self.armKp = wpilib.Preferences.getDouble(Constants.kArmPKey, self.armKp)
+            self.controller.setP(self.armKp)
+
+    def reachSetpoint(self):
+        pidOutput = self.controller.calculate(
+            self.encoder.getDistance(),
+            units.degreesToRadians(self.armSetpointDegrees),
+        )
+        self.motor.setVoltage(pidOutput)
+
+    def stop(self):
+        self.motor.set(0.0)