physics

  1"""
  2Swerve drive physics simulation for pyfrc.
  3
  4Applies a zero-delay plant model: each cycle, the drive encoder velocity
  5and steer encoder position are set directly from the module's last commanded
  6state.  This makes closed-loop PID control, odometry, and Mechanism2d
  7visualization all respond correctly in simulation without tuning a dynamic model.
  8
  9Robot field pose is integrated from swerve kinematics and fed back to:
 10  - The Field2d widget (visible in the sim GUI)
 11  - The NavX gyro sim device (so heading-based field-relative drive works)
 12"""
 13
 14import typing
 15
 16import rev
 17import wpilib
 18import wpilib.simulation
 19from pyfrc.physics.core import PhysicsInterface
 20from wpimath.geometry import Rotation2d, Twist2d
 21from wpimath.kinematics import SwerveModuleState
 22from wpimath.system.plant import DCMotor
 23
 24if typing.TYPE_CHECKING:
 25    from robot import MyRobot
 26
 27
 28class PhysicsEngine:
 29    """
 30    No-delay swerve drive physics model.
 31
 32    Drive encoder velocity and steer encoder position are written from
 33    _last_commanded_state each cycle, giving an instantaneous plant
 34    response suitable for basic simulation and visualization.
 35
 36    All encoders are seeded to 0 degrees (wheels forward) in __init__ to
 37    prevent the calibration-based baseline from interfering with optimize()
 38    before the first physics tick runs.
 39    """
 40
 41    def __init__(self, physics_controller: PhysicsInterface, robot: "MyRobot") -> None:
 42        self.physics_controller = physics_controller
 43
 44        modules = robot.container.drivetrain.swerve_modules
 45        self.modules = modules
 46        self.drive_sims = [
 47            rev.SparkSim(m.drive_motor, DCMotor.NEO()) for m in modules
 48        ]
 49        self.steer_sims = [
 50            rev.SparkSim(m.steer_motor, DCMotor.NEO()) for m in modules
 51        ]
 52        self.kinematics = robot.container.drivetrain.drive_kinematics
 53
 54        # Seed all encoders to 0° (wheels forward) immediately.
 55        #
 56        # baseline_relative_encoders() runs during module __init__ and sets each
 57        # steer relative encoder to the module's calibration angle (e.g. 241° for
 58        # FL).  The robot periodic runs before the first update_sim(), so if these
 59        # stale values are left in place, optimize() sees a large angle error on the
 60        # first set_state() call and begins oscillating.  Seeding here ensures the
 61        # first robot periodic sees a clean 0° reading on every module.
 62        for module, drive_sim, steer_sim in zip(self.modules, self.drive_sims, self.steer_sims):
 63            calibration = module.constants.encoder_calibration
 64
 65            steer_enc = steer_sim.getRelativeEncoderSim()
 66            steer_enc.setPosition(0.0)
 67            steer_enc.setVelocity(0.0)
 68
 69            # Set CANcoder raw = -calibration so absolute position reads 0.0 (forward).
 70            # current_raw_absolute_steer_position() uses is-not-None, so 0.0 is
 71            # accepted and returns 0°, consistent with the relative encoder above.
 72            module.absolute_encoder.sim_state.set_supply_voltage(12.0)
 73            module.absolute_encoder.sim_state.set_raw_position(-calibration)
 74
 75            drive_enc = drive_sim.getRelativeEncoderSim()
 76            drive_enc.setPosition(0.0)
 77            drive_enc.setVelocity(0.0)
 78
 79        # Seed pose from the drivetrain's configured default starting position.
 80        self._pose = robot.container.drivetrain.get_default_starting_pose()
 81
 82        # NavX yaw variable — may be unavailable depending on sim state.
 83        self._navx_yaw = None
 84        try:
 85            navx_sim = wpilib.simulation.SimDeviceSim("navX-Sensor[4]")
 86            self._navx_yaw = navx_sim.getDouble("Yaw")
 87        except Exception:
 88            # NavX sim device or Yaw entry may be unavailable in some sim setups;
 89            # ignore errors here and leave _navx_yaw as None.
 90            pass
 91
 92    def update_sim(self, now: float, tm_diff: float) -> None:
 93        """
 94        Called every simulation tick.
 95
 96        Reads _last_commanded_state from each module and writes those values
 97        directly to the encoder sims (no-delay model), then integrates robot
 98        pose from swerve kinematics.
 99
100        Using _last_commanded_state instead of SparkSim.getSetpoint() avoids
101        any uncertainty about whether the SparkSim API returns converted units
102        (degrees / m/s) or raw encoder units.
103
104        Args:
105            now: current timestamp in seconds
106            tm_diff: elapsed time since last call in seconds
107        """
108        module_states = []
109
110        for module, drive_sim, steer_sim in zip(
111            self.modules, self.drive_sims, self.steer_sims
112        ):
113            velocity = module._last_commanded_state.speed          # m/s (signed)
114            angle_deg = module._last_commanded_state.angle.degrees()  # degrees
115
116            # Normalize steer to [0°, 180°) so that the two equivalent wheel
117            # orientations — {+v, 180°} and {-v, 0°} — always write the same
118            # encoder value.  Without this, optimize() sees 0° one cycle and
119            # 180° the next, flip-flopping the commanded angle and oscillating
120            # the steer motor every cycle at the 0°/180° boundary.
121            steer_angle = angle_deg % 360.0
122            if steer_angle >= 180.0:
123                steer_angle -= 180.0
124
125            # Drive — set encoder velocity and integrate position.
126            # Uses the original velocity so drive odometry stays correct.
127            drive_enc = drive_sim.getRelativeEncoderSim()
128            drive_enc.setVelocity(velocity)
129            drive_enc.setPosition(drive_enc.getPosition() + velocity * tm_diff)
130
131            # Steer — set relative encoder to the normalized angle.
132            steer_enc = steer_sim.getRelativeEncoderSim()
133            steer_enc.setPosition(steer_angle)
134
135            # Update CANcoder with the normalized angle.
136            #   absolute_position = raw_position + calibration
137            #   raw_position = steer_angle/360 - calibration
138            calibration = module.constants.encoder_calibration
139            module.absolute_encoder.sim_state.set_supply_voltage(12.0)
140            module.absolute_encoder.sim_state.set_raw_position(
141                steer_angle / 360.0 - calibration
142            )
143
144            # Kinematics uses the original angle so chassis-speed integration
145            # (pose and NavX feedback) remains physically correct.
146            module_states.append(
147                SwerveModuleState(velocity, Rotation2d.fromDegrees(angle_deg))
148            )
149
150        # Integrate chassis speeds into robot pose.
151        speeds = self.kinematics.toChassisSpeeds(tuple(module_states))
152        self._pose = self._pose.exp(
153            Twist2d(
154                speeds.vx * tm_diff,
155                speeds.vy * tm_diff,
156                speeds.omega * tm_diff,
157            )
158        )
159
160        # Update Field2d widget with the integrated pose.
161        self.physics_controller.field.setRobotPose(self._pose)
162
163        # Feed integrated heading back to the NavX gyro sim device so that
164        # field-relative drive and pose estimation use the correct heading.
165        if self._navx_yaw is not None:
166            self._navx_yaw.set(self._pose.rotation().degrees())