522 lines
20 KiB
Python
522 lines
20 KiB
Python
# Copyright (c) 2020, salesforce.com, inc.
|
|
# All rights reserved.
|
|
# SPDX-License-Identifier: BSD-3-Clause
|
|
# For full license text, see the LICENSE file in the repo root
|
|
# or https://opensource.org/licenses/BSD-3-Clause
|
|
|
|
import numpy as np
|
|
|
|
from ai_economist.foundation.agents import agent_registry
|
|
from ai_economist.foundation.entities import landmark_registry, resource_registry
|
|
|
|
|
|
class Maps:
|
|
"""Manages the spatial configuration of the world as a set of entity maps.
|
|
|
|
A maps object is built during world construction, which is a part of environment
|
|
construction. The maps object is accessible through the world object. The maps
|
|
object maintains a map state for each of the spatial entities that are involved
|
|
in the constructed environment (which are determined by the "required_entities"
|
|
attributes of the Scenario and Component classes used to build the environment).
|
|
|
|
The Maps class also implements some of the basic spatial logic of the game,
|
|
such as which locations agents can occupy based on other agent locations and
|
|
locations of various landmarks.
|
|
|
|
Args:
|
|
size (list): A length-2 list specifying the dimensions of the 2D world.
|
|
Interpreted as [height, width].
|
|
n_agents (int): The number of mobile agents (does not include planner).
|
|
world_resources (list): The resources registered during environment
|
|
construction.
|
|
world_landmarks (list): The landmarks registered during environment
|
|
construction.
|
|
"""
|
|
|
|
def __init__(self, size, n_agents, world_resources, world_landmarks):
|
|
self.size = size
|
|
self.sz_h, self.sz_w = size
|
|
|
|
self.n_agents = n_agents
|
|
|
|
self.resources = world_resources
|
|
self.landmarks = world_landmarks
|
|
self.entities = world_resources + world_landmarks
|
|
|
|
self._maps = {} # All maps
|
|
self._blocked = [] # Solid objects that no agent can move through
|
|
self._private = [] # Solid objects that only permit movement for parent agents
|
|
self._public = [] # Non-solid objects that agents can move on top of
|
|
self._resources = [] # Non-solid objects that can be collected
|
|
|
|
self._private_landmark_types = []
|
|
self._resource_source_blocks = []
|
|
|
|
self._map_keys = []
|
|
|
|
self._accessibility_lookup = {}
|
|
|
|
for resource in self.resources:
|
|
resource_cls = resource_registry.get(resource)
|
|
if resource_cls.collectible:
|
|
self._maps[resource] = np.zeros(shape=self.size)
|
|
self._resources.append(resource)
|
|
self._map_keys.append(resource)
|
|
|
|
self.landmarks.append("{}SourceBlock".format(resource))
|
|
|
|
for landmark in self.landmarks:
|
|
dummy_landmark = landmark_registry.get(landmark)()
|
|
|
|
if dummy_landmark.public:
|
|
self._maps[landmark] = np.zeros(shape=self.size)
|
|
self._public.append(landmark)
|
|
self._map_keys.append(landmark)
|
|
|
|
elif dummy_landmark.blocking:
|
|
self._maps[landmark] = np.zeros(shape=self.size)
|
|
self._blocked.append(landmark)
|
|
self._map_keys.append(landmark)
|
|
self._accessibility_lookup[landmark] = len(self._accessibility_lookup)
|
|
|
|
elif dummy_landmark.private:
|
|
self._private_landmark_types.append(landmark)
|
|
self._maps[landmark] = dict(
|
|
owner=-np.ones(shape=self.size, dtype=np.int16),
|
|
health=np.zeros(shape=self.size),
|
|
)
|
|
self._private.append(landmark)
|
|
self._map_keys.append(landmark)
|
|
self._accessibility_lookup[landmark] = len(self._accessibility_lookup)
|
|
|
|
else:
|
|
raise NotImplementedError
|
|
self.reset_agent_maps(n_agents)
|
|
|
|
def reset_agent_maps(self,n_agents):
|
|
self.n_agents=n_agents
|
|
self._idx_map = np.stack(
|
|
[i * np.ones(shape=self.size) for i in range(self.n_agents)]
|
|
)
|
|
self._idx_array = np.arange(self.n_agents)
|
|
if self._accessibility_lookup:
|
|
self._accessibility = np.ones(
|
|
shape=[len(self._accessibility_lookup), self.n_agents] + self.size,
|
|
dtype=bool,
|
|
)
|
|
self._net_accessibility = None
|
|
else:
|
|
self._accessibility = None
|
|
self._net_accessibility = np.ones(
|
|
shape=[self.n_agents] + self.size, dtype=bool
|
|
)
|
|
|
|
self._agent_locs = [None for _ in range(self.n_agents)]
|
|
self._unoccupied = np.ones(self.size, dtype=bool)
|
|
|
|
def clear(self, entity_name=None):
|
|
"""Clear resource and landmark maps."""
|
|
if entity_name is not None:
|
|
assert entity_name in self._maps
|
|
if entity_name in self._private_landmark_types:
|
|
self._maps[entity_name] = dict(
|
|
owner=-np.ones(shape=self.size, dtype=np.int16),
|
|
health=np.zeros(shape=self.size),
|
|
)
|
|
else:
|
|
self._maps[entity_name] *= 0
|
|
|
|
else:
|
|
for name in self.keys():
|
|
self.clear(entity_name=name)
|
|
|
|
if self._accessibility is not None:
|
|
self._accessibility = np.ones_like(self._accessibility)
|
|
self._net_accessibility = None
|
|
|
|
def clear_agent_loc(self, agent=None):
|
|
"""Remove agents or agent from the world map."""
|
|
# Clear all agent locations
|
|
if agent is None:
|
|
self._agent_locs = [None for _ in range(self.n_agents)]
|
|
self._unoccupied[:, :] = 1
|
|
|
|
# Clear the location of the provided agent
|
|
else:
|
|
i = agent.idx
|
|
if self._agent_locs[i] is None:
|
|
return
|
|
r, c = self._agent_locs[i]
|
|
self._unoccupied[r, c] = 1
|
|
self._agent_locs[i] = None
|
|
|
|
def set_agent_loc(self, agent, r, c):
|
|
"""Set the location of agent to [r, c].
|
|
|
|
Note:
|
|
Things might break if you set the agent's location to somewhere it
|
|
cannot access. Don't do that.
|
|
"""
|
|
assert (0 <= r < self.size[0]) and (0 <= c < self.size[1])
|
|
i = agent.idx
|
|
# If the agent is currently on the board...
|
|
if self._agent_locs[i] is not None:
|
|
curr_r, curr_c = self._agent_locs[i]
|
|
# If the agent isn't actually moving, just return
|
|
if (curr_r, curr_c) == (r, c):
|
|
return
|
|
# Make the location the agent is currently at as unoccupied
|
|
# (since the agent is going to move)
|
|
self._unoccupied[curr_r, curr_c] = 1
|
|
|
|
# Set the agent location to the specified coordinates
|
|
# and update the occupation map
|
|
agent.state["loc"] = [r, c]
|
|
self._agent_locs[i] = [r, c]
|
|
self._unoccupied[r, c] = 0
|
|
|
|
def keys(self):
|
|
"""Return an iterable over map keys."""
|
|
return self._maps.keys()
|
|
|
|
def values(self):
|
|
"""Return an iterable over map values."""
|
|
return self._maps.values()
|
|
|
|
def items(self):
|
|
"""Return an iterable over map (key, value) pairs."""
|
|
return self._maps.items()
|
|
|
|
def get(self, entity_name, owner=False):
|
|
"""Return the map or ownership for entity_name."""
|
|
assert entity_name in self._maps
|
|
if entity_name in self._private_landmark_types:
|
|
sub_key = "owner" if owner else "health"
|
|
return self._maps[entity_name][sub_key]
|
|
return self._maps[entity_name]
|
|
|
|
def set(self, entity_name, map_state):
|
|
"""Set the map for entity_name."""
|
|
if entity_name in self._private_landmark_types:
|
|
assert "owner" in map_state
|
|
assert self.get(entity_name, owner=True).shape == map_state["owner"].shape
|
|
assert "health" in map_state
|
|
assert self.get(entity_name, owner=False).shape == map_state["health"].shape
|
|
|
|
h = np.maximum(0.0, map_state["health"])
|
|
o = map_state["owner"].astype(np.int16)
|
|
|
|
o[h <= 0] = -1
|
|
tmp = o[h > 0]
|
|
if len(tmp) > 0:
|
|
assert np.min(tmp) >= 0
|
|
|
|
self._maps[entity_name] = dict(owner=o, health=h)
|
|
|
|
owned_by_agent = o[None] == self._idx_map
|
|
owned_by_none = o[None] == -1
|
|
self._accessibility[
|
|
self._accessibility_lookup[entity_name]
|
|
] = np.logical_or(owned_by_agent, owned_by_none)
|
|
self._net_accessibility = None
|
|
|
|
else:
|
|
assert self.get(entity_name).shape == map_state.shape
|
|
self._maps[entity_name] = np.maximum(0, map_state)
|
|
|
|
if entity_name in self._blocked:
|
|
self._accessibility[
|
|
self._accessibility_lookup[entity_name]
|
|
] = np.repeat(map_state[None] == 0, self.n_agents, axis=0)
|
|
self._net_accessibility = None
|
|
|
|
def set_add(self, entity_name, map_state):
|
|
"""Add map_state to the existing map for entity_name."""
|
|
assert entity_name not in self._private_landmark_types
|
|
self.set(entity_name, self.get(entity_name) + map_state)
|
|
|
|
def get_point(self, entity_name, r, c, **kwargs):
|
|
"""Return the entity state at the specified coordinates."""
|
|
point_map = self.get(entity_name, **kwargs)
|
|
return point_map[r, c]
|
|
|
|
def set_point(self, entity_name, r, c, val, owner=None):
|
|
"""Set the entity state at the specified coordinates."""
|
|
if entity_name in self._private_landmark_types:
|
|
assert owner is not None
|
|
h = self._maps[entity_name]["health"]
|
|
o = self._maps[entity_name]["owner"]
|
|
assert o[r, c] == -1 or o[r, c] == int(owner)
|
|
h[r, c] = np.maximum(0, val)
|
|
if h[r, c] == 0:
|
|
o[r, c] = -1
|
|
else:
|
|
o[r, c] = int(owner)
|
|
|
|
self._maps[entity_name]["owner"] = o
|
|
self._maps[entity_name]["health"] = h
|
|
|
|
self._accessibility[
|
|
self._accessibility_lookup[entity_name], :, r, c
|
|
] = np.logical_or(o[r, c] == self._idx_array, o[r, c] == -1).astype(bool)
|
|
self._net_accessibility = None
|
|
|
|
else:
|
|
self._maps[entity_name][r, c] = np.maximum(0, val)
|
|
|
|
if entity_name in self._blocked:
|
|
self._accessibility[
|
|
self._accessibility_lookup[entity_name]
|
|
] = np.repeat(np.array([val]) == 0, self.n_agents, axis=0)
|
|
self._net_accessibility = None
|
|
|
|
def set_point_add(self, entity_name, r, c, value, **kwargs):
|
|
"""Add value to the existing entity state at the specified coordinates."""
|
|
self.set_point(
|
|
entity_name,
|
|
r,
|
|
c,
|
|
value + self.get_point(entity_name, r, c, **kwargs),
|
|
**kwargs
|
|
)
|
|
|
|
def is_accessible(self, r, c, agent_id):
|
|
"""Return True if agent with id agent_id can occupy the location [r, c]."""
|
|
return bool(self.accessibility[agent_id, r, c])
|
|
|
|
def location_resources(self, r, c):
|
|
"""Return {resource: health} dictionary for any resources at location [r, c]."""
|
|
return {
|
|
k: self._maps[k][r, c] for k in self._resources if self._maps[k][r, c] > 0
|
|
}
|
|
|
|
def location_landmarks(self, r, c):
|
|
"""Return {landmark: health} dictionary for any landmarks at location [r, c]."""
|
|
tmp = {k: self.get_point(k, r, c) for k in self.keys()}
|
|
return {k: v for k, v in tmp.items() if k not in self._resources and v > 0}
|
|
|
|
@property
|
|
def unoccupied(self):
|
|
"""Return a boolean map indicating which locations are unoccupied."""
|
|
return self._unoccupied
|
|
|
|
@property
|
|
def accessibility(self):
|
|
"""Return a boolean map indicating which locations are accessible."""
|
|
if self._net_accessibility is None:
|
|
self._net_accessibility = self._accessibility.prod(axis=0).astype(bool)
|
|
return self._net_accessibility
|
|
|
|
@property
|
|
def empty(self):
|
|
"""Return a boolean map indicating which locations are empty.
|
|
|
|
Empty locations have no landmarks or resources."""
|
|
return self.state.sum(axis=0) == 0
|
|
|
|
@property
|
|
def state(self):
|
|
"""Return the concatenated maps of landmark and resources."""
|
|
return np.stack([self.get(k) for k in self.keys()]).astype(np.float32)
|
|
|
|
@property
|
|
def owner_state(self):
|
|
"""Return the concatenated ownership maps of private landmarks."""
|
|
return np.stack(
|
|
[self.get(k, owner=True) for k in self._private_landmark_types]
|
|
).astype(np.int16)
|
|
|
|
@property
|
|
def state_dict(self):
|
|
"""Return a dictionary of the map states."""
|
|
return self._maps
|
|
|
|
|
|
class World:
|
|
"""Manages the environment's spatial- and agent-states.
|
|
|
|
The world object represents the state of the environment, minus whatever state
|
|
information is implicitly maintained by separate components. The world object
|
|
maintains the spatial state through an instance of the Maps class. Agent states
|
|
are maintained through instances of Agent classes (subclasses of BaseAgent),
|
|
with one such instance for each of the agents in the environment.
|
|
|
|
The world object is built during the environment construction, after the
|
|
required entities have been registered. As part of the world object construction,
|
|
it instantiates a map object and the agent objects.
|
|
|
|
The World class adds some functionality for interfacing with the spatial state
|
|
(the maps object) and setting/resetting agent locations. But its function is
|
|
mostly to wrap the stateful, non-component environment objects.
|
|
|
|
Args:
|
|
world_size (list): A length-2 list specifying the dimensions of the 2D world.
|
|
Interpreted as [height, width].
|
|
n_agents (int): The number of total agents (does not include planner).
|
|
agent_composition(dict): Dict of Agent Class names and amount
|
|
world_resources (list): The resources registered during environment
|
|
construction.
|
|
world_landmarks (list): The landmarks registered during environment
|
|
construction.
|
|
multi_action_mode_agents (bool): Whether "mobile" agents use multi action mode
|
|
(see BaseEnvironment in base_env.py).
|
|
multi_action_mode_planner (bool): Whether the planner agent uses multi action
|
|
mode (see BaseEnvironment in base_env.py).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
world_size,
|
|
agent_composition,
|
|
world_resources,
|
|
world_landmarks,
|
|
multi_action_mode_agents,
|
|
multi_action_mode_planner,
|
|
):
|
|
self.world_size = world_size
|
|
|
|
self.resources = world_resources
|
|
self.landmarks = world_landmarks
|
|
self.multi_action_mode_agents = bool(multi_action_mode_agents)
|
|
self.multi_action_mode_planner = bool(multi_action_mode_planner)
|
|
self._agent_class_idx_map={}
|
|
#create agents
|
|
self.create_agents(agent_composition)
|
|
|
|
self.maps = Maps(world_size, self.n_agents, world_resources, world_landmarks)
|
|
|
|
planner_class = agent_registry.get("BasicPlanner")
|
|
self._planner = planner_class(multi_action_mode=self.multi_action_mode_planner)
|
|
|
|
self.timestep = 0
|
|
|
|
# CUDA-related attributes (for GPU simulations).
|
|
# These will be set via the env_wrapper, if required.
|
|
self.use_cuda = False
|
|
self.cuda_function_manager = None
|
|
self.cuda_data_manager = None
|
|
|
|
def create_agents(self, agent_composition):
|
|
"""create_agents creates the world agent db with the given compostition."""
|
|
self.agent_composition=agent_composition
|
|
self.n_agents=0
|
|
self._agents = []
|
|
for k,v in agent_composition.items():
|
|
self._agent_class_idx_map[k]=[]
|
|
for offset in range(v):
|
|
agent_class=agent_registry.get(k)
|
|
agent=agent_class(self.n_agents,self.multi_action_mode_agents)
|
|
self._agents.append(agent)
|
|
self._agent_class_idx_map[k].append(str(self.n_agents))
|
|
self.n_agents+=1
|
|
|
|
def apply_agent_db(self):
|
|
"""Applys current agent db into lookup maps inside world and map itself. Enables insertion of new agents into existing env."""
|
|
self.n_agents=len(self._agents)
|
|
self._agent_class_idx_map={}
|
|
self.maps.reset_agent_maps(self.n_agents) # reset map lookups
|
|
#create mapping dict
|
|
for idx in range(self.n_agents):
|
|
cls=self.get_agent_class(idx)
|
|
agent=self._agents[idx]
|
|
if cls in self._agent_class_idx_map:
|
|
self._agent_class_idx_map[cls].append(idx)
|
|
else:
|
|
self._agent_class_idx_map[cls]=[idx]
|
|
# apply agent locs db to maps
|
|
if "loc" in agent.state:
|
|
self.maps.set_agent_loc(agent,*agent.loc)
|
|
|
|
@property
|
|
def agents(self):
|
|
"""Return a list of the agent objects in the world (sorted by index)."""
|
|
return self._agents
|
|
|
|
@property
|
|
def planner(self):
|
|
"""Return the planner agent object."""
|
|
return self._planner
|
|
|
|
@property
|
|
def loc_map(self):
|
|
"""Return a map indicating the agent index occupying each location.
|
|
|
|
Locations with a value of -1 are not occupied by an agent.
|
|
"""
|
|
idx_map = -np.ones(shape=self.world_size, dtype=np.int16)
|
|
for agent in self.agents:
|
|
r, c = agent.loc
|
|
idx_map[r, c] = int(agent.idx)
|
|
return idx_map
|
|
|
|
def get_random_order_agents(self):
|
|
"""The agent list in a randomized order."""
|
|
agent_order = np.random.permutation(self.n_agents)
|
|
agents = self.agents
|
|
return [agents[i] for i in agent_order]
|
|
|
|
def get_agent_class(self,idx):
|
|
"""Return class name of agent"""
|
|
return self.agents[idx].name
|
|
|
|
def is_valid(self, r, c):
|
|
"""Return True if the coordinates [r, c] are within the game boundaries."""
|
|
return (0 <= r < self.world_size[0]) and (0 <= c < self.world_size[1])
|
|
|
|
def is_location_accessible(self, r, c, agent):
|
|
"""Return True if location [r, c] is accessible to agent."""
|
|
if not self.is_valid(r, c):
|
|
return False
|
|
return self.maps.is_accessible(r, c, agent.idx)
|
|
|
|
def can_agent_occupy(self, r, c, agent):
|
|
"""Return True if location [r, c] is accessible to agent and unoccupied."""
|
|
if not self.is_location_accessible(r, c, agent):
|
|
return False
|
|
if self.maps.unoccupied[r, c]:
|
|
return True
|
|
return False
|
|
|
|
def clear_agent_locs(self):
|
|
"""Take all agents off the board. Useful for resetting."""
|
|
for agent in self.agents:
|
|
agent.state["loc"] = [-1, -1]
|
|
self.maps.clear_agent_loc()
|
|
|
|
def agent_locs_are_valid(self):
|
|
"""Returns True if all agent locations comply with world semantics."""
|
|
return all(
|
|
self.is_location_accessible(*agent.loc, agent) for agent in self.agents
|
|
)
|
|
|
|
def set_agent_loc(self, agent, r, c):
|
|
"""Set the agent's location to coordinates [r, c] if possible.
|
|
|
|
If agent cannot occupy [r, c], do nothing."""
|
|
if self.can_agent_occupy(r, c, agent):
|
|
self.maps.set_agent_loc(agent, r, c)
|
|
return [int(coord) for coord in agent.loc]
|
|
|
|
def location_resources(self, r, c):
|
|
"""Return {resource: health} dictionary for any resources at location [r, c]."""
|
|
if not self.is_valid(r, c):
|
|
return {}
|
|
return self.maps.location_resources(r, c)
|
|
|
|
def location_landmarks(self, r, c):
|
|
"""Return {landmark: health} dictionary for any landmarks at location [r, c]."""
|
|
if not self.is_valid(r, c):
|
|
return {}
|
|
return self.maps.location_landmarks(r, c)
|
|
|
|
def create_landmark(self, landmark_name, r, c, agent_idx=None):
|
|
"""Place a landmark on the world map.
|
|
|
|
Place landmark of type landmark_name at the given coordinates, indicating
|
|
agent ownership if applicable."""
|
|
self.maps.set_point(landmark_name, r, c, 1, owner=agent_idx)
|
|
|
|
def consume_resource(self, resource_name, r, c):
|
|
"""Consume a unit of resource_name from location [r, c]."""
|
|
self.maps.set_point_add(resource_name, r, c, -1)
|