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
 13macOS notes
 14-----------
 15Two platform-specific workarounds are applied during simulation on Mac:
 16
 171. **CANcoder sim frames unavailable at init** — CTRE phoenix6's sim backend
 18   on macOS does not deliver CAN frames in time for module initialization.
 19   ``baseline_relative_encoders()`` in ``swerve_module.py`` handles this by
 20   defaulting the steer baseline to 0° instead of raising.  The physics engine
 21   also guards against an empty module list (``len(module_states) != 4``).
 22
 232. **Xbox/PlayStation controller input** — Apple's Game Controller framework
 24   intercepts these devices, making them invisible to GLFW (the input backend
 25   used by the WPILib sim GUI).  An SDL2-based bridge
 26   (``utils/sim/sdl2_controller_bridge``) reads controllers via SDL2 and writes
 27   their state to the HAL sim.  See that module's docstring for setup details.
 28"""
 29
 30import typing
 31
 32import ntcore
 33import rev
 34import wpilib
 35import wpilib.simulation
 36from pyfrc.physics.core import PhysicsInterface
 37from wpimath.system.plant import LinearSystemId
 38from wpimath.geometry import Pose2d, Rotation2d, Twist2d
 39from wpimath.kinematics import SwerveModuleState
 40from wpimath.system.plant import DCMotor
 41
 42from utils.sim.sdl2_controller_bridge import start_controller_bridge
 43
 44if typing.TYPE_CHECKING:
 45    from robot import MyRobot
 46
 47
 48class PhysicsEngine:
 49    """
 50    No-delay swerve drive physics model.
 51
 52    Drive encoder velocity and steer encoder position are written from
 53    _last_commanded_state each cycle, giving an instantaneous plant
 54    response suitable for basic simulation and visualization.
 55
 56    All encoders are seeded to 0 degrees (wheels forward) in __init__ to
 57    prevent the calibration-based baseline from interfering with optimize()
 58    before the first physics tick runs.
 59    """
 60
 61    def __init__(self, physics_controller: PhysicsInterface, robot: "MyRobot") -> None:
 62        self.physics_controller = physics_controller
 63
 64        modules = robot.container.drivetrain.swerve_modules
 65        if len(modules) != 4:
 66            raise RuntimeError(
 67                f"PhysicsEngine requires exactly 4 swerve modules, got {len(modules)}. "
 68                "CANcoder initialization likely failed — check CAN IDs and phoenix6 sim setup."
 69            )
 70        self.modules = modules
 71        self.drive_sims = [
 72            rev.SparkSim(m.drive_motor, DCMotor.NEO()) for m in modules
 73        ]
 74        self.steer_sims = [
 75            rev.SparkSim(m.steer_motor, DCMotor.NEO()) for m in modules
 76        ]
 77        self.kinematics = robot.container.drivetrain.drive_kinematics
 78
 79        # Seed all encoders to 0° (wheels forward) immediately.
 80        for module, drive_sim, steer_sim in zip(self.modules, self.drive_sims, self.steer_sims):
 81            calibration = module.constants.encoder_calibration
 82
 83            steer_enc = steer_sim.getRelativeEncoderSim()
 84            steer_enc.setPosition(0.0)
 85            steer_enc.setVelocity(0.0)
 86
 87            module.absolute_encoder.sim_state.set_supply_voltage(12.0)
 88            module.absolute_encoder.sim_state.set_raw_position(-calibration)
 89
 90            drive_enc = drive_sim.getRelativeEncoderSim()
 91            drive_enc.setPosition(0.0)
 92            drive_enc.setVelocity(0.0)
 93
 94        # Shooter flywheel simulation using WPILib FlywheelSim
 95        shooter = robot.container.shooter
 96        self._has_shooter = shooter is not None
 97        if self._has_shooter:
 98            self._shooter_lead_sim = rev.SparkSim(shooter.leadMotor, DCMotor.neoVortex())
 99            self._shooter_follower_sim = rev.SparkSim(shooter.followerMotor, DCMotor.neoVortex())
100            # Two NEO Vortex motors, ~0.005 kg·m² moment of inertia (dual 4" wheels), 1:1 gearing
101            self._flywheel_sim = wpilib.simulation.FlywheelSim(
102                LinearSystemId.flywheelSystem(DCMotor.neoVortex(2), 0.005, 1.0),
103                DCMotor.neoVortex(2),
104            )
105
106        # Vision simulation — safe to skip if localization is disabled
107        self._localization = getattr(robot.container, 'localization', None)
108        if self._localization is not None:
109            self._localization.setup_sim()
110
111        # Seed pose from the drivetrain's configured default starting position.
112        self._pose = robot.container.drivetrain.get_default_starting_pose()
113
114        # Ground-truth pose struct publisher (for AdvantageScope ghost overlay)
115        self._gt_pose_pub = ntcore.NetworkTableInstance.getDefault().getStructTopic(
116            "SimGroundTruth/pose", Pose2d
117        ).publish()
118
119        # NavX yaw variable — may be unavailable depending on sim state.
120        self._navx_yaw = None
121        try:
122            navx_sim = wpilib.simulation.SimDeviceSim("navX-Sensor[4]")
123            self._navx_yaw = navx_sim.getDouble("Yaw")
124        except Exception:
125            # NavX sim device or Yaw entry may be unavailable in some sim setups;
126            # ignore errors here and leave _navx_yaw as None.
127            pass
128
129        # On macOS, start the SDL2 controller bridge (see module docstring).
130        start_controller_bridge()
131
132    def update_sim(self, now: float, tm_diff: float) -> None:
133        """
134        Called every simulation tick.
135
136        Reads _last_commanded_state from each module and writes those values
137        directly to the encoder sims (no-delay model), then integrates robot
138        pose from swerve kinematics.
139
140        Using _last_commanded_state instead of SparkSim.getSetpoint() avoids
141        any uncertainty about whether the SparkSim API returns converted units
142        (degrees / m/s) or raw encoder units.
143
144        Args:
145            now: current timestamp in seconds
146            tm_diff: elapsed time since last call in seconds
147        """
148        module_states = []
149
150        for module, drive_sim, steer_sim in zip(
151            self.modules, self.drive_sims, self.steer_sims
152        ):
153            velocity = module._last_commanded_state.speed
154            angle_deg = module._last_commanded_state.angle.degrees()
155
156            steer_angle = angle_deg % 360.0
157            if steer_angle >= 180.0:
158                steer_angle -= 180.0
159
160            drive_enc = drive_sim.getRelativeEncoderSim()
161            drive_enc.setVelocity(velocity)
162            drive_enc.setPosition(drive_enc.getPosition() + velocity * tm_diff)
163
164            steer_enc = steer_sim.getRelativeEncoderSim()
165            steer_enc.setPosition(steer_angle)
166
167            calibration = module.constants.encoder_calibration
168            module.absolute_encoder.sim_state.set_supply_voltage(12.0)
169            module.absolute_encoder.sim_state.set_raw_position(
170                steer_angle / 360.0 - calibration
171            )
172
173            module_states.append(
174                SwerveModuleState(velocity, Rotation2d.fromDegrees(angle_deg))
175            )
176
177        # Integrate chassis speeds into robot pose.
178        if len(module_states) != 4:
179            return
180        speeds = self.kinematics.toChassisSpeeds(tuple(module_states))
181        self._pose = self._pose.exp(
182            Twist2d(
183                speeds.vx * tm_diff,
184                speeds.vy * tm_diff,
185                speeds.omega * tm_diff,
186            )
187        )
188
189        # Feed ground-truth pose to vision sim so cameras see AprilTags
190        if self._localization is not None:
191            self._localization.update_sim(self._pose)
192
193        # Update Field2d widget and struct publisher with the integrated pose.
194        self.physics_controller.field.setRobotPose(self._pose)
195        self._gt_pose_pub.set(self._pose)
196
197        # Feed integrated heading back to the NavX gyro sim device so that
198        # field-relative drive and pose estimation use the correct heading.
199        if self._navx_yaw is not None:
200            self._navx_yaw.set(self._pose.rotation().degrees())
201
202        # Simulate shooter flywheel (if enabled)
203        if self._has_shooter:
204            flywheel_rpm = self._flywheel_sim.getAngularVelocity() * 60.0 / (2 * 3.14159)
205            self._shooter_lead_sim.iterate(flywheel_rpm, 12.0, tm_diff)
206            lead_output = self._shooter_lead_sim.getAppliedOutput()
207            self._flywheel_sim.setInputVoltage(lead_output * 12.0)
208            self._flywheel_sim.update(tm_diff)
209            flywheel_rpm = self._flywheel_sim.getAngularVelocity() * 60.0 / (2 * 3.14159)
210            lead_enc = self._shooter_lead_sim.getRelativeEncoderSim()
211            lead_enc.setVelocity(flywheel_rpm)
212            lead_enc.setPosition(lead_enc.getPosition() + flywheel_rpm / 60.0 * tm_diff)
213            follower_enc = self._shooter_follower_sim.getRelativeEncoderSim()
214            follower_enc.setVelocity(flywheel_rpm)
215            follower_enc.setPosition(lead_enc.getPosition())