Benchmark dev

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/warehouse_generator.py

@@ -0,0 +1,91 @@
+import numpy as np
+from pydantic import BaseModel, validator
+from pogema.utils import grid_to_str
+
+class WarehouseConfig(BaseModel):
+    wall_width: int = 5
+    wall_height: int = 2
+    walls_in_row: int = 5
+    walls_rows: int = 5
+    bottom_gap: int = 5
+    horizontal_gap: int = 1
+    vertical_gap: int = 3
+    wfi_instances: bool = True
+
+    @validator('wall_width', 'wall_height', 'walls_in_row', 'walls_rows', 'bottom_gap', 'horizontal_gap', 'vertical_gap')
+    def must_be_positive(cls, value, field):
+        if value <= 0:
+            raise ValueError(f'{field.name} must be a positive integer')
+        return value
+
+    @validator('walls_in_row', 'walls_rows')
+    def must_be_at_least_one(cls, value, field):
+        if value < 1:
+            raise ValueError(f'{field.name} must be at least 1')
+        return value
+
+def generate_warehouse(cfg: WarehouseConfig):
+    height = cfg.vertical_gap * (cfg.walls_rows + 1) + cfg.wall_height * cfg.walls_rows
+    width = cfg.bottom_gap * 2 + cfg.wall_width * cfg.walls_in_row + cfg.horizontal_gap * (cfg.walls_in_row - 1)
+
+    grid = np.zeros((height, width), dtype=int)
+
+    for row in range(cfg.walls_rows):
+        row_start = cfg.vertical_gap * (row + 1) + cfg.wall_height * row
+        for col in range(cfg.walls_in_row):
+            col_start = cfg.bottom_gap + col * (cfg.wall_width + cfg.horizontal_gap)
+            grid[row_start:row_start + cfg.wall_height, col_start:col_start + cfg.wall_width] = 1
+
+    return grid_to_str(grid)
+
+
+def generate_wfi_positions(grid_str, bottom_gap, vertical_gap):
+    if vertical_gap == 1:
+        raise ValueError("Cannot generate WFI instance with vertical_gap of 1.")
+
+    grid = [list(row) for row in grid_str.strip().split('\n')]
+    height = len(grid)
+    width = len(grid[0])
+
+    start_locations = []
+    goal_locations = []
+
+    for row in range(1, height - 1):
+        if row % 3 == 0:
+            continue
+        for col in range(bottom_gap - 1):
+            if grid[row][col] == '.':
+                start_locations.append((row, col))
+        for col in range(width - bottom_gap + 1, width):
+            if grid[row][col] == '.':
+                start_locations.append((row, col))
+
+    if vertical_gap == 2:
+        for row in range(1, height):
+            for col in range(width):
+                if grid[row][col] == '.' and grid[row - 1][col] == '#':
+                    goal_locations.append((row, col))
+    else:
+        for row in range(height):
+            for col in range(width):
+                if grid[row][col] == '.':
+                    if (row > 0 and grid[row - 1][col] == '#') or (row < height - 1 and grid[row + 1][col] == '#'):
+                        goal_locations.append((row, col))
+
+    return start_locations, goal_locations
+
+def generate_wfi_warehouse(cfg: WarehouseConfig):
+    grid = generate_warehouse(cfg)
+    start_locations, goal_locations = generate_wfi_positions(grid, cfg.bottom_gap, cfg.vertical_gap)
+    grid_list = [list(row) for row in grid.split('\n')]
+
+    for s in start_locations:
+        grid_list[s[0]][s[1]] = '$'
+    for s in goal_locations:
+        if grid_list[s[0]][s[1]] == '$':
+            grid_list[s[0]][s[1]] = '&'
+        else:
+            grid_list[s[0]][s[1]] = '@'
+    str_grid = '\n'.join([''.join(row) for row in grid_list])
+
+    return str_grid