hivebr 0.1.0__py3-none-any.whl

hivebr/__init__.py

@@ -0,0 +1,20 @@
+"""hivebr — OpenLab API for single- and multi-robot navigation experiments.
+
+The public API is backend-agnostic and intentionally avoids importing the
+pygame-based engine (``hivebr.engine``) at package import time, so the same
+user script can later target the physical ROS 2 arena.
+"""
+
+from .core import HiveBrAgent, SwarmController, SwarmExperiment
+from .state import Command, FleetState, RobotState
+
+__all__ = [
+    "HiveBrAgent",
+    "SwarmController",
+    "SwarmExperiment",
+    "Command",
+    "FleetState",
+    "RobotState",
+]
+
+__version__ = "0.1.0"

hivebr/bridge.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from .kinematics import normalize_command
+from .state import Command, FleetState, RobotState
+
+
+class Backend(ABC):
+    """Execution backend contract shared by the simulator and ROS 2 arena.
+
+    A backend is the only layer allowed to talk to a concrete engine, so that
+    the public API never imports pygame or ROS 2 directly.
+    """
+
+    @abstractmethod
+    def reset(self) -> None: ...
+
+    @abstractmethod
+    def get_states(self) -> dict[int, RobotState]: ...
+
+    @abstractmethod
+    def compute(self, states: dict[int, RobotState], dt: float) -> dict[int, Command]:
+        """Return one command per robot id from a frozen fleet snapshot."""
+
+    @abstractmethod
+    def step(self, commands: dict[int, Command], dt: float) -> None: ...
+
+    @abstractmethod
+    def run(self) -> None: ...
+
+    @abstractmethod
+    def close(self) -> None: ...
+
+
+class NavSim2DBackend(Backend):
+    """Backend for local simulation using the hivebr.engine renderer/physics."""
+
+    def __init__(self, experiment) -> None:
+        self.experiment = experiment
+        self._sim = None
+
+    def reset(self) -> None:
+        if self._sim is not None:
+            self._sim.reset()
+
+    def get_states(self) -> dict[int, RobotState]:
+        if self._sim is None:
+            return {}
+        robots = getattr(self._sim, "robots", [self._sim.robot])
+        states = {}
+        for i, robot in enumerate(robots):
+            goal = self._sim.robot_goal(i) if hasattr(self._sim, "robot_goal") else None
+            states[i] = _state_from_robot(i, robot, goal)
+        return states
+
+    def compute(self, states: dict[int, RobotState], dt: float) -> dict[int, Command]:
+        return self.experiment._dispatch(states, dt)
+
+    def step(self, commands: dict[int, Command], dt: float) -> None:
+        raise NotImplementedError(
+            "Headless stepping is not available yet; use run() which drives the "
+            "engine's interactive loop through the fleet controller."
+        )
+
+    def run(self) -> None:
+        from .engine import Simulator
+
+        config = self.experiment.build_engine_config()
+        self._sim = Simulator(config=config, fleet_controller=self._fleet_controller)
+        self._sim.run()
+
+    def close(self) -> None:
+        self._sim = None
+
+    def _fleet_controller(self, robots, goals, config, dt, lidar_readings=None):
+        states = {
+            i: _state_from_robot(i, r, goals[i] if goals is not None else None)
+            for i, r in enumerate(robots)
+        }
+        commands = self.experiment._dispatch(states, dt)
+        return [_engine_command(r, commands.get(i, Command())) for i, r in enumerate(robots)]
+
+
+class Ros2Backend(Backend):
+    """Placeholder for the physical arena: one ROS 2 node per agent or a single
+    overarching node for a SwarmController."""
+
+    def __init__(self, experiment) -> None:
+        self.experiment = experiment
+
+    def reset(self) -> None:
+        raise NotImplementedError("ROS 2 backend not implemented yet.")
+
+    def get_states(self) -> dict[int, RobotState]:
+        raise NotImplementedError("ROS 2 backend not implemented yet.")
+
+    def compute(self, states: dict[int, RobotState], dt: float) -> dict[int, Command]:
+        return self.experiment._dispatch(states, dt)
+
+    def step(self, commands: dict[int, Command], dt: float) -> None:
+        raise NotImplementedError("ROS 2 backend not implemented yet.")
+
+    def run(self) -> None:
+        raise NotImplementedError("ROS 2 backend not implemented yet.")
+
+    def close(self) -> None:
+        pass
+
+
+def _state_from_robot(rid: int, robot, goal=None) -> RobotState:
+    st = robot.state()
+    goal_tuple = tuple(float(v) for v in goal) if goal is not None else None
+    return RobotState(
+        id=rid,
+        x=float(st.get("x", getattr(robot, "x", 0.0))),
+        y=float(st.get("y", getattr(robot, "y", 0.0))),
+        theta=float(st.get("theta", getattr(robot, "theta", 0.0))),
+        v=float(st.get("v", 0.0)),
+        omega=float(st.get("omega", 0.0)),
+        type=str(st.get("type", "differential")),
+        radius=float(getattr(robot, "radius", 0.1375)),
+        battery=float(getattr(robot, "battery_percent", 100.0)),
+        goal=goal_tuple,
+    )
+
+
+def _engine_command(robot, command: Command) -> dict:
+    rtype = str(robot.state().get("type", "differential"))
+    if rtype == "holonomic":
+        return {"vx": command.v, "vy": command.omega}
+    if rtype == "ackermann":
+        return {"speed": command.v, "steer": command.omega}
+    return {"v": command.v, "omega": command.omega}
+
+
+def create_backend(name: str, experiment) -> Backend:
+    key = str(name).lower()
+    if key in ("navsim2d", "nav_sim2d", "local", "sim"):
+        return NavSim2DBackend(experiment)
+    if key in ("ros2", "ros", "arena"):
+        return Ros2Backend(experiment)
+    raise ValueError(f"Unknown backend: {name!r} (use 'navsim2d' or 'ros2')")

hivebr/core.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from copy import deepcopy
+import math
+
+from .bridge import create_backend
+from .kinematics import normalize_command
+from .state import Command, RobotState
+
+
+class HiveBrAgent(ABC):
+    """Decentralized (agent-centric) behavior for a single B-bot.
+
+    The engine instantiates this class once per robot and calls ``loop`` with a
+    frozen state snapshot plus the states of its neighbors.
+    """
+
+    def __init__(self, robot_id: int | None = None, **params):
+        self.id = robot_id
+        self.params = params
+
+    def setup(self) -> None:
+        """Optional one-time initialization."""
+
+    @abstractmethod
+    def loop(
+        self,
+        state: RobotState,
+        neighbors: dict[int, RobotState],
+    ) -> "Command | tuple[float, float]":
+        """Return the desired (v, omega) for this agent."""
+
+    def reset(self) -> None:
+        """Optional per-episode reset."""
+
+
+class SwarmController(ABC):
+    """Centralized (global-centric) control for the whole fleet.
+
+    The engine instantiates this class exactly once and passes the states of
+    all robots, expecting one command per robot id back.
+    """
+
+    def setup(self) -> None:
+        """Optional one-time initialization."""
+
+    @abstractmethod
+    def loop(self, states: dict[int, RobotState]) -> dict[int, "Command | tuple[float, float]"]:
+        """Return a mapping ``{robot_id: command}``."""
+
+    def reset(self) -> None:
+        """Optional per-episode reset."""
+
+
+class SwarmExperiment:
+    """Entry point that binds user code to a backend (navsim2d or ROS 2)."""
+
+    def __init__(
+        self,
+        robots=1,
+        backend: str = "navsim2d",
+        config=None,
+        dt: float = 0.03,
+        orientation: bool | None = None,
+        goal_tolerance: float | None = None,
+        comm_radius: float | None = None,
+    ):
+        self.dt = float(dt)
+        self.backend_name = backend
+        self.config = config
+        self.orientation = orientation
+        self.goal_tolerance = goal_tolerance
+        self.comm_radius = comm_radius
+        self.agent_classes: dict[int, type[HiveBrAgent]] = {}
+        self.agent_params: dict[int, dict] = {}
+        self.routes: dict[int, list] = {}
+        self.controller: SwarmController | None = None
+        self.agents: dict[int, HiveBrAgent] = {}
+        self._mode: str | None = None
+        self._backend = None
+        self._num_robots, self._fleet = _parse_robots(robots)
+
+    def assign_agent(self, robot_id=None, robot_ids=None, agent_class=None, **params) -> "SwarmExperiment":
+        if agent_class is None:
+            raise ValueError("assign_agent requires an agent_class.")
+        ids = []
+        if robot_id is not None:
+            ids.append(int(robot_id))
+        if robot_ids is not None:
+            ids.extend(int(i) for i in robot_ids)
+        for rid in ids:
+            self.agent_classes[rid] = agent_class
+            self.agent_params[rid] = dict(params)
+        return self
+
+    def set_controller(self, controller: SwarmController) -> "SwarmExperiment":
+        self.controller = controller
+        return self
+
+    def assign_route(self, robot_id, waypoints) -> "SwarmExperiment":
+        """Give a B-bot its waypoint route (each item ``(x, y)`` or ``(x, y, theta)``)."""
+        self.routes[int(robot_id)] = [list(wp) for wp in waypoints]
+        return self
+
+    def _build(self) -> None:
+        if self.controller is not None:
+            self._mode = "centralized"
+            self.controller.setup()
+            return
+
+        self._mode = "decentralized"
+        self.agents = {
+            rid: cls(robot_id=rid, **self.agent_params.get(rid, {}))
+            for rid, cls in self.agent_classes.items()
+        }
+        for agent in self.agents.values():
+            agent.setup()
+
+    def _dispatch(self, states: dict[int, RobotState], dt: float) -> dict[int, Command]:
+        if self._mode is None:
+            self._build()
+
+        if self._mode == "centralized":
+            raw = self.controller.loop(states)
+            return {int(rid): normalize_command(cmd) for rid, cmd in raw.items()}
+
+        comm_radius = self.comm_radius
+        if comm_radius is None and isinstance(self.config, dict):
+            comm_radius = self.config.get("interaction", {}).get("comm_radius")
+
+        neighbors = {}
+        for rid, st in states.items():
+            visible = {}
+            for j, other in states.items():
+                if j == rid:
+                    continue
+                if comm_radius is None or math.hypot(st.x - other.x, st.y - other.y) <= comm_radius:
+                    visible[j] = other
+            neighbors[rid] = visible
+
+        commands: dict[int, Command] = {}
+        for rid, st in states.items():
+            agent = self.agents.get(rid)
+            commands[rid] = normalize_command(agent.loop(st, neighbors[rid])) if agent else Command()
+        return commands
+
+    def reset(self) -> None:
+        for agent in self.agents.values():
+            agent.reset()
+        if self.controller is not None:
+            self.controller.reset()
+        if self._backend is not None:
+            self._backend.reset()
+
+    def step(self) -> None:
+        backend = self._ensure_backend()
+        states = backend.get_states()
+        commands = self._dispatch(states, self.dt)
+        backend.step(commands, self.dt)
+
+    def run(self) -> None:
+        self._build()
+        self._ensure_backend().run()
+
+    def close(self) -> None:
+        if self._backend is not None:
+            self._backend.close()
+            self._backend = None
+
+    def build_engine_config(self) -> dict:
+        cfg = deepcopy(self.config) if self.config is not None else None
+        if cfg is None:
+            from .engine.simulator import DEFAULT_CONFIG
+
+            cfg = deepcopy(DEFAULT_CONFIG)
+        cfg["robots"] = [deepcopy(meta) for meta in self._fleet]
+        for index, meta in enumerate(cfg["robots"]):
+            meta.setdefault("initial_pose", list(_spread_pose(index)))
+            if index in self.routes:
+                meta["waypoints"] = [list(wp) for wp in self.routes[index]]
+
+        cfg.setdefault("waypoints", {})
+        if self.orientation is not None:
+            cfg["waypoints"]["define_orientation"] = bool(self.orientation)
+        if self.goal_tolerance is not None:
+            cfg["waypoints"]["goal_tolerance"] = float(self.goal_tolerance)
+        return cfg
+
+    def _ensure_backend(self):
+        if self._backend is None:
+            self._backend = create_backend(self.backend_name, self)
+        return self._backend
+
+
+def _spread_pose(index: int, spacing: float = 1.0) -> tuple[float, float, float]:
+    """Default spawn layout: a row starting near the world origin."""
+    return (1.0 + spacing * index, 1.0, 0.0)
+
+
+def _parse_robots(robots) -> tuple[int, list[dict]]:
+    if isinstance(robots, int):
+        return robots, [{"id": i} for i in range(robots)]
+    fleet = [dict(meta) for meta in robots]
+    for i, meta in enumerate(fleet):
+        meta.setdefault("id", i)
+    return len(fleet), fleet

hivebr/engine/__init__.py

@@ -0,0 +1,15 @@
+from .config import load_config
+from .simulator import Simulator
+from .data_logger import DataLogger, SimLog
+from .robots import create_fleet, create_robot
+from . import interaction
+
+__all__ = [
+    "Simulator",
+    "load_config",
+    "DataLogger",
+    "SimLog",
+    "create_fleet",
+    "create_robot",
+    "interaction",
+]

hivebr/engine/camera.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+from dataclasses import dataclass
+import numpy as np
+
+
+@dataclass
+class Camera2D:
+    world_width: float
+    world_height: float
+    viewport_width: int
+    viewport_height: int
+    initial_zoom: float = 1.0
+    min_zoom: float = 0.25
+    max_zoom: float = 8.0
+    zoom_step: float = 1.1
+    margin_x: float = 1.0
+    margin_y: float = 1.0
+
+    def __post_init__(self):
+        self.zoom = self.initial_zoom
+        self.offset = np.array([-self.margin_x, -self.margin_y], dtype=float)
+        self.compute_scale()
+
+    def compute_scale(self):
+        scale_x = self.viewport_width / self.world_width
+        scale_y = self.viewport_height / self.world_height
+        self.base_scale = min(scale_x, scale_y)
+        self.scale = self.base_scale * self.zoom
+
+    def resize(self, viewport_width: int, viewport_height: int):
+        self.viewport_width = max(1, int(viewport_width))
+        self.viewport_height = max(1, int(viewport_height))
+        self.compute_scale()
+
+    def reset(self):
+        self.zoom = self.initial_zoom
+        self.offset = np.array([-self.margin_x, -self.margin_y], dtype=float)
+        self.compute_scale()
+
+    def center_on_origin_with_margin(self):
+        self.offset = np.array([-self.margin_x, -self.margin_y], dtype=float)
+
+    def zoom_at(self, mouse_pos, wheel_y: int):
+        old_world = self.screen_to_world(*mouse_pos)
+        factor = self.zoom_step if wheel_y > 0 else 1.0 / self.zoom_step
+        self.zoom = float(np.clip(self.zoom * factor, self.min_zoom, self.max_zoom))
+        self.compute_scale()
+        new_world = self.screen_to_world(*mouse_pos)
+        self.offset += old_world - new_world
+
+    def pan_pixels(self, dx: float, dy: float):
+        self.offset[0] -= dx / self.scale
+        self.offset[1] += dy / self.scale
+
+    def _display_offsets(self):
+        ox = (self.viewport_width - self.world_width * self.scale) / 2
+        oy = (self.viewport_height - self.world_height * self.scale) / 2
+        return ox, oy
+
+    def world_to_screen(self, px: float, py: float):
+        ox, oy = self._display_offsets()
+        px_cam = px - self.offset[0]
+        py_cam = py - self.offset[1]
+        sx = int(px_cam * self.scale + ox)
+        sy = int(self.viewport_height - (py_cam * self.scale + oy))
+        return sx, sy
+
+    def screen_to_world(self, sx: float, sy: float):
+        ox, oy = self._display_offsets()
+        px = (sx - ox) / self.scale + self.offset[0]
+        py = ((self.viewport_height - sy) - oy) / self.scale + self.offset[1]
+        return np.array([px, py], dtype=float)

hivebr/engine/config.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+from pathlib import Path
+from typing import Any
+import yaml
+
+
+def load_config(path: str | Path) -> dict[str, Any]:
+    path = Path(path)
+    with path.open("r", encoding="utf-8") as f:
+        return yaml.safe_load(f)