idea174-catopt-swarm/catopt_swarm/adapters.py

77 lines
2.3 KiB
Python

from __future__ import annotations
from dataclasses import dataclass
from typing import Any
from .models import LocalProblem, PlanDelta
@dataclass(frozen=True)
class _AdapterDefaults:
max_step: float
energy_budget: float
class DroneControllerAdapter:
name = "drone-controller"
domain = "aerial"
defaults = _AdapterDefaults(max_step=12.0, energy_budget=80.0)
def to_local_problem(self, robot_id: str, mission: dict[str, Any]) -> LocalProblem:
return LocalProblem(
robot_id=robot_id,
target_pos=float(mission["target_pos"]),
initial_pos=float(mission.get("initial_pos", 0.0)),
max_step=float(mission.get("max_step", self.defaults.max_step)),
energy_budget=float(mission.get("energy_budget", self.defaults.energy_budget)),
metadata={"domain": self.domain, **mission.get("metadata", {})},
)
def to_plan_delta(
self,
robot_id: str,
contract_id: str,
updates: dict[str, Any],
sequence: int,
author: str = "drone-controller",
) -> PlanDelta:
return PlanDelta(
author=author,
contract_id=contract_id,
target_robot_id=robot_id,
sequence=sequence,
updates=updates,
)
class GroundRoverControllerAdapter:
name = "ground-rover-controller"
domain = "ground"
defaults = _AdapterDefaults(max_step=8.0, energy_budget=120.0)
def to_local_problem(self, robot_id: str, mission: dict[str, Any]) -> LocalProblem:
return LocalProblem(
robot_id=robot_id,
target_pos=float(mission["target_pos"]),
initial_pos=float(mission.get("initial_pos", 0.0)),
max_step=float(mission.get("max_step", self.defaults.max_step)),
energy_budget=float(mission.get("energy_budget", self.defaults.energy_budget)),
metadata={"domain": self.domain, **mission.get("metadata", {})},
)
def to_plan_delta(
self,
robot_id: str,
contract_id: str,
updates: dict[str, Any],
sequence: int,
author: str = "ground-rover-controller",
) -> PlanDelta:
return PlanDelta(
author=author,
contract_id=contract_id,
target_robot_id=robot_id,
sequence=sequence,
updates=updates,
)