Split shadowing and hijacking for consistent vehicles in bubble

smarts/core/bubble_manager.py

@@ -18,6 +18,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 from functools import lru_cache
+from sys import maxsize
 import logging
 from collections import defaultdict
 from copy import deepcopy
@@ -36,7 +37,7 @@
 from smarts.core.utils.string import truncate
 from smarts.core.vehicle import Vehicle, VehicleState
 from smarts.core.vehicle_index import VehicleIndex
-from smarts.sstudio.types import BoidAgentActor
+from smarts.sstudio.types import BoidAgentActor, BubbleLimits
 from smarts.sstudio.types import Bubble as SSBubble
 from smarts.sstudio.types import SocialAgentActor
 from smarts.zoo.registry import make as make_social_agent
@@ -64,17 +65,33 @@ class Bubble:
     def __init__(self, bubble: SSBubble, sumo_road_network: SumoRoadNetwork):
         geometry = bubble.zone.to_geometry(sumo_road_network)
 
-        self._bubble_heading = 0
-        self._bubble = bubble
-        self._limit = bubble.limit
-        self._cached_inner_geometry = geometry
-        self._exclusion_prefixes = bubble.exclusion_prefixes
+        bubble_limit = (
+            bubble.limit or BubbleLimits()
+            if bubble.limit is None or isinstance(bubble.limit, BubbleLimits)
+            else BubbleLimits(bubble.limit, bubble.limit + 1)
+        )
 
         if isinstance(bubble.actor, BoidAgentActor):
+
+            def safe_min(a, b):
+                return min(a or maxsize, b or maxsize)
+
             if bubble.limit is None:
-                self._limit = bubble.actor.capacity
+                bubble_limit = bubble.actor.capacity
             elif bubble.actor.capacity is not None:
-                self._limit = min(bubble.limit, bubble.actor.capacity)
+                hijack_limit = safe_min(
+                    bubble.limit.hijack_limit, bubble.actor.capacity
+                )
+                shadow_limit = safe_min(
+                    bubble.limit.shadow_limit, bubble.actor.capacity.shadow_limit
+                )
+                bubble_limit = BubbleLimits(hijack_limit, shadow_limit)
+
+        self._bubble_heading = 0
+        self._bubble = bubble
+        self._limit = bubble_limit
+        self._cached_inner_geometry = geometry
+        self._exclusion_prefixes = bubble.exclusion_prefixes
 
         self._cached_airlock_geometry = self._cached_inner_geometry.buffer(
             bubble.margin, cap_style=CAP_STYLE.square, join_style=JOIN_STYLE.mitre,
@@ -118,7 +135,7 @@ def geometry(self) -> Polygon:
     def airlock_geometry(self) -> Polygon:
         return self._cached_airlock_geometry
 
-    def is_admissible(
+    def admissibility(
         self,
         vehicle_id: str,
         index: VehicleIndex,
@@ -130,29 +147,43 @@ def is_admissible(
         """
         for prefix in self.exclusion_prefixes:
             if vehicle_id.startswith(prefix):
-                return False
+                return False, False
 
+        hijackable, shadowable = True, True
         if self._limit is not None:
             # Already hijacked (according to VehicleIndex) + to be hijacked (running cursors)
-            current_hijacked_or_shadowed_vehicle_ids = {
+            current_hijacked_vehicle_ids = {
+                v_id
+                for v_id in vehicle_ids_in_bubbles[self]
+                if index.vehicle_is_hijacked(v_id)
+            }
+            current_shadowed_vehicle_ids = {
                 v_id
                 for v_id in vehicle_ids_in_bubbles[self]
-                if index.vehicle_is_hijacked(v_id) or index.vehicle_is_shadowed(v_id)
+                if index.vehicle_is_shadowed(v_id)
             }
-            per_bubble_veh_ids = BubbleManager.vehicle_ids_per_bubble(
+            vehicle_ids_by_bubble_state = BubbleManager._vehicle_ids_divided_by_bubble_state(
                 frozenset(running_cursors)
             )
-            running_hijacked_or_shadowed_vehicle_ids = per_bubble_veh_ids[self]
 
-            all_hijacked_or_shadowed_vehicle_ids = (
-                current_hijacked_or_shadowed_vehicle_ids
-                | running_hijacked_or_shadowed_vehicle_ids
+            all_hijacked_vehicle_ids = (
+                current_hijacked_vehicle_ids
+                | vehicle_ids_by_bubble_state[BubbleState.InAirlock][self]
+            ) - {vehicle_id}
+
+            all_shadowed_vehicle_ids = (
+                current_shadowed_vehicle_ids
+                | vehicle_ids_by_bubble_state[BubbleState.InBubble][self]
             ) - {vehicle_id}
 
-            if len(all_hijacked_or_shadowed_vehicle_ids) >= self._limit:
-                return False
+            hijackable = len(all_hijacked_vehicle_ids) < (
+                self._limit.hijack_limit or maxsize
+            )
+            shadowable = len(all_shadowed_vehicle_ids) + len(
+                all_hijacked_vehicle_ids
+            ) < (self._limit.shadow_limit or maxsize)
 
-        return True
+        return hijackable, shadowable
 
     def in_bubble_or_airlock(self, position):
         if not isinstance(position, Point):
@@ -227,10 +258,10 @@ def from_pos(
         vehicle_ids_per_bubble: Dict[Bubble, Set[str]],
         running_cursors: Set["Cursor"],
     ):
-        in_bubble, in_airlock = bubble.in_bubble_or_airlock(pos)
+        in_bubble_zone, in_airlock_zone = bubble.in_bubble_or_airlock(pos)
         is_social = vehicle.id in index.social_vehicle_ids()
         is_hijacked, is_shadowed = index.vehicle_is_hijacked_or_shadowed(vehicle.id)
-        is_admissible = bubble.is_admissible(
+        is_hijack_admissible, is_airlock_admissible = bubble.admissibility(
             vehicle.id, index, vehicle_ids_per_bubble, running_cursors
         )
         was_in_this_bubble = vehicle.id in vehicle_ids_per_bubble[bubble]
@@ -242,25 +273,25 @@ def from_pos(
         #       time-based airlocking. For robust code we'll want to handle these
         #       scenarios (e.g. hijacking if didn't airlock first)
         transition = None
-        if is_social and not is_shadowed and is_admissible and in_airlock:
+        if is_social and not is_shadowed and is_airlock_admissible and in_airlock_zone:
             transition = BubbleTransition.AirlockEntered
-        elif is_shadowed and is_admissible and in_bubble:
+        elif is_shadowed and is_hijack_admissible and in_bubble_zone:
             transition = BubbleTransition.Entered
-        elif is_hijacked and in_airlock:
+        elif is_hijacked and in_airlock_zone:
             # XXX: This may get called repeatedly because we don't actually change
             #      any state when this happens.
             transition = BubbleTransition.Exited
         elif (
             was_in_this_bubble
             and (is_shadowed or is_hijacked)
-            and not (in_airlock or in_bubble)
+            and not (in_airlock_zone or in_bubble_zone)
         ):
             transition = BubbleTransition.AirlockExited
 
         state = None
-        if is_admissible and in_bubble:
+        if is_hijack_admissible and in_bubble_zone:
             state = BubbleState.InBubble
-        elif is_admissible and in_airlock:
+        elif is_airlock_admissible and in_airlock_zone:
             state = BubbleState.InAirlock
 
         return cls(
@@ -296,17 +327,25 @@ def is_active(bubble):
         return [bubble for bubble in self._bubbles if is_active(bubble)]
 
     @staticmethod
-    @lru_cache(maxsize=2)
-    def vehicle_ids_per_bubble(cursors: FrozenSet[Cursor]) -> Dict[Bubble, Set[str]]:
-        grouped = defaultdict(set)
+    def _vehicle_ids_divided_by_bubble_state(
+        cursors: FrozenSet[Cursor],
+    ) -> Dict[Bubble, Set[str]]:
+        vehicle_ids_grouped_by_cursor = defaultdict(lambda: defaultdict(set))
         for cursor in cursors:
-            if (
-                cursor.state == BubbleState.InBubble
-                or cursor.state == BubbleState.InAirlock
-            ):
-                grouped[cursor.bubble].add(cursor.vehicle_id)
+            vehicle_ids_grouped_by_cursor[cursor.state][cursor.bubble].add(
+                cursor.vehicle_id
+            )
+        return vehicle_ids_grouped_by_cursor
 
-        return grouped
+    @classmethod
+    @lru_cache(maxsize=2)
+    def vehicle_ids_per_bubble(
+        cls, cursors: FrozenSet[Cursor],
+    ) -> Dict[Bubble, Set[str]]:
+        vid = cls._vehicle_ids_divided_by_bubble_state(cursors)
+        return defaultdict(
+            set, {**vid[BubbleState.InBubble], **vid[BubbleState.InAirlock]}
+        )
 
     def step(self, sim):
         self._move_travelling_bubbles(sim)

smarts/core/tests/test_bubble_manager.py

@@ -35,6 +35,13 @@
 # TODO: Add test for travelling bubbles
 
 
+@pytest.fixture
+def time_resolution(request):
+    tr = getattr(request, "param", 0.1)
+    assert tr >= 1e-10, "Should be a non-negative non-zero real number"
+    return tr
+
+
 @pytest.fixture
 def bubble(request):
     """
@@ -43,7 +50,7 @@ def bubble(request):
     return t.Bubble(
         zone=t.PositionalZone(pos=(100, 0), size=(10, 10)),
         margin=2,
-        limit=getattr(request, "param", 10),
+        limit=getattr(request, "param", t.BubbleLimits(10, 11)),
         actor=t.SocialAgentActor(
             name="zoo-car", agent_locator="zoo.policies:keep-lane-agent-v0"
         ),
@@ -65,9 +72,10 @@ def mock_provider():
 
 
 @pytest.fixture
-def smarts(scenarios, mock_provider):
+def smarts(scenarios, mock_provider, time_resolution):
     smarts_ = SMARTS(
-        agent_interfaces={}, traffic_sim=SumoTrafficSimulation(time_resolution=0.1),
+        agent_interfaces={},
+        traffic_sim=SumoTrafficSimulation(time_resolution=time_resolution,),
     )
     smarts_.add_provider(mock_provider)
     smarts_.reset(next(scenarios))
@@ -122,9 +130,11 @@ def test_bubble_manager_state_change(smarts, mock_provider):
         assert got_hijacked == hijacked, assert_msg
 
 
-@pytest.mark.parametrize("bubble", [1], indirect=True)
-def test_bubble_manager_limit(smarts, mock_provider):
+@pytest.mark.parametrize("bubble", [t.BubbleLimits(1, 1)], indirect=True)
+def test_bubble_manager_limit(smarts, mock_provider, time_resolution):
     vehicle_ids = ["vehicle-1", "vehicle-2", "vehicle-3"]
+    speed = 2.5
+    distance_per_step = speed * time_resolution
     for x in range(200):
         vehicle_ids = {
             v_id
@@ -136,9 +146,10 @@ def test_bubble_manager_limit(smarts, mock_provider):
             (
                 v_id,
                 Pose.from_center(
-                    (80 + y * 0.5 + x * 0.25, y * 4 - 4, 0), Heading(math.pi / 2)
+                    (80 + y * 0.5 + x * distance_per_step, y * 4 - 4, 0),
+                    Heading(-math.pi / 2),
                 ),
-                10,
+                speed,  # speed
             )
             for y, v_id in enumerate(vehicle_ids)
         ]
@@ -161,7 +172,7 @@ def test_vehicle_spawned_in_bubble_is_not_captured(smarts, mock_provider):
                 (
                     vehicle_id,
                     Pose.from_center((100 + x, 0, 0), Heading(-math.pi / 2)),
-                    10,
+                    10,  # speed
                 )
             ]
         )
@@ -179,7 +190,7 @@ def test_vehicle_spawned_outside_bubble_is_captured(smarts, mock_provider):
                 (
                     vehicle_id,
                     Pose.from_center((90 + x, 0, 0), Heading(-math.pi / 2)),
-                    10,
+                    10,  # speed
                 )
             ]
         )

smarts/sstudio/types.py

@@ -21,6 +21,7 @@
 import random
 import collections.abc as collections_abc
 from dataclasses import dataclass, field
+from sys import maxsize
 from typing import Any, Dict, Optional, Sequence, Tuple, Union
 
 from shapely.geometry import GeometryCollection, MultiPolygon, Polygon, Point
@@ -212,7 +213,7 @@ class BoidAgentActor(SocialAgentActor):
 
     # The max number of vehicles that this agent will control at a time. This value is
     # honored when using a bubble for boid dynamic assignment.
-    capacity: int = None
+    capacity: "BubbleLimits" = None
     """The capacity of the boid agent to take over vehicles."""
 
 
@@ -609,6 +610,20 @@ def to_geometry(self, road_network: SumoRoadNetwork = None) -> Polygon:
         return Polygon([p0, (p0[0], p1[1]), p1, (p1[0], p0[1])])
 
 
+@dataclass(frozen=True)
+class BubbleLimits:
+    hijack_limit: int = maxsize
+    """The maximum number of vehicles the bubble can hijack"""
+    shadow_limit: int = maxsize
+    """The maximum number of vehicles the bubble can shadow"""
+
+    def __post_init__(self):
+        if self.shadow_limit is None:
+            raise ValueError("Shadow limit must be a non-negative real number")
+        if self.hijack_limit is None or self.shadow_limit < self.hijack_limit:
+            raise ValueError("Shadow limit must be >= hijack limit")
+
+
 @dataclass(frozen=True)
 class Bubble:
     """A descriptor that defines a capture bubble for social agents."""
@@ -622,7 +637,7 @@ class Bubble:
     # If limit != None it will only allow that specified number of vehicles to be
     # hijacked. N.B. when actor = BoidAgentActor the lesser of the actor capacity
     # and bubble limit will be used.
-    limit: int = None
+    limit: BubbleLimits = None
     """The maximum number of actors that could be captured."""
     exclusion_prefixes: Tuple[str, ...] = field(default_factory=tuple)
     """Used to exclude social actors from capture."""