import time from collections import deque, UserList from enum import Enum from typing import List, Union, NamedTuple, Dict import numpy as np from environments.factory.base.base_factory import BaseFactory from environments.helpers import Constants as c from environments import helpers as h from environments.factory.base.objects import Agent, Entity, Action, Tile, MoveableEntity from environments.factory.base.registers import Entities, EntityObjectRegister, ObjectRegister, \ MovingEntityObjectRegister, Register from environments.factory.renderer import RenderEntity NO_ITEM = 0 ITEM_DROP_OFF = 1 def inventory_slice_name(agent_i): if isinstance(agent_i, int): return f'{c.INVENTORY.name}_{c.AGENT.value}#{agent_i}' else: return f'{c.INVENTORY.name}_{agent_i}' class Item(MoveableEntity): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._auto_despawn = -1 @property def auto_despawn(self): return self._auto_despawn @property def can_collide(self): return False @property def encoding(self): # Edit this if you want items to be drawn in the ops differently return 1 def set_auto_despawn(self, auto_despawn): self._auto_despawn = auto_despawn class ItemRegister(MovingEntityObjectRegister): def as_array(self): self._array[:] = c.FREE_CELL.value for item in self: if item.pos != c.NO_POS.value: self._array[0, item.x, item.y] = item.encoding return self._array _accepted_objects = Item def spawn_items(self, tiles: List[Tile]): items = [Item(tile) for tile in tiles] self.register_additional_items(items) def despawn_items(self, items: List[Item]): items = [items] if isinstance(items, Item) else items for item in items: del self[item] class Inventory(UserList): @property def is_blocking_light(self): return False @property def name(self): return self.agent.name def __init__(self, pomdp_r: int, level_shape: (int, int), agent: Agent, capacity: int): super(Inventory, self).__init__() self.agent = agent self.pomdp_r = pomdp_r self._level_shape = level_shape if self.pomdp_r: self._array = np.zeros((1, pomdp_r * 2 + 1, pomdp_r * 2 + 1)) else: self._array = np.zeros((1, *self._level_shape)) self.capacity = min(capacity, self._array.size) def as_array(self): self._array[:] = c.FREE_CELL.value for item_idx, item in enumerate(self): x_diff, y_diff = divmod(item_idx, self._array.shape[1]) self._array[0, int(x_diff), int(y_diff)] = item.encoding return self._array def __repr__(self): return f'{self.__class__.__name__}[{self.agent.name}]({self.data})' def append(self, item) -> None: if len(self) < self.capacity: super(Inventory, self).append(item) else: raise RuntimeError('Inventory is full') def belongs_to_entity(self, entity): return self.agent == entity def summarize_state(self): return {val.name: val.summarize_state() for val in self} class Inventories(ObjectRegister): _accepted_objects = Inventory is_blocking_light = False can_be_shadowed = False hide_from_obs_builder = True def __init__(self, *args, pomdp_r=0, **kwargs): super(Inventories, self).__init__(*args, is_per_agent=True, individual_slices=True, **kwargs) self.is_observable = True def as_array(self): # self._array[:] = c.FREE_CELL.value for inv_idx, inventory in enumerate(self): self._array[inv_idx] = inventory.as_array() return self._array def spawn_inventories(self, agents, pomdp_r, capacity): inventories = [self._accepted_objects(pomdp_r, self._level_shape, agent, capacity) for _, agent in enumerate(agents)] self.register_additional_items(inventories) def idx_by_entity(self, entity): try: return next((idx for idx, inv in enumerate(self) if inv.belongs_to_entity(entity))) except StopIteration: return None def by_entity(self, entity): try: return next((inv for inv in self if inv.belongs_to_entity(entity))) except StopIteration: return None class DropOffLocation(Entity): @property def can_collide(self): return False @property def encoding(self): return ITEM_DROP_OFF def __init__(self, *args, storage_size_until_full: int = 5, auto_item_despawn_interval: int = 5, **kwargs): super(DropOffLocation, self).__init__(*args, **kwargs) self.auto_item_despawn_interval = auto_item_despawn_interval self.storage = deque(maxlen=storage_size_until_full or None) def place_item(self, item: Item): if self.is_full: raise RuntimeWarning("There is currently no way to clear the storage or make it unfull.") return c.NOT_VALID else: self.storage.append(item) item.set_auto_despawn(self.auto_item_despawn_interval) return c.VALID @property def is_full(self): return False if not self.storage.maxlen else self.storage.maxlen == len(self.storage) class DropOffLocations(EntityObjectRegister): _accepted_objects = DropOffLocation def as_array(self): self._array[:] = c.FREE_CELL.value for item in self: if item.pos != c.NO_POS.value: self._array[0, item.x, item.y] = item.encoding return self._array class ItemProperties(NamedTuple): n_items: int = 5 # How many items are there at the same time spawn_frequency: int = 10 # Spawn Frequency in Steps n_drop_off_locations: int = 5 # How many DropOff locations are there at the same time max_dropoff_storage_size: int = 0 # How many items are needed until the drop off is full max_agent_inventory_capacity: int = 5 # How many items are needed until the agent inventory is full agent_can_interact: bool = True # Whether agents have the possibility to interact with the domain items # noinspection PyAttributeOutsideInit, PyAbstractClass class ItemFactory(BaseFactory): # noinspection PyMissingConstructor def __init__(self, *args, item_properties: ItemProperties = ItemProperties(), env_seed=time.time_ns(), **kwargs): if isinstance(item_properties, dict): item_properties = ItemProperties(**item_properties) self.item_properties = item_properties kwargs.update(env_seed=env_seed) self._item_rng = np.random.default_rng(env_seed) assert (item_properties.n_items <= ((1 + kwargs.get('pomdp_r', 0) * 2) ** 2)) or not kwargs.get('pomdp_r', 0) super().__init__(*args, **kwargs) @property def additional_actions(self) -> Union[Action, List[Action]]: # noinspection PyUnresolvedReferences super_actions = super().additional_actions super_actions.append(Action(enum_ident=h.EnvActions.ITEM_ACTION)) return super_actions @property def additional_entities(self) -> Dict[(Enum, Entities)]: # noinspection PyUnresolvedReferences super_entities = super().additional_entities empty_tiles = self[c.FLOOR].empty_tiles[:self.item_properties.n_drop_off_locations] drop_offs = DropOffLocations.from_tiles(empty_tiles, self._level_shape, storage_size_until_full=self.item_properties.max_dropoff_storage_size) item_register = ItemRegister(self._level_shape) empty_tiles = self[c.FLOOR].empty_tiles[:self.item_properties.n_items] item_register.spawn_items(empty_tiles) inventories = Inventories(self._level_shape if not self.pomdp_r else ((self.pomdp_diameter,) * 2)) inventories.spawn_inventories(self[c.AGENT], self.pomdp_r, self.item_properties.max_agent_inventory_capacity) super_entities.update({c.DROP_OFF: drop_offs, c.ITEM: item_register, c.INVENTORY: inventories}) return super_entities def additional_per_agent_obs_build(self, agent) -> List[np.ndarray]: additional_per_agent_obs_build = super().additional_per_agent_obs_build(agent) additional_per_agent_obs_build.append(self[c.INVENTORY].by_entity(agent).as_array()) return additional_per_agent_obs_build def do_item_action(self, agent: Agent): inventory = self[c.INVENTORY].by_entity(agent) if drop_off := self[c.DROP_OFF].by_pos(agent.pos): if inventory: valid = drop_off.place_item(inventory.pop(0)) return valid else: return c.NOT_VALID elif item := self[c.ITEM].by_pos(agent.pos): try: inventory.append(item) item.move(self._NO_POS_TILE) return c.VALID except RuntimeError: return c.NOT_VALID else: return c.NOT_VALID def do_additional_actions(self, agent: Agent, action: Action) -> Union[None, c]: # noinspection PyUnresolvedReferences valid = super().do_additional_actions(agent, action) if valid is None: if action == h.EnvActions.ITEM_ACTION: if self.item_properties.agent_can_interact: valid = self.do_item_action(agent) return valid else: return c.NOT_VALID else: return None else: return valid def do_additional_reset(self) -> None: # noinspection PyUnresolvedReferences super().do_additional_reset() self._next_item_spawn = self.item_properties.spawn_frequency self.trigger_item_spawn() def trigger_item_spawn(self): if item_to_spawns := max(0, (self.item_properties.n_items - len(self[c.ITEM]))): empty_tiles = self[c.FLOOR].empty_tiles[:item_to_spawns] self[c.ITEM].spawn_items(empty_tiles) self._next_item_spawn = self.item_properties.spawn_frequency self.print(f'{item_to_spawns} new items have been spawned; next spawn in {self._next_item_spawn}') else: self.print('No Items are spawning, limit is reached.') def do_additional_step(self) -> dict: # noinspection PyUnresolvedReferences info_dict = super().do_additional_step() for item in list(self[c.ITEM].values()): if item.auto_despawn >= 1: item.set_auto_despawn(item.auto_despawn-1) elif not item.auto_despawn: self[c.ITEM].delete_item(item) else: pass if not self._next_item_spawn: self.trigger_item_spawn() else: self._next_item_spawn = max(0, self._next_item_spawn-1) return info_dict def calculate_additional_reward(self, agent: Agent) -> (int, dict): # noinspection PyUnresolvedReferences reward, info_dict = super().calculate_additional_reward(agent) if h.EnvActions.ITEM_ACTION == agent.temp_action: if agent.temp_valid: if drop_off := self[c.DROP_OFF].by_pos(agent.pos): info_dict.update({f'{agent.name}_item_dropoff': 1}) self.print(f'{agent.name} just dropped of an item at {drop_off.pos}.') reward += 0.5 else: info_dict.update({f'{agent.name}_item_pickup': 1}) self.print(f'{agent.name} just picked up an item at {agent.pos}') reward += 0.1 else: info_dict.update({f'{agent.name}_failed_item_action': 1}) self.print(f'{agent.name} just tried to pick up an item at {agent.pos}, but failed.') reward -= 0.1 return reward, info_dict def render_additional_assets(self, mode='human'): # noinspection PyUnresolvedReferences additional_assets = super().render_additional_assets() items = [RenderEntity(c.ITEM.value, item.tile.pos) for item in self[c.ITEM]] additional_assets.extend(items) drop_offs = [RenderEntity(c.DROP_OFF.value, drop_off.tile.pos) for drop_off in self[c.DROP_OFF]] additional_assets.extend(drop_offs) return additional_assets if __name__ == '__main__': import random render = False item_props = ItemProperties() factory = ItemFactory(item_properties=item_props, n_agents=3, done_at_collision=False, frames_to_stack=0, level_name='rooms', max_steps=4000, omit_agent_in_obs=True, parse_doors=True, pomdp_r=3, record_episodes=False, verbose=False ) # noinspection DuplicatedCode n_actions = factory.action_space.n - 1 _ = factory.observation_space for epoch in range(100): random_actions = [[random.randint(0, n_actions) for _ in range(factory.n_agents)] for _ in range(200)] env_state = factory.reset() rew = 0 for agent_i_action in random_actions: env_state, step_r, done_bool, info_obj = factory.step(agent_i_action) rew += step_r if render: factory.render() if done_bool: break print(f'Factory run {epoch} done, reward is:\n {rew}')