144 lines
4.2 KiB
Python
144 lines
4.2 KiB
Python
import numpy as np
|
|
from enum import Enum, auto
|
|
import matplotlib.pyplot as plt
|
|
|
|
class Direction(Enum):
|
|
NORTH = (0, 1)
|
|
SOUTH = (0, -1)
|
|
WEST = (-1, 0)
|
|
EAST = (1, 0)
|
|
|
|
class CellState(Enum):
|
|
OPEN = auto()
|
|
WALL = auto()
|
|
|
|
class Colours(Enum):
|
|
WALL = [30,30,30]
|
|
STOCKED = [200,100,0]
|
|
DEPLETED = [0,50,50]
|
|
UNDEFINED = [255,0,255]
|
|
AGENT = [0,255,0]
|
|
|
|
class Cell():
|
|
def __init__(self, state, resources):
|
|
self.state = state
|
|
self.resources = resources
|
|
|
|
def get_colour(self):
|
|
if self.state is CellState.WALL:
|
|
return Colours.WALL
|
|
elif self.state is CellState.OPEN:
|
|
if self.resources > 0:
|
|
return Colours.STOCKED
|
|
elif self.resources == 0:
|
|
return Colours.DEPLETED
|
|
else:
|
|
return Colours.UNDEFINED
|
|
else:
|
|
return Colours.UNDEFINED
|
|
|
|
class Environment:
|
|
def __init__(self, shape):
|
|
self.shape = shape
|
|
self.agents = []
|
|
|
|
self.resource_map = np.round(
|
|
np.random.normal(loc=50, scale=10, size=shape)
|
|
)
|
|
|
|
cols, rows = shape
|
|
self.cells = [[] for _ in range(rows)]
|
|
for y in range(rows):
|
|
for x in range(cols):
|
|
if (x == 0 or y == 0 or x+1==cols or y+1==rows):
|
|
self.cells[y].append(Cell(CellState.WALL, 0))
|
|
else:
|
|
self.cells[y].append(Cell(CellState.OPEN, self.resource_map[x,y]))
|
|
|
|
def step(self):
|
|
for agent in self.agents:
|
|
agent.step()
|
|
|
|
def render(self):
|
|
cols, rows = self.shape
|
|
pixel_data = [[] for _ in range(rows)]
|
|
|
|
for y in range(self.shape[1]):
|
|
for x in range(self.shape[0]):
|
|
#print((x, y, self.cell((x,y)).state, self.cell((x,y)).get_colour().value))
|
|
pixel_data[y].append(self.cell((x,y)).get_colour().value)
|
|
|
|
for agent in self.agents:
|
|
x, y = agent.position
|
|
pixel_data[y][x] = Colours.AGENT.value
|
|
|
|
return np.array(pixel_data)
|
|
|
|
def cell(self, position):
|
|
x, y = tuple(position)
|
|
return self.cells[y][x]
|
|
|
|
def is_wall(self, position):
|
|
x, y = tuple(position)
|
|
return self.cells[y][x].state == CellState.WALL
|
|
|
|
def register_agent(self, agent):
|
|
self.agents.append(agent)
|
|
|
|
def unregister_agent(self, agent):
|
|
self.agents.remove(agent)
|
|
|
|
class Agent:
|
|
def __init__(self, environment, position):
|
|
self.position = np.array(position)
|
|
self.environment = environment
|
|
self.environment.register_agent(self)
|
|
|
|
def step(self):
|
|
pass
|
|
|
|
def move(self, direction, respect_walls=True):
|
|
new_position = self.position + direction.value
|
|
if respect_walls and self.environment.is_wall(new_position):
|
|
return False
|
|
else:
|
|
self.position = new_position
|
|
return True
|
|
|
|
def die(self):
|
|
self.environment.unregister_agent(self)
|
|
|
|
class DirectionalAgent(Agent):
|
|
def __init__(self, environment, position, initial_direction):
|
|
super().__init__(environment, position)
|
|
self.direction = initial_direction
|
|
|
|
def get_cell_in_front(self):
|
|
position_in_front = self.position + self.direction.value
|
|
return self.environment.cell(position_in_front)
|
|
|
|
def wall_in_front(self):
|
|
position_in_front = self.position + self.direction.value
|
|
return self.environment.is_wall(position_in_front)
|
|
|
|
def turn_left(self):
|
|
raise NotImplementedError
|
|
|
|
def turn_right(self):
|
|
raise NotImplementedError
|
|
|
|
def reverse_direction(self):
|
|
new_vector = np.array(self.direction.value) * -1
|
|
self.direction = Direction(tuple(new_vector))
|
|
|
|
class CAMSReverseAndSidestepAgent(DirectionalAgent):
|
|
def __init__(self, environment, position, initial_direction, required_resources = 100):
|
|
super().__init__(environment, position, initial_direction)
|
|
self.resources = 0
|
|
self.required_resources = required_resources
|
|
|
|
def step(self):
|
|
if self.wall_in_front():
|
|
self.reverse_direction()
|
|
self.move(self.direction)
|