Repository: wookayin/expt

Issue:
Feature Request: Plotting the median
Hi @wookayin, I have been using the `expt` package to do plotting, and the experience has been great. 

Any chance you could add support for plotting the **median**? I am trying to plot the human normalized score used in classical Atari literature, so it would be useful to plot the median instead of mean. As an example, the snippet below plots the mean

```python
# Costa: Note the data is not the same as Mnih et al., 2015
# Note the random agent score on Video Pinball is sometimes greater than the
# human score under other evaluation methods.
atari_human_normalized_scores = {
    'Alien-v5': (227.8, 7127.7),
    'Amidar-v5': (5.8, 1719.5),
    'Assault-v5': (222.4, 742.0),
    'Asterix-v5': (210.0, 8503.3),
    'Asteroids-v5': (719.1, 47388.7),
    'Atlantis-v5': (12850.0, 29028.1),
    'BankHeist-v5': (14.2, 753.1),
    'BattleZone-v5': (2360.0, 37187.5),
    'BeamRider-v5': (363.9, 16926.5),
    'Berzerk-v5': (123.7, 2630.4),
    'Bowling-v5': (23.1, 160.7),
    'Boxing-v5': (0.1, 12.1),
    'Breakout-v5': (1.7, 30.5),
    'Centipede-v5': (2090.9, 12017.0),
    'ChopperCommand-v5': (811.0, 7387.8),
    'CrazyClimber-v5': (10780.5, 35829.4),
    'Defender-v5': (2874.5, 18688.9),
    # 'DemonAttack-v5': (152.1, 1971.0),
    # 'DoubleDunk-v5': (-18.6, -16.4),
    # 'Enduro-v5': (0.0, 860.5),
    # 'FishingDerby-v5': (-91.7, -38.7),
    # 'Freeway-v5': (0.0, 29.6),
    # 'Frostbite-v5': (65.2, 4334.7),
    # 'Gopher-v5': (257.6, 2412.5),
    # 'Gravitar-v5': (173.0, 3351.4),
    # 'Hero-v5': (1027.0, 30826.4),
    # 'IceHockey-v5': (-11.2, 0.9),
    # 'Jamesbond-v5': (29.0, 302.8),
    # 'Kangaroo-v5': (52.0, 3035.0),
    # 'Krull-v5': (1598.0, 2665.5),
    # 'KungFuMaster-v5': (258.5, 22736.3),
    # 'MontezumaRevenge-v5': (0.0, 4753.3),
    # 'MsPacman-v5': (307.3, 6951.6),
    # 'NameThisGame-v5': (2292.3, 8049.0),
    # 'Phoenix-v5': (761.4, 7242.6),
    # 'Pitfall-v5': (-229.4, 6463.7),
    # 'Pong-v5': (-20.7, 14.6),
    # 'PrivateEye-v5': (24.9, 69571.3),
    # 'Qbert-v5': (163.9, 13455.0),
    # 'Riverraid-v5': (1338.5, 17118.0),
    # 'RoadRunner-v5': (11.5, 7845.0),
    # 'Robotank-v5': (2.2, 11.9),
    # 'Seaquest-v5': (68.4, 42054.7),
    # 'Skiing-v5': (-17098.1, -4336.9),
    # 'Solaris-v5': (1236.3, 12326.7),
    # 'SpaceInvaders-v5': (148.0, 1668.7),
    # 'StarGunner-v5': (664.0, 10250.0),
    # 'Surround-v5': (-10.0, 6.5),
    # 'Tennis-v5': (-23.8, -8.3),
    # 'TimePilot-v5': (3568.0, 5229.2),
    # 'Tutankham-v5': (11.4, 167.6),
    # 'UpNDown-v5': (533.4, 11693.2),
    # 'Venture-v5': (0.0, 1187.5),
    # 'VideoPinball-v5': (16256.9, 17667.9),
    # 'WizardOfWor-v5': (563.5, 4756.5),
    # 'YarsRevenge-v5': (3092.9, 54576.9),
    # 'Zaxxon-v5': (32.5, 9173.3),
}

import enum
from matplotlib import axis
import numpy as np
import expt
from expt import Run, Hypothesis, Experiment
import pandas as pd
import matplotlib.pyplot as plt
import wandb
import wandb.apis.reports as wb  # noqa
from expt.plot import GridPlot


def create_expt_runs(wandb_runs):
    runs = []
    for idx, run in enumerate(wandb_runs):
        wandb_run = run.history()
        if 'videos' in wandb_run:
            wandb_run = wandb_run.drop(columns=['videos'], axis=1)
        runs += [Run(f"seed{idx}", wandb_run)]
    return runs

api = wandb.Api()

env_ids = atari_human_normalized_scores.keys()

NUM_FRAME_STACK = 4
runss = []
for env_id in env_ids:
    api = wandb.Api()
    wandb_runs = api.runs(
        path="costa-huang/envpool-atari",
        filters={'$and': [{'config.env_id.value': env_id}, {'config.exp_name.value': 'ppo_atari_envpool_xla_jax'}]}
    )
    expt_runs = create_expt_runs(wandb_runs)

    # normalize scores and adjust x-axis from steps to frames
    for expt_run in expt_runs:
        expt_run.df["charts/avg_episodic_return"] = (
            expt_run.df["charts/avg_episodic_return"] - atari_human_normalized_scores[env_id][0]) / \
                (atari_human_normalized_scores[env_id][1] - atari_human_normalized_scores[env_id][0]
        )
        expt_run.df["global_step"] *= NUM_FRAME_STACK
    runss.extend(expt_runs)


h = Hypothesis("CleanRL's PPO + Envpool", runss)
fig, ax = plt.subplots(figsize=(4, 4))
g = h.plot(x='global_step', y="charts/avg_episodic_return", rolling=50, n_samples=400, legend=False, err_fn=lambda h: h.grouped.sem(), err_style="fill",
 suptitle="", ax=ax,)
ax.set_title("CleanRL's PPO + Envpool")
ax.yaxis.set_label_text("Human normalized score")
ax.xaxis.set_label_text("Frames")
plt.savefig("test.png", bbox_inches='tight')
```

<img width="401" alt="image" src="https://user-images.githubusercontent.com/5555347/179360171-1acfa6c0-55ea-48d5-8031-f03ae867ee55.png">


Comments:

Patch:
diff --git a/expt/plot.py b/expt/plot.py
index 657d7e5..5a09b36 100644
--- a/expt/plot.py
+++ b/expt/plot.py
@@ -21,6 +21,11 @@ warnings.filterwarnings("ignore", category=UserWarning,
                         message='Creating legend with loc="best"')
 # yapf: enable
 
+HypothesisSummaryFn = Callable[  # see HypothesisPlotter
+    [Hypothesis], pd.DataFrame]
+HypothesisSummaryErrFn = Callable[  # see HypothesisPlotter
+    [Hypothesis], Union[pd.DataFrame, Tuple[pd.DataFrame, pd.DataFrame]]]
+
 
 class GridPlot:
   """Multi-plot grid subplots.
@@ -269,7 +274,8 @@ class HypothesisPlotter:
                *args,
                subplots=True,
                err_style="runs",
-               err_fn: Optional[Callable[[Hypothesis], pd.DataFrame]] = None,
+               err_fn: Optional[HypothesisSummaryFn] = None,
+               representative_fn: Optional[HypothesisSummaryErrFn] = None,
                std_alpha=0.2,
                runs_alpha=0.2,
                n_samples=None,
@@ -309,13 +315,31 @@ class HypothesisPlotter:
             (i) runs, unit_traces: Show individual runs/traces (see runs_alpha)
             (ii) band, fill: Show as shaded area (see std_alpha)
             (iii) None or False: do not display any errors
-      - err_fn (Callable: Hypothesis -> pd.DataFrame):
-          A strategy to compute the standard error or deviation. This should
-          return the standard err results as a DataFrame, having the same
-          column and index as the hypothesis.
-          Defaults to "standard deviation.", i.e. `hypothoses.grouped.std()`.
-          To use standard error, use `err_fn=lambda h: h.grouped.sem()`.
-      - std_alpha (float): If not None, will show the 1-std range as a
+      - err_fn (Callable: Hypothesis -> pd.DataFrame | Tuple):
+          A strategy to compute the error range when err_style is band or fill.
+          Defaults to "standard deviation.", i.e. `hypothosis.grouped.std()`.
+          This function may return either:
+            (i) a single DataFrame, representing the standard error,
+              which must have the same column and index as the hypothesis; or
+            (ii) a tuple of two DataFrames, representing the error range
+              (lower, upper). Both DataFrames must also have the same
+              column and index as the hypothesis.
+          In the case of (i), we assume that a custom `representative_fn` is
+          NOT being used, but the representative value of the hypothesis is
+          the grouped mean of the Hypothesis, i.e., `hypothesis.mean()`.
+          (Example) To use standard error for the bands, you can use either
+          `err_fn=lambda h: h.grouped.sem()` or
+          `err_fn=lambda h: (h.grouped.mean() - h.grouped.sem(),
+                             h.grouped.mean() + h.grouped.sem())`.
+      - representative_fn (Callable: Hypothesis -> pd.DataFrame):
+          A strategy to compute the representative value (usually drawn
+          in a thicker line) when plotting.
+          This function should return a DataFrame that has the same column
+          and index as the hypothesis.
+          Defaults to "sample mean.", i.e., `hypothesis.mean()`
+          For instance, to use median instead of mean, use
+          `representative_fn=lambda h: h.grouped.median()`
+      - std_alpha (float): If not None, will show the error band as a
           shaded area. Defaults 0.2,
       - runs_alpha (float): If not None, will draw an individual line
           for each run. Defaults 0.2.
@@ -339,17 +363,50 @@ class HypothesisPlotter:
       # nothing to draw (no rows)
       raise ValueError("No data to plot, all runs have empty DataFrame.")
 
-    mean, std = None, None
-    _h_interpolated = None
+    def _representative_and_err(h: Hypothesis) -> Tuple[
+        pd.DataFrame,  # representative (mean)
+        Tuple[pd.DataFrame, pd.DataFrame]  # error band range (stderr)
+    ]:  # yapf: disable
+      """Evaluate representative_fn and err_fn."""
 
-    def _mean_and_err(h: Hypothesis):  # type: ignore
-      mean = h.grouped.mean()
+      representative = representative_fn(h) if representative_fn \
+                       else h.grouped.mean()   # noqa: E127
+      err_range: Tuple[pd.DataFrame, pd.DataFrame]
       std = err_fn(h) if err_fn else h.grouped.std()
-      return mean, std
+
+      # Condition check: when representative_fn is given,
+      # err_fn should return a range (i.e., tuple)
+      if representative_fn and err_fn and not isinstance(std, tuple):
+        raise ValueError(
+            "When representative_fn is given, err_fn must return a range "
+            "(tuple of pd.DataFrame) representing the lower and upper value "
+            "of the error band. Pass err_fn=None to use the default one, "
+            "or try: lambda h: (h.mean() + h.std(), h.mean() - h.std()). "
+            f"err_fn returned: {std}")
+
+      if isinstance(std, pd.DataFrame):
+        mean = h.grouped.mean()
+        err_range = (mean - std, mean + std)
+        return representative, err_range
+
+      elif (isinstance(std, tuple) and len(std) == 2 and
+            isinstance(std[0], pd.DataFrame) and
+            isinstance(std[1], pd.DataFrame)):
+        err_range = (std[0], std[1])
+        return representative, err_range  # type: ignore
+
+      raise ValueError("err_fn must return either a tuple of "
+                       "two DataFrames or a single DataFrame, but "
+                       f"got {type(std)}")
+
+    NULL = pd.DataFrame()
+    representative = NULL
+    err = (NULL, NULL)
+    _h_interpolated = None
 
     if 'x' not in kwargs:
       # index (same across runs) being x value, so we can simply average
-      mean, std = _mean_and_err(self._parent)
+      representative, err = _representative_and_err(self._parent)
     else:
       # might have different x values --- we need to interpolate.
       # (i) check if the x-column is consistent?
@@ -363,31 +420,33 @@ class HypothesisPlotter:
             "recommended.", UserWarning)
         n_samples = 10000
       else:
-        mean, std = _mean_and_err(self._parent)
+        representative, err = _representative_and_err(self._parent)
 
     if n_samples is not None:
       # subsample by interpolation, then average.
       _h_interpolated = self._parent.interpolate(
           x_column=kwargs.get('x', None), n_samples=n_samples)
-      mean, std = _mean_and_err(_h_interpolated)
+      representative, err = _representative_and_err(_h_interpolated)
 
       # Now that the index of group-averaged dataframes are the x samples
       # we interpolated on, we can let DataFrame.plot use them as index
       if 'x' in kwargs:
         del kwargs['x']
 
-    if not isinstance(std, pd.DataFrame):
-      raise TypeError(f"err_fn should return a pd.DataFrame, got {type(std)}")
+    if not isinstance(representative, pd.DataFrame):
+      raise TypeError("representative_fn should return a pd.DataFrame, "
+                      f"but got {type(err)}")
 
     # there might be many NaN values if each column is being logged
     # at a different period. We fill in the missing values.
-    mean = mean.interpolate()  # type: ignore
-    std = std.interpolate()  # type: ignore
-    assert mean is not None and std is not None
+    representative = representative.interpolate()
+    assert representative is not None
+    err = (err[0].interpolate(), err[1].interpolate())
+    assert err[0] is not None and err[1] is not None
 
     # determine which columns to draw (i.e. y) before smoothing.
     # should only include numerical values
-    y: Iterable[str] = kwargs.get('y', None) or mean.columns
+    y: Iterable[str] = kwargs.get('y', None) or representative.columns
     if isinstance(y, str):
       y = [y]
     if 'x' in kwargs:
@@ -397,24 +456,25 @@ class HypothesisPlotter:
       # TODO(remove): this is hack to handle homogeneous column names
       # over different hypotheses in a single of experiment, because it
       # will end up adding dummy columns instead of ignoring unknowns.
-      extra_y = set(y) - set(mean.columns)
+      extra_y = set(y) - set(representative.columns)
       for yi in extra_y:
-        mean[yi] = np.nan
+        representative[yi] = np.nan
 
     def _should_include_column(col_name: str) -> bool:
       if not col_name:  # empty name
         return False
 
       # unknown column in the DataFrame
-      assert mean is not None
-      dtypes = mean.dtypes.to_dict()  # type: ignore
+      assert representative is not None
+      dtypes = representative.dtypes.to_dict()  # type: ignore
       if col_name not in dtypes:
         if ignore_unknown:
           return False  # just ignore, no error
         else:
-          raise ValueError(f"Unknown column name '{col_name}'. " +
-                           f"Available columns: {list(mean.columns)}; " +
-                           "Use ignore_unknown=True to ignore unknown columns.")
+          raise ValueError(
+              f"Unknown column name '{col_name}'. " +
+              f"Available columns: {list(representative.columns)}; " +
+              "Use ignore_unknown=True to ignore unknown columns.")
 
       # include only numeric values (integer or float)
       if not (dtypes[col_name].kind in ('i', 'f')):
@@ -424,8 +484,10 @@ class HypothesisPlotter:
     y = [yi for yi in y if _should_include_column(yi)]
 
     if rolling:
-      mean = mean.rolling(rolling, min_periods=1, center=True).mean()
-      std = std.rolling(rolling, min_periods=1, center=True).mean()
+      representative = representative.rolling(
+          rolling, min_periods=1, center=True).mean()
+      err = (err[0].rolling(rolling, min_periods=1, center=True).mean(),
+             err[1].rolling(rolling, min_periods=1, center=True).mean())
 
     # suptitle: defaults to hypothesis name if ax/grid was not given
     if suptitle is None and (ax is None and grid is None):
@@ -433,8 +495,8 @@ class HypothesisPlotter:
 
     return self._do_plot(
         y,
-        mean,  # type: ignore
-        std,  # type: ignore
+        representative,  # type: ignore
+        err,  # type: ignore
         _h_interpolated=_h_interpolated,
         n_samples=n_samples,
         subplots=subplots,
@@ -465,8 +527,8 @@ class HypothesisPlotter:
   def _do_plot(
       self,
       y: List[str],
-      mean: pd.DataFrame,
-      std: pd.DataFrame,
+      representative: pd.DataFrame,  # usually mean
+      err_range: Tuple[pd.DataFrame, pd.DataFrame],  # usually mean ± stderr
       *,
       _h_interpolated: Optional[Hypothesis] = None,  # type: ignore
       n_samples: Optional[int],
@@ -544,7 +606,7 @@ class HypothesisPlotter:
     else:
       kwargs['legend'] = bool(legend)
 
-    axes = mean.plot(*args, subplots=subplots, ax=ax, **kwargs)
+    axes = representative.plot(*args, subplots=subplots, ax=ax, **kwargs)
 
     if err_style not in self.KNOWN_ERR_STYLES:
       raise ValueError(f"Unknown err_style '{err_style}', "
@@ -556,10 +618,10 @@ class HypothesisPlotter:
         ax = cast(Axes, ax)
         mean_line = ax.get_lines()[-1]
         x = kwargs.get('x', None)
-        x_values = mean[x].values if x else mean[yi].index
+        x_values = representative[x].values if x else representative[yi].index
         ax.fill_between(x_values,
-                        (mean - std)[yi].values,
-                        (mean + std)[yi].values,
+                        err_range[0][yi].values,
+                        err_range[1][yi].values,
                         color=mean_line.get_color(),
                         alpha=std_alpha)  # yapf: disable
 
@@ -623,8 +685,8 @@ class HypothesisHvPlotter(HypothesisPlotter):
   def _do_plot(
       self,
       y: List[str],
-      mean: pd.DataFrame,
-      std: pd.DataFrame,
+      representative: pd.DataFrame,
+      err_range: Tuple[pd.DataFrame, pd.DataFrame],  # usually mean ± stderr
       *,
       _h_interpolated: Optional[Hypothesis] = None,
       n_samples: Optional[int],
@@ -642,7 +704,7 @@ class HypothesisHvPlotter(HypothesisPlotter):
       args: List,
       kwargs: Dict,
   ):
-    if not hasattr(mean, 'hvplot'):
+    if not hasattr(representative, 'hvplot'):
       import hvplot.pandas
 
     if subplots:
@@ -650,7 +712,7 @@ class HypothesisHvPlotter(HypothesisPlotter):
 
       # TODO implement various options for hvplot.
       kwargs.update(dict(y=y))
-      p = mean.hvplot(shared_axes=False, subplots=True, **kwargs)
+      p = representative.hvplot(shared_axes=False, subplots=True, **kwargs)
 
       # Display a single legend without duplication
       if legend and isinstance(p.data, dict):
@@ -674,9 +736,9 @@ class HypothesisHvPlotter(HypothesisPlotter):
       raise NotImplementedError
 
     if err_style in ('band', 'fill') and std_alpha:
-      band_lower = mean - std
+      # TODO
+      band_lower, band_upper = err_range
       band_lower['_facet'] = 'lower'
-      band_upper = mean + std
       band_upper['_facet'] = 'upper'
       band = pd.concat([band_lower.add_suffix('.min'),
                         band_upper.add_suffix('.max')], axis=1)  # yapf: disable


Test patch:
diff --git a/expt/plot_test.py b/expt/plot_test.py
index f10c603..985f08d 100644
--- a/expt/plot_test.py
+++ b/expt/plot_test.py
@@ -1,7 +1,7 @@
 """Tests for expt.plot"""
 import contextlib
 import sys
-from typing import List, cast
+from typing import List, Tuple, cast
 
 import matplotlib
 import matplotlib.pyplot as plt
@@ -43,6 +43,7 @@ def matplotlib_rcparams(kwargs: dict):
 
 # -----------------------------------------------------------------------------
 # Fixtures
+# pylint: disable=redefined-outer-name
 
 
 @pytest.fixture
@@ -247,8 +248,12 @@ class TestHypothesisPlot:
   def test_error_range_custom_fn(self, hypothesis: Hypothesis):
     """Tests plot(err_fn=...)"""
 
+    # Case 1: err_fn returns a single DataFrame.
+    # ------------------------------------------
     def err_fn(h: Hypothesis) -> pd.DataFrame:
-      return cast(pd.DataFrame, h.grouped.std()).applymap(lambda x: 5000)
+      df: pd.DataFrame = h.grouped.mean()
+      df['loss'][:] = 5000
+      return df
 
     # without interpolation
     g = hypothesis.plot(x='step', y='loss', err_style='fill', err_fn=err_fn)
@@ -264,6 +269,52 @@ class TestHypothesisPlot:
         x='step', y='loss', err_style='runs', n_samples=100,
         err_fn=err_fn)  # Note: with err_style='runs', err_fn is not useful..?
 
+    # Case 2: err_fn returns a tuple of two DataFrames.
+    # -------------------------------------------------
+    def err_fn2(h: Hypothesis) -> Tuple[pd.DataFrame, pd.DataFrame]:
+      df: pd.DataFrame = h.grouped.mean()
+      std: pd.DataFrame = h.grouped.sem()
+      return (df - std, df + std * 100000)
+
+    # without interpolation
+    g = hypothesis.plot(x='step', y='loss', err_style='fill', err_fn=err_fn2)
+    band = g['loss'].collections[0].get_paths()[0].vertices
+
+    # std is approximately 0.25 (0.25 * 100_000 ~= 25000)
+    assert -1 <= np.min(band[:, 1]) <= 0
+    assert 20000 <= np.max(band[:, 1]) <= 30000
+
+  def test_representative_custom_fn(self, hypothesis: Hypothesis):
+    """Tests plot(representative_fn=...)"""
+
+    def repr_fn(h: Hypothesis) -> pd.DataFrame:
+      # A dummy function that manipulates the representative value ('mean')
+      df: pd.DataFrame = h.grouped.mean()
+      df['loss'] = np.asarray(df.reset_index()['step']) * -1.0
+      return df
+
+    def _ensure_representative_curve(line):
+      assert line.get_alpha() is None
+      return line
+
+    # without interpolation
+    g = hypothesis.plot(x='step', y='loss', representative_fn=repr_fn)
+    line = _ensure_representative_curve(g['loss'].get_lines()[0])
+    np.testing.assert_array_equal(line.get_xdata() * -1, line.get_ydata())
+
+    # with interpolation
+    # yapf: disable
+    g = hypothesis.plot(x='step', y='loss', n_samples=100,
+                        representative_fn=repr_fn, err_style='fill')  # fill
+    line = _ensure_representative_curve(g['loss'].get_lines()[0])
+    np.testing.assert_array_equal(line.get_xdata() * -1, line.get_ydata())
+
+    g = hypothesis.plot(x='step', y='loss', n_samples=100,
+                        representative_fn=repr_fn, err_style='runs')  # runs
+    line = _ensure_representative_curve(g['loss'].get_lines()[0])
+    np.testing.assert_array_equal(line.get_xdata() * -1, line.get_ydata())
+    # yapf: enable
+
 
 class TestExperimentPlot: