def mplu_arrow(ax, x0, y0, x1, y1, **kwargs):
_color = kwargs.get('color', 'black')
_scale = kwargs.get('scale', 25)
arrow = FancyArrowPatch((x0, y0), (x1, y1),
arrowstyle=ArrowStyle('Simple',
head_width=1,
head_length=1),
mutation_scale=_scale)
arrow.set_color(_color)
ax.add_artist(arrow)
def initialise():
"""Sets up all the nodes and arrows on the window"""
for state_name, center in positions["state_nodes"].items():
center = tuple(np.array(center) + (self.x_shift, self.y_shift))
state_node = Circle(
center,
1.5,
color=self.visual_constants["state_color_light"],
zorder=1)
self.state_nodes[state_name] = state_node
ax.add_artist(state_node)
for state_action_tuple, center in positions["action_nodes"].items(
):
center = tuple(np.array(center) + (self.x_shift, self.y_shift))
action_node = Circle(
center,
0.8,
color=self.visual_constants["action_color_light"])
self.action_nodes[state_action_tuple] = action_node
ax.add_artist(action_node)
for state_name in structure:
for action_name in structure[state_name]:
arrow_start = self.state_nodes[state_name].center
arrow_end = self.action_nodes[(state_name,
action_name)].center
arrow_style = "simple, tail_width=0.1, head_width=4, head_length=8"
arrow = FancyArrowPatch(arrow_start,
arrow_end,
connectionstyle="arc3, rad=-.4",
arrowstyle=arrow_style,
color="k",
shrinkA=30,
shrinkB=16)
arrow.set_color("#66ff66")
self.agent_arrows[(state_name, action_name)] = arrow
ax.add_artist(arrow)
for state_name in structure:
for action_name in structure[state_name]:
for next_state_name in structure[state_name][action_name]:
arrow_start = self.action_nodes[(state_name,
action_name)].center
arrow_end = self.state_nodes[next_state_name].center
arrow_style = "simple, tail_width=0.1, head_width=4, head_length=8"
arrow = FancyArrowPatch(
arrow_start,
arrow_end,
linestyle=(0, (2, 5)),
connectionstyle="arc3, rad=-.4",
arrowstyle=arrow_style,
color="k",
shrinkA=16,
shrinkB=30)
arrow.set_color("#66ff66")
self.env_arrows[(state_name, action_name,
next_state_name)] = arrow
ax.add_artist(arrow)
for state_name, node in self.state_nodes.items():
self.state_values[state_name] = ax.text(
node.center[0],
node.center[1],
state_name,
horizontalalignment='center',
verticalalignment='center',
zorder=5,
color="w")
for state_action_tuple, node in self.action_nodes.items():
self.action_values[state_action_tuple] = ax.text(
node.center[0],
node.center[1],
state_action_tuple[1],
horizontalalignment='center',
verticalalignment='center',
zorder=5,
color="w")
self.status = ax.text(7,
0,
"start",
horizontalalignment='left',
verticalalignment='center',
zorder=5,
color="k",
fontsize="large",
fontweight="bold")
return tuple(
list(self.state_nodes.values()) +
list(self.state_values.values()) +
list(self.agent_arrows.values()) +
list(self.env_arrows.values()) + [self.status])
