rlshiftlab 0.2.0__py3-none-any.whl

rlshiftlab/runner.py

@@ -0,0 +1,458 @@
+"""Paired episode evaluation and suite orchestration."""
+
+from __future__ import annotations
+
+import json
+import platform
+import zlib
+from collections.abc import Callable, Mapping, Sequence
+from copy import deepcopy
+from numbers import Integral
+from typing import TYPE_CHECKING, Any
+
+import gymnasium as gym
+import numpy as np
+
+from rlshiftlab import __version__
+from rlshiftlab.records import EpisodeRecord, EpisodeResult, Transition
+from rlshiftlab.results import Results
+from rlshiftlab.shifts import ActionSlip, ObservationNoise, StickyActions
+from rlshiftlab.stats import bootstrap_mean_ci, paired_diff_ci, wilson_interval
+
+if TYPE_CHECKING:
+    from rlshiftlab.groups import GroupResults
+
+StepMetric = Callable[[Transition], float]
+EpisodeMetric = Callable[[EpisodeResult], float]
+
+
+def _metric_maps(
+    step_metrics: Mapping[str, StepMetric] | None,
+    episode_metrics: Mapping[str, EpisodeMetric] | None,
+) -> tuple[dict[str, StepMetric], dict[str, EpisodeMetric]]:
+    """Validate metric callbacks while preserving their insertion order."""
+    checked: list[dict[str, Any]] = []
+    for kind, metrics in (("step_metrics", step_metrics), ("episode_metrics", episode_metrics)):
+        if metrics is None:
+            checked.append({})
+            continue
+        if not isinstance(metrics, Mapping):
+            raise TypeError(f"{kind} must be a mapping of metric names to callbacks")
+        selected: dict[str, Any] = {}
+        for name, callback in metrics.items():
+            if not isinstance(name, str) or not name.strip():
+                raise ValueError(f"{kind} contains an invalid metric name {name!r}")
+            if name in {"return", "length", "success"}:
+                raise ValueError(f"metric name {name!r} is reserved")
+            if not callable(callback):
+                raise TypeError(f"metric {name!r} callback must be callable")
+            selected[name] = callback
+        checked.append(selected)
+    duplicate = set(checked[0]).intersection(checked[1])
+    if duplicate:
+        raise ValueError(f"metric name {min(duplicate)!r} appears in both callback maps")
+    return checked[0], checked[1]
+
+
+def _metric_value(value: Any, name: str) -> float:
+    """Convert a finite numeric scalar without accepting strings or vectors."""
+    try:
+        array = np.asarray(value)
+        if array.ndim != 0 or array.dtype.kind not in "bifu":
+            raise ValueError("expected a numeric scalar")
+        result = float(array)
+        if not np.isfinite(result):
+            raise ValueError("expected a finite numeric scalar")
+        return result
+    except (TypeError, ValueError, OverflowError) as exc:
+        raise ValueError(f"metric {name!r} must return a finite numeric scalar") from exc
+
+
+def _positive_int(value: int, name: str, minimum: int = 1) -> None:
+    if isinstance(value, bool) or not isinstance(value, Integral) or value < minimum:
+        raise ValueError(f"{name} must be an integer >= {minimum}")
+
+
+def _has_time_limit(env: gym.Env) -> bool:
+    current = env
+    while isinstance(current, gym.Wrapper):
+        if isinstance(current, gym.wrappers.TimeLimit):
+            return True
+        current = current.env
+    return getattr(getattr(env, "spec", None), "max_episode_steps", None) is not None
+
+
+def run_episodes(
+    env: gym.Env,
+    policy: Callable[[Any], Any],
+    n_episodes: int,
+    seed: int = 0,
+    max_steps: int = 1000,
+    *,
+    step_metrics: Mapping[str, StepMetric] | None = None,
+    episode_metrics: Mapping[str, EpisodeMetric] | None = None,
+) -> list[EpisodeResult]:
+    """Run a frozen policy with paired reset seeds; the caller owns the env."""
+    _positive_int(n_episodes, "n_episodes")
+    _positive_int(max_steps, "max_steps")
+    if isinstance(seed, bool) or not isinstance(seed, Integral) or seed < 0:
+        raise ValueError("seed must be a nonnegative integer")
+    if not callable(policy):
+        raise TypeError("policy must be callable")
+    selected_step_metrics, selected_episode_metrics = _metric_maps(step_metrics, episode_metrics)
+    metrics_enabled = bool(selected_step_metrics or selected_episode_metrics)
+
+    has_time_limit = _has_time_limit(env)
+    reset_policy = getattr(policy, "reset", None)
+    episodes = []
+    for episode_index in range(n_episodes):
+        if callable(reset_policy):
+            reset_policy()
+        episode_seed = int(seed + episode_index)
+        observation, _ = env.reset(seed=episode_seed)
+        episode_return = 0.0
+        length = 0
+        terminated = truncated = False
+        metric_totals = {name: 0.0 for name in selected_step_metrics}
+        final_info: dict[str, Any] = {}
+        while not (terminated or truncated):
+            # Capture the policy's input before it or the environment can mutate it.
+            metric_observation = deepcopy(observation) if metrics_enabled else None
+            action = policy(observation)
+            metric_action = deepcopy(action) if metrics_enabled else None
+            next_observation, reward, terminated, truncated, info = env.step(action)
+            episode_return += float(reward)
+            length += 1
+            if not has_time_limit and length >= max_steps and not terminated and not truncated:
+                truncated = True
+            if metrics_enabled:
+                # One snapshot per step; each callback receives an independent copy.
+                transition = Transition(
+                    observation=metric_observation,
+                    requested_action=metric_action,
+                    next_observation=deepcopy(next_observation),
+                    reward=float(reward),
+                    terminated=bool(terminated),
+                    truncated=bool(truncated),
+                    info=deepcopy(info),
+                    step_index=length - 1,
+                    episode_seed=episode_seed,
+                )
+                for name, callback in selected_step_metrics.items():
+                    try:
+                        contribution = _metric_value(callback(deepcopy(transition)), name)
+                    except Exception as exc:
+                        raise ValueError(
+                            f"step metric {name!r} failed in episode {episode_index}, "
+                            f"step {length - 1}"
+                        ) from exc
+                    metric_totals[name] += contribution
+                    if not np.isfinite(metric_totals[name]):
+                        raise ValueError(
+                            f"step metric {name!r} accumulated a nonfinite value "
+                            f"in episode {episode_index}"
+                        )
+                final_info = deepcopy(transition.info)
+            observation = next_observation
+        result = EpisodeResult(
+            episode_return, length, bool(terminated), bool(truncated), metric_totals, final_info
+        )
+        if selected_episode_metrics:
+            base_result = deepcopy(result)
+            completed_metrics: dict[str, float] = {}
+            for name, callback in selected_episode_metrics.items():
+                try:
+                    completed_metrics[name] = _metric_value(callback(deepcopy(base_result)), name)
+                except Exception as exc:
+                    raise ValueError(f"episode metric {name!r} failed in episode {episode_index}") from exc
+            result.metrics.update(completed_metrics)
+        episodes.append(result)
+    return episodes
+
+
+class Suite:
+    """Evaluate policies under one shift at a time with paired episode seeds."""
+
+    def __init__(
+        self,
+        env_fn: Callable[[], gym.Env],
+        shifts: Sequence[ObservationNoise | ActionSlip | StickyActions],
+        n_episodes: int,
+        seed: int = 0,
+        success_fn: Callable[[EpisodeResult], bool] | None = None,
+        max_steps: int = 1000,
+        *,
+        step_metrics: Mapping[str, StepMetric] | None = None,
+        episode_metrics: Mapping[str, EpisodeMetric] | None = None,
+    ) -> None:
+        if not callable(env_fn):
+            raise TypeError("env_fn must be callable")
+        _positive_int(n_episodes, "n_episodes", minimum=2)
+        _positive_int(max_steps, "max_steps")
+        if isinstance(seed, bool) or not isinstance(seed, Integral) or seed < 0:
+            raise ValueError("seed must be a nonnegative integer")
+        if success_fn is not None and not callable(success_fn):
+            raise TypeError("success_fn must be callable")
+        selected_step_metrics, selected_episode_metrics = _metric_maps(
+            step_metrics, episode_metrics
+        )
+        selected_shifts = list(shifts)
+        if not selected_shifts:
+            raise ValueError("shifts must contain at least one shift")
+        if any(not isinstance(shift, (ObservationNoise, ActionSlip, StickyActions)) for shift in selected_shifts):
+            raise TypeError("shifts must contain only ObservationNoise, ActionSlip, or StickyActions")
+        names = [shift.name for shift in selected_shifts]
+        if len(set(names)) != len(names):
+            raise ValueError("shift names must be unique")
+
+        validation_env = env_fn()
+        try:
+            self._validate_spaces(validation_env)
+            for shift in selected_shifts:
+                for level in shift.levels:
+                    shift.wrap(validation_env, level)
+        finally:
+            validation_env.close()
+
+        self.env_fn = env_fn
+        self.shifts = selected_shifts
+        self.n_episodes = int(n_episodes)
+        self.seed = int(seed)
+        self.success_fn = success_fn
+        self.max_steps = int(max_steps)
+        self.step_metrics = selected_step_metrics
+        self.episode_metrics = selected_episode_metrics
+
+    @staticmethod
+    def _validate_spaces(env: gym.Env) -> None:
+        observation_space = env.observation_space
+        action_space = env.action_space
+        if not isinstance(observation_space, (gym.spaces.Box, gym.spaces.Discrete)):
+            raise TypeError(
+                f"unsupported observation space {observation_space!r}; expected Box or Discrete"
+            )
+        if isinstance(action_space, gym.spaces.Discrete):
+            return
+        if isinstance(action_space, gym.spaces.Box) and (
+            np.all(np.isfinite(action_space.low)) and np.all(np.isfinite(action_space.high))
+        ):
+            return
+        raise TypeError(
+            f"unsupported action space {action_space!r}; expected Discrete or finite Box"
+        )
+
+    def _bootstrap_seed(self, policy: str, shift: str, level: float, metric: str) -> int:
+        key = json.dumps([policy, shift, level, metric], ensure_ascii=False, separators=(",", ":"))
+        stream = np.random.default_rng([self.seed, zlib.crc32(key.encode("utf-8"))])
+        return int(stream.integers(0, 2**32))
+
+    def _run_condition(
+        self,
+        policy: Callable[[Any], Any],
+        shift: ObservationNoise | ActionSlip | StickyActions | None,
+        level: float,
+    ) -> list[EpisodeResult]:
+        env = self.env_fn()
+        condition_env = env
+        try:
+            if shift is not None:
+                condition_env = shift.wrap(env, level)
+            return run_episodes(
+                condition_env,
+                policy,
+                self.n_episodes,
+                self.seed,
+                self.max_steps,
+                step_metrics=self.step_metrics,
+                episode_metrics=self.episode_metrics,
+            )
+        finally:
+            condition_env.close()
+
+    def _row(
+        self,
+        policy_name: str,
+        shift_name: str,
+        level: float,
+        episodes: list[EpisodeResult],
+        nominal_returns: list[float] | None,
+        successes: list[bool] | None,
+    ) -> dict[str, Any]:
+        returns = [episode.episode_return for episode in episodes]
+        mean, low, high = bootstrap_mean_ci(
+            returns, seed=self._bootstrap_seed(policy_name, shift_name, level, "return")
+        )
+        if nominal_returns is None:
+            drop = drop_low = drop_high = None
+        else:
+            drop, drop_low, drop_high = paired_diff_ci(
+                nominal_returns,
+                returns,
+                seed=self._bootstrap_seed(policy_name, shift_name, level, "drop"),
+            )
+        if successes is None:
+            success = success_low = success_high = None
+        else:
+            count = sum(successes)
+            success, success_low, success_high = wilson_interval(count, len(episodes))
+        return {
+            "policy": policy_name,
+            "shift": shift_name,
+            "level": level,
+            "n_episodes": self.n_episodes,
+            "mean_return": mean,
+            "return_ci_low": low,
+            "return_ci_high": high,
+            "drop": drop,
+            "drop_ci_low": drop_low,
+            "drop_ci_high": drop_high,
+            "mean_length": float(np.mean([episode.length for episode in episodes])),
+            "success_rate": success,
+            "success_ci_low": success_low,
+            "success_ci_high": success_high,
+        }
+
+    def _condition_records(
+        self,
+        policy_name: str,
+        shift_name: str,
+        level: float,
+        episodes: list[EpisodeResult],
+    ) -> tuple[list[EpisodeRecord], list[bool] | None]:
+        successes: list[bool] | None = [] if self.success_fn is not None else None
+        records = []
+        for index, episode in enumerate(episodes):
+            success = None
+            if self.success_fn is not None:
+                success = bool(self.success_fn(episode))
+                assert successes is not None
+                successes.append(success)
+            records.append(
+                EpisodeRecord(
+                    policy=policy_name,
+                    member=None,
+                    shift=shift_name,
+                    level=level,
+                    episode_index=index,
+                    episode_seed=self.seed + index,
+                    episode=episode,
+                    success=success,
+                )
+            )
+        return records, successes
+
+    def _metric_rows(
+        self,
+        policy_name: str,
+        shift_name: str,
+        level: float,
+        episodes: list[EpisodeResult],
+        nominal_episodes: list[EpisodeResult] | None,
+    ) -> list[dict[str, Any]]:
+        metric_rows = []
+        for name in (*self.step_metrics, *self.episode_metrics):
+            values = [episode.metrics[name] for episode in episodes]
+            mean, low, high = bootstrap_mean_ci(
+                values,
+                seed=self._bootstrap_seed(policy_name, shift_name, level, f"metric:{name}"),
+            )
+            if nominal_episodes is None:
+                drop = drop_low = drop_high = None
+            else:
+                nominal_values = [episode.metrics[name] for episode in nominal_episodes]
+                drop, drop_low, drop_high = paired_diff_ci(
+                    nominal_values,
+                    values,
+                    seed=self._bootstrap_seed(
+                        policy_name, shift_name, level, f"metric:{name}:drop"
+                    ),
+                )
+            metric_rows.append(
+                {
+                    "policy": policy_name,
+                    "metric": name,
+                    "shift": shift_name,
+                    "level": level,
+                    "n_episodes": self.n_episodes,
+                    "mean": mean,
+                    "ci_low": low,
+                    "ci_high": high,
+                    "drop": drop,
+                    "drop_ci_low": drop_low,
+                    "drop_ci_high": drop_high,
+                }
+            )
+        return metric_rows
+
+    def evaluate(self, policies: Mapping[str, Callable[[Any], Any]]) -> Results:
+        """Evaluate every policy nominally and at each configured shift level."""
+        if not isinstance(policies, Mapping) or not policies:
+            raise ValueError("policies must be a nonempty mapping of names to callables")
+        for name, policy in policies.items():
+            if not isinstance(name, str):
+                raise TypeError("policy names must be strings")
+            if not callable(policy):
+                raise TypeError(f"policy {name!r} must be callable")
+
+        rows = []
+        records: list[EpisodeRecord] = []
+        metric_rows: list[dict[str, Any]] = []
+        for policy_name, policy in policies.items():
+            nominal_episodes = self._run_condition(policy, None, 0.0)
+            nominal_returns = [episode.episode_return for episode in nominal_episodes]
+            condition_records, successes = self._condition_records(
+                policy_name, "nominal", 0.0, nominal_episodes
+            )
+            records.extend(condition_records)
+            rows.append(self._row(policy_name, "nominal", 0.0, nominal_episodes, None, successes))
+            metric_rows.extend(
+                self._metric_rows(policy_name, "nominal", 0.0, nominal_episodes, None)
+            )
+            for shift in self.shifts:
+                for level in shift.levels:
+                    episodes = self._run_condition(policy, shift, level)
+                    condition_records, successes = self._condition_records(
+                        policy_name, shift.name, level, episodes
+                    )
+                    records.extend(condition_records)
+                    rows.append(
+                        self._row(
+                            policy_name, shift.name, level, episodes, nominal_returns, successes
+                        )
+                    )
+                    metric_rows.extend(
+                        self._metric_rows(
+                            policy_name, shift.name, level, episodes, nominal_episodes
+                        )
+                    )
+
+        meta = {
+            "rlshiftlab": __version__,
+            "gymnasium": gym.__version__,
+            "numpy": np.__version__,
+            "python": platform.python_version(),
+            "seed": self.seed,
+            "n_episodes": self.n_episodes,
+            "shifts": {shift.name: list(shift.levels) for shift in self.shifts},
+            "policies": list(policies),
+        }
+        if self.step_metrics or self.episode_metrics:
+            meta["step_metrics"] = list(self.step_metrics)
+            meta["episode_metrics"] = list(self.episode_metrics)
+        return Results(
+            rows=rows,
+            meta=meta,
+            records=records,
+            metric_rows=metric_rows,
+        )
+
+    def evaluate_groups(
+        self, groups: Mapping[str, Mapping[str, Callable[[Any], Any]]]
+    ) -> GroupResults:
+        """Evaluate independently trained frozen members in named groups."""
+        from rlshiftlab.groups import evaluate_groups
+
+        return evaluate_groups(self, groups)
+
+
+__all__ = ["EpisodeRecord", "EpisodeResult", "Suite", "Transition", "run_episodes"]

rlshiftlab/shifts.py

@@ -0,0 +1,170 @@
+"""Observation and action shifts for frozen-policy evaluation."""
+
+from __future__ import annotations
+
+import zlib
+from collections.abc import Sequence
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+
+
+def _level(value: float, name: str, probability: bool = False) -> float:
+    try:
+        level = float(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"{name} level must be a finite number greater than 0") from exc
+    if not np.isfinite(level) or level <= 0:
+        raise ValueError(f"{name} level must be a finite number greater than 0")
+    if probability and level > 1:
+        raise ValueError(f"{name} probability level must be at most 1")
+    return level
+
+
+def _levels(values: Sequence[float], name: str, probability: bool = False) -> list[float]:
+    levels = sorted({_level(value, name, probability) for value in values})
+    if not levels:
+        raise ValueError(f"{name} requires at least one level")
+    return levels
+
+
+def _check_action_space(space: gym.Space[Any], name: str) -> None:
+    if isinstance(space, gym.spaces.Discrete):
+        return
+    if isinstance(space, gym.spaces.Box) and (
+        np.all(np.isfinite(space.low)) and np.all(np.isfinite(space.high))
+    ):
+        return
+    raise TypeError(f"{name} does not support action space {space!r}; expected Discrete or finite Box")
+
+
+def _generator(seed: int, name: str) -> np.random.Generator:
+    return np.random.default_rng([seed, zlib.crc32(name.encode("utf-8"))])
+
+
+class _ShiftWrapper(gym.Wrapper):
+    """Base wrapper with one reproducible random stream per seeded episode."""
+
+    def __init__(self, env: gym.Env, level: float, name: str) -> None:
+        super().__init__(env)
+        self.level = level
+        self._name = name
+        self._rng = np.random.default_rng()
+
+    def reset(
+        self, *, seed: int | None = None, options: dict[str, Any] | None = None
+    ) -> tuple[Any, dict[str, Any]]:
+        if seed is not None:
+            self._rng = _generator(seed, self._name)
+        return self.env.reset(seed=seed, options=options)
+
+
+class _ObservationNoiseWrapper(_ShiftWrapper):
+    def _corrupt(self, observation: Any) -> Any:
+        space = self.observation_space
+        if isinstance(space, gym.spaces.Box):
+            noise = self._rng.normal(0.0, self.level, size=space.shape)
+            return (np.asarray(observation) + noise).astype(space.dtype)
+        if isinstance(space, gym.spaces.Discrete):
+            replace = self._rng.random() < self.level
+            candidate = int(space.start + self._rng.integers(space.n))
+            return candidate if replace else observation
+        raise AssertionError("observation space was validated by ObservationNoise.wrap")
+
+    def reset(
+        self, *, seed: int | None = None, options: dict[str, Any] | None = None
+    ) -> tuple[Any, dict[str, Any]]:
+        observation, info = super().reset(seed=seed, options=options)
+        return self._corrupt(observation), info
+
+    def step(self, action: Any) -> tuple[Any, float, bool, bool, dict[str, Any]]:
+        observation, reward, terminated, truncated, info = self.env.step(action)
+        return self._corrupt(observation), reward, terminated, truncated, info
+
+
+class _ActionSlipWrapper(_ShiftWrapper):
+    def step(self, action: Any) -> tuple[Any, float, bool, bool, dict[str, Any]]:
+        space = self.action_space
+        slip = self._rng.random() < self.level
+        if isinstance(space, gym.spaces.Discrete):
+            candidate = int(space.start + self._rng.integers(space.n))
+        elif isinstance(space, gym.spaces.Box):
+            candidate = self._rng.uniform(space.low, space.high).astype(space.dtype)
+        else:
+            raise TypeError(f"{self._name} does not support action space {space!r}")
+        return self.env.step(candidate if slip else action)
+
+
+class _StickyActionsWrapper(_ShiftWrapper):
+    def __init__(self, env: gym.Env, level: float, name: str) -> None:
+        super().__init__(env, level, name)
+        self._previous_action: Any = None
+        self._has_previous = False
+
+    def reset(
+        self, *, seed: int | None = None, options: dict[str, Any] | None = None
+    ) -> tuple[Any, dict[str, Any]]:
+        self._previous_action = None
+        self._has_previous = False
+        return super().reset(seed=seed, options=options)
+
+    def step(self, action: Any) -> tuple[Any, float, bool, bool, dict[str, Any]]:
+        sticky = self._rng.random() < self.level
+        executed = self._previous_action if sticky and self._has_previous else action
+        if isinstance(self.action_space, gym.spaces.Box):
+            # Keep a separate copy because environments can mutate action arrays.
+            self._previous_action = np.array(executed, copy=True)
+            executed = np.array(executed, copy=True)
+        else:
+            self._previous_action = executed
+        self._has_previous = True
+        return self.env.step(executed)
+
+
+class ObservationNoise:
+    """Add Box noise or randomly replace Discrete observations."""
+
+    name = "observation_noise"
+
+    def __init__(self, levels: Sequence[float]) -> None:
+        self.levels = _levels(levels, self.name)
+
+    def wrap(self, env: gym.Env, level: float) -> gym.Wrapper:
+        """Wrap an environment at a positive observation-noise level."""
+        space = env.observation_space
+        if not isinstance(space, (gym.spaces.Box, gym.spaces.Discrete)):
+            raise TypeError(f"{self.name} does not support observation space {space!r}; expected Box or Discrete")
+        validated = _level(level, self.name, isinstance(space, gym.spaces.Discrete))
+        return _ObservationNoiseWrapper(env, validated, self.name)
+
+
+class ActionSlip:
+    """Randomly replace actions with samples from the action space."""
+
+    name = "action_slip"
+
+    def __init__(self, levels: Sequence[float]) -> None:
+        self.levels = _levels(levels, self.name, probability=True)
+
+    def wrap(self, env: gym.Env, level: float) -> gym.Wrapper:
+        """Wrap an environment at a positive action-slip probability."""
+        _check_action_space(env.action_space, self.name)
+        return _ActionSlipWrapper(env, _level(level, self.name, probability=True), self.name)
+
+
+class StickyActions:
+    """Repeat the previous executed action with a given probability."""
+
+    name = "sticky_actions"
+
+    def __init__(self, levels: Sequence[float]) -> None:
+        self.levels = _levels(levels, self.name, probability=True)
+
+    def wrap(self, env: gym.Env, level: float) -> gym.Wrapper:
+        """Wrap an environment at a positive sticky-action probability."""
+        _check_action_space(env.action_space, self.name)
+        return _StickyActionsWrapper(env, _level(level, self.name, probability=True), self.name)
+
+
+__all__ = ["ActionSlip", "ObservationNoise", "StickyActions"]