idea174-catopt-swarm/catopt_swarm/solver.py

from __future__ import annotations

from math import fsum
from typing import Iterable, Sequence

from .models import (
    AuditLogEntry,
    ConvergenceCertificate,
    LocalProblem,
    PlanDelta,
    SafetyPolicy,
    SharedVariables,
    SwarmSolution,
)
from .verification import verify_swarm_solution


class ADMMSolver:
    def __init__(self, rho: float = 1.0, max_iter: int = 50, epsilon: float = 1e-3, staleness_bound: int = 2):
        self.rho = rho
        self.max_iter = max_iter
        self.epsilon = epsilon
        self.staleness_bound = staleness_bound

    def _ordered_deltas(self, deltas: Iterable[PlanDelta]) -> list[PlanDelta]:
        return sorted(
            deltas,
            key=lambda delta: (delta.timestamp, delta.author, delta.contract_id, delta.sequence),
        )

    def _apply_deltas(
        self,
        robots: Sequence[LocalProblem],
        shared: SharedVariables,
        deltas: Iterable[PlanDelta],
    ) -> list[str]:
        applied: list[str] = []
        indexed = {robot.robot_id: robot for robot in robots}
        for delta in self._ordered_deltas(deltas):
            if delta.base_version > shared.version:
                continue
            if shared.version - delta.base_version > self.staleness_bound:
                continue
            if delta.target_robot_id is None or delta.target_robot_id not in indexed:
                continue
            last_sequence = shared.last_applied_sequences.get(delta.target_robot_id, -1)
            if delta.sequence <= last_sequence:
                continue
            indexed[delta.target_robot_id].apply_updates(delta.updates)
            indexed[delta.target_robot_id].version = max(indexed[delta.target_robot_id].version, delta.sequence)
            shared.last_applied_sequences[delta.target_robot_id] = delta.sequence
            shared.version = max(shared.version, delta.base_version, delta.sequence)
            applied.append(f"{delta.contract_id}:{delta.author}:{delta.sequence}")
        if applied:
            shared.delta_clock += len(applied)
        return applied

    def _project_for_separation(self, candidates: list[float], minimum: float) -> list[float]:
        if minimum <= 0 or len(candidates) <= 1:
            return candidates
        center = fsum(candidates) / len(candidates)
        half_span = minimum * (len(candidates) - 1) / 2.0
        return [center - half_span + index * minimum for index in range(len(candidates))]

    def _current_position(self, robot: LocalProblem) -> float:
        current_pos = robot.current_pos
        if current_pos is None:
            current_pos = robot.initial_pos
        return float(current_pos)

    def solve(
        self,
        robots: Sequence[LocalProblem],
        safety_policy: SafetyPolicy,
        incoming_deltas: Iterable[PlanDelta] = (),
    ) -> SwarmSolution:
        if not robots:
            raise ValueError("at least one robot is required")

        working = [robot.model_copy(deep=True) for robot in robots]
        shared = SharedVariables(
            version=max(robot.version for robot in working),
            global_consensus=fsum(self._current_position(robot) for robot in working) / len(working),
        )
        audit_log: list[AuditLogEntry] = []
        applied_deltas = self._apply_deltas(working, shared, incoming_deltas)

        final_primal = 0.0
        final_dual = 0.0
        converged = False

        for iteration in range(self.max_iter):
            old_consensus = shared.global_consensus
            candidate_positions: list[float] = []
            previous_positions: list[float] = []

            for robot in working:
                current_pos = self._current_position(robot)
                previous_positions.append(current_pos)
                raw_position = (
                    2.0 * robot.target_pos + self.rho * (shared.global_consensus - robot.dual_var)
                ) / (2.0 + self.rho)
                step = max(-robot.max_step, min(robot.max_step, raw_position - current_pos))
                candidate_positions.append(current_pos + step)

            projected_positions = self._project_for_separation(candidate_positions, safety_policy.min_separation)

            for robot, new_position in zip(working, projected_positions):
                robot.current_pos = new_position
                if not safety_policy.verify(robot.initial_pos, new_position):
                    raise ValueError(f"Safety violation for {robot.robot_id}: travel bound exceeded")

            shared.global_consensus = fsum(projected_positions) / len(projected_positions)
            shared.version += 1

            final_primal = 0.0
            for robot, previous_position in zip(working, previous_positions):
                current_pos = self._current_position(robot)
                robot.dual_var += 0.5 * (current_pos - shared.global_consensus)
                final_primal = max(final_primal, abs(current_pos - previous_position))

            final_dual = abs(shared.global_consensus - old_consensus) * self.rho
            audit_log.append(
                AuditLogEntry(
                    iteration=iteration,
                    shared_version=shared.version,
                    consensus=shared.global_consensus,
                    primal_residual=final_primal,
                    dual_residual=final_dual,
                    applied_deltas=applied_deltas if iteration == 0 else [],
                )
            )

            if final_primal < self.epsilon and final_dual < self.epsilon:
                converged = True
                break

        certificate = ConvergenceCertificate(
            converged=converged,
            iterations=len(audit_log),
            epsilon=self.epsilon,
            final_primal_residual=final_primal,
            final_dual_residual=final_dual,
        )
        solution = SwarmSolution(
            consensus=shared.global_consensus,
            robots=working,
            shared=shared,
            audit_log=audit_log,
            certificate=certificate,
        )
        report = verify_swarm_solution(solution, safety_policy)
        if not report.passed:
            raise ValueError("; ".join(report.violations))
        return solution

    def run(
        self,
        robots: Sequence[LocalProblem],
        safety_policy: SafetyPolicy,
        incoming_deltas: Iterable[PlanDelta] = (),
    ) -> float:
        return self.solve(robots, safety_policy, incoming_deltas=incoming_deltas).consensus


def reference_scenario() -> SwarmSolution:
    robots = [
        LocalProblem(robot_id="drone-1", target_pos=0.0, current_pos=0.0),
        LocalProblem(robot_id="drone-2", target_pos=10.0, current_pos=10.0),
    ]
    policy = SafetyPolicy(max_travel_distance=15.0, min_separation=1.0, energy_budget=20.0)
    return ADMMSolver(rho=1.5, max_iter=40, epsilon=1e-4).solve(robots, policy)