def draw_arena_segment(self, top, bottom, left, right, rgb_tuple):
"""
Helper function that creates the geometry for a segment of the arena. Intended to be used by the viewer
:param top:
:param bottom:
:param rgb_tuple:
:return: A FilledPolygon object that can be added to the viewer using add_geom
"""
l, r, t, b = (left + self.x_shift) * self.scale, \
(right + self.x_shift) * self.scale, \
(top + self.y_shift) * self.scale, \
(bottom + self.y_shift) * self.scale
arena_segment = rendering.FilledPolygon([(l, b), (l, t), (r, t),
(r, b)])
arena_segment.add_attr(rendering.Transform(translation=(0, 0)))
arena_transform = rendering.Transform()
arena_segment.add_attr(arena_transform)
arena_segment.set_color(rgb_tuple[0], rgb_tuple[1], rgb_tuple[2])
return arena_segment
def draw_arena_segment(self, top, bottom, rgb_tuple):
"""
Helper function that creates the geometry for a segment of the arena. Intended to be used by the viewer
:param top:
:param bottom:
:param rgb_tuple:
:return: A FilledPolygon object that can be added to the viewer using add_geom
"""
l, r, t, b = self.arena_constraints["x_min"] * self.scale, \
self.arena_constraints["x_max"] * self.scale, \
top * self.scale, \
bottom * self.scale
arena_segment = rendering.FilledPolygon([(l, b), (l, t), (r, t),
(r, b)])
arena_segment.add_attr(rendering.Transform(translation=(0, 0)))
arena_transform = rendering.Transform()
arena_segment.add_attr(arena_transform)
arena_segment.set_color(rgb_tuple[0], rgb_tuple[1], rgb_tuple[2])
return arena_segment
def render(self, mode='human'):
"""
Renders the environment, placing all agents and boxes in appropriate positions
:param mode:
:return:
"""
screen_width = self.arena_constraints["x_max"] * self.scale
screen_height = self.arena_constraints["y_max"] * self.scale
if self.viewer is None:
self.viewer = rendering.Viewer(screen_width, screen_height)
home_top = self.arena_constraints["y_min"] + self.home_length
main_top = self.arena_constraints["y_max"] - self.goal_length
goal_top = self.arena_constraints["y_max"]
# Draw home
home = self.draw_arena_segment(home_top,
self.arena_constraints["y_min"],
self.home_colour)
self.viewer.add_geom(home)
# Draw main
main_area = self.draw_arena_segment(main_top, home_top,
self.main_colour)
self.viewer.add_geom(main_area)
# Draw goal
goal = self.draw_arena_segment(goal_top, main_top,
self.goal_colour)
self.viewer.add_geom(goal)
# Draw grid
grid_lines = self.draw_grid()
for line in grid_lines:
self.viewer.add_geom(line)
# Draw agent(s)
for i in range(self.num_agents):
agent = rendering.FilledPolygon([
(0, 0), (self.agent_width * self.scale, 0),
(self.agent_width / 2 * self.scale,
self.agent_height * self.scale)
])
agent.set_color(self.agent_colour[0], self.agent_colour[1],
self.agent_colour[2])
agent.add_attr(
rendering.Transform(
translation=(-self.agent_width / 2 * self.scale,
-self.agent_height / 2 * self.scale)))
agent.add_attr(self.agent_transforms[i])
self.viewer.add_geom(agent)
# Draw box(es)
for i in range(self.num_small_boxes + self.num_medium_boxes +
self.num_large_boxes):
if i < self.num_small_boxes and self.num_small_boxes > 0:
l, r, t, b = -self.small_box_width / 2 * self.scale, self.small_box_width / 2 * self.scale, self.small_box_width * self.scale, 0
elif i >= self.num_small_boxes and self.num_medium_boxes > 0:
l, r, t, b = -self.medium_box_width / 2 * self.scale, self.medium_box_width / 2 * self.scale, self.small_box_width * self.scale, 0
elif i >= self.num_small_boxes and self.num_large_boxes > 0:
l, r, t, b = -self.large_box_width / 2 * self.scale, self.large_box_width / 2 * self.scale, self.small_box_width * self.scale, 0
box = rendering.FilledPolygon([(l, b), (l, t), (r, t), (r, b)])
box.set_color(self.box_colour[0], self.box_colour[1],
self.box_colour[2])
box.add_attr(rendering.Transform(translation=(0, 0)))
box.add_attr(self.box_transforms[i])
self.viewer.add_geom(box)
# Set position of agent(s)
for i in range(self.num_agents):
self.agent_transforms[i].set_translation(
(self.agent_positions[i][0] - self.arena_constraints["x_min"] +
0.5) * self.scale,
(self.agent_positions[i][1] - self.arena_constraints["y_min"] +
0.5) * self.scale)
# Set position of box(es)
for box_id, box_details in self.boxes_in_arena.items():
if box_details[2] == "small":
self.box_transforms[box_id].set_translation(
(box_details[0] - self.arena_constraints["x_min"] + 0.5) *
self.scale,
(box_details[1] - self.arena_constraints["y_min"] + 0.125)
* self.scale)
elif box_details[2] == "medium":
self.box_transforms[box_id].set_translation(
(box_details[0] - self.arena_constraints["x_min"] +
(0.5 * self.medium_box_size)) * self.scale,
(box_details[1] - self.arena_constraints["y_min"] + 0.125)
* self.scale)
elif box_details[2] == "large":
self.box_transforms[box_id].set_translation(
(box_details[0] - self.arena_constraints["x_min"] +
(0.5 * self.num_agents)) * self.scale,
(box_details[1] - self.arena_constraints["y_min"] + 0.125)
* self.scale)
return self.viewer.render(return_rgb_array=mode == 'rgb_array')