<NME> LearningSession.py
<BEF> ADDFILE
<MSG> LearningSession

<DFF> @@ -0,0 +1,560 @@
+# This file is part of DmpBbo, a set of libraries and programs for the
+# black-box optimization of dynamical movement primitives.
+# Copyright (C) 2022 Freek Stulp
+#
+# DmpBbo is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published by
+# the Free Software Foundation, either version 2 of the License, or
+# (at your option) any later version.
+#
+# DmpBbo is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with DmpBbo.  If not, see <http://www.gnu.org/licenses/>.
+
+
+import numpy as np
+import math
+import os
+import sys
+import copy
+import argparse
+import inspect
+from glob import glob
+import matplotlib.pyplot as plt
+from pylab import mean
+from matplotlib.patches import Ellipse
+import warnings
+
+# Avoid warnings when plotting very narrow covariance matrices
+warnings.simplefilter("ignore", np.ComplexWarning)
+
+
+lib_path = os.path.abspath("../../python/")
+sys.path.append(lib_path)
+
+from to_jsonpickle import *
+from bbo.DistributionGaussian import DistributionGaussian
+
+
+def plot_error_ellipse(mu, cov, ax=None, **kwargs):
+    """
+Plot the error ellipse at a point given it's covariance matrix
+
+Parameters
+----------
+mu : array (2,)
+The center of the ellipse
+
+cov : array (2,2)
+The covariance matrix for the point
+
+ax : matplotlib.Axes, optional
+The axis to overplot on
+
+**kwargs : dict
+These keywords are passed to matplotlib.patches.Ellipse
+
+From https://github.com/dfm/dfmplot/blob/master/dfmplot/ellipse.py
+
+"""
+    # some sane defaults
+    facecolor = kwargs.pop("facecolor", "none")
+    edgecolor = kwargs.pop("edgecolor", "k")
+
+    x, y = mu
+    U, S, V = np.linalg.svd(cov)
+    theta = np.degrees(np.arctan2(U[1, 0], U[0, 0]))
+    ellipsePlot = Ellipse(
+        xy=[x, y],
+        width=2 * np.sqrt(S[0]),
+        height=2 * np.sqrt(S[1]),
+        angle=theta,
+        facecolor=facecolor,
+        edgecolor=edgecolor,
+        **kwargs,
+    )
+
+    if not ax:
+        ax = plt.gca()
+    lines = ax.add_patch(ellipsePlot)
+    return lines
+
+
+def plotUpdateLines(n_samples_per_update, ax):
+    """ Plot vertical lines when an parameter update occured during the optimization.
+    \param[in] n_samples_per_update Vector specifying how many samples were used between updates.
+    \param[in] ax Axis object to plot the lines in.
+    """
+
+    # Find good number of horizontal update lines to plot
+    updates = np.arange(0, len(n_samples_per_update))
+    while len(n_samples_per_update) > 20:
+        n_samples_per_update = n_samples_per_update[0:-1:5]
+        updates = updates[0:-1:5]
+
+    y_limits = ax.get_ylim()
+
+    n = n_samples_per_update
+    h = ax.plot([n, n], y_limits)
+    plt.setp(h, color="#bbbbbb", linewidth=0.5, zorder=0)
+        
+    for ii in range(len(n_samples_per_update) - 1):
+        y = y_limits[0] + 0.9 * (y_limits[1] - y_limits[0])
+        ax.text(
+            n_samples_per_update[ii + 1],
+            y,
+            str(updates[ii + 1]),
+            horizontalalignment="center",
+            verticalalignment="top",
+            rotation="vertical",
+        )
+
+    y = y_limits[0] + 0.95 * (y_limits[1] - y_limits[0])
+    ax.text(
+        mean(ax.get_xlim()),
+        y,
+        "number of updates",
+        horizontalalignment="center",
+        verticalalignment="top",
+    )
+    ax.set_ylim(y_limits)
+
+
+def plotLearningCurve(learning_curve, **kwargs):
+    """ Plot a learning curve.
+    \param[in] learing_curve A learning curve that has the following format
+        #rows is number of optimization updates
+        column 0: Number of samples at which the cost was evaluated
+        column 1: The total cost 
+        column 2...: Individual cost components (column 1 is their sum)
+    \param[in] ax Axis to plot the learning curve on.
+    \param[in] costs_all Vector of costs of each sample (default=[])
+    \param[in] cost_labels Vector of strings for the different cost components  (default=[]).
+    """
+
+    cost_labels = kwargs.get("cost_labels", None)
+    ax = kwargs.get("ax") or plt.axes()
+
+    # Plot costs at evaluations
+    learning_curve = np.array(learning_curve)
+    samples_eval = learning_curve[:, 0]
+    costs_eval = learning_curve[:, 1:]
+    # Sum of cost components
+    lines = ax.plot(samples_eval, costs_eval[:, 0], "-", color="black", linewidth=2)
+    # Individual cost components
+    if costs_eval.shape[1] > 1:
+        lines.extend(ax.plot(samples_eval, costs_eval[:, 1:], "-", linewidth=1))
+
+    # Annotation
+    ax.set_xlabel("number of evaluations")
+    ax.set_ylabel("cost")
+    ax.set_title("Learning curve")
+
+    if cost_labels:
+        cost_labels.insert(0, "total cost")
+        plt.legend(lines, cost_labels)
+
+    plotUpdateLines(samples_eval, ax)
+
+    return (lines, ax)
+
+
+def plotExplorationCurve(exploration_curve, **kwargs):
+    """ Plot an exploration curve.
+    \param[in] exploration An exploration curve that has the following format
+        #rows is number of optimization updates
+        column 0: Number of samples at which the cost was evaluated
+        column 1: The exploration at that update 
+    \param[in] ax Axis to plot the learning curve on.
+    """
+    ax = kwargs.get("ax") or plt.axes()
+
+    exploration_curve = np.array(exploration_curve)
+    samples_eval = exploration_curve[:, 0]
+    explo = exploration_curve[:, 1]
+
+    line = ax.plot(samples_eval, explo, "-", color="green", linewidth=2)
+    plotUpdateLines(samples_eval, ax)
+    ax.set_xlabel("number of evaluations")
+    ax.set_ylabel("sqrt of max. eigval of covar")
+    ax.set_title("Exploration magnitude")
+
+    return (line, ax)
+
+
+def setColor(handle, i_update, n_updates):
+    """ Set the color of an object, according to how far the optimization has proceeded.
+    \param[in] handle Handle to the object
+    \param[in] i_update Which update is the optimization at now?
+    \param[in] n_update Which is the number of updates the optimization will run?
+    """
+    if i_update == 0:
+        plt.setp(handle, color="r", linewidth=2)
+    else:
+        c = 1.0 * i_update / n_updates
+        cur_color = [0.0 - 0.0 * c, 0.0 + 1.0 * c, 0.0 - 0.0 * c]
+        plt.setp(handle, color=cur_color)
+
+
+def plotUpdate(
+    distribution,
+    cost_eval,
+    samples,
+    costs,
+    weights,
+    distribution_new,
+    **kwargs):
+    """ Save an optimization update to a directory.
+    \param[in] distribution Gaussian distribution before the update
+    \param[in] cost_eval Cost of the mean of the distribution 
+    \param[in] samples The samples in the search space
+    \param[in] costs The cost of each sample 
+    \param[in] weights The weight of each sample 
+    \param[in] distribution_new Gaussian distribution after the update
+    \param[in] ax Axis to plot the update on.
+    \param[in] highlight Wether to highlight this update (default=False)
+    \param[in] plot_samples Whether to plot the individual samples (default=False)
+    """
+    ax = kwargs.get("ax") or plt.axes()
+    highlight = kwargs.get("highlight",False)
+    plot_samples = kwargs.get("plot_samples",False)
+
+    if samples is None:
+        plot_samples = False
+
+    n_dims = len(distribution.mean)
+    if n_dims == 1:
+        raise ValueError("Sorry, only know how to plot for n_dims==2, but you provided n_dims==1")
+
+    # ZZZ Take into consideration block_covar_sizes to plot sub blocks in
+    # different subplots
+
+    if n_dims >= 2:
+        distr_mean = distribution.mean[0:2]
+        distr_covar = distribution.covar[0:2, 0:2]
+        distr_new_mean = distribution_new.mean[0:2]
+        distr_new_covar = distribution_new.covar[0:2, 0:2]
+        if samples is not None:
+            samples = samples[:, 0:2]
+
+        if plot_samples:
+            max_marker_size = 80
+            for ii in range(len(weights)):
+                cur_marker_size = max_marker_size * weights[ii]
+                sample_handle = ax.plot(
+                    samples[ii, 0], samples[ii, 1], "o", color="green"
+                )
+                plt.setp(
+                    sample_handle,
+                    markersize=cur_marker_size,
+                    markerfacecolor=(0.5, 0.8, 0.5),
+                    markeredgecolor="none",
+                )
+
+                ax.plot(samples[:, 0], samples[:, 1], ".", color="black")
+            ax.plot(
+                (distr_mean[0], distr_new_mean[0]),
+                (distr_mean[1], distr_new_mean[1]),
+                "-",
+                color="blue",
+            )
+
+        mean_handle = ax.plot(distr_mean[0], distr_mean[1], "o", label="old")
+        mean_handle_new = ax.plot(
+            distr_new_mean[0], distr_new_mean[1], "o", label="new"
+        )
+        mean_handle_link = ax.plot(
+            [distr_mean[0], distr_new_mean[0]], [distr_mean[1], distr_new_mean[1]], "-"
+        )
+        patch = plot_error_ellipse(distr_mean, distr_covar, ax)
+        patch_new = plot_error_ellipse(distr_new_mean, distr_new_covar, ax)
+        if highlight:
+            plt.setp(mean_handle, color="red")
+            plt.setp(mean_handle_new, color="blue")
+            plt.setp(patch, edgecolor="red")
+            plt.setp(patch_new, edgecolor="blue")
+        else:
+            plt.setp(mean_handle, color="gray")
+            plt.setp(mean_handle_new, color="gray")
+            plt.setp(patch, edgecolor="gray")
+            plt.setp(patch_new, edgecolor="gray")
+        plt.setp(mean_handle_link, color="gray")
+    ax.set_aspect("equal")
+
+    # plt.rcParams['text.usetex']=True
+    # ax.set_xlabel(r'$\theta_1$')
+    # ax.set_ylabel(r'$\theta_2$')
+    ax.set_xlabel("dim 1 (of " + str(n_dims) + ")")
+    ax.set_ylabel("dim 2 (of " + str(n_dims) + ")")
+    ax.set_title("Search space")
+    return mean_handle, mean_handle_new, patch, patch_new
+
+
+class LearningSession:
+    def __init__(self, n_samples_per_update, root_directory=None, **kwargs):
+        self._n_samples_per_update = n_samples_per_update
+        self._root_dir = root_directory
+        self._cache = {}
+        self._last_update_added = None
+
+        self.tell(n_samples_per_update, "n_samples_per_update")
+        for name, value in kwargs.items():
+            self.tell(value, name)
+            
+        if root_directory and 'cost_function' in kwargs:
+            cost_function = kwargs['cost_function']
+            src = inspect.getsourcelines(cost_function.__class__)
+            src = " ".join(src[0])
+            src = src.replace("(CostFunction)", "")
+            filename = os.path.join(root_directory, "cost_function.py")
+            with open(filename, "w") as f:
+                f.write(src)
+                
+    @classmethod
+    def from_dir(cls, directory):
+        filename = os.path.join(directory, "n_samples_per_update.txt")
+        n_samples_per_update = np.atleast_1d(np.loadtxt(filename))
+        n_samples_per_update = int(n_samples_per_update[0])
+        return cls(n_samples_per_update, directory)
+
+    def addEval(self, i_update, eval_sample, eval_cost):
+        self.tell(eval_cost, "eval_cost", i_update)
+        self.tell(eval_sample, "eval_sample", i_update)
+
+    def addUpdate(
+        self, i_update, distribution, samples, costs, weights, distribution_new=None
+    ):
+        self.tell(distribution, "distribution", i_update)
+        self.tell(samples, "samples", i_update)
+        self.tell(costs, "costs", i_update)
+        self.tell(weights, "weights", i_update)
+        if distribution_new:
+            self.tell(distribution_new, "distribution_new", i_update)
+        self._last_update_added = i_update
+
+    def getNUpdates(self):
+        if not self._root_dir:
+            return self._last_update_added
+        else:
+            update_dirs = sorted(glob(os.path.join(self._root_dir, "update*/")))
+            last_dir = update_dirs[-1]
+            last_dir = last_dir.replace("update", "")
+            last_dir = last_dir.replace(self._root_dir, "")
+            last_dir = last_dir.replace("/", "")
+            return int(last_dir)
+
+    def getBaseName(self, name, i_update=None, i_sample=None):
+
+        if isinstance(i_sample, int):
+            basename = f"{i_sample:03d}_{name}"  # e.g. 0 => "000_name"
+        elif i_sample:
+            basename = f"{i_sample}_{name}"  # e.g. "eval" => "eval_name"
+        else:
+            basename = name
+
+        if isinstance(i_update, int):
+            # e.g. "update00023/000_name"
+            basename = os.path.join(f"update{i_update:05d}", basename)
+
+        return basename
+
+    def exists(self, name, i_update=None, i_sample=None):
+        try:
+            obj = self.ask(name, i_update, i_sample)
+        except Exception:
+            return False
+        return True
+
+    def ask(self, name, i_update=None, i_sample=None):
+
+        basename = self.getBaseName(name, i_update, i_sample)
+
+        if basename in self._cache:
+            return copy.deepcopy(self._cache[basename])
+
+        if not self._root_dir:
+            # Not in cache, and no directory given: Error
+            raise KeyError("Could not find basename in cache: " + basename)
+
+        else:
+            abs_basename = os.path.join(self._root_dir, basename)
+            if os.path.isfile(abs_basename + ".json"):
+                with open(abs_basename + ".json", "r") as f:
+                    obj = jsonpickle.decode(f.read())
+
+            elif os.path.isfile(abs_basename + ".txt"):
+                obj = np.loadtxt(abs_basename + ".txt")
+
+            else:
+                raise IOError("Could not find file with basename: " + abs_basename)
+
+            self._cache[basename] = obj
+
+            return obj
+
+    def tell(self, obj, name, i_update=None, i_sample=None):
+
+        basename = self.getBaseName(name, i_update, i_sample)
+
+        self._cache[basename] = copy.deepcopy(obj)
+
+        if self._root_dir:
+            abs_basename = os.path.join(self._root_dir, basename)
+
+            # Make sure directory exists
+            os.makedirs(os.path.dirname(abs_basename), exist_ok=True)
+
+            if isinstance(obj, list) or isinstance(obj, int):
+                # Try to convert list to numpy array
+                obj = np.atleast_1d(np.array(obj))
+
+            if isinstance(obj, np.ndarray):
+                filename = abs_basename + ".txt"
+                np.savetxt(filename, obj)
+            else:
+                filename = abs_basename + ".json"
+                with open(filename, "w") as out_file:
+                    out_file.write(to_jsonpickle(obj))
+
+            return filename
+
+    def getEvalCosts(self, i_update):
+        has_eval = self.exists("cost", 0, "eval")
+        if has_eval:
+            cost_eval = self.ask("cost", i_update, "eval")
+            # Evaluation at the beginning of an update
+            i_sample = i_update * self._n_samples_per_update
+        else:
+            update_costs = self.getSampleCosts(i_update)
+            cost_eval = np.mean(update_costs, axis=0)
+            # Evaluation between two updates (mean of all samples)
+            i_sample = (i_update + 0.5) * self._n_samples_per_update
+        return cost_eval, i_sample
+
+    def getSampleCosts(self, i_update):
+        if self.exists("costs", i_update):
+            # Everything in one file
+            sample_costs = self.ask("costs", i_update)
+        else:
+            # One file for each sample
+            n = self._n_samples_per_update
+            sample_costs = [self.ask("costs", i_update, s) for s in range(n)]
+        return np.array(sample_costs)
+
+    def getLearningCurve(self):
+        learning_curve = []
+        n_updates = self.getNUpdates()
+        for i_update in range(n_updates):
+            costs_eval, i_sample = self.getEvalCosts(i_update)
+            costs_eval = np.atleast_1d(costs_eval)
+            learning_curve.append(np.concatenate(([i_sample], costs_eval)))
+
+        cost_labels = []
+        if self.exists('cost_function'):
+            cost_function = self.ask('cost_function')
+            cost_labels = cost_function.costLabels();
+        elif self.exists('task'):
+            cost_function = self.ask('task')
+            cost_labels = cost_function.costLabels();
+
+        return (learning_curve, cost_labels)
+
+    def plotLearningCurve(self, ax=None):
+        """
+        has eval or not
+        has cost components or not
+        saved in costs or 000_cost...
+        """
+        if not ax:
+            ax = plt.axes()
+
+        # Plot costs of individual samples
+        costs = []
+        for i_update in range(self.getNUpdates()):
+            sample_costs = self.getSampleCosts(i_update)
+            costs.extend([c[0] for c in sample_costs]) # Only include sum
+        ax.plot(costs, ".", color="gray")
+
+        # Plot learning curve itself
+        learning_curve, cost_labels = self.getLearningCurve()
+        return plotLearningCurve(learning_curve, ax=ax, cost_labels=cost_labels)
+
+    def plotExplorationCurve(self, ax=None):
+        n_updates = self.getNUpdates()
+        curve = np.zeros((n_updates, 2))
+        for i_update in range(n_updates):
+            distribution = self.ask("distribution", i_update)
+            cur_exploration = np.sqrt(distribution.maxEigenValue())
+            curve[i_update, 0] = i_update * self._n_samples_per_update
+            curve[i_update, 1] = cur_exploration
+
+        if not ax:
+            ax = plt.axes()
+        return plotExplorationCurve(curve, ax=ax)
+
+    def plotDistributionUpdates(self, ax=None):
+        if not ax:
+            ax = plt.axes()
+
+        n_updates = self.getNUpdates()
+        for i_update in range(n_updates):
+            highlight = i_update == 0
+            self.plotDistributionUpdate(i_update, ax, highlight=highlight)
+        return ax # zzz
+
+    def plotDistributionUpdate(self, i_update, ax=None, **kwargs):
+        #ax = kwargs.get("ax") or plt.axes()
+        highlight = kwargs.get("highlight",False)
+        plot_samples = kwargs.get("plot_samples",False)
+        
+        distribution = self.ask("distribution", i_update)
+        costs_eval = self.getEvalCosts(i_update)
+        costs = self.getSampleCosts(i_update)
+        samples = self.ask("samples", i_update)
+        weights = self.ask("weights", i_update)
+        distribution_new = self.ask("distribution_new", i_update)
+
+        if not ax:
+            ax = plt.axes()
+            
+        return plotUpdate(
+            distribution,
+            costs_eval,
+            samples,
+            costs,
+            weights,
+            distribution_new,
+            ax=ax,
+            highlight=highlight,
+            plot_samples=plot_samples,
+        )
+
+    def plot(self, fig=None):
+        if not fig:
+            fig = plt.figure(figsize=(15, 5))
+        axs = [fig.add_subplot(131 + sp) for sp in range(3)]
+        self.plotDistributionUpdates(axs[0])
+        self.plotExplorationCurve(axs[1])
+        self.plotLearningCurve(axs[2])
+        return fig
+
+    @staticmethod
+    def plotMultiple(sessions):
+        pass
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "directory", type=str, help="directory from which to read the learning session"
+    )
+    args = parser.parse_args()
+
+    session = LearningSession.from_dir(args.directory)
+    session.plot()
+    plt.show()

{'commit': 'diff --git a/demos/python/TaskViapoint.py b/demos/python/TaskViapoint.py\nindex d1c6c3a..a11475c 100644\n--- a/demos/python/TaskViapoint.py\n+++ b/demos/python/TaskViapoint.py\n@@ -140,40 +140,4 @@ class TaskViapoint(Task):\n         else:\n             line_handles = []\n             \n-        return (line_handles, ax)\n-        \n-    def saveToFile(self,directory,filename):\n-        if not os.path.exists(directory):\n-            os.makedirs(directory)\n-            \n-        my_file = open(directory+"/"+filename, \'w\')\n-        for x in self.viapoint_:\n-            my_file.write(str(x)+" ")\n-        if self.viapoint_time_:\n-            my_file.write(str(self.viapoint_time_)+" ")\n-        else:\n-            my_file.write(str(-1.0)+" ")\n-        my_file.write(str(self.viapoint_radius_)+" ")\n-        \n-        for x in self.goal_:\n-            my_file.write(str(x)+" ")\n-        if self.goal_time_:\n-            my_file.write(str(self.goal_time_)+" ")\n-        else:\n-            my_file.write(str(-1.0)+" ")\n-            \n-        my_file.write(str(self.viapoint_weight_)+" ")\n-        my_file.write(str(self.acceleration_weight_)+" ")\n-        my_file.write(str(self.goal_weight_)+" ")\n-        my_file.close()\n-\n-if __name__==\'__main__\':\n-    counter = 0\n-    for n_dims in [1,2,5]:\n-        for viapoint_time in [0.5, None]:\n-            viapoint = np.linspace(0.0,1.0,n_dims)\n-            task = TaskViapoint(viapoint, viapoint_time)\n-            task.saveToFile("/tmp/demoTaskViapoint/","viapoint"+str(counter)+".txt")\n-            counter += 1\n-        \n-        \n+        return (line_handles, ax)\n\\ No newline at end of file\ndiff --git a/python/bbo/LearningSession.py b/python/bbo/LearningSession.py\nnew file mode 100644\nindex 0000000..60094bf\n--- /dev/null\n+++ b/python/bbo/LearningSession.py\n@@ -0,0 +1,560 @@\n+# This file is part of DmpBbo, a set of libraries and programs for the\n+# black-box optimization of dynamical movement primitives.\n+# Copyright (C) 2022 Freek Stulp\n+#\n+# DmpBbo is free software: you can redistribute it and/or modify\n+# it under the terms of the GNU Lesser General Public License as published by\n+# the Free Software Foundation, either version 2 of the License, or\n+# (at your option) any later version.\n+#\n+# DmpBbo is distributed in the hope that it will be useful,\n+# but WITHOUT ANY WARRANTY; without even the implied warranty of\n+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the\n+# GNU Lesser General Public License for more details.\n+#\n+# You should have received a copy of the GNU Lesser General Public License\n+# along with DmpBbo.  If not, see <http://www.gnu.org/licenses/>.\n+\n+\n+import numpy as np\n+import math\n+import os\n+import sys\n+import copy\n+import argparse\n+import inspect\n+from glob import glob\n+import matplotlib.pyplot as plt\n+from pylab import mean\n+from matplotlib.patches import Ellipse\n+import warnings\n+\n+# Avoid warnings when plotting very narrow covariance matrices\n+warnings.simplefilter("ignore", np.ComplexWarning)\n+\n+\n+lib_path = os.path.abspath("../../python/")\n+sys.path.append(lib_path)\n+\n+from to_jsonpickle import *\n+from bbo.DistributionGaussian import DistributionGaussian\n+\n+\n+def plot_error_ellipse(mu, cov, ax=None, **kwargs):\n+    """\n+Plot the error ellipse at a point given it\'s covariance matrix\n+\n+Parameters\n+----------\n+mu : array (2,)\n+The center of the ellipse\n+\n+cov : array (2,2)\n+The covariance matrix for the point\n+\n+ax : matplotlib.Axes, optional\n+The axis to overplot on\n+\n+**kwargs : dict\n+These keywords are passed to matplotlib.patches.Ellipse\n+\n+From https://github.com/dfm/dfmplot/blob/master/dfmplot/ellipse.py\n+\n+"""\n+    # some sane defaults\n+    facecolor = kwargs.pop("facecolor", "none")\n+    edgecolor = kwargs.pop("edgecolor", "k")\n+\n+    x, y = mu\n+    U, S, V = np.linalg.svd(cov)\n+    theta = np.degrees(np.arctan2(U[1, 0], U[0, 0]))\n+    ellipsePlot = Ellipse(\n+        xy=[x, y],\n+        width=2 * np.sqrt(S[0]),\n+        height=2 * np.sqrt(S[1]),\n+        angle=theta,\n+        facecolor=facecolor,\n+        edgecolor=edgecolor,\n+        **kwargs,\n+    )\n+\n+    if not ax:\n+        ax = plt.gca()\n+    lines = ax.add_patch(ellipsePlot)\n+    return lines\n+\n+\n+def plotUpdateLines(n_samples_per_update, ax):\n+    """ Plot vertical lines when an parameter update occured during the optimization.\n+    \\param[in] n_samples_per_update Vector specifying how many samples were used between updates.\n+    \\param[in] ax Axis object to plot the lines in.\n+    """\n+\n+    # Find good number of horizontal update lines to plot\n+    updates = np.arange(0, len(n_samples_per_update))\n+    while len(n_samples_per_update) > 20:\n+        n_samples_per_update = n_samples_per_update[0:-1:5]\n+        updates = updates[0:-1:5]\n+\n+    y_limits = ax.get_ylim()\n+\n+    n = n_samples_per_update\n+    h = ax.plot([n, n], y_limits)\n+    plt.setp(h, color="#bbbbbb", linewidth=0.5, zorder=0)\n+        \n+    for ii in range(len(n_samples_per_update) - 1):\n+        y = y_limits[0] + 0.9 * (y_limits[1] - y_limits[0])\n+        ax.text(\n+            n_samples_per_update[ii + 1],\n+            y,\n+            str(updates[ii + 1]),\n+            horizontalalignment="center",\n+            verticalalignment="top",\n+            rotation="vertical",\n+        )\n+\n+    y = y_limits[0] + 0.95 * (y_limits[1] - y_limits[0])\n+    ax.text(\n+        mean(ax.get_xlim()),\n+        y,\n+        "number of updates",\n+        horizontalalignment="center",\n+        verticalalignment="top",\n+    )\n+    ax.set_ylim(y_limits)\n+\n+\n+def plotLearningCurve(learning_curve, **kwargs):\n+    """ Plot a learning curve.\n+    \\param[in] learing_curve A learning curve that has the following format\n+        #rows is number of optimization updates\n+        column 0: Number of samples at which the cost was evaluated\n+        column 1: The total cost \n+        column 2...: Individual cost components (column 1 is their sum)\n+    \\param[in] ax Axis to plot the learning curve on.\n+    \\param[in] costs_all Vector of costs of each sample (default=[])\n+    \\param[in] cost_labels Vector of strings for the different cost components  (default=[]).\n+    """\n+\n+    cost_labels = kwargs.get("cost_labels", None)\n+    ax = kwargs.get("ax") or plt.axes()\n+\n+    # Plot costs at evaluations\n+    learning_curve = np.array(learning_curve)\n+    samples_eval = learning_curve[:, 0]\n+    costs_eval = learning_curve[:, 1:]\n+    # Sum of cost components\n+    lines = ax.plot(samples_eval, costs_eval[:, 0], "-", color="black", linewidth=2)\n+    # Individual cost components\n+    if costs_eval.shape[1] > 1:\n+        lines.extend(ax.plot(samples_eval, costs_eval[:, 1:], "-", linewidth=1))\n+\n+    # Annotation\n+    ax.set_xlabel("number of evaluations")\n+    ax.set_ylabel("cost")\n+    ax.set_title("Learning curve")\n+\n+    if cost_labels:\n+        cost_labels.insert(0, "total cost")\n+        plt.legend(lines, cost_labels)\n+\n+    plotUpdateLines(samples_eval, ax)\n+\n+    return (lines, ax)\n+\n+\n+def plotExplorationCurve(exploration_curve, **kwargs):\n+    """ Plot an exploration curve.\n+    \\param[in] exploration An exploration curve that has the following format\n+        #rows is number of optimization updates\n+        column 0: Number of samples at which the cost was evaluated\n+        column 1: The exploration at that update \n+    \\param[in] ax Axis to plot the learning curve on.\n+    """\n+    ax = kwargs.get("ax") or plt.axes()\n+\n+    exploration_curve = np.array(exploration_curve)\n+    samples_eval = exploration_curve[:, 0]\n+    explo = exploration_curve[:, 1]\n+\n+    line = ax.plot(samples_eval, explo, "-", color="green", linewidth=2)\n+    plotUpdateLines(samples_eval, ax)\n+    ax.set_xlabel("number of evaluations")\n+    ax.set_ylabel("sqrt of max. eigval of covar")\n+    ax.set_title("Exploration magnitude")\n+\n+    return (line, ax)\n+\n+\n+def setColor(handle, i_update, n_updates):\n+    """ Set the color of an object, according to how far the optimization has proceeded.\n+    \\param[in] handle Handle to the object\n+    \\param[in] i_update Which update is the optimization at now?\n+    \\param[in] n_update Which is the number of updates the optimization will run?\n+    """\n+    if i_update == 0:\n+        plt.setp(handle, color="r", linewidth=2)\n+    else:\n+        c = 1.0 * i_update / n_updates\n+        cur_color = [0.0 - 0.0 * c, 0.0 + 1.0 * c, 0.0 - 0.0 * c]\n+        plt.setp(handle, color=cur_color)\n+\n+\n+def plotUpdate(\n+    distribution,\n+    cost_eval,\n+    samples,\n+    costs,\n+    weights,\n+    distribution_new,\n+    **kwargs):\n+    """ Save an optimization update to a directory.\n+    \\param[in] distribution Gaussian distribution before the update\n+    \\param[in] cost_eval Cost of the mean of the distribution \n+    \\param[in] samples The samples in the search space\n+    \\param[in] costs The cost of each sample \n+    \\param[in] weights The weight of each sample \n+    \\param[in] distribution_new Gaussian distribution after the update\n+    \\param[in] ax Axis to plot the update on.\n+    \\param[in] highlight Wether to highlight this update (default=False)\n+    \\param[in] plot_samples Whether to plot the individual samples (default=False)\n+    """\n+    ax = kwargs.get("ax") or plt.axes()\n+    highlight = kwargs.get("highlight",False)\n+    plot_samples = kwargs.get("plot_samples",False)\n+\n+    if samples is None:\n+        plot_samples = False\n+\n+    n_dims = len(distribution.mean)\n+    if n_dims == 1:\n+        raise ValueError("Sorry, only know how to plot for n_dims==2, but you provided n_dims==1")\n+\n+    # ZZZ Take into consideration block_covar_sizes to plot sub blocks in\n+    # different subplots\n+\n+    if n_dims >= 2:\n+        distr_mean = distribution.mean[0:2]\n+        distr_covar = distribution.covar[0:2, 0:2]\n+        distr_new_mean = distribution_new.mean[0:2]\n+        distr_new_covar = distribution_new.covar[0:2, 0:2]\n+        if samples is not None:\n+            samples = samples[:, 0:2]\n+\n+        if plot_samples:\n+            max_marker_size = 80\n+            for ii in range(len(weights)):\n+                cur_marker_size = max_marker_size * weights[ii]\n+                sample_handle = ax.plot(\n+                    samples[ii, 0], samples[ii, 1], "o", color="green"\n+                )\n+                plt.setp(\n+                    sample_handle,\n+                    markersize=cur_marker_size,\n+                    markerfacecolor=(0.5, 0.8, 0.5),\n+                    markeredgecolor="none",\n+                )\n+\n+                ax.plot(samples[:, 0], samples[:, 1], ".", color="black")\n+            ax.plot(\n+                (distr_mean[0], distr_new_mean[0]),\n+                (distr_mean[1], distr_new_mean[1]),\n+                "-",\n+                color="blue",\n+            )\n+\n+        mean_handle = ax.plot(distr_mean[0], distr_mean[1], "o", label="old")\n+        mean_handle_new = ax.plot(\n+            distr_new_mean[0], distr_new_mean[1], "o", label="new"\n+        )\n+        mean_handle_link = ax.plot(\n+            [distr_mean[0], distr_new_mean[0]], [distr_mean[1], distr_new_mean[1]], "-"\n+        )\n+        patch = plot_error_ellipse(distr_mean, distr_covar, ax)\n+        patch_new = plot_error_ellipse(distr_new_mean, distr_new_covar, ax)\n+        if highlight:\n+            plt.setp(mean_handle, color="red")\n+            plt.setp(mean_handle_new, color="blue")\n+            plt.setp(patch, edgecolor="red")\n+            plt.setp(patch_new, edgecolor="blue")\n+        else:\n+            plt.setp(mean_handle, color="gray")\n+            plt.setp(mean_handle_new, color="gray")\n+            plt.setp(patch, edgecolor="gray")\n+            plt.setp(patch_new, edgecolor="gray")\n+        plt.setp(mean_handle_link, color="gray")\n+    ax.set_aspect("equal")\n+\n+    # plt.rcParams[\'text.usetex\']=True\n+    # ax.set_xlabel(r\'$\\theta_1$\')\n+    # ax.set_ylabel(r\'$\\theta_2$\')\n+    ax.set_xlabel("dim 1 (of " + str(n_dims) + ")")\n+    ax.set_ylabel("dim 2 (of " + str(n_dims) + ")")\n+    ax.set_title("Search space")\n+    return mean_handle, mean_handle_new, patch, patch_new\n+\n+\n+class LearningSession:\n+    def __init__(self, n_samples_per_update, root_directory=None, **kwargs):\n+        self._n_samples_per_update = n_samples_per_update\n+        self._root_dir = root_directory\n+        self._cache = {}\n+        self._last_update_added = None\n+\n+        self.tell(n_samples_per_update, "n_samples_per_update")\n+        for name, value in kwargs.items():\n+            self.tell(value, name)\n+            \n+        if root_directory and \'cost_function\' in kwargs:\n+            cost_function = kwargs[\'cost_function\']\n+            src = inspect.getsourcelines(cost_function.__class__)\n+            src = " ".join(src[0])\n+            src = src.replace("(CostFunction)", "")\n+            filename = os.path.join(root_directory, "cost_function.py")\n+            with open(filename, "w") as f:\n+                f.write(src)\n+                \n+    @classmethod\n+    def from_dir(cls, directory):\n+        filename = os.path.join(directory, "n_samples_per_update.txt")\n+        n_samples_per_update = np.atleast_1d(np.loadtxt(filename))\n+        n_samples_per_update = int(n_samples_per_update[0])\n+        return cls(n_samples_per_update, directory)\n+\n+    def addEval(self, i_update, eval_sample, eval_cost):\n+        self.tell(eval_cost, "eval_cost", i_update)\n+        self.tell(eval_sample, "eval_sample", i_update)\n+\n+    def addUpdate(\n+        self, i_update, distribution, samples, costs, weights, distribution_new=None\n+    ):\n+        self.tell(distribution, "distribution", i_update)\n+        self.tell(samples, "samples", i_update)\n+        self.tell(costs, "costs", i_update)\n+        self.tell(weights, "weights", i_update)\n+        if distribution_new:\n+            self.tell(distribution_new, "distribution_new", i_update)\n+        self._last_update_added = i_update\n+\n+    def getNUpdates(self):\n+        if not self._root_dir:\n+            return self._last_update_added\n+        else:\n+            update_dirs = sorted(glob(os.path.join(self._root_dir, "update*/")))\n+            last_dir = update_dirs[-1]\n+            last_dir = last_dir.replace("update", "")\n+            last_dir = last_dir.replace(self._root_dir, "")\n+            last_dir = last_dir.replace("/", "")\n+            return int(last_dir)\n+\n+    def getBaseName(self, name, i_update=None, i_sample=None):\n+\n+        if isinstance(i_sample, int):\n+            basename = f"{i_sample:03d}_{name}"  # e.g. 0 => "000_name"\n+        elif i_sample:\n+            basename = f"{i_sample}_{name}"  # e.g. "eval" => "eval_name"\n+        else:\n+            basename = name\n+\n+        if isinstance(i_update, int):\n+            # e.g. "update00023/000_name"\n+            basename = os.path.join(f"update{i_update:05d}", basename)\n+\n+        return basename\n+\n+    def exists(self, name, i_update=None, i_sample=None):\n+        try:\n+            obj = self.ask(name, i_update, i_sample)\n+        except Exception:\n+            return False\n+        return True\n+\n+    def ask(self, name, i_update=None, i_sample=None):\n+\n+        basename = self.getBaseName(name, i_update, i_sample)\n+\n+        if basename in self._cache:\n+            return copy.deepcopy(self._cache[basename])\n+\n+        if not self._root_dir:\n+            # Not in cache, and no directory given: Error\n+            raise KeyError("Could not find basename in cache: " + basename)\n+\n+        else:\n+            abs_basename = os.path.join(self._root_dir, basename)\n+            if os.path.isfile(abs_basename + ".json"):\n+                with open(abs_basename + ".json", "r") as f:\n+                    obj = jsonpickle.decode(f.read())\n+\n+            elif os.path.isfile(abs_basename + ".txt"):\n+                obj = np.loadtxt(abs_basename + ".txt")\n+\n+            else:\n+                raise IOError("Could not find file with basename: " + abs_basename)\n+\n+            self._cache[basename] = obj\n+\n+            return obj\n+\n+    def tell(self, obj, name, i_update=None, i_sample=None):\n+\n+        basename = self.getBaseName(name, i_update, i_sample)\n+\n+        self._cache[basename] = copy.deepcopy(obj)\n+\n+        if self._root_dir:\n+            abs_basename = os.path.join(self._root_dir, basename)\n+\n+            # Make sure directory exists\n+            os.makedirs(os.path.dirname(abs_basename), exist_ok=True)\n+\n+            if isinstance(obj, list) or isinstance(obj, int):\n+                # Try to convert list to numpy array\n+                obj = np.atleast_1d(np.array(obj))\n+\n+            if isinstance(obj, np.ndarray):\n+                filename = abs_basename + ".txt"\n+                np.savetxt(filename, obj)\n+            else:\n+                filename = abs_basename + ".json"\n+                with open(filename, "w") as out_file:\n+                    out_file.write(to_jsonpickle(obj))\n+\n+            return filename\n+\n+    def getEvalCosts(self, i_update):\n+        has_eval = self.exists("cost", 0, "eval")\n+        if has_eval:\n+            cost_eval = self.ask("cost", i_update, "eval")\n+            # Evaluation at the beginning of an update\n+            i_sample = i_update * self._n_samples_per_update\n+        else:\n+            update_costs = self.getSampleCosts(i_update)\n+            cost_eval = np.mean(update_costs, axis=0)\n+            # Evaluation between two updates (mean of all samples)\n+            i_sample = (i_update + 0.5) * self._n_samples_per_update\n+        return cost_eval, i_sample\n+\n+    def getSampleCosts(self, i_update):\n+        if self.exists("costs", i_update):\n+            # Everything in one file\n+            sample_costs = self.ask("costs", i_update)\n+        else:\n+            # One file for each sample\n+            n = self._n_samples_per_update\n+            sample_costs = [self.ask("costs", i_update, s) for s in range(n)]\n+        return np.array(sample_costs)\n+\n+    def getLearningCurve(self):\n+        learning_curve = []\n+        n_updates = self.getNUpdates()\n+        for i_update in range(n_updates):\n+            costs_eval, i_sample = self.getEvalCosts(i_update)\n+            costs_eval = np.atleast_1d(costs_eval)\n+            learning_curve.append(np.concatenate(([i_sample], costs_eval)))\n+\n+        cost_labels = []\n+        if self.exists(\'cost_function\'):\n+            cost_function = self.ask(\'cost_function\')\n+            cost_labels = cost_function.costLabels();\n+        elif self.exists(\'task\'):\n+            cost_function = self.ask(\'task\')\n+            cost_labels = cost_function.costLabels();\n+\n+        return (learning_curve, cost_labels)\n+\n+    def plotLearningCurve(self, ax=None):\n+        """\n+        has eval or not\n+        has cost components or not\n+        saved in costs or 000_cost...\n+        """\n+        if not ax:\n+            ax = plt.axes()\n+\n+        # Plot costs of individual samples\n+        costs = []\n+        for i_update in range(self.getNUpdates()):\n+            sample_costs = self.getSampleCosts(i_update)\n+            costs.extend([c[0] for c in sample_costs]) # Only include sum\n+        ax.plot(costs, ".", color="gray")\n+\n+        # Plot learning curve itself\n+        learning_curve, cost_labels = self.getLearningCurve()\n+        return plotLearningCurve(learning_curve, ax=ax, cost_labels=cost_labels)\n+\n+    def plotExplorationCurve(self, ax=None):\n+        n_updates = self.getNUpdates()\n+        curve = np.zeros((n_updates, 2))\n+        for i_update in range(n_updates):\n+            distribution = self.ask("distribution", i_update)\n+            cur_exploration = np.sqrt(distribution.maxEigenValue())\n+            curve[i_update, 0] = i_update * self._n_samples_per_update\n+            curve[i_update, 1] = cur_exploration\n+\n+        if not ax:\n+            ax = plt.axes()\n+        return plotExplorationCurve(curve, ax=ax)\n+\n+    def plotDistributionUpdates(self, ax=None):\n+        if not ax:\n+            ax = plt.axes()\n+\n+        n_updates = self.getNUpdates()\n+        for i_update in range(n_updates):\n+            highlight = i_update == 0\n+            self.plotDistributionUpdate(i_update, ax, highlight=highlight)\n+        return ax # zzz\n+\n+    def plotDistributionUpdate(self, i_update, ax=None, **kwargs):\n+        #ax = kwargs.get("ax") or plt.axes()\n+        highlight = kwargs.get("highlight",False)\n+        plot_samples = kwargs.get("plot_samples",False)\n+        \n+        distribution = self.ask("distribution", i_update)\n+        costs_eval = self.getEvalCosts(i_update)\n+        costs = self.getSampleCosts(i_update)\n+        samples = self.ask("samples", i_update)\n+        weights = self.ask("weights", i_update)\n+        distribution_new = self.ask("distribution_new", i_update)\n+\n+        if not ax:\n+            ax = plt.axes()\n+            \n+        return plotUpdate(\n+            distribution,\n+            costs_eval,\n+            samples,\n+            costs,\n+            weights,\n+            distribution_new,\n+            ax=ax,\n+            highlight=highlight,\n+            plot_samples=plot_samples,\n+        )\n+\n+    def plot(self, fig=None):\n+        if not fig:\n+            fig = plt.figure(figsize=(15, 5))\n+        axs = [fig.add_subplot(131 + sp) for sp in range(3)]\n+        self.plotDistributionUpdates(axs[0])\n+        self.plotExplorationCurve(axs[1])\n+        self.plotLearningCurve(axs[2])\n+        return fig\n+\n+    @staticmethod\n+    def plotMultiple(sessions):\n+        pass\n+\n+\n+if __name__ == "__main__":\n+\n+    parser = argparse.ArgumentParser()\n+    parser.add_argument(\n+        "directory", type=str, help="directory from which to read the learning session"\n+    )\n+    args = parser.parse_args()\n+\n+    session = LearningSession.from_dir(args.directory)\n+    session.plot()\n+    plt.show()\ndiff --git a/python/dmp_bbo/LearningSessionTask.py b/python/dmp_bbo/LearningSessionTask.py\nindex 335f392..6d103c4 100644\n--- a/python/dmp_bbo/LearningSessionTask.py\n+++ b/python/dmp_bbo/LearningSessionTask.py\n@@ -25,6 +25,7 @@ from glob import glob\n import matplotlib.pyplot as plt\n from pylab import mean\n from matplotlib.patches import Ellipse\n+import inspect\n \n \n lib_path = os.path.abspath("../../python/")\n@@ -39,6 +40,14 @@ class LearningSessionTask(LearningSession):\n         super().__init__(n_samples_per_update, directory, **kwargs)\n         self.task_ = kwargs.get("task", None)\n         self.task_solver_ = kwargs.get("task_solver", None)\n+        \n+        if directory:\n+            src = inspect.getsourcelines(self.task_.__class__)\n+            src = " ".join(src[0])\n+            src = src.replace("(Task)", "")\n+            filename = os.path.join(directory, "task.py")\n+            with open(filename, "w") as f:\n+                f.write(src)\n \n     def addRollout(self, i_update, i_sample, sample, cost_vars, cost):\n         self.tell(sample, "sample", i_update, i_sample)\n', 'commit_subject': 'LearningSession', 'repo_name': 'stulp/dmpbbo', 'committer': {'name': 'Freek Stulp', 'email': '68f227b90bfee9d6edcfa839ec768d905c9d219d@dlr.de', 'time_sec': '1650805510', 'tz_offset': '120', 'date': {'seconds': '1650805510'}}, 'commit_tree': '20c7784ff564587746447aeab2e09f7a9d338be7', 'commit_message': 'LearningSession\n', 'language_name': 'Shell', 'license': 'lgpl-2.1', 'hunks': '@@ -0,0 +1,560 @@\n+# This file is part of DmpBbo, a set of libraries and programs for the\n+# black-box optimization of dynamical movement primitives.\n+# Copyright (C) 2022 Freek Stulp\n+#\n+# DmpBbo is free software: you can redistribute it and/or modify\n+# it under the terms of the GNU Lesser General Public License as published by\n+# the Free Software Foundation, either version 2 of the License, or\n+# (at your option) any later version.\n+#\n+# DmpBbo is distributed in the hope that it will be useful,\n+# but WITHOUT ANY WARRANTY; without even the implied warranty of\n+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the\n+# GNU Lesser General Public License for more details.\n+#\n+# You should have received a copy of the GNU Lesser General Public License\n+# along with DmpBbo.  If not, see <http://www.gnu.org/licenses/>.\n+\n+\n+import numpy as np\n+import math\n+import os\n+import sys\n+import copy\n+import argparse\n+import inspect\n+from glob import glob\n+import matplotlib.pyplot as plt\n+from pylab import mean\n+from matplotlib.patches import Ellipse\n+import warnings\n+\n+# Avoid warnings when plotting very narrow covariance matrices\n+warnings.simplefilter("ignore", np.ComplexWarning)\n+\n+\n+lib_path = os.path.abspath("../../python/")\n+sys.path.append(lib_path)\n+\n+from to_jsonpickle import *\n+from bbo.DistributionGaussian import DistributionGaussian\n+\n+\n+def plot_error_ellipse(mu, cov, ax=None, **kwargs):\n+    """\n+Plot the error ellipse at a point given it\'s covariance matrix\n+\n+Parameters\n+----------\n+mu : array (2,)\n+The center of the ellipse\n+\n+cov : array (2,2)\n+The covariance matrix for the point\n+\n+ax : matplotlib.Axes, optional\n+The axis to overplot on\n+\n+**kwargs : dict\n+These keywords are passed to matplotlib.patches.Ellipse\n+\n+From https://github.com/dfm/dfmplot/blob/master/dfmplot/ellipse.py\n+\n+"""\n+    # some sane defaults\n+    facecolor = kwargs.pop("facecolor", "none")\n+    edgecolor = kwargs.pop("edgecolor", "k")\n+\n+    x, y = mu\n+    U, S, V = np.linalg.svd(cov)\n+    theta = np.degrees(np.arctan2(U[1, 0], U[0, 0]))\n+    ellipsePlot = Ellipse(\n+        xy=[x, y],\n+        width=2 * np.sqrt(S[0]),\n+        height=2 * np.sqrt(S[1]),\n+        angle=theta,\n+        facecolor=facecolor,\n+        edgecolor=edgecolor,\n+        **kwargs,\n+    )\n+\n+    if not ax:\n+        ax = plt.gca()\n+    lines = ax.add_patch(ellipsePlot)\n+    return lines\n+\n+\n+def plotUpdateLines(n_samples_per_update, ax):\n+    """ Plot vertical lines when an parameter update occured during the optimization.\n+    \\param[in] n_samples_per_update Vector specifying how many samples were used between updates.\n+    \\param[in] ax Axis object to plot the lines in.\n+    """\n+\n+    # Find good number of horizontal update lines to plot\n+    updates = np.arange(0, len(n_samples_per_update))\n+    while len(n_samples_per_update) > 20:\n+        n_samples_per_update = n_samples_per_update[0:-1:5]\n+        updates = updates[0:-1:5]\n+\n+    y_limits = ax.get_ylim()\n+\n+    n = n_samples_per_update\n+    h = ax.plot([n, n], y_limits)\n+    plt.setp(h, color="#bbbbbb", linewidth=0.5, zorder=0)\n+        \n+    for ii in range(len(n_samples_per_update) - 1):\n+        y = y_limits[0] + 0.9 * (y_limits[1] - y_limits[0])\n+        ax.text(\n+            n_samples_per_update[ii + 1],\n+            y,\n+            str(updates[ii + 1]),\n+            horizontalalignment="center",\n+            verticalalignment="top",\n+            rotation="vertical",\n+        )\n+\n+    y = y_limits[0] + 0.95 * (y_limits[1] - y_limits[0])\n+    ax.text(\n+        mean(ax.get_xlim()),\n+        y,\n+        "number of updates",\n+        horizontalalignment="center",\n+        verticalalignment="top",\n+    )\n+    ax.set_ylim(y_limits)\n+\n+\n+def plotLearningCurve(learning_curve, **kwargs):\n+    """ Plot a learning curve.\n+    \\param[in] learing_curve A learning curve that has the following format\n+        #rows is number of optimization updates\n+        column 0: Number of samples at which the cost was evaluated\n+        column 1: The total cost \n+        column 2...: Individual cost components (column 1 is their sum)\n+    \\param[in] ax Axis to plot the learning curve on.\n+    \\param[in] costs_all Vector of costs of each sample (default=[])\n+    \\param[in] cost_labels Vector of strings for the different cost components  (default=[]).\n+    """\n+\n+    cost_labels = kwargs.get("cost_labels", None)\n+    ax = kwargs.get("ax") or plt.axes()\n+\n+    # Plot costs at evaluations\n+    learning_curve = np.array(learning_curve)\n+    samples_eval = learning_curve[:, 0]\n+    costs_eval = learning_curve[:, 1:]\n+    # Sum of cost components\n+    lines = ax.plot(samples_eval, costs_eval[:, 0], "-", color="black", linewidth=2)\n+    # Individual cost components\n+    if costs_eval.shape[1] > 1:\n+        lines.extend(ax.plot(samples_eval, costs_eval[:, 1:], "-", linewidth=1))\n+\n+    # Annotation\n+    ax.set_xlabel("number of evaluations")\n+    ax.set_ylabel("cost")\n+    ax.set_title("Learning curve")\n+\n+    if cost_labels:\n+        cost_labels.insert(0, "total cost")\n+        plt.legend(lines, cost_labels)\n+\n+    plotUpdateLines(samples_eval, ax)\n+\n+    return (lines, ax)\n+\n+\n+def plotExplorationCurve(exploration_curve, **kwargs):\n+    """ Plot an exploration curve.\n+    \\param[in] exploration An exploration curve that has the following format\n+        #rows is number of optimization updates\n+        column 0: Number of samples at which the cost was evaluated\n+        column 1: The exploration at that update \n+    \\param[in] ax Axis to plot the learning curve on.\n+    """\n+    ax = kwargs.get("ax") or plt.axes()\n+\n+    exploration_curve = np.array(exploration_curve)\n+    samples_eval = exploration_curve[:, 0]\n+    explo = exploration_curve[:, 1]\n+\n+    line = ax.plot(samples_eval, explo, "-", color="green", linewidth=2)\n+    plotUpdateLines(samples_eval, ax)\n+    ax.set_xlabel("number of evaluations")\n+    ax.set_ylabel("sqrt of max. eigval of covar")\n+    ax.set_title("Exploration magnitude")\n+\n+    return (line, ax)\n+\n+\n+def setColor(handle, i_update, n_updates):\n+    """ Set the color of an object, according to how far the optimization has proceeded.\n+    \\param[in] handle Handle to the object\n+    \\param[in] i_update Which update is the optimization at now?\n+    \\param[in] n_update Which is the number of updates the optimization will run?\n+    """\n+    if i_update == 0:\n+        plt.setp(handle, color="r", linewidth=2)\n+    else:\n+        c = 1.0 * i_update / n_updates\n+        cur_color = [0.0 - 0.0 * c, 0.0 + 1.0 * c, 0.0 - 0.0 * c]\n+        plt.setp(handle, color=cur_color)\n+\n+\n+def plotUpdate(\n+    distribution,\n+    cost_eval,\n+    samples,\n+    costs,\n+    weights,\n+    distribution_new,\n+    **kwargs):\n+    """ Save an optimization update to a directory.\n+    \\param[in] distribution Gaussian distribution before the update\n+    \\param[in] cost_eval Cost of the mean of the distribution \n+    \\param[in] samples The samples in the search space\n+    \\param[in] costs The cost of each sample \n+    \\param[in] weights The weight of each sample \n+    \\param[in] distribution_new Gaussian distribution after the update\n+    \\param[in] ax Axis to plot the update on.\n+    \\param[in] highlight Wether to highlight this update (default=False)\n+    \\param[in] plot_samples Whether to plot the individual samples (default=False)\n+    """\n+    ax = kwargs.get("ax") or plt.axes()\n+    highlight = kwargs.get("highlight",False)\n+    plot_samples = kwargs.get("plot_samples",False)\n+\n+    if samples is None:\n+        plot_samples = False\n+\n+    n_dims = len(distribution.mean)\n+    if n_dims == 1:\n+        raise ValueError("Sorry, only know how to plot for n_dims==2, but you provided n_dims==1")\n+\n+    # ZZZ Take into consideration block_covar_sizes to plot sub blocks in\n+    # different subplots\n+\n+    if n_dims >= 2:\n+        distr_mean = distribution.mean[0:2]\n+        distr_covar = distribution.covar[0:2, 0:2]\n+        distr_new_mean = distribution_new.mean[0:2]\n+        distr_new_covar = distribution_new.covar[0:2, 0:2]\n+        if samples is not None:\n+            samples = samples[:, 0:2]\n+\n+        if plot_samples:\n+            max_marker_size = 80\n+            for ii in range(len(weights)):\n+                cur_marker_size = max_marker_size * weights[ii]\n+                sample_handle = ax.plot(\n+                    samples[ii, 0], samples[ii, 1], "o", color="green"\n+                )\n+                plt.setp(\n+                    sample_handle,\n+                    markersize=cur_marker_size,\n+                    markerfacecolor=(0.5, 0.8, 0.5),\n+                    markeredgecolor="none",\n+                )\n+\n+                ax.plot(samples[:, 0], samples[:, 1], ".", color="black")\n+            ax.plot(\n+                (distr_mean[0], distr_new_mean[0]),\n+                (distr_mean[1], distr_new_mean[1]),\n+                "-",\n+                color="blue",\n+            )\n+\n+        mean_handle = ax.plot(distr_mean[0], distr_mean[1], "o", label="old")\n+        mean_handle_new = ax.plot(\n+            distr_new_mean[0], distr_new_mean[1], "o", label="new"\n+        )\n+        mean_handle_link = ax.plot(\n+            [distr_mean[0], distr_new_mean[0]], [distr_mean[1], distr_new_mean[1]], "-"\n+        )\n+        patch = plot_error_ellipse(distr_mean, distr_covar, ax)\n+        patch_new = plot_error_ellipse(distr_new_mean, distr_new_covar, ax)\n+        if highlight:\n+            plt.setp(mean_handle, color="red")\n+            plt.setp(mean_handle_new, color="blue")\n+            plt.setp(patch, edgecolor="red")\n+            plt.setp(patch_new, edgecolor="blue")\n+        else:\n+            plt.setp(mean_handle, color="gray")\n+            plt.setp(mean_handle_new, color="gray")\n+            plt.setp(patch, edgecolor="gray")\n+            plt.setp(patch_new, edgecolor="gray")\n+        plt.setp(mean_handle_link, color="gray")\n+    ax.set_aspect("equal")\n+\n+    # plt.rcParams[\'text.usetex\']=True\n+    # ax.set_xlabel(r\'$\\theta_1$\')\n+    # ax.set_ylabel(r\'$\\theta_2$\')\n+    ax.set_xlabel("dim 1 (of " + str(n_dims) + ")")\n+    ax.set_ylabel("dim 2 (of " + str(n_dims) + ")")\n+    ax.set_title("Search space")\n+    return mean_handle, mean_handle_new, patch, patch_new\n+\n+\n+class LearningSession:\n+    def __init__(self, n_samples_per_update, root_directory=None, **kwargs):\n+        self._n_samples_per_update = n_samples_per_update\n+        self._root_dir = root_directory\n+        self._cache = {}\n+        self._last_update_added = None\n+\n+        self.tell(n_samples_per_update, "n_samples_per_update")\n+        for name, value in kwargs.items():\n+            self.tell(value, name)\n+            \n+        if root_directory and \'cost_function\' in kwargs:\n+            cost_function = kwargs[\'cost_function\']\n+            src = inspect.getsourcelines(cost_function.__class__)\n+            src = " ".join(src[0])\n+            src = src.replace("(CostFunction)", "")\n+            filename = os.path.join(root_directory, "cost_function.py")\n+            with open(filename, "w") as f:\n+                f.write(src)\n+                \n+    @classmethod\n+    def from_dir(cls, directory):\n+        filename = os.path.join(directory, "n_samples_per_update.txt")\n+        n_samples_per_update = np.atleast_1d(np.loadtxt(filename))\n+        n_samples_per_update = int(n_samples_per_update[0])\n+        return cls(n_samples_per_update, directory)\n+\n+    def addEval(self, i_update, eval_sample, eval_cost):\n+        self.tell(eval_cost, "eval_cost", i_update)\n+        self.tell(eval_sample, "eval_sample", i_update)\n+\n+    def addUpdate(\n+        self, i_update, distribution, samples, costs, weights, distribution_new=None\n+    ):\n+        self.tell(distribution, "distribution", i_update)\n+        self.tell(samples, "samples", i_update)\n+        self.tell(costs, "costs", i_update)\n+        self.tell(weights, "weights", i_update)\n+        if distribution_new:\n+            self.tell(distribution_new, "distribution_new", i_update)\n+        self._last_update_added = i_update\n+\n+    def getNUpdates(self):\n+        if not self._root_dir:\n+            return self._last_update_added\n+        else:\n+            update_dirs = sorted(glob(os.path.join(self._root_dir, "update*/")))\n+            last_dir = update_dirs[-1]\n+            last_dir = last_dir.replace("update", "")\n+            last_dir = last_dir.replace(self._root_dir, "")\n+            last_dir = last_dir.replace("/", "")\n+            return int(last_dir)\n+\n+    def getBaseName(self, name, i_update=None, i_sample=None):\n+\n+        if isinstance(i_sample, int):\n+            basename = f"{i_sample:03d}_{name}"  # e.g. 0 => "000_name"\n+        elif i_sample:\n+            basename = f"{i_sample}_{name}"  # e.g. "eval" => "eval_name"\n+        else:\n+            basename = name\n+\n+        if isinstance(i_update, int):\n+            # e.g. "update00023/000_name"\n+            basename = os.path.join(f"update{i_update:05d}", basename)\n+\n+        return basename\n+\n+    def exists(self, name, i_update=None, i_sample=None):\n+        try:\n+            obj = self.ask(name, i_update, i_sample)\n+        except Exception:\n+            return False\n+        return True\n+\n+    def ask(self, name, i_update=None, i_sample=None):\n+\n+        basename = self.getBaseName(name, i_update, i_sample)\n+\n+        if basename in self._cache:\n+            return copy.deepcopy(self._cache[basename])\n+\n+        if not self._root_dir:\n+            # Not in cache, and no directory given: Error\n+            raise KeyError("Could not find basename in cache: " + basename)\n+\n+        else:\n+            abs_basename = os.path.join(self._root_dir, basename)\n+            if os.path.isfile(abs_basename + ".json"):\n+                with open(abs_basename + ".json", "r") as f:\n+                    obj = jsonpickle.decode(f.read())\n+\n+            elif os.path.isfile(abs_basename + ".txt"):\n+                obj = np.loadtxt(abs_basename + ".txt")\n+\n+            else:\n+                raise IOError("Could not find file with basename: " + abs_basename)\n+\n+            self._cache[basename] = obj\n+\n+            return obj\n+\n+    def tell(self, obj, name, i_update=None, i_sample=None):\n+\n+        basename = self.getBaseName(name, i_update, i_sample)\n+\n+        self._cache[basename] = copy.deepcopy(obj)\n+\n+        if self._root_dir:\n+            abs_basename = os.path.join(self._root_dir, basename)\n+\n+            # Make sure directory exists\n+            os.makedirs(os.path.dirname(abs_basename), exist_ok=True)\n+\n+            if isinstance(obj, list) or isinstance(obj, int):\n+                # Try to convert list to numpy array\n+                obj = np.atleast_1d(np.array(obj))\n+\n+            if isinstance(obj, np.ndarray):\n+                filename = abs_basename + ".txt"\n+                np.savetxt(filename, obj)\n+            else:\n+                filename = abs_basename + ".json"\n+                with open(filename, "w") as out_file:\n+                    out_file.write(to_jsonpickle(obj))\n+\n+            return filename\n+\n+    def getEvalCosts(self, i_update):\n+        has_eval = self.exists("cost", 0, "eval")\n+        if has_eval:\n+            cost_eval = self.ask("cost", i_update, "eval")\n+            # Evaluation at the beginning of an update\n+            i_sample = i_update * self._n_samples_per_update\n+        else:\n+            update_costs = self.getSampleCosts(i_update)\n+            cost_eval = np.mean(update_costs, axis=0)\n+            # Evaluation between two updates (mean of all samples)\n+            i_sample = (i_update + 0.5) * self._n_samples_per_update\n+        return cost_eval, i_sample\n+\n+    def getSampleCosts(self, i_update):\n+        if self.exists("costs", i_update):\n+            # Everything in one file\n+            sample_costs = self.ask("costs", i_update)\n+        else:\n+            # One file for each sample\n+            n = self._n_samples_per_update\n+            sample_costs = [self.ask("costs", i_update, s) for s in range(n)]\n+        return np.array(sample_costs)\n+\n+    def getLearningCurve(self):\n+        learning_curve = []\n+        n_updates = self.getNUpdates()\n+        for i_update in range(n_updates):\n+            costs_eval, i_sample = self.getEvalCosts(i_update)\n+            costs_eval = np.atleast_1d(costs_eval)\n+            learning_curve.append(np.concatenate(([i_sample], costs_eval)))\n+\n+        cost_labels = []\n+        if self.exists(\'cost_function\'):\n+            cost_function = self.ask(\'cost_function\')\n+            cost_labels = cost_function.costLabels();\n+        elif self.exists(\'task\'):\n+            cost_function = self.ask(\'task\')\n+            cost_labels = cost_function.costLabels();\n+\n+        return (learning_curve, cost_labels)\n+\n+    def plotLearningCurve(self, ax=None):\n+        """\n+        has eval or not\n+        has cost components or not\n+        saved in costs or 000_cost...\n+        """\n+        if not ax:\n+            ax = plt.axes()\n+\n+        # Plot costs of individual samples\n+        costs = []\n+        for i_update in range(self.getNUpdates()):\n+            sample_costs = self.getSampleCosts(i_update)\n+            costs.extend([c[0] for c in sample_costs]) # Only include sum\n+        ax.plot(costs, ".", color="gray")\n+\n+        # Plot learning curve itself\n+        learning_curve, cost_labels = self.getLearningCurve()\n+        return plotLearningCurve(learning_curve, ax=ax, cost_labels=cost_labels)\n+\n+    def plotExplorationCurve(self, ax=None):\n+        n_updates = self.getNUpdates()\n+        curve = np.zeros((n_updates, 2))\n+        for i_update in range(n_updates):\n+            distribution = self.ask("distribution", i_update)\n+            cur_exploration = np.sqrt(distribution.maxEigenValue())\n+            curve[i_update, 0] = i_update * self._n_samples_per_update\n+            curve[i_update, 1] = cur_exploration\n+\n+        if not ax:\n+            ax = plt.axes()\n+        return plotExplorationCurve(curve, ax=ax)\n+\n+    def plotDistributionUpdates(self, ax=None):\n+        if not ax:\n+            ax = plt.axes()\n+\n+        n_updates = self.getNUpdates()\n+        for i_update in range(n_updates):\n+            highlight = i_update == 0\n+            self.plotDistributionUpdate(i_update, ax, highlight=highlight)\n+        return ax # zzz\n+\n+    def plotDistributionUpdate(self, i_update, ax=None, **kwargs):\n+        #ax = kwargs.get("ax") or plt.axes()\n+        highlight = kwargs.get("highlight",False)\n+        plot_samples = kwargs.get("plot_samples",False)\n+        \n+        distribution = self.ask("distribution", i_update)\n+        costs_eval = self.getEvalCosts(i_update)\n+        costs = self.getSampleCosts(i_update)\n+        samples = self.ask("samples", i_update)\n+        weights = self.ask("weights", i_update)\n+        distribution_new = self.ask("distribution_new", i_update)\n+\n+        if not ax:\n+            ax = plt.axes()\n+            \n+        return plotUpdate(\n+            distribution,\n+            costs_eval,\n+            samples,\n+            costs,\n+            weights,\n+            distribution_new,\n+            ax=ax,\n+            highlight=highlight,\n+            plot_samples=plot_samples,\n+        )\n+\n+    def plot(self, fig=None):\n+        if not fig:\n+            fig = plt.figure(figsize=(15, 5))\n+        axs = [fig.add_subplot(131 + sp) for sp in range(3)]\n+        self.plotDistributionUpdates(axs[0])\n+        self.plotExplorationCurve(axs[1])\n+        self.plotLearningCurve(axs[2])\n+        return fig\n+\n+    @staticmethod\n+    def plotMultiple(sessions):\n+        pass\n+\n+\n+if __name__ == "__main__":\n+\n+    parser = argparse.ArgumentParser()\n+    parser.add_argument(\n+        "directory", type=str, help="directory from which to read the learning session"\n+    )\n+    args = parser.parse_args()\n+\n+    session = LearningSession.from_dir(args.directory)\n+    session.plot()\n+    plt.show()\n', 'addition_count': 560, 'deletion_count': 0, 'src_file': '/dev/null', 'tgt_file': 'b/python/bbo/LearningSession.py', 'file_extension': '.py', 'before_file': 'ADDFILE', 'message': 'LearningSession\n', 'author_name': 'Freek Stulp', 'lang': 'py'}