import time from typing import List, Union, NamedTuple import random import numpy as np from environments.helpers import Constants as c from environments import helpers as h from environments.factory.base.base_factory import BaseFactory from environments.factory.base.objects import Agent, Action, Object, Slice from environments.factory.base.registers import Entities from environments.factory.renderer import Renderer, Entity from environments.utility_classes import MovementProperties DIRT = "dirt" CLEAN_UP_ACTION = 'clean_up' class DirtProperties(NamedTuple): clean_amount: int = 1 # How much does the robot clean with one actions. max_spawn_ratio: float = 0.2 # On max how much tiles does the dirt spawn in percent. gain_amount: float = 0.3 # How much dirt does spawn per tile spawn_frequency: int = 5 # Spawn Frequency in Steps max_local_amount: int = 2 # Max dirt amount per tile. max_global_amount: int = 20 # Max dirt amount in the whole environment. dirt_smear_amount: float = 0.2 # Agents smear dirt, when not cleaning up in place def softmax(x): """Compute softmax values for each sets of scores in x.""" e_x = np.exp(x - np.max(x)) return e_x / e_x.sum() def entropy(x): return -(x * np.log(x + 1e-8)).sum() # noinspection PyAttributeOutsideInit class SimpleFactory(BaseFactory): @property def additional_actions(self) -> List[Object]: return [Action(CLEAN_UP_ACTION)] @property def additional_entities(self) -> Union[Entities, List[Entities]]: return [] @property def additional_slices(self) -> List[Slice]: return [Slice('dirt', np.zeros(self._level_shape))] def _is_clean_up_action(self, action: Union[str, int]): if isinstance(action, str): action = self._actions.by_name(action) return self._actions[action].name == CLEAN_UP_ACTION def __init__(self, *args, dirt_properties: DirtProperties = DirtProperties(), **kwargs): self.dirt_properties = dirt_properties self._renderer = None # expensive - don't use it when not required ! self._dirt_rng = np.random.default_rng(kwargs.get('seed', default=time.time_ns())) super(SimpleFactory, self).__init__(*args, **kwargs) def _flush_state(self): super(SimpleFactory, self)._flush_state() self._obs_cube[self._slices.get_idx_by_name(DIRT)] = self._slices.by_name(DIRT).slice def render(self, mode='human'): if not self._renderer: # lazy init height, width = self._obs_cube.shape[1:] self._renderer = Renderer(width, height, view_radius=self.pomdp_r, fps=5) dirt_slice = self._slices.by_name(DIRT).slice dirt = [Entity('dirt', tile.pos, min(0.15 + dirt_slice[tile.pos], 1.5), 'scale') for tile in [tile for tile in self._tiles if dirt_slice[tile.pos]]] walls = [Entity('wall', pos) for pos in np.argwhere(self._slices.by_enum(c.LEVEL).slice == c.OCCUPIED_CELL.value)] def asset_str(agent): # What does this abonimation do? # if any([x is None for x in [self._slices[j] for j in agent.collisions]]): # print('error') col_names = [x.name for x in agent.temp_collisions] if c.AGENT.value in col_names: return 'agent_collision', 'blank' elif not agent.temp_valid or c.LEVEL.name in col_names or c.AGENT.name in col_names: return c.AGENT.value, 'invalid' elif self._is_clean_up_action(agent.temp_action): return c.AGENT.value, 'valid' else: return c.AGENT.value, 'idle' agents = [] for i, agent in enumerate(self._agents): name, state = asset_str(agent) agents.append(Entity(name, agent.pos, 1, 'none', state, i+1, agent.temp_light_map)) doors = [] if self.parse_doors: for i, door in enumerate(self._doors): name, state = 'door_open' if door.is_open else 'door_closed', 'blank' agents.append(Entity(name, door.pos, 1, 'none', state, i+1)) self._renderer.render(dirt+walls+agents+doors) def spawn_dirt(self) -> None: dirt_slice = self._slices.by_name(DIRT).slice # dirty_tiles = [tile for tile in self._tiles if dirt_slice[tile.pos]] curr_dirt_amount = dirt_slice.sum() if not curr_dirt_amount > self.dirt_properties.max_global_amount: free_for_dirt = self._tiles.empty_tiles # randomly distribute dirt across the grid new_spawn = self._dirt_rng.uniform(0, self.dirt_properties.max_spawn_ratio) n_dirt_tiles = max(0, int(new_spawn * len(free_for_dirt))) for tile in free_for_dirt[:n_dirt_tiles]: new_value = dirt_slice[tile.pos] + self.dirt_properties.gain_amount dirt_slice[tile.pos] = min(new_value, self.dirt_properties.max_local_amount) else: pass def clean_up(self, agent: Agent) -> bool: dirt_slice = self._slices.by_name(DIRT).slice if old_dirt_amount := dirt_slice[agent.pos]: new_dirt_amount = old_dirt_amount - self.dirt_properties.clean_amount dirt_slice[agent.pos] = max(new_dirt_amount, c.FREE_CELL.value) return True else: return False def do_additional_step(self) -> dict: if smear_amount := self.dirt_properties.dirt_smear_amount: dirt_slice = self._slices.by_name(DIRT).slice for agent in self._agents: if agent.temp_valid and agent.last_pos != h.NO_POS: if dirt := dirt_slice[agent.last_pos]: if smeared_dirt := round(dirt * smear_amount, 2): dirt_slice[agent.last_pos] = max(0, dirt_slice[agent.last_pos]-smeared_dirt) dirt_slice[agent.pos] = min((self.dirt_properties.max_local_amount, dirt_slice[agent.pos] + smeared_dirt) ) if not self._next_dirt_spawn: self.spawn_dirt() self._next_dirt_spawn = self.dirt_properties.spawn_frequency else: self._next_dirt_spawn -= 1 return {} def do_additional_actions(self, agent: Agent, action: int) -> bool: if self._is_clean_up_action(action): valid = self.clean_up(agent) return valid else: return c.NOT_VALID.value def do_additional_reset(self) -> None: self.spawn_dirt() self._next_dirt_spawn = self.dirt_properties.spawn_frequency def calculate_reward(self) -> (int, dict): info_dict = dict() dirt_slice = self._slices.by_name(DIRT).slice dirty_tiles = [dirt_slice[tile.pos] for tile in self._tiles if dirt_slice[tile.pos]] current_dirt_amount = sum(dirty_tiles) dirty_tile_count = len(dirty_tiles) if dirty_tile_count: dirt_distribution_score = entropy(softmax(dirt_slice)) / dirty_tile_count else: dirt_distribution_score = 0 info_dict.update(dirt_amount=current_dirt_amount) info_dict.update(dirty_tile_count=dirty_tile_count) info_dict.update(dirt_distribution_score=dirt_distribution_score) try: # penalty = current_dirt_amount reward = 0 except (ZeroDivisionError, RuntimeWarning): reward = 0 for agent in self._agents: if agent.temp_collisions: self.print(f't = {self._steps}\t{agent.name} has collisions with {agent.temp_collisions}') if self._is_clean_up_action(agent.temp_action): if agent.temp_valid: reward += 0.5 self.print(f'{agent.name} did just clean up some dirt at {agent.pos}.') info_dict.update(dirt_cleaned=1) else: reward -= 0.01 self.print(f'{agent.name} just tried to clean up some dirt at {agent.pos}, but failed.') info_dict.update({f'{agent.name}_failed_action': 1}) info_dict.update({f'{agent.name}_failed_action': 1}) info_dict.update({f'{agent.name}_failed_dirt_cleanup': 1}) elif self._actions.is_moving_action(agent.temp_action): if agent.temp_valid: # info_dict.update(movement=1) reward -= 0.00 else: # self.print('collision') reward -= 0.01 self.print(f'{agent.name} just hit the wall at {agent.pos}.') info_dict.update({f'{agent.name}_vs_LEVEL': 1}) elif self._actions.is_door_usage(agent.temp_action): if agent.temp_valid: self.print(f'{agent.name} did just use the door at {agent.pos}.') info_dict.update(door_used=1) else: reward -= 0.01 self.print(f'{agent.name} just tried to use a door at {agent.pos}, but failed.') info_dict.update({f'{agent.name}_failed_action': 1}) info_dict.update({f'{agent.name}_failed_door_open': 1}) else: info_dict.update(no_op=1) reward -= 0.00 for other_agent in agent.temp_collisions: info_dict.update({f'{agent.name}_vs_{other_agent.name}': 1}) self.print(f"reward is {reward}") # Potential based rewards -> # track the last reward , minus the current reward = potential return reward, info_dict if __name__ == '__main__': render = True dirt_props = DirtProperties(1, 0.05, 0.1, 3, 1, 20, 0.0) move_props = MovementProperties(True, True, False) factory = SimpleFactory(n_agents=1, done_at_collision=False, frames_to_stack=0, level_name='rooms', max_steps=400, omit_agent_slice_in_obs=True, parse_doors=True, pomdp_r=3, record_episodes=False, verbose=False ) 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() r = 0 for agent_i_action in random_actions: env_state, step_r, done_bool, info_obj = factory.step(agent_i_action) r += step_r if render: factory.render() if done_bool: break print(f'Factory run {epoch} done, reward is:\n {r}')