rlshiftlab 0.2.0__py3-none-any.whl

rlshiftlab/__init__.py

@@ -0,0 +1,23 @@
+"""Stress-test frozen reinforcement learning policies under controlled shifts."""
+
+__version__ = "0.2.0"
+
+from rlshiftlab.comparisons import ComparisonResults
+from rlshiftlab.groups import GroupResults
+from rlshiftlab.results import Results
+from rlshiftlab.runner import EpisodeRecord, EpisodeResult, Suite, Transition
+from rlshiftlab.shifts import ActionSlip, ObservationNoise, StickyActions
+
+__all__ = [
+    "ActionSlip",
+    "ComparisonResults",
+    "EpisodeRecord",
+    "EpisodeResult",
+    "GroupResults",
+    "ObservationNoise",
+    "Results",
+    "StickyActions",
+    "Suite",
+    "Transition",
+    "__version__",
+]

rlshiftlab/comparisons.py

@@ -0,0 +1,301 @@
+"""Direct differences between frozen policies on matched evaluation episodes."""
+
+from __future__ import annotations
+
+import json
+import zlib
+from collections.abc import Mapping
+from dataclasses import dataclass
+from math import ceil
+from numbers import Integral
+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.stats import crossed_independent_group_diff_ci, paired_diff_ci
+
+if TYPE_CHECKING:
+    from rlshiftlab.groups import GroupResults
+    from rlshiftlab.results import Results
+
+
+@dataclass
+class ComparisonResults:
+    """Pointwise differences; positive means the left metric is larger."""
+
+    rows: list[dict[str, Any]]
+    meta: dict[str, Any]
+
+    def to_csv(self, path: str | PathLike[str]) -> None:
+        """Write condition contrasts with a stable comparison-only schema."""
+        from rlshiftlab.exports import COMPARISON_COLUMNS, write_rows
+
+        write_rows(path, self.rows, COMPARISON_COLUMNS)
+
+    def summary(self) -> str:
+        """Summarize pointwise left-minus-right differences by condition."""
+        left = str(self.meta.get("left", "left"))
+        right = str(self.meta.get("right", "right"))
+        metric = str(self.meta.get("metric", "return"))
+        lines = [
+            f"{left} - {right}: {metric} (positive means a higher measured value for {left})",
+            "shift  level  mean difference  pointwise 95% CI",
+        ]
+        for row in self.rows:
+            lines.append(
+                f"{row['shift']}  {float(row['level']):g}  {float(row['mean_diff']):.3f}  "
+                f"[{float(row['ci_low']):.3f}, {float(row['ci_high']):.3f}]"
+            )
+        return "\n".join(lines)
+
+    def plot(self, path: str | PathLike[str] | None = None) -> Figure:
+        """Plot pointwise contrast intervals with a zero reference per shift."""
+        shifts = list(self.meta.get("shifts", {}))
+        if not shifts:
+            shifts = list(dict.fromkeys(row["shift"] for row in self.rows if row["shift"] != "nominal"))
+        if not shifts:
+            raise ValueError("comparison contains no shifted conditions to plot")
+        nominal = next((row for row in self.rows if row["shift"] == "nominal"), None)
+        if nominal is None:
+            raise ValueError("comparison is missing its nominal condition")
+        columns = min(3, len(shifts))
+        figure = Figure(figsize=(5.5 * columns, 4 * ceil(len(shifts) / columns)))
+        FigureCanvasAgg(figure)
+        left = str(self.meta.get("left", "left"))
+        right = str(self.meta.get("right", "right"))
+        metric = str(self.meta.get("metric", "return"))
+        for index, shift in enumerate(shifts, start=1):
+            axis = figure.add_subplot(ceil(len(shifts) / columns), columns, index)
+            shifted = sorted(
+                (row for row in self.rows if row["shift"] == shift),
+                key=lambda row: float(row["level"]),
+            )
+            points = [nominal, *shifted]
+            x = [float(point["level"]) for point in points]
+            mean = [float(point["mean_diff"]) for point in points]
+            low = [float(point["ci_low"]) for point in points]
+            high = [float(point["ci_high"]) for point in points]
+            axis.axhline(0.0, color="black", linewidth=1, linestyle="--", alpha=0.6)
+            line = axis.plot(x, mean, marker="o", label=f"{left} - {right}")[0]
+            axis.fill_between(x, low, high, color=line.get_color(), alpha=0.18)
+            axis.set_title(shift.replace("_", " "))
+            axis.set_xlabel("shift strength")
+            axis.set_ylabel(f"{metric}: {left} - {right} (pointwise 95% CI)")
+            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]) -> None:
+        """Write a standalone HTML comparison report."""
+        from rlshiftlab.report import write_html
+
+        write_html(self, path)
+
+
+def _setup(
+    results: Results | GroupResults,
+    left: str,
+    right: str,
+    metric: str,
+    grouped: bool,
+) -> tuple[list[tuple[str, float]], list[int], dict[str, list[str | None]]]:
+    if not isinstance(left, str) or not left or not isinstance(right, str) or not right:
+        raise ValueError("left and right must be nonempty policy or group names")
+    if left == right:
+        raise ValueError("left and right must be distinct")
+    labels = results.meta.get("groups" if grouped else "policies")
+    if grouped:
+        if not isinstance(labels, Mapping):
+            raise ValueError("group metadata is required for a group comparison")
+        if left not in labels or right not in labels:
+            raise ValueError("left and right must be existing groups")
+        members = {name: list(labels[name]) for name in (left, right)}
+        if any(len(ids) < 2 or len(ids) != len(set(ids)) for ids in members.values()):
+            raise ValueError("groups must declare at least two distinct members")
+    else:
+        if not isinstance(labels, (list, tuple)) or left not in labels or right not in labels:
+            raise ValueError("left and right must be existing policies")
+        members = {left: [None], right: [None]}
+
+    configured_metrics = {
+        "return", "length", "success", *results.meta.get("step_metrics", []),
+        *results.meta.get("episode_metrics", []),
+    }
+    if not isinstance(metric, str) or metric not in configured_metrics:
+        raise ValueError(f"metric {metric!r} is not configured")
+    shifts = results.meta.get("shifts")
+    if not isinstance(shifts, Mapping):
+        raise TypeError("shift metadata must be a mapping for a comparison")
+    conditions = [("nominal", 0.0)] + [
+        (name, float(level)) for name, levels in shifts.items() for level in levels
+    ]
+    if len(conditions) != len(set(conditions)):
+        raise ValueError("condition metadata contains duplicate shift levels")
+    count = results.meta.get("n_episodes")
+    seed = results.meta.get("seed")
+    if (
+        isinstance(count, bool) or not isinstance(count, Integral) or count < 2
+        or isinstance(seed, bool) or not isinstance(seed, Integral) or seed < 0
+    ):
+        raise ValueError("comparison requires valid evaluation seed and episode count")
+    return conditions, [int(seed + index) for index in range(count)], members
+
+
+def _numeric_value(record: EpisodeRecord, metric: str) -> float:
+    if metric == "return":
+        value = record.episode.episode_return
+    elif metric == "length":
+        value = record.episode.length
+    elif metric == "success":
+        value = record.success
+        if value is None:
+            raise ValueError("success metric requires recorded success values")
+    else:
+        if metric not in record.episode.metrics:
+            raise ValueError(f"record lacks configured metric {metric!r}")
+        value = record.episode.metrics[metric]
+    try:
+        array = np.asarray(value)
+        if array.ndim != 0 or array.dtype.kind not in "bifu":
+            raise ValueError("not a numeric scalar")
+        number = float(array)
+    except (TypeError, ValueError, OverflowError) as exc:
+        raise ValueError(f"record has invalid {metric!r} value") from exc
+    if not np.isfinite(number):
+        raise ValueError(f"record has nonfinite {metric!r} value")
+    return number
+
+
+def _panels(
+    records: list[EpisodeRecord],
+    entities: dict[str, list[str | None]],
+    conditions: list[tuple[str, float]],
+    episode_seeds: list[int],
+    metric: str,
+) -> dict[tuple[str, str, float], np.ndarray]:
+    """Validate every expected key, then align cells by declared member and seed."""
+    expected_conditions = set(conditions)
+    keyed: dict[tuple[str, str | None, str, float, int], float] = {}
+    for record in records:
+        if record.policy not in entities:
+            continue
+        if not isinstance(record.shift, str):
+            raise TypeError(f"invalid condition shift for {record.policy!r}")
+        if (
+            isinstance(record.level, bool)
+            or not isinstance(record.level, (int, float))
+            or not np.isfinite(record.level)
+        ):
+            raise ValueError(f"invalid condition level for {record.policy!r}")
+        condition = (record.shift, record.level)
+        if condition not in expected_conditions:
+            raise ValueError(f"unexpected condition for {record.policy!r}: {condition!r}")
+        if record.member not in entities[record.policy]:
+            raise ValueError(f"unexpected member {record.member!r} for {record.policy!r}")
+        index = record.episode_index
+        if isinstance(index, bool) or not isinstance(index, Integral) or not 0 <= index < len(episode_seeds):
+            raise ValueError(f"invalid episode index for {record.policy!r}")
+        if (
+            isinstance(record.episode_seed, bool)
+            or not isinstance(record.episode_seed, Integral)
+            or record.episode_seed != episode_seeds[index]
+        ):
+            raise ValueError(f"mismatched episode seed for {record.policy!r}")
+        key = (record.policy, record.member, record.shift, record.level, int(index))
+        if key in keyed:
+            raise ValueError(f"duplicate condition episode key for {record.policy!r}")
+        keyed[key] = _numeric_value(record, metric)
+
+    expected_count = sum(len(ids) for ids in entities.values()) * len(conditions) * len(episode_seeds)
+    if len(keyed) != expected_count:
+        raise ValueError("incomplete condition and evaluation episode grid")
+
+    panels: dict[tuple[str, str, float], np.ndarray] = {}
+    for name, members in entities.items():
+        for shift, level in conditions:
+            try:
+                panel = [
+                    [keyed[(name, member, shift, level, index)] for index in range(len(episode_seeds))]
+                    for member in members
+                ]
+            except KeyError as exc:
+                raise ValueError("incomplete condition and evaluation episode grid") from exc
+            panels[(name, shift, level)] = np.asarray(panel, dtype=float)
+    return panels
+
+
+def _comparison_seed(
+    suite_seed: int, left: str, right: str, shift: str, level: float, metric: str
+) -> int:
+    key = json.dumps([left, right, shift, level, metric], ensure_ascii=False, separators=(",", ":"))
+    rng = np.random.default_rng([suite_seed, zlib.crc32(key.encode("utf-8"))])
+    return int(rng.integers(0, 2**32))
+
+
+def compare_results(
+    results: Results | GroupResults,
+    left: str,
+    right: str,
+    metric: str = "return",
+    *,
+    grouped: bool = False,
+) -> ComparisonResults:
+    """Compare complete raw panels with paired evaluation episode keys."""
+    conditions, episode_seeds, members = _setup(results, left, right, metric, grouped)
+    panels = _panels(results.records, members, conditions, episode_seeds, metric)
+    canonical_left, canonical_right = sorted((left, right))
+    reverse = left != canonical_left
+    method = "crossed_member_episode_bootstrap" if grouped else "paired_episode_bootstrap"
+    rows: list[dict[str, Any]] = []
+    for shift, level in conditions:
+        first = panels[(canonical_left, shift, level)]
+        second = panels[(canonical_right, shift, level)]
+        seed = _comparison_seed(
+            int(results.meta["seed"]), canonical_left, canonical_right, shift, level, metric
+        )
+        if grouped:
+            mean, low, high = crossed_independent_group_diff_ci(first, second, seed=seed)
+        else:
+            mean, low, high = paired_diff_ci(first[0], second[0], seed=seed)
+        if reverse:
+            mean, low, high = -mean, -high, -low
+        rows.append(
+            {
+                "left": left,
+                "right": right,
+                "shift": shift,
+                "level": level,
+                "metric": metric,
+                "n_episodes": len(episode_seeds),
+                "n_members_left": len(members[left]) if grouped else None,
+                "n_members_right": len(members[right]) if grouped else None,
+                "method": method,
+                "mean_diff": mean,
+                "ci_low": low,
+                "ci_high": high,
+            }
+        )
+    return ComparisonResults(
+        rows=rows,
+        meta={
+            **results.meta,
+            "left": left,
+            "right": right,
+            "metric": metric,
+            "comparison_method": method,
+            "ci_confidence": 0.95,
+            "ci_replicates": 2000,
+            "comparison_direction": "left minus right; positive means left has a higher metric",
+            "ci_scope": "pointwise, not simultaneous or multiple-comparison adjusted",
+        },
+    )
+
+
+__all__ = ["ComparisonResults", "compare_results"]

rlshiftlab/exports.py

@@ -0,0 +1,98 @@
+"""Stable CSV schemas for v0.2 evaluation and audit exports."""
+
+from __future__ import annotations
+
+import csv
+from collections.abc import Sequence
+from os import PathLike
+from pathlib import Path
+from typing import Any
+
+from rlshiftlab.records import EpisodeRecord
+from rlshiftlab.results import CSV_COLUMNS
+
+METRIC_COLUMNS = (
+    "policy", "metric", "shift", "level", "n_episodes", "mean", "ci_low", "ci_high",
+    "drop", "drop_ci_low", "drop_ci_high",
+)
+
+GROUP_COLUMNS = (
+    "group", "shift", "level", "n_members", "n_episodes_per_member", "total_episodes",
+    "mean_return", "return_ci_low", "return_ci_high", "member_std", "member_min",
+    "member_max", "drop", "drop_ci_low", "drop_ci_high", "mean_length",
+    "success_rate", "success_ci_low", "success_ci_high",
+)
+
+GROUP_METRIC_COLUMNS = (
+    "group", "metric", "shift", "level", "n_members", "n_episodes_per_member",
+    "total_episodes", "mean", "ci_low", "ci_high", "member_std", "member_min",
+    "member_max", "drop", "drop_ci_low", "drop_ci_high",
+)
+
+MEMBER_COLUMNS = ("group", "member", *CSV_COLUMNS[1:])
+
+EPISODE_COLUMNS = (
+    "policy", "member", "shift", "level", "episode_index", "episode_seed",
+    "episode_return", "length", "terminated", "truncated", "success",
+)
+
+COMPARISON_COLUMNS = (
+    "left", "right", "shift", "level", "metric", "n_episodes", "n_members_left",
+    "n_members_right", "method", "mean_diff", "ci_low", "ci_high",
+)
+
+
+def write_rows(
+    path: str | PathLike[str], rows: Sequence[dict[str, Any]], columns: Sequence[str]
+) -> None:
+    """Write one deterministic UTF-8 CSV with a fixed field order."""
+    with Path(path).open("w", encoding="utf-8", newline="") as output:
+        writer = csv.DictWriter(output, fieldnames=columns, lineterminator="\n")
+        writer.writeheader()
+        for row in rows:
+            writer.writerow(
+                {
+                    column: (
+                        "" if row.get(column) is None else repr(row[column])
+                        if isinstance(row[column], float)
+                        else row[column]
+                    )
+                    for column in columns
+                }
+            )
+
+
+def episode_columns(meta: dict[str, Any]) -> tuple[str, ...]:
+    """Append custom metric names in their configured order."""
+    names = [*meta.get("step_metrics", ()), *meta.get("episode_metrics", ())]
+    return (*EPISODE_COLUMNS, *(f"metric:{name}" for name in names))
+
+
+def episode_rows(records: Sequence[EpisodeRecord], meta: dict[str, Any]) -> list[dict[str, Any]]:
+    """Convert episode records to scalar audit rows without Gym info or observations."""
+    names = [*meta.get("step_metrics", ()), *meta.get("episode_metrics", ())]
+    rows = []
+    for record in records:
+        episode = record.episode
+        row = {
+            "policy": record.policy,
+            "member": record.member,
+            "shift": record.shift,
+            "level": record.level,
+            "episode_index": record.episode_index,
+            "episode_seed": record.episode_seed,
+            "episode_return": episode.episode_return,
+            "length": episode.length,
+            "terminated": episode.terminated,
+            "truncated": episode.truncated,
+            "success": record.success,
+        }
+        row.update({f"metric:{name}": episode.metrics[name] for name in names})
+        rows.append(row)
+    return rows
+
+
+__all__ = [
+    "COMPARISON_COLUMNS", "EPISODE_COLUMNS", "GROUP_COLUMNS", "GROUP_METRIC_COLUMNS",
+    "MEMBER_COLUMNS", "METRIC_COLUMNS", "episode_columns", "episode_rows", "write_rows",
+]