Skip to content

Visualisation

rle.reusable.visualisation provides small matplotlib helpers for experiment figures. It is metric-spec driven: the same API handles one metric, two side-by-side metrics, and 2x2 or larger grids.

The sandbox notebook lives at:

notebooks/reusable/generic/visualisation_mo.py

Standalone legend export lives at:

notebooks/reusable/generic/legend_mo.py

Basic Use

from rle.reusable.visualisation import plot_metric_pair
from rle.reusable.visualisation.models import MetricPlotSpec, PlotStyle

fig, axes = plot_metric_pair(
    histories_by_algorithm,
    MetricPlotSpec("reward", title="Reward", ylabel="Reward"),
    MetricPlotSpec("safety", title="Safety", ylabel="Safety", lower_clip=0.0),
    style=PlotStyle(band="ci95", smooth=True),
)

data_by_algorithm is a mapping from algorithm label to runs. Each run can be:

  • a list of per-episode row dicts;
  • a dict of array-like metric series such as trainer logs;
  • a list of either of those for multiple seeds/runs.

X Axes

Episodic row data uses episodes on the x-axis by default. If rows have an episode field, that value is used. Otherwise the row order is plotted as 1-based episode index. This keeps episode charts readable even when training itself is step-count based.

Use an explicit x_key for step axes:

MetricPlotSpec("reward", x_key="t_end", x_label="Step")
MetricPlotSpec("reward_mean", x_key="global_step")

t_end, step, and global_step are styled as step axes with compact scientific tick formatting. Event markers use the same x units as the plot: use {"x": 12, "label": "L1"} for episode plots, and {"step": 20_000, "label": "WM"} for step plots.

For paper/report curves where sample efficiency matters, use the explicit completed-episode helpers. They keep the generic default above unchanged, but default their own x-axis to t_end:

from rle.reusable.visualisation import plot_episode_metric_series

fig, ax = plot_episode_metric_series(
    episode_histories_by_alg,
    y_key="cum_reward",
    ylabel="Episode Return",
)

Truncation

Runs often finish at different horizons. PlotStyle.truncate controls how the alignment grid is cut:

  • None: keep the full union of x values. After a shorter run ends, it no longer contributes to later points.
  • "shared": cut to the largest horizon reached by every nonempty run.
  • a number: drop points with x greater than that value.

This is separate from smoothing and bands. Bands are computed after alignment and truncation.

Moving-average smoothing normalizes the truncated window at the two curve boundaries. Constant series therefore remain constant instead of acquiring zero-padding dips at the first and last points. EMA smoothing remains causal.

Metrics

MetricPlotSpec.cumulative=True converts per-row values into a running total and prepends a (0, 0) baseline. This is useful for cumulative reward or violation counts over t_end.

MetricPlotSpec.normalize_by="ep_len" divides the metric by another row or series field before aggregation. This is the usual way to plot reward per step or safety per step while keeping episodes on the x-axis.

MetricPlotSpec.lower_clip=0.0 clips the plotted mean and confidence band, and also clamps the visible y-axis floor to 0.0. This keeps nonnegative safety or count metrics anchored on the bottom spine instead of leaving Matplotlib's default autoscale padding below the zero gridline.

PlotStyle.band supports:

  • ci95: 95% t-interval by seed/run count;
  • std: mean plus/minus one standard deviation;
  • minmax: min/max envelope.

Paper Panels

Use PlotStyle(paper_panel=True) when a single metric plot will be embedded as a small panel in a paper grid:

fig, ax = plot_metric(
    metrics,
    MetricPlotSpec("reward_mean", x_key="global_step", title="Mean Step Reward"),
    style=PlotStyle(paper_panel=True, legend=False),
    save_path="reward_paper.png",
)

Paper-panel plots keep the normal reusable plotting semantics, but use a compact 3.2 x 2.1 inch single-panel canvas with larger relative axis labels, tick labels, offset text, line strokes, grid lines, and spines. Existing plots are unchanged unless this option is set. Multi-panel layouts keep their normal default canvas sizes; pass figure_size explicitly when a paper needs a custom grid size.

Reusable Exports

The visualisation run archive format is the csh-compatible .np archive:

from rle.reusable.visualisation import (
    load_experiment_run_episode_histories,
    load_experiment_run_metrics,
    load_reusable_episode_histories,
    load_reusable_metric_runs,
    plot_metric,
    save_run_export,
)

save_run_export(
    "exports/stub_env/runs/shielded_ippo_0.np",
    algo="Shielded-IPPO",
    run_number=0,
    seed=0,
    env_name="stub_env",
    formula_name="safety",
    formula="G safe",
    env_config={},
    metrics={"global_step": [100, 200], "reward_mean": [1.0, 1.2]},
    episode_history=[
        {"episode": 1, "t_end": 100, "ep_len": 100, "cum_reward": 1.0, "cum_violations": 0}
    ],
)

metrics = load_reusable_metric_runs("exports/stub_env", algorithms="all")
episode_histories = load_reusable_episode_histories("exports/stub_env")
fig, ax = plot_metric(
    metrics,
    MetricPlotSpec("reward_mean", x_key="global_step", title="Mean Step Reward"),
)

The reusable visualisation loaders read <export_dir>/runs/*.np archives. The canonical filename layout is:

<export_dir>/runs/<algorithm_slug>_<run_idx>.np

load_reusable_metric_runs(...) returns algorithm_name -> list[metrics]; load_reusable_episode_histories(...) returns algorithm_name -> list[episode_history_per_run]. This keeps per-update metric arrays separate from completed-episode rows while leaving both easy to plot. load_experiment_run_metrics(...) and load_experiment_run_episode_histories(...) provide the same data in explicit algorithm/run order for experiment scripts that already know the requested algorithm list and run count. By default, auto-loaded algorithms from complete_export_algorithms are included only when all requested num_runs archives load cleanly. Pass min_complete_runs_for_auto_load=1 or another positive integer to scan the same 0..num_runs-1 window while keeping auto-loaded algorithms with at least that many valid archives. Every archive must contain the complete canonical metadata and an explicit episode_history list.

export_episode_graph_variants(...) writes reusable completed-episode graph sets directly under the algorithm scopes all/ and shielded/: reward, reward_legend, and papers/reward as both .png and .pdf files, plus matching safety variants. The shielded scope includes canonical Shielded-* and Contract-* labels. Exporting again replaces the canonical reward and safety variants in those scopes.

Legend Exports

Standalone legend images live in the same reusable visualisation package. They are useful when a paper figure grid shares one algorithm legend across many plots.

from rle.reusable.visualisation import export_graph_legend_variants

export_graph_legend_variants(
    "exports/legends",
    algorithms=("IPPO", "Shielded-IPPO", "Contract-IPPO", "ICPO"),
)

This writes the three standard PNG variants by default:

  • legend_algorithms.png: compact/default multi-row legend.
  • longer_legend_algorithms.png: single-row legend for wide figures.
  • two_line_legend_algorithms.png: centered two-line legend.

Pass formats=("png", "pdf") to also export matching PDF files. The reusable legend notebook does this by default for paper workflows.

Paper styles use TrueType fonts in PDF/PS output and preserve text in SVG output, so labels remain editable in common vector-graphics tools.

The generic legend notebook lives at:

notebooks/reusable/generic/legend_mo.py

It discovers algorithms from reusable experiment exports. Both grouped and direct suite exports are supported:

exports/<group>/<algorithm_slug>/timings.json
exports/<group>/runs/<algorithm_slug>_<run_idx>.np

<export_dir>/<algorithm_slug>/timings.json
<export_dir>/runs/<algorithm_slug>_<run_idx>.np

Run archives provide display labels through algo. canonical_algorithm_label(...) maps the canonical display names and their canonical slugs, such as Shielded-IPPO/shielded_ippo and Contract-IQL/contract_iql. Other spellings are treated as distinct labels. When no matching run archive is available, the algorithm slug is mapped when canonical and otherwise prettified. Archived run records under _archived_complete_runs are ignored. Sorting uses ALGORITHM_LEGEND_ORDER first, then appends unknown algorithms alphabetically.

Run it as a script:

uv run notebooks/reusable/generic/legend_mo.py \
  --export-root exports \
  --groups all \
  --output-dir exports/legends

Use --groups csh,omsh, --algorithms ippo,icpo, or the interactive grouped checkbox grid when you want a subset of discovered groups or algorithms.

Adapting To Your Project

Copy src/rle/reusable/visualisation and the generic notebooks if the project uses matplotlib figures. Add matplotlib and numpy to the project dependencies. If the project keeps the reusable experiment export layout, the metric loader, timing notebook, and legend notebook can discover outputs without custom adapters.

The first project-specific block is usually metric naming. Choose the keys that algorithms will return in ExperimentRunResult.metrics, then define plot specs around those keys:

reward = MetricPlotSpec("reward_mean", title="Mean Step Reward", ylabel="Mean Step Reward")
safety = MetricPlotSpec("safety_mean", title="Safety", ylabel="Violations")

Generic MetricPlotSpec plots still default row histories to episode/order. Set x_key="t_end" for completed-episode sample-efficiency curves. If a project needs a different axis, make that explicit with x_key="episode" or x_key="global_step" in the project plotting code.

The second project-specific block is style:

style = PlotStyle(
    band="ci95",
    smooth=True,
    truncate="shared",
    show_seed_traces=False,
)

Use truncate=None when late points from longer seeds should remain visible, truncate="shared" when every plotted point should include every nonempty run, and a numeric truncation when the paper needs a fixed horizon. Bands are computed after truncation and alignment, so changing truncation changes the seed count available at each x value.

The third project-specific block is algorithm identity. Prefer the csh display labels used by the reusable palette and legend order, for example Shielded-IQL, Shielded-IPPO, Contract-IPPO, Contract-IQL, and Joint PPO. Extend DEFAULT_PALETTE and ALGORITHM_LEGEND_ORDER when a project introduces new algorithms that should have stable colors and legend order:

project_palette = {
    **DEFAULT_PALETTE,
    "My Algorithm": "#8c564b",
}

export_graph_legend_variants(
    "exports/legends",
    algorithms=("IPPO", "ICPO", "My Algorithm"),
    palette=project_palette,
)

The functions most likely to be used or lightly wrapped in a project are:

  • load_reusable_metric_runs(...) for records written by the reusable runner;
  • load_reusable_episode_histories(...) for completed-episode histories;
  • plot_metric(...) for one-panel inspection;
  • plot_episode_metric_series(...) for completed-episode sample-efficiency curves by t_end;
  • plot_episode_metric_and_metric_side_by_side(...) for mixed figures such as episodic return beside cumulative step-based safety;
  • export_episode_graph_variants(...) for paper/report completed-episode PNGs and PDFs;
  • plot_metric_pair(...) for the common reward/safety side-by-side case;
  • plot_metric_grid(...) for 2x2 or larger metric grids;
  • discover_reusable_legend_algorithms(...) for export-driven legends;
  • export_graph_legend_variants(...) for the three standard legend PNGs, with optional matching PDFs.

If the trainer emits metrics in a different shape, add a small project loader that returns the same in-memory structure expected by the plotting functions: algorithm label -> list of per-run metric mappings or per-episode row lists. Keep the reusable plotting layer focused on generic alignment, aggregation, and styling rather than teaching it every trainer's native log format.