colosseum.analysis.visualization

  1from typing import Dict, Tuple, Union, TYPE_CHECKING, List
  2
  3import networkx as nx
  4import numpy as np
  5import seaborn as sns
  6import toolz
  7from matplotlib import pyplot as plt
  8from matplotlib.colors import LinearSegmentedColormap
  9
 10from colosseum.mdp.utils import MDPCommunicationClass
 11
 12if TYPE_CHECKING:
 13    from colosseum.mdp import ContinuousMDP, EpisodicMDP, NODE_TYPE
 14
 15custom_dark = sns.color_palette("dark")
 16p = custom_dark.pop(1)
 17
 18
 19def plot_MDP_graph(
 20    mdp: Union["ContinuousMDP", "EpisodicMDP"],
 21    node_palette: List[Tuple[float, float, float]] = sns.color_palette("bright"),
 22    action_palette: List[Tuple[float, float, float]] = custom_dark,
 23    save_file: str = None,
 24    ax: plt.Axes = None,
 25    figsize: Tuple[int, int] = None,
 26    node_labels: Union[bool, Dict["NODE_TYPE", Union[float, str]]] = None,
 27    action_labels: Union[bool, Dict[Tuple["NODE_TYPE", int], Union[float, str]]] = None,
 28    int_labels_offset_x: int = 10,
 29    int_labels_offset_y: int = 10,
 30    continuous_form: bool = True,
 31    prog: str = "neato",
 32    ncol: int = 4,
 33    title: str = None,
 34    legend_fontsize: int = None,
 35    font_color_state_labels: str = "k",
 36    font_color_state_actions_labels: str = "k",
 37    cm_state_labels: LinearSegmentedColormap = None,
 38    cm_state_actions_labels: LinearSegmentedColormap = None,
 39    no_written_state_labels=True,
 40    no_written_state_action_labels=True,
 41    node_size=150,
 42):
 43    """
 44    In the MDP graph representation, states are associated to round nodes and actions to square nodes. Each state is
 45    connected to action nodes for all the available actions. Each action node is connected to the state node such that
 46    the probability of transitioning to them when playing the corresponding action from the corresponding state is
 47    greater than one.
 48
 49    Parameters
 50    ----------
 51    mdp : Union["ContinuousMDP", "EpisodicMDP"]
 52        The MDP to be visualised.
 53    node_palette : List[Tuple[float, float, float]]
 54        The color palette for the nodes corresponding to states. By default, the seaborn 'bright' palette is used.
 55    action_palette : List[Tuple[float, float, float]]
 56        The color palette for the nodes corresponding to actions. By default, the seaborn 'dark' palette is used.
 57    save_file : str
 58        The path file where the visualization will be store. Be default, the visualization will not be stored.
 59    ax : plt.Axes
 60        The ax object where the plot will be put. By default, a new axis is created.
 61    figsize : Tuple[int, int]
 62        The size of the figures in the grid of plots. By default, None is provided.
 63    node_labels : Union[bool, Dict["NODE_TYPE", Union[float, str]]]
 64        The dictionary mapping state nodes to labels (either a float or a str). If provided as bool, then the default labels
 65        are used. By default, no label is plotted.
 66    action_labels : Union[bool, Dict[Tuple["NODE_TYPE", int], Union[float, str]]]
 67        The dictionary mapping action nodes to labels (either a float or a str). If provided as bool, then the default labels
 68        are used. By default, no label is plotted.
 69    int_labels_offset_x : int
 70        x-axis offset for the node labels. By default, the offset is set to 10.
 71    int_labels_offset_y : int
 72        y-axis offset for the node labels. By default, the offset is set to 10.
 73    continuous_form : bool
 74        If True and the MDP is episodic, the continuous form of the MDP is used. By default, the continuous form is
 75        used.
 76    prog : str
 77        The default program to create the graph layout to be passed to the networkx library. By default, 'neato' is used.
 78    ncol : int
 79        The number of columns in the legend. By default, four columns are used.
 80    title : str
 81        The title to be given to the plot. By default, no title is used.
 82    legend_fontsize : int
 83        The font size of the legend. Vy default, it is not specified.
 84    font_color_state_labels : str
 85        The color code for the labels of the state nodes. By default, it is set to 'k'.
 86    font_color_state_actions_labels : str
 87        The color code for the labels of the action nodes. By default, it is set to 'k'.
 88    cm_state_labels : LinearSegmentedColormap
 89        The matplotlib color map for the labels of the state nodes. By default, it is not specified.
 90    cm_state_actions_labels : LinearSegmentedColormap
 91        The matplotlib color map for the labels of the action nodes. By default, it is not specified.
 92    no_written_state_labels : bool
 93        If True, the labels for the node states are not shown. By default, it is set to True.
 94    no_written_state_action_labels : bool
 95        If True, the labels for the node actions are not shown. By default, it is set to True.
 96    node_size : int
 97        The size of the node in the plot. By default, it is set to 100.
 98    """
 99    show = ax is None
100
101    sns.reset_defaults()
102    G, probs = (
103        _create_epi_MDP_graph(mdp)
104        if mdp.is_episodic() and not continuous_form
105        else _create_MDP_graph(mdp)
106    )
107    T, R = mdp.transition_matrix_and_rewards
108
109    layout = nx.nx_agraph.graphviz_layout(G, prog=prog)
110
111    if ax is None:
112        ax = _create_ax(layout, figsize)
113
114    if mdp.is_episodic() and not continuous_form:
115        if node_labels is not None and cm_state_labels is not None:
116            node_color = [
117                cm_state_labels(node_labels[node] / max(node_labels.values()))
118                for node in mdp.get_episodic_graph(False).nodes
119            ]
120        else:
121            node_color = [
122                node_palette[5]  # brown
123                if node[1] in mdp.starting_nodes and node[0] == 0
124                else node_palette[2]  # green
125                if R[mdp.node_to_index(node[1])].max() == R.max()
126                else node_palette[-2]  # yellow
127                if node[1] in mdp.recurrent_nodes_set
128                else node_palette[-3]  # grey
129                for node in mdp.get_episodic_graph(False).nodes
130            ]
131    else:
132        if node_labels is not None and cm_state_labels is not None:
133            node_color = [
134                cm_state_labels(node_labels[node] / max(node_labels.values()))
135                for node in mdp.G.nodes
136            ]
137        else:
138            node_color = [
139                node_palette[5]  # brown
140                if node in mdp.starting_nodes
141                else node_palette[2]  # green
142                if R[mdp.node_to_index[node]].max() == R.max()
143                else node_palette[-2]  # yellow
144                if node in mdp.recurrent_nodes_set
145                else node_palette[-3]  # grey
146                for node in mdp.G.nodes
147            ]
148
149    # Lazy way to create nice legend handles
150    x, y = list(layout.values())[0]
151    if cm_state_labels is None:
152        ax.scatter(x, y, color=node_palette[2], label="Highly rewarding state")
153        ax.scatter(x, y, color=node_palette[-2], label="State")
154        if mdp.communication_class == MDPCommunicationClass.WEAKLY_COMMUNICATING:
155            ax.scatter(x, y, color=node_palette[-3], label="Transient state")
156        ax.scatter(x, y, color=node_palette[5], label="Starting state")
157    ax.plot(x, y, color=node_palette[-3], label="Transition probability")
158    if cm_state_actions_labels is None:
159        for a in range(mdp.n_actions):
160            ax.plot(x, y, color=action_palette[a], label=f"Action: {a}", marker="s")
161
162    G_nodes = (
163        mdp.get_episodic_graph(False).nodes
164        if mdp.is_episodic() and not continuous_form
165        else mdp.G.nodes
166    )
167    nx.draw_networkx_nodes(
168        G,
169        layout,
170        G_nodes,
171        ax=ax,
172        node_color=node_color,
173        edgecolors="black",
174        node_size=node_size,
175    )
176    for a in range(mdp.n_actions):
177        na_nodes = (
178            [an for an in G.nodes if type(an[1]) == int and an[-1] == a]
179            if mdp.is_episodic() and not continuous_form
180            else [an for an in G.nodes if type(an) == tuple and an[-1] == a]
181        )
182        nx.draw_networkx_nodes(
183            G,
184            layout,
185            na_nodes,
186            node_shape="s",
187            ax=ax,
188            node_size=node_size,
189            node_color=[action_palette[a]]
190            if cm_state_actions_labels is None
191            else [
192                cm_state_actions_labels(action_labels[an] / max(action_labels.values()))
193                for an in na_nodes
194            ],
195            edgecolors="black",
196        )
197        nx.draw_networkx_edges(
198            G,
199            layout,
200            edgelist=[e for e in G.edges if type(e[0][0]) != tuple and e[1][1] == a]
201            if mdp.is_episodic() and not continuous_form
202            else [e for e in G.edges if type(e[0]) != tuple and e[1][1] == a],
203            ax=ax,
204            edge_color=action_palette[a],
205        )
206    nx.draw_networkx_edges(
207        G,
208        layout,
209        edgelist=[e for e in G.edges if type(e[0][0]) == tuple]
210        if mdp.is_episodic() and not continuous_form
211        else [e for e in G.edges if type(e[0]) == tuple],
212        ax=ax,
213        edge_color=action_palette[-3],
214        width=[probs[e] for e in G.edges if type(e[0][0]) == tuple]
215        if mdp.is_episodic() and not continuous_form
216        else [probs[e] for e in G.edges if type(e[0]) == tuple],
217    )
218    ax.legend(ncol=ncol, fontsize=legend_fontsize)
219    if node_labels is not None and not no_written_state_labels:
220        if type(node_labels) == bool and node_labels:
221            node_labels = {
222                n: (
223                    f"h={n[0]},{n[1]}"
224                    if mdp.is_episodic() and not continuous_form
225                    else str(n)
226                )
227                for n in G_nodes
228            }
229        assert all(n in G.nodes for n in node_labels)
230        nx.draw_networkx_labels(
231            G,
232            toolz.valmap(
233                lambda x: [x[0] + int_labels_offset_x, x[1] + int_labels_offset_y],
234                layout,
235            ),
236            node_labels,
237            font_color=font_color_state_labels,
238            ax=ax,
239            verticalalignment="center_baseline",
240        )
241    if action_labels is not None and not no_written_state_action_labels:
242        if type(action_labels) == bool and action_labels:
243            action_labels = {
244                n: str(n[1])
245                for n in (
246                    (an for an in G.nodes if type(an[1]) == int)
247                    if mdp.is_episodic() and not continuous_form
248                    else (an for an in G.nodes if type(an) == tuple)
249                )
250            }
251        assert all(n in G.nodes and type(n[1]) == int for n in action_labels)
252        nx.draw_networkx_labels(
253            G,
254            toolz.valmap(
255                lambda x: [x[0] + int_labels_offset_x, x[1] + int_labels_offset_y],
256                layout,
257            ),
258            action_labels,
259            font_color=font_color_state_actions_labels,
260            ax=ax,
261            verticalalignment="center_baseline",
262        )
263
264    ax.axis("off")
265    if title is not None:
266        ax.set_title(title)
267    if save_file is not None:
268        plt.savefig(save_file, bbox_inches="tight")
269    if show:
270        plt.show()
271
272
273def plot_MCGraph(
274    mdp: Union["ContinuousMDP", "EpisodicMDP"],
275    node_palette: List[Tuple[float, float, float]] = sns.color_palette("deep"),
276    labels: Union[bool, Dict["NODE_TYPE", Union[float, str]]] = {},
277    font_color_labels="k",
278    save_file: str = None,
279    ax: plt.Axes = None,
280    figsize: Tuple[int, int] = None,
281    prog: str = None,
282    legend_fontsize: int = None,
283    node_size=100,
284    cm_state_labels=None,
285    no_written_state_labels=True,
286):
287    """
288    In the Markov chain-based representation, states are associated to nodes and actions are not visualized. Each state
289    is connected to states such that there exists at least one action with corresponding transition probability greater
290    than zero.
291
292    Parameters
293    ----------
294    mdp : Union["ContinuousMDP", "EpisodicMDP"]
295        The MDP to be visualized.
296    node_palette : List[Tuple[float, float, float]]
297        The color palette for the nodes. By default, the seaborn 'bright' palette is used.
298    labels : Union[bool, Dict["NODE_TYPE", Union[float, str]]]
299        The dictionary mapping nodes to labels (either a float or a str). If provided as bool, then the default labels
300        are used. By default, no label is plotted.
301    font_color_labels
302    save_file : str
303        The path file where the visualization will be store. Be default, the visualization will not be stored.
304    ax : plt.Axes
305        The ax object where the plot will be put. By default, a new axis is created.
306    figsize : Tuple[int, int]
307        The size of the figures in the grid of plots. By default, None is provided.
308    prog : str
309        The default program to create the graph layout to be passed to the networkx library. By default, it is not
310        specified.
311    legend_fontsize : int
312        The font size of the legend. Vy default, it is not specified.
313    node_size : int
314        The size of the node in the plot. By default, it is set to 100.
315    cm_state_labels : LinearSegmentedColormap
316        The matplotlib color map for the labels of the state nodes. By default, it is not specified.
317    no_written_state_labels : bool
318        If True, the labels for the nodes are not shown. By default, it is set to True.
319
320    """
321
322    show = ax is None
323
324    _, R = mdp.transition_matrix_and_rewards
325
326    if cm_state_labels is not None:
327        node_color = [
328            cm_state_labels(labels[n] / max(labels.values())) for n in mdp.G.nodes
329        ]
330    else:
331        node_color = [
332            node_palette[0]  # brown
333            if node in mdp.starting_nodes
334            else node_palette[2]  # green
335            if R[mdp.node_to_index[node]].max() == R.max()
336            else node_palette[1]  # yellow
337            if node in mdp.recurrent_nodes_set
338            else node_palette[-1]  # grey
339            for node in mdp.G.nodes
340        ]
341
342    if ax is None:
343        ax = _create_ax(mdp.graph_layout, figsize)
344
345    if cm_state_labels is None:
346        x, y = list(mdp.graph_layout.values())[0]
347        ax.scatter(x, y, color=node_palette[2], label="Highly rewarding state")
348        ax.scatter(x, y, color=node_palette[1], label="State")
349        if mdp.communication_class == MDPCommunicationClass.WEAKLY_COMMUNICATING:
350            ax.scatter(x, y, color=node_palette[-1], label="Transient state")
351        ax.scatter(x, y, color=node_palette[0], label="Starting state")
352
353    nx.draw(
354        mdp.G,
355        mdp.graph_layout
356        if prog is None
357        else nx.nx_agraph.graphviz_layout(mdp.G, prog=prog),
358        node_color=node_color,
359        node_size=node_size,
360        edgecolors="black",
361        edge_color=node_palette[-3],
362        labels={}
363        if cm_state_labels is not None and no_written_state_labels
364        else labels,
365        font_color=font_color_labels,
366        ax=ax,
367    )
368
369    if cm_state_labels is None:
370        ax.legend(fontsize=legend_fontsize)
371    if save_file is not None:
372        plt.savefig(save_file)
373    if show:
374        plt.show()
375
376
377def _create_ax(layout, figsize: Tuple[int, int] = None):
378    if figsize is None:
379        positions = np.array(list(layout.values()))
380        max_distance = max(
381            np.sqrt(np.sum((positions[i] - positions[j]) ** 2))
382            for i in range(len(layout))
383            for j in range(i + 1, len(layout))
384        )
385        figsize = max(6, min(20, int(max_distance / 70)))
386        figsize = (figsize, figsize)
387    plt.figure(None, figsize)
388    cf = plt.gcf()
389    cf.set_facecolor("w")
390    ax = cf.add_axes((0, 0, 1, 1)) if cf._axstack() is None else cf.gca()
391    ax.spines["top"].set_visible(False)
392    ax.spines["right"].set_visible(False)
393    ax.spines["bottom"].set_visible(False)
394    ax.spines["left"].set_visible(False)
395    return ax
396
397
398def _create_MDP_graph(mdp: Union["ContinuousMDP", "EpisodicMDP"]):
399    T, R = mdp.transition_matrix_and_rewards
400
401    probs = dict()
402    G = nx.DiGraph()
403    for s in range(mdp.n_states):
404        n = mdp.index_to_node[s]
405        if s not in G.nodes:
406            G.add_node(n)
407        for a in range(mdp.n_actions):
408            an = tuple((n, a))
409            G.add_edge(n, an)
410            for nn in np.where(T[s, a] > 0)[0]:
411                G.add_edge(an, mdp.index_to_node[nn])
412                probs[an, mdp.index_to_node[nn]] = T[s, a, nn]
413
414    return G, probs
415
416
417def _create_epi_MDP_graph(mdp: "EpisodicMDP"):
418    G_epi = mdp.get_episodic_graph(False)
419    T, R = mdp.episodic_transition_matrix_and_rewards
420
421    probs = dict()
422    G = nx.DiGraph()
423    for n in G_epi.nodes:
424        if n not in G.nodes:
425            G.add_node(n)
426
427        for a in range(mdp.n_actions):
428            an = tuple((n, a))
429            G.add_edge(n, an)
430            for nn in G_epi.successors(n):
431                G.add_edge(an, nn)
432                probs[an, nn] = T[
433                    n[0], mdp.node_to_index(n[1]), a, mdp.node_to_index(nn[1])
434                ]
435
436    return G, probs