wehub-resource-sync
315 lines
11 KiB
Python
315 lines
11 KiB
Python
import gymnasium as gym
|
|
|
|
from ray.rllib.env.multi_agent_env import MultiAgentEnv
|
|
|
|
# Map representation: Always six rooms (as the name suggests) with doors in between.
|
|
MAPS = {
|
|
"small": [
|
|
"WWWWWWWWWWWWW",
|
|
"W W W W",
|
|
"W W W",
|
|
"W W W",
|
|
"W WWWW WWWW W",
|
|
"W W W W",
|
|
"W W W",
|
|
"W W GW",
|
|
"WWWWWWWWWWWWW",
|
|
],
|
|
"medium": [
|
|
"WWWWWWWWWWWWWWWWWWW",
|
|
"W W W W",
|
|
"W W W",
|
|
"W W W",
|
|
"W WWWWWWW WWWWWWW W",
|
|
"W W W W",
|
|
"W W W",
|
|
"W W GW",
|
|
"WWWWWWWWWWWWWWWWWWW",
|
|
],
|
|
"large": [
|
|
"WWWWWWWWWWWWWWWWWWWWWWWWW",
|
|
"W W W W",
|
|
"W W W W",
|
|
"W W W",
|
|
"W W W",
|
|
"W W W W",
|
|
"WW WWWWWWWWW WWWWWWWWWW W",
|
|
"W W W W",
|
|
"W W W",
|
|
"W W W W",
|
|
"W W W",
|
|
"W W W GW",
|
|
"WWWWWWWWWWWWWWWWWWWWWWWWW",
|
|
],
|
|
}
|
|
|
|
|
|
class SixRoomEnv(gym.Env):
|
|
"""A grid-world with six rooms (arranged as 2x3), which are connected by doors.
|
|
|
|
The agent starts in the upper left room and has to reach a designated goal state
|
|
in one of the rooms using primitive actions up, left, down, and right.
|
|
|
|
The agent receives a small penalty of -0.01 on each step and a reward of +10.0 when
|
|
reaching the goal state.
|
|
"""
|
|
|
|
def __init__(self, config=None):
|
|
super().__init__()
|
|
|
|
# User can provide a custom map or a recognized map name (small, medium, large).
|
|
self.map = config.get("custom_map", MAPS.get(config.get("map"), MAPS["small"]))
|
|
self.time_limit = config.get("time_limit", 50)
|
|
|
|
# Define observation space: Discrete, index fields.
|
|
self.observation_space = gym.spaces.Discrete(len(self.map) * len(self.map[0]))
|
|
# Primitive actions: up, down, left, right.
|
|
self.action_space = gym.spaces.Discrete(4)
|
|
|
|
# Initialize environment state.
|
|
self.reset()
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
self._agent_pos = (1, 1)
|
|
self._ts = 0
|
|
# Return high-level observation.
|
|
return self._agent_discrete_pos, {}
|
|
|
|
def step(self, action):
|
|
next_pos = _get_next_pos(action, self._agent_pos)
|
|
|
|
self._ts += 1
|
|
|
|
# Check if the move ends up in a wall. If so -> Ignore the move and stay
|
|
# where we are right now.
|
|
if self.map[next_pos[0]][next_pos[1]] != "W":
|
|
self._agent_pos = next_pos
|
|
|
|
# Check if the agent has reached the global goal state.
|
|
if self.map[self._agent_pos[0]][self._agent_pos[1]] == "G":
|
|
return self._agent_discrete_pos, 10.0, True, False, {}
|
|
|
|
# Small step penalty.
|
|
return self._agent_discrete_pos, -0.01, False, self._ts >= self.time_limit, {}
|
|
|
|
@property
|
|
def _agent_discrete_pos(self):
|
|
x = self._agent_pos[0]
|
|
y = self._agent_pos[1]
|
|
# discrete position = row idx * columns + col idx
|
|
return x * len(self.map[0]) + y
|
|
|
|
|
|
class HierarchicalSixRoomEnv(MultiAgentEnv):
|
|
def __init__(self, config=None):
|
|
super().__init__()
|
|
|
|
# User can provide a custom map or a recognized map name (small, medium, large).
|
|
self.map = config.get("custom_map", MAPS.get(config.get("map"), MAPS["small"]))
|
|
self.max_steps_low_level = config.get("max_steps_low_level", 15)
|
|
self.time_limit = config.get("time_limit", 50)
|
|
self.num_low_level_agents = config.get("num_low_level_agents", 3)
|
|
|
|
self.agents = self.possible_agents = ["high_level_agent"] + [
|
|
f"low_level_agent_{i}" for i in range(self.num_low_level_agents)
|
|
]
|
|
|
|
# Define basic observation space: Discrete, index fields.
|
|
observation_space = gym.spaces.Discrete(len(self.map) * len(self.map[0]))
|
|
# Low level agents always see where they are right now and what the target
|
|
# state should be.
|
|
low_level_observation_space = gym.spaces.Tuple(
|
|
(observation_space, observation_space)
|
|
)
|
|
# Primitive actions: up, down, left, right.
|
|
low_level_action_space = gym.spaces.Discrete(4)
|
|
|
|
self.observation_spaces = {"high_level_agent": observation_space}
|
|
self.observation_spaces.update(
|
|
{
|
|
f"low_level_agent_{i}": low_level_observation_space
|
|
for i in range(self.num_low_level_agents)
|
|
}
|
|
)
|
|
self.action_spaces = {
|
|
"high_level_agent": gym.spaces.Tuple(
|
|
(
|
|
# The new target observation.
|
|
observation_space,
|
|
# Low-level policy that should get us to the new target observation.
|
|
gym.spaces.Discrete(self.num_low_level_agents),
|
|
)
|
|
)
|
|
}
|
|
self.action_spaces.update(
|
|
{
|
|
f"low_level_agent_{i}": low_level_action_space
|
|
for i in range(self.num_low_level_agents)
|
|
}
|
|
)
|
|
|
|
# Initialize environment state.
|
|
self.reset()
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
self._agent_pos = (1, 1)
|
|
self._low_level_steps = 0
|
|
self._high_level_action = None
|
|
# Number of times the low-level agent reached the given target (by the high
|
|
# level agent).
|
|
self._num_targets_reached = 0
|
|
|
|
self._ts = 0
|
|
|
|
# Return high-level observation.
|
|
return {
|
|
"high_level_agent": self._agent_discrete_pos,
|
|
}, {}
|
|
|
|
def step(self, action_dict):
|
|
self._ts += 1
|
|
|
|
terminateds = {"__all__": self._ts >= self.time_limit}
|
|
truncateds = {"__all__": False}
|
|
|
|
# High-level agent acted: Set next goal and next low-level policy to use.
|
|
# Note that the agent does not move in this case and stays at its current
|
|
# location.
|
|
if "high_level_agent" in action_dict:
|
|
self._high_level_action = action_dict["high_level_agent"]
|
|
low_level_agent = f"low_level_agent_{self._high_level_action[1]}"
|
|
self._low_level_steps = 0
|
|
# Return next low-level observation for the now-active agent.
|
|
# We want this agent to act next.
|
|
return (
|
|
{
|
|
low_level_agent: (
|
|
self._agent_discrete_pos, # current
|
|
self._high_level_action[0], # target
|
|
)
|
|
},
|
|
# Penalty for a target state that's close to the current state.
|
|
{
|
|
"high_level_agent": (
|
|
self.eucl_dist(
|
|
self._agent_discrete_pos,
|
|
self._high_level_action[0],
|
|
self.map,
|
|
)
|
|
/ (len(self.map) ** 2 + len(self.map[0]) ** 2) ** 0.5
|
|
)
|
|
- 1.0,
|
|
},
|
|
terminateds,
|
|
truncateds,
|
|
{},
|
|
)
|
|
# Low-level agent made a move (primitive action).
|
|
else:
|
|
assert len(action_dict) == 1
|
|
|
|
# Increment low-level step counter.
|
|
self._low_level_steps += 1
|
|
|
|
target_discrete_pos, low_level_agent = self._high_level_action
|
|
low_level_agent = f"low_level_agent_{low_level_agent}"
|
|
next_pos = _get_next_pos(action_dict[low_level_agent], self._agent_pos)
|
|
|
|
# Check if the move ends up in a wall. If so -> Ignore the move and stay
|
|
# where we are right now.
|
|
if self.map[next_pos[0]][next_pos[1]] != "W":
|
|
self._agent_pos = next_pos
|
|
|
|
# Check if the agent has reached the global goal state.
|
|
if self.map[self._agent_pos[0]][self._agent_pos[1]] == "G":
|
|
rewards = {
|
|
"high_level_agent": 10.0,
|
|
# +1.0 if the goal position was also the target position for the
|
|
# low level agent.
|
|
low_level_agent: float(
|
|
self._agent_discrete_pos == target_discrete_pos
|
|
),
|
|
}
|
|
terminateds["__all__"] = True
|
|
return (
|
|
{"high_level_agent": self._agent_discrete_pos},
|
|
rewards,
|
|
terminateds,
|
|
truncateds,
|
|
{},
|
|
)
|
|
|
|
# Low-level agent has reached its target location (given by the high-level):
|
|
# - Hand back control to high-level agent.
|
|
# - Reward low level agent and high-level agent with small rewards.
|
|
elif self._agent_discrete_pos == target_discrete_pos:
|
|
self._num_targets_reached += 1
|
|
rewards = {
|
|
"high_level_agent": 1.0,
|
|
low_level_agent: 1.0,
|
|
}
|
|
return (
|
|
{"high_level_agent": self._agent_discrete_pos},
|
|
rewards,
|
|
terminateds,
|
|
truncateds,
|
|
{},
|
|
)
|
|
|
|
# Low-level agent has not reached anything.
|
|
else:
|
|
# Small step penalty for low-level agent.
|
|
rewards = {low_level_agent: -0.01}
|
|
# Reached time budget -> Hand back control to high level agent.
|
|
if self._low_level_steps >= self.max_steps_low_level:
|
|
rewards["high_level_agent"] = -0.01
|
|
return (
|
|
{"high_level_agent": self._agent_discrete_pos},
|
|
rewards,
|
|
terminateds,
|
|
truncateds,
|
|
{},
|
|
)
|
|
else:
|
|
return (
|
|
{
|
|
low_level_agent: (
|
|
self._agent_discrete_pos, # current
|
|
target_discrete_pos, # target
|
|
),
|
|
},
|
|
rewards,
|
|
terminateds,
|
|
truncateds,
|
|
{},
|
|
)
|
|
|
|
@property
|
|
def _agent_discrete_pos(self):
|
|
x = self._agent_pos[0]
|
|
y = self._agent_pos[1]
|
|
# discrete position = row idx * columns + col idx
|
|
return x * len(self.map[0]) + y
|
|
|
|
@staticmethod
|
|
def eucl_dist(pos1, pos2, map):
|
|
x1, y1 = pos1 % len(map[0]), pos1 // len(map)
|
|
x2, y2 = pos2 % len(map[0]), pos2 // len(map)
|
|
return ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** 0.5
|
|
|
|
|
|
def _get_next_pos(action, pos):
|
|
x, y = pos
|
|
# Up.
|
|
if action == 0:
|
|
return x - 1, y
|
|
# Down.
|
|
elif action == 1:
|
|
return x + 1, y
|
|
# Left.
|
|
elif action == 2:
|
|
return x, y - 1
|
|
# Right.
|
|
else:
|
|
return x, y + 1
|