Read the docstring at the top 🙏

mcl.py

from math import cos, sin, exp, pi, sqrt, radians
from random import random # returns in [0, 1)

"""
A minimal reference implementation of Monte Carlo Localization for a VEX robot.
**DO NOT USE THIS. Here are 10 reasons why.**
1. Terrible performance, there are a lot of optimizations you can do.
2. I don't cut off gaussian after a maximum error (which most impls do).
3. Probably doesn't even compile lol
4. We are assuming that distance sensors all start from the same point,
    and that that point is the robot tracking center.
    In a real implementation, you would have offsets set for each sensor.
5. No raycasting (for perf), so the distance to wall is approximated.
6. This was written for easy readability, not performance.
7. Just the summing weights is O(n^2), which is terrible.
8. Strays from MCL theory in the error calculation, which is not ideal.
   this **only** works because the robot is in a square field, if you want
   to account for obstacles (i.e. Push Back goals) you'll need to switch
   to a raycast.
9. Strays from MCL theory in the weight calculation, which is also not ideal.
10. Provides no way to set the initial particle positions.
"""

N_PARTICLES = 100 # tune this
GAUSSIAN_STDEV = 1 # tune this
GAUSSIAN_FACTOR = 1 # tune this

THETA_NOISE = radians(2) # radians converts from degrees to radians
XY_NOISE = 0.1 # inches, this is the noise in the particle's position

class Beam:
    angle: float # radians
    distance: float # inches

class Particle:
    def __init__(self, x: float = 0, y: float = 0, theta: float = 0, weight: float = -1):
        self.x = x
        self.y = y
        self.theta = theta
        self.weight = weight

    def expected_point(self, beam: Beam) -> tuple[float, float, float]:
        global_theta = self.theta + beam.angle
        return (
            self.x + beam.distance * cos(global_theta),
            self.y + beam.distance * sin(global_theta),
            global_theta
        )

    @staticmethod
    def distance_to_wall(point: tuple[float, float, float]) -> float:
        # this approximates distance from point to wall
        # **not as accurate as raycasting**
        # but works fine in my experience
        return min([
            abs((point[0] - 72) / cos(point[2])),
            abs((point[0] + 72) / cos(point[2])),
            abs((point[1] - 72) / sin(point[2])),
            abs((point[1] + 72) / sin(point[2]))
        ])

    @staticmethod
    def gaussian(x: float) -> float:
        return GAUSSIAN_FACTOR * exp(-0.5 * (x / GAUSSIAN_STDEV) ** 2) / (GAUSSIAN_STDEV * sqrt(2 * pi))

    def update_weight(self, beams: list[Beam]):
        """
        This weighting function **does not** follow MCL theory!
        It calculates the weight of the particle based on the distance
        to the wall from the expected point of each beam.
        In reality, you would want to calculate the expected distance of the
        sensor for each particle and compare to the actual distance.
        This requires a raycast, however, which is more expensive.
        I use this method in my implementation because it is simpler,
        less computationally expensive, and works well enough.
        """

        # might want to do something like taking to exp of the sum to
        # have faster convergence
        #
        # to follow theory, might want to switch to a product of gaussians.
        # I use sum b/c I found it easier to tune.
        self.weight = sum([
            self.gaussian(self.distance_to_wall(self.expected_point(beam)))
            for beam in beams
        ])

    def update_delta_noise(self, delta: tuple[float, float, float]):
        self.x += XY_NOISE * 2 * (random() - 0.5) + delta[0]
        self.y += XY_NOISE * 2 * (random() - 0.5) + delta[1]
        self.theta += THETA_NOISE * 2 * (random() - 0.5) + delta[2]

    def copy(self) -> 'Particle':
        return Particle(self.x, self.y, self.theta, self.weight)

class MCL:
    def __init__(self):
        self.particles = [Particle() for _ in range(N_PARTICLES)]
        self.average_pose = (0, 0, 0)

    def update_step(self, delta: tuple[float, float, float]) -> None:
        # delta = [delta x, delta y, delta theta]
        for particle in self.particles:
            particle.update_delta_noise(delta)
            

    def resample_step(self, beams: list[Beam]) -> None:
        for particle in self.particles:
            particle.update_weight(beams)
        # resampling step
        weights = [particle.weight for particle in self.particles]
        # DO NOT use this sum implementation, it is O(n^2).
        # Use prefix sums when actually implementing this.
        sum_weights = [
            sum(weights[0:i+1]) for i in range(N_PARTICLES)
        ]
        random_offset = random() / N_PARTICLES
        offsets = [
            random_offset + i / N_PARTICLES for i in range(N_PARTICLES)
        ]
        offsets = [
            offset * sum(weights) for offset in offsets
        ]
        new_particles = []
        for offset in offsets:
            for particle, cumulative_weight in zip(self.particles, sum_weights):
                if cumulative_weight >= offset:
                    new_particles.append(particle.copy())
                    break
        self.particles = new_particles
        self.average_pose = (
            sum(p.x for p in self.particles) / N_PARTICLES,
            sum(p.y for p in self.particles) / N_PARTICLES,
            sum(p.theta for p in self.particles) / N_PARTICLES
        )

    def run(self, beams: list[Beam], delta: tuple[float, float, float]) -> tuple[float, float, float]:
        self.update_step(delta)
        self.resample_step(beams)
        return self.average_pose