rlshiftlab 0.2.0__py3-none-any.whl

rlshiftlab/groups.py

@@ -0,0 +1,440 @@
+"""Evaluation summaries for groups of independently trained frozen policies."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass, field
+from math import ceil
+from os import PathLike
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+from matplotlib.backends.backend_agg import FigureCanvasAgg
+from matplotlib.figure import Figure
+
+from rlshiftlab.records import EpisodeRecord
+from rlshiftlab.results import CSV_COLUMNS
+from rlshiftlab.stats import crossed_mean_ci, crossed_paired_diff_ci
+
+if TYPE_CHECKING:
+    from rlshiftlab.comparisons import ComparisonResults
+    from rlshiftlab.runner import Suite
+
+
+@dataclass
+class GroupResults:
+    """Group condition summaries, member summaries, and raw episode records."""
+
+    rows: list[dict[str, Any]]
+    meta: dict[str, Any]
+    member_rows: list[dict[str, Any]] = field(default_factory=list)
+    records: list[EpisodeRecord] = field(default_factory=list)
+    metric_rows: list[dict[str, Any]] = field(default_factory=list)
+
+    def to_csv(self, path: str | PathLike[str]) -> None:
+        """Write group-level condition summaries."""
+        from rlshiftlab.exports import GROUP_COLUMNS, write_rows
+
+        write_rows(path, self.rows, GROUP_COLUMNS)
+
+    def to_member_csv(self, path: str | PathLike[str]) -> None:
+        """Write each member's condition summaries."""
+        from rlshiftlab.exports import MEMBER_COLUMNS, write_rows
+
+        write_rows(path, self.member_rows, MEMBER_COLUMNS)
+
+    def to_episode_csv(self, path: str | PathLike[str]) -> None:
+        """Write one scalar audit row per member and evaluation episode."""
+        from rlshiftlab.exports import episode_columns, episode_rows, write_rows
+
+        write_rows(path, episode_rows(self.records, self.meta), episode_columns(self.meta))
+
+    def to_metrics_csv(self, path: str | PathLike[str]) -> None:
+        """Write custom metric summaries with training-member spread."""
+        from rlshiftlab.exports import GROUP_METRIC_COLUMNS, write_rows
+
+        write_rows(path, self.metric_rows, GROUP_METRIC_COLUMNS)
+
+    def summary(self) -> str:
+        """Summarize measured group means, member spread, and paired drops."""
+        headings = (
+            "group", "shift", "level", "members", "episodes/member", "mean return (95% CI)",
+            "member SD", "drop (95% CI)",
+        )
+        lines = []
+        for row in self.rows:
+            drop = (
+                "" if row["drop"] is None else
+                f"{row['drop']:.3f} [{row['drop_ci_low']:.3f}, {row['drop_ci_high']:.3f}]"
+            )
+            lines.append(
+                (
+                    str(row["group"]),
+                    str(row["shift"]),
+                    f"{row['level']:g}",
+                    str(row["n_members"]),
+                    str(row["n_episodes_per_member"]),
+                    (
+                        f"{row['mean_return']:.3f} "
+                        f"[{row['return_ci_low']:.3f}, {row['return_ci_high']:.3f}]"
+                    ),
+                    f"{row['member_std']:.3f}",
+                    drop,
+                )
+            )
+        widths = [
+            max([len(heading), *(len(line[index]) for line in lines)])
+            for index, heading in enumerate(headings)
+        ]
+        header = "  ".join(heading.ljust(widths[index]) for index, heading in enumerate(headings))
+        rule = "  ".join("-" * width for width in widths)
+        body = [
+            "  ".join(value.ljust(widths[index]) for index, value in enumerate(line))
+            for line in lines
+        ]
+        return "\n".join([header, rule, *body])
+
+    def plot(
+        self,
+        path: str | PathLike[str] | None = None,
+        *,
+        metric: str = "return",
+        band: str = "ci",
+    ) -> Figure:
+        """Plot group means with CI or descriptive member-SD bands."""
+        if band not in {"ci", "std"}:
+            raise ValueError("band must be 'ci' or 'std'")
+        if metric == "return":
+            plot_rows = self.rows
+            mean_key, low_key, high_key = "mean_return", "return_ci_low", "return_ci_high"
+            measure_label = "mean return"
+        else:
+            configured = [
+                *self.meta.get("step_metrics", ()), *self.meta.get("episode_metrics", ())
+            ]
+            if metric not in configured:
+                raise ValueError(f"custom metric {metric!r} is not configured")
+            plot_rows = [row for row in self.metric_rows if row["metric"] == metric]
+            mean_key, low_key, high_key = "mean", "ci_low", "ci_high"
+            measure_label = f"mean {metric}"
+        shifts = list(self.meta.get("shifts") or dict.fromkeys(
+            row["shift"] for row in plot_rows if row["shift"] != "nominal"
+        ))
+        if not shifts:
+            raise ValueError("results contain no shifted conditions to plot")
+        groups = self.meta.get("groups") or dict.fromkeys(row["group"] for row in plot_rows)
+        columns = min(3, len(shifts))
+        rows = ceil(len(shifts) / columns)
+        figure = Figure(figsize=(5.5 * columns, 4 * rows))
+        FigureCanvasAgg(figure)
+        for index, shift in enumerate(shifts, start=1):
+            axis = figure.add_subplot(rows, columns, index)
+            for group in groups:
+                nominal = next(
+                    (row for row in plot_rows if row["group"] == group and row["shift"] == "nominal"),
+                    None,
+                )
+                shifted = sorted(
+                    (row for row in plot_rows if row["group"] == group and row["shift"] == shift),
+                    key=lambda row: float(row["level"]),
+                )
+                if nominal is None or not shifted:
+                    continue
+                points = [nominal, *shifted]
+                x = [float(point["level"]) for point in points]
+                means = [float(point[mean_key]) for point in points]
+                if band == "ci":
+                    lows = [float(point[low_key]) for point in points]
+                    highs = [float(point[high_key]) for point in points]
+                else:
+                    lows = [float(point[mean_key] - point["member_std"]) for point in points]
+                    highs = [float(point[mean_key] + point["member_std"]) for point in points]
+                line = axis.plot(x, means, marker="o", label=str(group))[0]
+                axis.fill_between(x, lows, highs, color=line.get_color(), alpha=0.18)
+            axis.set_title(shift.replace("_", " "))
+            axis.set_xlabel("shift strength")
+            axis.set_ylabel(
+                f"{measure_label} (95% pointwise CI)" if band == "ci"
+                else f"{measure_label} ± 1 member SD (descriptive)"
+            )
+            axis.grid(alpha=0.2)
+            axis.legend()
+        figure.tight_layout()
+        if path is not None:
+            figure.savefig(path, dpi=150)
+        return figure
+
+    def to_html(
+        self, path: str | PathLike[str], *, comparisons: Any = None
+    ) -> None:
+        """Write a standalone report for group evaluation."""
+        from rlshiftlab.report import write_html
+
+        write_html(self, path, comparisons=comparisons)
+
+    def compare(
+        self, left: str, right: str, metric: str = "return"
+    ) -> ComparisonResults:
+        """Compare group means with shared episode and independent member draws."""
+        from rlshiftlab.comparisons import compare_results
+
+        return compare_results(self, left, right, metric=metric, grouped=True)
+
+
+def _validated_groups(
+    groups: Mapping[str, Mapping[str, Callable[[Any], Any]]],
+) -> dict[str, dict[str, Callable[[Any], Any]]]:
+    if not isinstance(groups, Mapping) or not groups:
+        raise ValueError("groups must be a nonempty mapping of group names to member mappings")
+    selected: dict[str, dict[str, Callable[[Any], Any]]] = {}
+    for group_name, members in groups.items():
+        if not isinstance(group_name, str) or not group_name.strip():
+            raise ValueError(f"invalid group name {group_name!r}; expected a nonempty string")
+        if group_name in selected:
+            raise ValueError(f"duplicate group name {group_name!r}")
+        if not isinstance(members, Mapping) or len(members) < 2:
+            raise ValueError(f"group {group_name!r} must have at least two members")
+        selected_members: dict[str, Callable[[Any], Any]] = {}
+        for member_id, policy in members.items():
+            if not isinstance(member_id, str) or not member_id.strip():
+                raise ValueError(
+                    f"group {group_name!r} has invalid member ID {member_id!r}"
+                )
+            if member_id in selected_members:
+                raise ValueError(f"group {group_name!r} has duplicate member ID {member_id!r}")
+            if not callable(policy):
+                raise TypeError(
+                    f"group {group_name!r} member {member_id!r} policy must be callable"
+                )
+            selected_members[member_id] = policy
+        selected[group_name] = selected_members
+    return selected
+
+
+def _condition_matrix(
+    records: list[EpisodeRecord],
+    group: str,
+    members: list[str],
+    shift: str,
+    level: float,
+    episode_seeds: list[int],
+    metric: str,
+) -> np.ndarray:
+    """Build and validate one complete member by episode panel from raw records."""
+    cells: dict[tuple[str, int], float] = {}
+    expected_members = set(members)
+    expected_seeds = set(episode_seeds)
+    for record in records:
+        if record.policy != group or record.shift != shift or record.level != level:
+            continue
+        if record.member not in expected_members or record.episode_seed not in expected_seeds:
+            raise ValueError(f"group {group!r} has an unexpected member or episode seed")
+        if record.episode_seed != episode_seeds[record.episode_index]:
+            raise ValueError(f"group {group!r} has a mismatched episode index and seed")
+        key = (record.member, record.episode_seed)
+        if key in cells:
+            raise ValueError(f"group {group!r} has duplicate condition episode key {key!r}")
+        if metric == "return":
+            value = record.episode.episode_return
+        elif metric == "length":
+            value = record.episode.length
+        elif metric == "success":
+            if record.success is None:
+                raise ValueError("success metric requires a configured success_fn")
+            value = float(record.success)
+        else:
+            value = record.episode.metrics[metric]
+        cells[key] = float(value)
+    expected_count = len(members) * len(episode_seeds)
+    if len(cells) != expected_count:
+        raise ValueError(
+            f"group {group!r} {shift!r} level {level} has incomplete episode records"
+        )
+    return np.asarray(
+        [[cells[(member, seed)] for seed in episode_seeds] for member in members],
+        dtype=float,
+    )
+
+
+def _spread(values: np.ndarray) -> tuple[float, float, float]:
+    member_means = values.mean(axis=1)
+    return (
+        float(np.std(member_means, ddof=1)),
+        float(np.min(member_means)),
+        float(np.max(member_means)),
+    )
+
+
+def evaluate_groups(
+    suite: Suite,
+    groups: Mapping[str, Mapping[str, Callable[[Any], Any]]],
+) -> GroupResults:
+    """Evaluate frozen member policies and aggregate without flattening seed runs."""
+    selected = _validated_groups(groups)
+    conditions = [("nominal", 0.0)] + [
+        (shift.name, level) for shift in suite.shifts for level in shift.levels
+    ]
+    episode_seeds = [suite.seed + index for index in range(suite.n_episodes)]
+    rows: list[dict[str, Any]] = []
+    member_rows: list[dict[str, Any]] = []
+    records: list[EpisodeRecord] = []
+    metric_rows: list[dict[str, Any]] = []
+    metric_names = [*suite.step_metrics, *suite.episode_metrics]
+    base_meta: dict[str, Any] = {}
+
+    for group_name, members in selected.items():
+        evaluated = suite.evaluate(members)
+        if not base_meta:
+            base_meta = dict(evaluated.meta)
+        for member_row in evaluated.rows:
+            member_rows.append(
+                {
+                    "group": group_name,
+                    "member": member_row["policy"],
+                    **{
+                        column: member_row[column]
+                        for column in CSV_COLUMNS
+                        if column != "policy"
+                    },
+                }
+            )
+        group_records = [
+            EpisodeRecord(
+                policy=group_name,
+                member=record.policy,
+                shift=record.shift,
+                level=record.level,
+                episode_index=record.episode_index,
+                episode_seed=record.episode_seed,
+                episode=record.episode,
+                success=record.success,
+            )
+            for record in evaluated.records
+        ]
+        records.extend(group_records)
+        member_ids = list(members)
+        nominal_return = _condition_matrix(
+            group_records, group_name, member_ids, "nominal", 0.0, episode_seeds, "return"
+        )
+        nominal_metrics = {
+            metric: _condition_matrix(
+                group_records, group_name, member_ids, "nominal", 0.0, episode_seeds, metric
+            )
+            for metric in metric_names
+        }
+
+        for shift_name, level in conditions:
+            common = {
+                "group": group_name,
+                "shift": shift_name,
+                "level": level,
+                "n_members": len(member_ids),
+                "n_episodes_per_member": suite.n_episodes,
+                "total_episodes": len(member_ids) * suite.n_episodes,
+            }
+            returns = _condition_matrix(
+                group_records, group_name, member_ids, shift_name, level, episode_seeds, "return"
+            )
+            mean_return, return_low, return_high = crossed_mean_ci(
+                returns,
+                seed=suite._bootstrap_seed(group_name, shift_name, level, "group:return"),
+            )
+            member_std, member_min, member_max = _spread(returns)
+            if shift_name == "nominal":
+                drop = drop_low = drop_high = None
+            else:
+                drop, drop_low, drop_high = crossed_paired_diff_ci(
+                    nominal_return,
+                    returns,
+                    seed=suite._bootstrap_seed(group_name, shift_name, level, "group:return:drop"),
+                )
+            lengths = _condition_matrix(
+                group_records, group_name, member_ids, shift_name, level, episode_seeds, "length"
+            )
+            if suite.success_fn is None:
+                success = success_low = success_high = None
+            else:
+                successes = _condition_matrix(
+                    group_records, group_name, member_ids, shift_name, level, episode_seeds, "success"
+                )
+                success, success_low, success_high = crossed_mean_ci(
+                    successes,
+                    seed=suite._bootstrap_seed(group_name, shift_name, level, "group:success"),
+                )
+            rows.append(
+                {
+                    **common,
+                    "mean_return": mean_return,
+                    "return_ci_low": return_low,
+                    "return_ci_high": return_high,
+                    "member_std": member_std,
+                    "member_min": member_min,
+                    "member_max": member_max,
+                    "drop": drop,
+                    "drop_ci_low": drop_low,
+                    "drop_ci_high": drop_high,
+                    "mean_length": float(np.mean(lengths)),
+                    "success_rate": success,
+                    "success_ci_low": success_low,
+                    "success_ci_high": success_high,
+                }
+            )
+            for metric in metric_names:
+                values = _condition_matrix(
+                    group_records,
+                    group_name,
+                    member_ids,
+                    shift_name,
+                    level,
+                    episode_seeds,
+                    metric,
+                )
+                mean, low, high = crossed_mean_ci(
+                    values,
+                    seed=suite._bootstrap_seed(group_name, shift_name, level, f"group:{metric}"),
+                )
+                metric_std, metric_min, metric_max = _spread(values)
+                if shift_name == "nominal":
+                    metric_drop = metric_drop_low = metric_drop_high = None
+                else:
+                    metric_drop, metric_drop_low, metric_drop_high = crossed_paired_diff_ci(
+                        nominal_metrics[metric],
+                        values,
+                        seed=suite._bootstrap_seed(
+                            group_name, shift_name, level, f"group:{metric}:drop"
+                        ),
+                    )
+                metric_rows.append(
+                    {
+                        **common,
+                        "metric": metric,
+                        "mean": mean,
+                        "ci_low": low,
+                        "ci_high": high,
+                        "member_std": metric_std,
+                        "member_min": metric_min,
+                        "member_max": metric_max,
+                        "drop": metric_drop,
+                        "drop_ci_low": metric_drop_low,
+                        "drop_ci_high": metric_drop_high,
+                    }
+                )
+
+    meta = {
+        key: value for key, value in base_meta.items() if key != "policies"
+    }
+    meta.update(
+        {
+            "groups": {name: list(members) for name, members in selected.items()},
+            "evaluation_seeds": episode_seeds,
+            "n_members": {name: len(members) for name, members in selected.items()},
+            "ci_method": "crossed member and evaluation episode percentile bootstrap",
+            "ci_confidence": 0.95,
+            "ci_replicates": 2000,
+            "member_spread": "sample SD and observed range of member means; descriptive, not a CI",
+        }
+    )
+    return GroupResults(rows, meta, member_rows, records, metric_rows)
+
+
+__all__ = ["GroupResults", "evaluate_groups"]

rlshiftlab/records.py

@@ -0,0 +1,50 @@
+"""Episode and transition records for frozen-policy evaluations."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class EpisodeResult:
+    """The return, stopping state, and optional metrics of one episode."""
+
+    episode_return: float
+    length: int
+    terminated: bool
+    truncated: bool
+    metrics: dict[str, float] = field(default_factory=dict)
+    final_info: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class Transition:
+    """One policy-visible step supplied to a custom metric callback."""
+
+    observation: Any
+    requested_action: Any
+    next_observation: Any
+    reward: float
+    terminated: bool
+    truncated: bool
+    info: dict[str, Any]
+    step_index: int
+    episode_seed: int
+
+
+@dataclass
+class EpisodeRecord:
+    """An auditable episode keyed by condition and evaluation seed."""
+
+    policy: str
+    member: str | None
+    shift: str
+    level: float
+    episode_index: int
+    episode_seed: int
+    episode: EpisodeResult
+    success: bool | None = None
+
+
+__all__ = ["EpisodeRecord", "EpisodeResult", "Transition"]