Benchmark dev (#36)

* New way to generate instances

* SOC_Makespan Metric

* Fixed WFI generator for LMAPF

* Warehouse generator

* grid_to_str in utils

* Refactoring for custom map to list instances creation

* Refactored wrappers

* Removed warehouse_generator.py

* Bump version to 1.2.3a3

---------

Co-authored-by: Anton Andreychuk <[email protected]>

pogema/__init__.py

@@ -9,7 +9,8 @@
 from pogema.grid_config import Easy32x32, Normal32x32, Hard32x32, ExtraHard32x32
 from pogema.grid_config import Easy64x64, Normal64x64, Hard64x64, ExtraHard64x64
 
-__version__ = '1.2.3a2'
+__version__ = '1.2.3a3'
+
 
 __all__ = [
     'GridConfig',

pogema/envs.py

@@ -7,9 +7,9 @@
 from pogema.grid import Grid, GridLifeLong
 from pogema.grid_config import GridConfig
 from pogema.wrappers.metrics import LifeLongAverageThroughputMetric, NonDisappearEpLengthMetric, \
-    NonDisappearCSRMetric, NonDisappearISRMetric, EpLengthMetric, ISRMetric, CSRMetric
+    NonDisappearCSRMetric, NonDisappearISRMetric, EpLengthMetric, ISRMetric, CSRMetric, SumOfCostsAndMakespanMetric
 from pogema.wrappers.multi_time_limit import MultiTimeLimit
-from pogema.generator import generate_new_target
+from pogema.generator import generate_new_target, generate_from_possible_targets
 from pogema.wrappers.persistence import PersistentWrapper
 
 
@@ -300,6 +300,18 @@ def _initialize_grid(self):
         seeds = main_rng.integers(np.iinfo(np.int32).max, size=self.grid_config.num_agents)
         self.random_generators = [np.random.default_rng(seed) for seed in seeds]
 
+    def _generate_new_target(self, agent_idx):
+        if self.grid_config.possible_targets_xy is not None:
+            new_goal = generate_from_possible_targets(self.random_generators[agent_idx], 
+                                                     self.grid_config.possible_targets_xy, 
+                                                     self.grid.positions_xy[agent_idx])
+            return (new_goal[0] + self.grid_config.obs_radius, new_goal[1] + self.grid_config.obs_radius)
+        else:
+            return generate_new_target(self.random_generators[agent_idx],
+                                    self.grid.point_to_component,
+                                    self.grid.component_to_points,
+                                    self.grid.positions_xy[agent_idx])
+
     def step(self, action: list):
         assert len(action) == self.grid_config.num_agents
         rewards = []
@@ -317,10 +329,7 @@ def step(self, action: list):
                 rewards.append(0.0)
 
             if self.grid.on_goal(agent_idx):
-                self.grid.finishes_xy[agent_idx] = generate_new_target(self.random_generators[agent_idx],
-                                                                       self.grid.point_to_component,
-                                                                       self.grid.component_to_points,
-                                                                       self.grid.positions_xy[agent_idx])
+                self.grid.finishes_xy[agent_idx] = self._generate_new_target(agent_idx)
 
         for agent_idx in range(self.grid_config.num_agents):
             infos[agent_idx]['is_active'] = self.grid.is_active[agent_idx]
@@ -382,6 +391,7 @@ def _make_pogema(grid_config):
             env = NonDisappearISRMetric(env)
             env = NonDisappearCSRMetric(env)
             env = NonDisappearEpLengthMetric(env)
+            env = SumOfCostsAndMakespanMetric(env)
         elif grid_config.on_target == 'finish':
             env = ISRMetric(env)
             env = CSRMetric(env)

pogema/generator.py

@@ -103,6 +103,14 @@ def placing(order, components, grid, start_id, num_agents):
 
     return positions_xy, finishes_xy
 
+def generate_from_possible_positions(grid_config: GridConfig):
+    if len(grid_config.possible_agents_xy) < grid_config.num_agents or len(grid_config.possible_targets_xy) < grid_config.num_agents:
+        raise OverflowError(f"Can't create task. Not enough possible positions for {grid_config.num_agents} agents.")
+    rng = np.random.default_rng(grid_config.seed)
+    rng.shuffle(grid_config.possible_agents_xy)
+    rng.shuffle(grid_config.possible_targets_xy)
+    return grid_config.possible_agents_xy[:grid_config.num_agents], grid_config.possible_targets_xy[:grid_config.num_agents]
+
 
 def generate_positions_and_targets_fast(obstacles, grid_config):
     c = grid_config
@@ -117,6 +125,11 @@ def generate_positions_and_targets_fast(obstacles, grid_config):
 
     return placing(order=order, components=components, grid=grid, start_id=start_id, num_agents=c.num_agents)
 
+def generate_from_possible_targets(rnd_generator, possible_positions, position):
+    new_target = tuple(rnd_generator.choice(possible_positions, 1)[0])
+    while new_target == position:
+        new_target = tuple(rnd_generator.choice(possible_positions, 1)[0])
+    return new_target
 
 def generate_new_target(rnd_generator, point_to_component, component_to_points, position):
     component_id = point_to_component[position]

pogema/grid.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 from pogema.generator import generate_obstacles, generate_positions_and_targets_fast, \
-    get_components
+    get_components, generate_from_possible_positions
 from .grid_config import GridConfig
 from .grid_registry import in_registry, get_grid
 from .utils import render_grid
@@ -16,7 +16,6 @@ def __init__(self, grid_config: GridConfig, add_artificial_border: bool = True,
 
         self.config = grid_config
         self.rnd = np.random.default_rng(grid_config.seed)
-
         if self.config.map is None:
             self.obstacles = generate_obstacles(self.config)
         else:
@@ -41,6 +40,8 @@ def __init__(self, grid_config: GridConfig, add_artificial_border: bool = True,
                     warnings.warn(f"There is an obstacle on a finish point ({s_x}, {s_y}), replacing with free cell",
                                   Warning, stacklevel=2)
                 self.obstacles[f_x, f_y] = grid_config.FREE
+        elif grid_config.possible_agents_xy and grid_config.possible_targets_xy:
+            self.starts_xy, self.finishes_xy = generate_from_possible_positions(self.config)
         else:
             self.starts_xy, self.finishes_xy = generate_positions_and_targets_fast(self.obstacles, self.config)
 

pogema/grid_config.py

@@ -16,6 +16,8 @@ class GridConfig(CommonSettings, ):
     obs_radius: int = 5
     agents_xy: Optional[list] = None
     targets_xy: Optional[list] = None
+    possible_agents_xy: Optional[list] = None
+    possible_targets_xy: Optional[list] = None
     collision_system: Literal['block_both', 'priority', 'soft'] = 'priority'
     persistent: bool = False
     observation_type: Literal['POMAPF', 'MAPF', 'default'] = 'default'
@@ -53,18 +55,26 @@ def obs_radius_must_be_positive(cls, v):
         return v
 
     @validator('map', always=True)
-    def map_validation(cls, v, values, ):
+    def map_validation(cls, v, values):
         if v is None:
             return None
         if isinstance(v, str):
-            v, agents_xy, targets_xy = cls.str_map_to_list(v, values['FREE'], values['OBSTACLE'])
-            if agents_xy and targets_xy and values['agents_xy'] is not None and values['targets_xy'] is not None:
-                raise KeyError("""Can't create task. Please provide agents_xy and targets_xy only ones.
+            v, agents_xy, targets_xy, possible_agents_xy, possible_targets_xy = cls.str_map_to_list(v, values['FREE'],
+                                                                                                    values['OBSTACLE'])
+            if agents_xy and targets_xy and values.get('agents_xy') is not None and values.get(
+                    'targets_xy') is not None:
+                raise KeyError("""Can't create task. Please provide agents_xy and targets_xy only once.
                 Either with parameters or with a map.""")
+            if (agents_xy or targets_xy) and (possible_agents_xy or possible_targets_xy):
+                raise KeyError("""Can't create task. Mark either possible locations or precise ones.""")
             elif agents_xy and targets_xy:
                 values['agents_xy'] = agents_xy
                 values['targets_xy'] = targets_xy
                 values['num_agents'] = len(agents_xy)
+            elif (values.get('agents_xy') is None or values.get(
+                    'targets_xy') is None) and possible_agents_xy and possible_targets_xy:
+                values['possible_agents_xy'] = possible_agents_xy
+                values['possible_targets_xy'] = possible_targets_xy
         size = len(v)
         area = 0
         for line in v:
@@ -88,6 +98,14 @@ def targets_xy_validation(cls, v, values):
             values['num_agents'] = len(v)
         return v
 
+    @validator('possible_agents_xy')
+    def possible_agents_xy_validation(cls, v):
+        return v
+
+    @validator('possible_targets_xy')
+    def possible_targets_xy_validation(cls, v):
+        return v
+
     @staticmethod
     def check_positions(v, size):
         for position in v:
@@ -100,35 +118,53 @@ def str_map_to_list(str_map, free, obstacle):
         obstacles = []
         agents = {}
         targets = {}
-        for idx, line in enumerate(str_map.split()):
+        possible_agents_xy = []
+        possible_targets_xy = []
+        special_chars = {'@', '$', '!'}
+
+        for row_idx, line in enumerate(str_map.split()):
             row = []
-            for char in line:
+            for col_idx, char in enumerate(line):
+                position = (row_idx, col_idx)
+
                 if char == '.':
                     row.append(free)
+                    possible_agents_xy.append(position)
+                    possible_targets_xy.append(position)
                 elif char == '#':
                     row.append(obstacle)
+                elif char in special_chars:
+                    row.append(free)
+                    if char == '@':
+                        possible_agents_xy.append(position)
+                    elif char == '$':
+                        possible_targets_xy.append(position)
                 elif 'A' <= char <= 'Z':
-                    targets[char.lower()] = len(obstacles), len(row)
+                    targets[char.lower()] = position
                     row.append(free)
+                    possible_agents_xy.append(position)
+                    possible_targets_xy.append(position)
                 elif 'a' <= char <= 'z':
-                    agents[char.lower()] = len(obstacles), len(row)
+                    agents[char.lower()] = position
                     row.append(free)
+                    possible_agents_xy.append(position)
+                    possible_targets_xy.append(position)
                 else:
-                    raise KeyError(f"Unsupported symbol '{char}' at line {idx}")
+                    raise KeyError(f"Unsupported symbol '{char}' at line {row_idx}")
+
             if row:
-                if obstacles:
-                    assert len(obstacles[-1]) == len(row), f"Wrong string size for row {idx};"
+                assert len(obstacles[-1]) == len(row) if obstacles else True, f"Wrong string size for row {row_idx};"
                 obstacles.append(row)
 
-        targets_xy = []
-        agents_xy = []
-        for _, (x, y) in sorted(agents.items()):
-            agents_xy.append([x, y])
-        for _, (x, y) in sorted(targets.items()):
-            targets_xy.append([x, y])
+        agents_xy = [[x, y] for _, (x, y) in sorted(agents.items())]
+        targets_xy = [[x, y] for _, (x, y) in sorted(targets.items())]
+
+        assert len(targets_xy) == len(agents_xy), "Mismatch in number of agents and targets."
+
+        if not any(char in special_chars for char in str_map):
+            possible_agents_xy, possible_targets_xy = None, None
 
-        assert len(targets_xy) == len(agents_xy)
-        return obstacles, agents_xy, targets_xy
+        return obstacles, agents_xy, targets_xy, possible_agents_xy, possible_targets_xy
 
 
 class PredefinedDifficultyConfig(GridConfig):

pogema/utils.py

@@ -9,6 +9,10 @@ class AgentsTargetsSizeError(Exception):
     pass
 
 
+def grid_to_str(grid):
+    return '\n'.join(''.join('.' if cell == 0 else '#' for cell in row) for row in grid)
+
+
 def check_grid(obstacles, agents_xy, targets_xy):
     if bool(agents_xy) != bool(targets_xy):
         raise AgentsTargetsSizeError("Agents and targets must be defined together/undefined together!")

pogema/wrappers/metrics.py

@@ -1,3 +1,4 @@
+import numpy as np
 from gymnasium import Wrapper
 
 
@@ -42,26 +43,6 @@ def _compute_stats(self, step, is_on_goal, finished):
             return result
 
 
-class LifeLongAttritionMetric(AbstractMetric):
-
-    def __init__(self, env):
-        super().__init__(env)
-        self._attrition_steps = 0
-        self._on_goal_steps = 0
-
-    def _compute_stats(self, step, is_on_goal, finished):
-        for agent_idx, on_goal in enumerate(is_on_goal):
-            if not on_goal:
-                self._attrition_steps += 1
-            else:
-                self._on_goal_steps += 1
-        if finished:
-            result = {
-                'attrition': self._attrition_steps / max(1, self._on_goal_steps)}
-            self._solved_instances = 0
-            return result
-
-
 class NonDisappearCSRMetric(AbstractMetric):
 
     def _compute_stats(self, step, is_on_goal, finished):
@@ -124,3 +105,48 @@ def _compute_stats(self, step, is_on_goal, finished):
             results = {'ISR': self._solved_instances / self.get_num_agents()}
             self._solved_instances = 0
             return results
+
+
+class SumOfCostsAndMakespanMetric(AbstractMetric):
+    def __init__(self, env):
+        super().__init__(env)
+        self._solve_time = [None for _ in range(self.get_num_agents())]
+
+    def _compute_stats(self, step, is_on_goal, finished):
+        for idx, on_goal in enumerate(is_on_goal):
+            if self._solve_time[idx] is None and (on_goal or finished):
+                self._solve_time[idx] = step
+            if not on_goal and not finished:
+                self._solve_time[idx] = None
+
+        if finished:
+            result = {'SoC': sum(self._solve_time) + self.get_num_agents(), 'makespan': max(self._solve_time) + 1}
+            self._solve_time = [None for _ in range(self.get_num_agents())]
+            return result
+
+
+class AgentsInObsWrapper(Wrapper):
+    def __init__(self, env):
+        super().__init__(env)
+        self._avg_num_agents = None
+
+    def count_agents(self, observations):
+        avg_num_agents = []
+        for obs in observations:
+            avg_num_agents.append(obs['agents'].sum().sum())
+        self._avg_num_agents.append(np.mean(avg_num_agents))
+
+    def step(self, actions):
+        observations, rewards, terminated, truncated, infos = self.env.step(actions)
+        self.count_agents(observations)
+        if all(terminated) or all(truncated):
+            if 'metrics' not in infos[0]:
+                infos[0]['metrics'] = {}
+            infos[0]['metrics'].update(avg_num_agents_in_obs=float(np.mean(self._avg_num_agents)))
+        return observations, rewards, terminated, truncated, infos
+
+    def reset(self, **kwargs):
+        self._avg_num_agents = []
+        observations, info = self.env.reset(**kwargs)
+        self.count_agents(observations)
+        return observations, info

tests/test_grid.py

@@ -207,3 +207,45 @@ def test_duplicated_params():
 def test_custom_grid_with_empty_agents_and_targets():
     grid_map = """...."""
     Grid(GridConfig(agents_xy=None, targets_xy=None, map=grid_map, num_agents=1))
+
+
+def test_custom_grid_with_specific_positions():
+    grid_map = """
+        !!!!!!!!!!!!!!!!!!
+        !@@!@@!$$$$$$$$$$!
+        !@@!@@!##########!
+        !@@!@@!$$$$$$$$$$!
+        !!!!!!!!!!!!!!!!!!
+        !@@!@@!$$$$$$$$$$!
+        !@@!@@!##########!
+        !@@!@@!$$$$$$$$$$!
+        !!!!!!!!!!!!!!!!!!
+    """
+    Grid(GridConfig(obs_radius=2, size=4, num_agents=24, map=grid_map))
+    with pytest.raises(OverflowError):
+        Grid(GridConfig(obs_radius=2, size=4, num_agents=25, map=grid_map))
+
+    grid_map = """
+        !!!!!!!!!!!
+        !@@!@@!$$$$
+        !@@!@@!####
+        !@@!@@!$$$$
+        !!!!!!!!!!!
+        !@@!@@!$$$$
+        !@@!@@!####
+        !@@!@@!$$$$
+        !!!!!!!!!!!
+    """
+    Grid(GridConfig(obs_radius=2, num_agents=16, map=grid_map))
+    with pytest.raises(OverflowError):
+        Grid(GridConfig(obs_radius=2, num_agents=17, map=grid_map))
+
+    grid_map = """
+            !!!!!!!!!!!
+            !@@!@@!.Ab.
+            !@@!@@!####
+            !@@!@@!.aB.
+
+        """
+    with pytest.raises(KeyError):
+        Grid(GridConfig(obs_radius=2, map=grid_map))