cotterbot 0.1.0__py3-none-any.whl

cotter/config.py

@@ -0,0 +1,201 @@
+"""YAML config schema for CLI-driven test runs.
+
+A config file declares which test categories to run and their
+parameters; categories are optional and only the ones present execute.
+Example::
+
+    env: InvertedPendulum-v5
+    algo: PPO
+    base_seed: 0
+    success:
+      type: min_length
+      value: 1000
+    performance:
+      p0: 0.80
+      p1: 0.95
+      alpha: 0.05
+      beta: 0.05
+      n_max: 50
+    safety:
+      n_episodes: 20
+      limits:
+        cotter/joint_velocities: 5.0
+        cotter/actuator_forces: 2.5
+    regression:
+      baseline: artifacts/victim_ppo_inverted_pendulum.zip
+      n_pairs: 30
+      alpha: 0.05
+    adversarial:
+      epsilon: 0.07
+      n_episodes: 20
+      min_success_rate: 0.5
+      train: true
+      timesteps: 150000
+    report: report.json
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import yaml
+
+from cotter.success import make_success_fn
+from cotter.tests.safety import SafetyLimit
+
+
+class ConfigError(ValueError):
+    """The config file is invalid; the message says where and why."""
+
+
+@dataclass
+class PerformanceConfig:
+    p0: float = 0.80
+    p1: float = 0.95
+    alpha: float = 0.05
+    beta: float = 0.05
+    n_max: int = 50
+
+
+@dataclass
+class SafetyConfig:
+    limits: list[SafetyLimit] = field(default_factory=list)
+    n_episodes: int = 20
+    n_workers: int = 1  # parallel rollout workers; 1 = serial
+
+
+@dataclass
+class RegressionConfig:
+    baseline: Path = Path()
+    n_pairs: int = 30
+    alpha: float = 0.05
+    n_workers: int = 1  # parallel rollout workers; 1 = serial
+    metric: str = "return"  # continuous metric for Wilcoxon: "return" or an info key
+
+
+@dataclass
+class AdversarialConfig:
+    epsilon: float = 0.05
+    n_episodes: int = 20
+    min_success_rate: float = 0.5
+    train: bool = True
+    timesteps: int = 150_000
+    use_zoo: bool = False  # reuse/cache the trained adversary in the zoo
+    zoo_root: str | None = None  # override the default ~/.cotter/zoo
+    n_workers: int = 1  # parallel rollout workers for the fixed-N eval; 1 = serial
+    max_seconds: float | None = None  # wall-clock time-box on adversary training
+
+
+@dataclass
+class RunConfig:
+    env: str
+    success: dict
+    algo: str = "PPO"
+    base_seed: int = 0
+    backend: str = "gymnasium"
+    performance: PerformanceConfig | None = None
+    safety: SafetyConfig | None = None
+    regression: RegressionConfig | None = None
+    adversarial: AdversarialConfig | None = None
+    report: Path | None = None
+
+    def success_fn(self):
+        return make_success_fn(self.success)
+
+
+_KNOWN_TOP_KEYS = {
+    "env", "algo", "base_seed", "backend", "success",
+    "performance", "safety", "regression", "adversarial", "report",
+}
+
+
+def _section(data: dict, name: str, cls, **transforms):
+    raw = data.get(name)
+    if raw is None:
+        return None
+    if not isinstance(raw, dict):
+        raise ConfigError(f"'{name}' must be a mapping; got {type(raw).__name__}")
+    valid = {f.name for f in cls.__dataclass_fields__.values()}
+    unknown = set(raw) - valid
+    if unknown:
+        raise ConfigError(f"unknown keys in '{name}': {sorted(unknown)} (expected {sorted(valid)})")
+    kwargs = dict(raw)
+    for key, fn in transforms.items():
+        if key in kwargs:
+            kwargs[key] = fn(kwargs[key])
+    try:
+        return cls(**kwargs)
+    except (TypeError, ValueError) as exc:
+        raise ConfigError(f"invalid '{name}' section: {exc}") from exc
+
+
+def _parse_limits(raw) -> list[SafetyLimit]:
+    if not isinstance(raw, dict) or not raw:
+        raise ConfigError("'safety.limits' must be a non-empty mapping of quantity -> max_abs")
+    limits = []
+    for quantity, max_abs in raw.items():
+        if not isinstance(max_abs, (int, float)) or isinstance(max_abs, bool):
+            raise ConfigError(f"safety limit '{quantity}' must be numeric; got {max_abs!r}")
+        try:
+            limits.append(SafetyLimit(str(quantity), float(max_abs)))
+        except ValueError as exc:
+            raise ConfigError(f"safety limit '{quantity}': {exc}") from exc
+    return limits
+
+
+def parse_config(data: dict, config_dir: Path | None = None) -> RunConfig:
+    """Validate a parsed YAML mapping into a RunConfig.
+
+    Relative paths (regression baseline, report) resolve against
+    ``config_dir`` when given, so configs work from any cwd.
+    """
+    if not isinstance(data, dict):
+        raise ConfigError(f"config root must be a mapping; got {type(data).__name__}")
+    unknown = set(data) - _KNOWN_TOP_KEYS
+    if unknown:
+        raise ConfigError(f"unknown top-level keys: {sorted(unknown)} (expected {sorted(_KNOWN_TOP_KEYS)})")
+    if "env" not in data:
+        raise ConfigError("config must set 'env' (a Gymnasium env id)")
+    if "success" not in data:
+        raise ConfigError("config must set 'success' (a success criterion mapping)")
+
+    def resolve(p) -> Path:
+        path = Path(p)
+        if config_dir is not None and not path.is_absolute():
+            path = config_dir / path
+        return path
+
+    cfg = RunConfig(
+        env=str(data["env"]),
+        success=dict(data["success"]),
+        algo=str(data.get("algo", "PPO")),
+        base_seed=int(data.get("base_seed", 0)),
+        backend=str(data.get("backend", "gymnasium")),
+        performance=_section(data, "performance", PerformanceConfig),
+        safety=_section(data, "safety", SafetyConfig, limits=_parse_limits),
+        regression=_section(data, "regression", RegressionConfig, baseline=resolve),
+        adversarial=_section(data, "adversarial", AdversarialConfig),
+        report=resolve(data["report"]) if "report" in data else None,
+    )
+    if cfg.safety is not None and not cfg.safety.limits:
+        raise ConfigError("'safety' section requires a 'limits' mapping")
+    if cfg.regression is not None and cfg.regression.baseline == Path():
+        raise ConfigError("'regression' section requires a 'baseline' policy path")
+    make_success_fn(cfg.success)  # fail at load time, not mid-run
+    return cfg
+
+
+def load_config(path: str | Path) -> RunConfig:
+    """Load and validate a YAML config file."""
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(f"config file not found: {path}")
+    try:
+        data = yaml.safe_load(path.read_text())
+    except yaml.YAMLError as exc:
+        raise ConfigError(f"{path} is not valid YAML: {exc}") from exc
+    try:
+        return parse_config(data, config_dir=path.resolve().parent)
+    except ConfigError as exc:
+        raise ConfigError(f"{path}: {exc}") from exc

cotter/envs/factory.py

@@ -0,0 +1,31 @@
+"""Picklable environment factory for parallel rollouts.
+
+``AsyncVectorEnv`` constructs its sub-environments inside worker
+processes, so it needs a factory that survives pickling under the
+``spawn`` start method used on macOS. A module-level callable class does;
+a local closure does not. :class:`WrappedEnvFactory` rebuilds the same
+instrumented env that :func:`cotter.pipeline.make_env` produces.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import gymnasium as gym
+
+from cotter.envs.registry import make_env_by_id
+from cotter.envs.wrapper import CotterWrapper
+
+
+@dataclass(frozen=True)
+class WrappedEnvFactory:
+    """Zero-argument callable that builds one wrapped env by id."""
+
+    env_id: str
+
+    def __call__(self) -> gym.Env:
+        env = make_env_by_id(self.env_id)
+        try:
+            return CotterWrapper(env)
+        except TypeError:
+            return env

cotter/envs/registry.py

@@ -0,0 +1,37 @@
+"""Environment creation with optional extension registration.
+
+gymnasium-robotics environments (Fetch, Shadow Hand, ...) are not
+registered until the package is imported. :func:`make_env_by_id` retries
+a failed ``gym.make`` after registering installed extension packages, so
+config files can name any installed env without boilerplate.
+"""
+
+from __future__ import annotations
+
+import gymnasium as gym
+from gymnasium.error import NameNotFound
+
+_registered = False
+
+
+def register_extension_envs() -> None:
+    """Register env packages that require an explicit import (idempotent)."""
+    global _registered
+    if _registered:
+        return
+    try:
+        import gymnasium_robotics
+
+        gym.register_envs(gymnasium_robotics)
+    except ImportError:
+        pass
+    _registered = True
+
+
+def make_env_by_id(env_id: str, **kwargs) -> gym.Env:
+    """``gym.make`` that knows about installed extension packages."""
+    try:
+        return gym.make(env_id, **kwargs)
+    except NameNotFound:
+        register_extension_envs()
+        return gym.make(env_id, **kwargs)

cotter/envs/wrapper.py

@@ -0,0 +1,73 @@
+"""Gymnasium wrapper that surfaces safety-relevant MuJoCo quantities.
+
+:class:`CotterWrapper` guarantees that every ``step()`` info dict contains
+the physical quantities the safety tests read, copied out of the MuJoCo
+``data`` structure of the wrapped environment:
+
+==============================  =========================================
+info key                        contents
+==============================  =========================================
+``cotter/joint_velocities``      ``data.qvel`` — generalized joint
+                                 velocities (rad/s or m/s per DoF)
+``cotter/actuator_forces``       ``data.actuator_force`` — scalar force/
+                                 torque produced by each actuator (N or Nm)
+``cotter/contact_count``         ``data.ncon`` — number of active contact
+                                 points this step
+``cotter/contact_forces``        ``data.cfrc_ext`` L2 norms — magnitude of
+                                 the external contact wrench on each body
+                                 (one value per body; N-scale)
+==============================  =========================================
+
+The wrapper works with any Gymnasium MuJoCo environment (anything whose
+unwrapped env exposes a ``data`` attribute with these fields) and fails
+loudly at construction time otherwise.
+"""
+
+from __future__ import annotations
+
+import gymnasium as gym
+import numpy as np
+
+JOINT_VELOCITIES = "cotter/joint_velocities"
+ACTUATOR_FORCES = "cotter/actuator_forces"
+CONTACT_COUNT = "cotter/contact_count"
+CONTACT_FORCES = "cotter/contact_forces"
+
+INSTRUMENTED_KEYS = (JOINT_VELOCITIES, ACTUATOR_FORCES, CONTACT_COUNT, CONTACT_FORCES)
+
+
+class CotterWrapper(gym.Wrapper):
+    """Instrument a MuJoCo env's info dict with safety-relevant quantities."""
+
+    def __init__(self, env: gym.Env) -> None:
+        super().__init__(env)
+        data = getattr(env.unwrapped, "data", None)
+        if data is None or not all(
+            hasattr(data, attr)
+            for attr in ("qvel", "actuator_force", "ncon", "cfrc_ext")
+        ):
+            raise TypeError(
+                f"CotterWrapper requires a MuJoCo-backed environment exposing "
+                f"`unwrapped.data` with qvel/actuator_force/ncon; "
+                f"got {type(env.unwrapped).__name__}. Wrap a Gymnasium MuJoCo "
+                "env (e.g. gym.make('InvertedPendulum-v5'))."
+            )
+
+    def _instrument(self, info: dict) -> dict:
+        data = self.env.unwrapped.data
+        info[JOINT_VELOCITIES] = np.array(data.qvel, dtype=float, copy=True)
+        info[ACTUATOR_FORCES] = np.array(data.actuator_force, dtype=float, copy=True)
+        info[CONTACT_COUNT] = int(data.ncon)
+        # one L2 wrench magnitude per body; body 0 is the world and always 0
+        info[CONTACT_FORCES] = np.linalg.norm(
+            np.asarray(data.cfrc_ext, dtype=float), axis=1
+        )
+        return info
+
+    def reset(self, **kwargs):
+        obs, info = self.env.reset(**kwargs)
+        return obs, self._instrument(info)
+
+    def step(self, action):
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        return obs, reward, terminated, truncated, self._instrument(info)

cotter/pipeline.py

@@ -0,0 +1,214 @@
+"""Execute a RunConfig against a policy: the engine behind `cotter run`.
+
+Runs whichever test categories the config declares and returns the
+aggregated :class:`~cotter.report.TestReport`. Categories run in a fixed
+order (performance, safety, regression, adversarial) with seeds derived
+from ``base_seed`` per category, so adding or removing one category does
+not change another's rollouts.
+"""
+
+from __future__ import annotations
+
+import importlib
+from pathlib import Path
+from typing import Callable
+
+import gymnasium as gym
+
+from cotter.backends import BackendFactory
+from cotter.config import RunConfig
+from cotter.envs.factory import WrappedEnvFactory
+from cotter.envs.wrapper import CotterWrapper
+from cotter.policy import Policy, load_policy
+from cotter.report import TestReport
+from cotter.runner import (
+    make_seed_sequence,
+    rollout_one,
+    run_rollouts,
+    run_rollouts_parallel,
+)
+from cotter.tests.adversarial import get_adversary, run_adversarial_test
+from cotter.tests.regression import mcnemar_exact, wilcoxon_regression
+from cotter.tests.safety import evaluate_safety
+from cotter.tests.sprt import run_sprt
+
+# per-category seed offsets, kept stable so results are comparable across
+# configs that enable different category subsets
+_PERF_SEED, _SAFETY_SEED, _REGRESSION_SEED, _ADV_SEED = 100, 200, 300, 400
+
+
+def resolve_algo(name: str):
+    """Look up an SB3 algorithm class by name (PPO, SAC, TD3, ...)."""
+    sb3 = importlib.import_module("stable_baselines3")
+    algo = getattr(sb3, name, None)
+    if algo is None:
+        available = [a for a in ("A2C", "DDPG", "DQN", "PPO", "SAC", "TD3") if hasattr(sb3, a)]
+        raise ValueError(f"unknown SB3 algorithm '{name}' (available: {available})")
+    return algo
+
+
+def make_env(env_id: str, backend: str = "gymnasium", log: Callable[[str], None] = print) -> gym.Env:
+    """Create the test environment via the named backend.
+
+    The env is instrumented with :class:`CotterWrapper` when the backend
+    produces a MuJoCo-backed env; otherwise safety quantities are
+    unavailable and a note is logged.
+    """
+    env = BackendFactory.from_name(backend).make_env(env_id)
+    if not isinstance(env, CotterWrapper):
+        log(f"[cotter] {env_id} is not MuJoCo-backed; safety quantities unavailable")
+    return env
+
+
+def _obtain_adversary(policy_path, policy, env, cfg, adv_cfg, log):
+    """Load a cached adversary from the zoo, or train one and cache it.
+
+    Falls back to :func:`get_adversary` (train-with-random-fallback) when
+    the zoo is disabled or the victim cannot be hashed.
+    """
+    if not adv_cfg.use_zoo:
+        return get_adversary(
+            policy, env, adv_cfg.epsilon, timesteps=adv_cfg.timesteps,
+            seed=cfg.base_seed, log=log, max_seconds=adv_cfg.max_seconds,
+        )
+
+    from cotter.tests.adversarial import train_adversary
+    from cotter.zoo import AdversaryZoo
+
+    zoo = AdversaryZoo(adv_cfg.zoo_root) if adv_cfg.zoo_root else AdversaryZoo()
+    # the artifact path hashes identically however it is later referenced
+    victim_ref = Path(policy_path)
+
+    cached = zoo.load(cfg.env, victim_ref, adv_cfg.epsilon)
+    if cached is not None:
+        log(f"[cotter]   reusing cached adversary from zoo ({zoo.root})")
+        return cached, "reused cached zoo adversary"
+
+    log(f"[cotter]   no cached adversary; training and storing in zoo ({zoo.root})")
+    adversary = train_adversary(
+        policy, env, adv_cfg.epsilon, timesteps=adv_cfg.timesteps,
+        seed=cfg.base_seed, max_seconds=adv_cfg.max_seconds,
+    )
+    entry = zoo.save(adversary, cfg.env, victim_ref, notes=f"trained via cotter run, {cfg.env}")
+    log(f"[cotter]   stored adversary at {entry.path}")
+    return adversary, "trained PPO adversary (cached in zoo)"
+
+
+def run_from_config(
+    policy_path: str | Path,
+    cfg: RunConfig,
+    log: Callable[[str], None] = print,
+) -> TestReport:
+    env = make_env(cfg.env, backend=cfg.backend, log=log)
+    if cfg.safety is not None and not isinstance(env, CotterWrapper):
+        raise ValueError(
+            f"config declares safety limits but {cfg.env} is not MuJoCo-backed, "
+            "so the instrumented quantities do not exist"
+        )
+
+    algo = resolve_algo(cfg.algo)
+    policy = load_policy(Path(policy_path), env, algo=algo)
+    success_fn = cfg.success_fn()
+    log(f"[cotter] loaded policy '{policy.name}' onto {cfg.env}")
+
+    report = TestReport(
+        policy_name=policy.name,
+        env_id=cfg.env,
+        metadata={"base_seed": cfg.base_seed, "success": cfg.success, "algo": cfg.algo},
+    )
+
+    if cfg.performance is not None:
+        p = cfg.performance
+        log(f"[cotter] performance: SPRT p0={p.p0} p1={p.p1} n_max={p.n_max}")
+        seeds = make_seed_sequence(p.n_max, cfg.base_seed + _PERF_SEED)
+
+        def trial(i: int) -> bool:
+            rec = rollout_one(policy, env, seeds[i], success_fn, record_infos=False)
+            log(f"[cotter]   trial {i + 1}: length={rec.length} "
+                f"return={rec.total_reward:.1f} success={rec.success}")
+            return rec.success
+
+        result = run_sprt(trial, p0=p.p0, p1=p.p1, alpha=p.alpha, beta=p.beta, n_max=p.n_max)
+        report.add_sprt(result)
+        log(f"[cotter]   => {result.decision.value} after {result.n_trials} trials")
+
+    factory = WrappedEnvFactory(cfg.env)
+
+    def dispatch(pol, n, *, seeds=None, base=0, record_infos, n_workers):
+        if n_workers > 1:
+            return run_rollouts_parallel(
+                pol, factory, n, success_fn, seeds=seeds, base_seed=base,
+                record_infos=record_infos, n_workers=n_workers,
+            )
+        return run_rollouts(
+            pol, env, n, success_fn, seeds=seeds, base_seed=base,
+            record_infos=record_infos,
+        )
+
+    if cfg.safety is not None:
+        s = cfg.safety
+        log(f"[cotter] safety: {len(s.limits)} limit(s) over {s.n_episodes} episodes"
+            + (f" ({s.n_workers} workers)" if s.n_workers > 1 else ""))
+        rollouts = dispatch(
+            policy, s.n_episodes, base=cfg.base_seed + _SAFETY_SEED,
+            record_infos=True, n_workers=s.n_workers,
+        )
+        result = evaluate_safety(rollouts.episode_infos, s.limits)
+        report.add_safety(result)
+        for quantity, worst in result.worst_observed.items():
+            limit = next(l.max_abs for l in s.limits if l.quantity == quantity)
+            log(f"[cotter]   worst |{quantity}| = {worst:.4f} (limit {limit})")
+        log(f"[cotter]   => {result.decision.value}")
+
+    if cfg.regression is not None:
+        r = cfg.regression
+        log(f"[cotter] regression: vs baseline {r.baseline} on {r.n_pairs} paired seeds")
+        baseline = load_policy(r.baseline, env, algo=algo, name=f"baseline:{r.baseline.stem}")
+        seeds = make_seed_sequence(r.n_pairs, cfg.base_seed + _REGRESSION_SEED)
+        # info-based metrics need per-episode final info recorded
+        record = r.metric != "return"
+        base_rs = dispatch(baseline, r.n_pairs, seeds=seeds, record_infos=record, n_workers=r.n_workers)
+        cand_rs = dispatch(policy, r.n_pairs, seeds=seeds, record_infos=record, n_workers=r.n_workers)
+        # baseline vs candidate: the CLI policy is the candidate
+        mcnemar = mcnemar_exact(base_rs.successes, cand_rs.successes, alpha=r.alpha)
+        report.add_regression(mcnemar, name="success_mcnemar")
+        wilcoxon = wilcoxon_regression(
+            base_rs.metric_values(r.metric), cand_rs.metric_values(r.metric), alpha=r.alpha
+        )
+        report.add_regression(wilcoxon, name=f"{r.metric}_wilcoxon")
+        log(f"[cotter]   baseline {base_rs.success_rate:.0%} vs candidate "
+            f"{cand_rs.success_rate:.0%} (metric: {r.metric})")
+        log(f"[cotter]   => McNemar {mcnemar.decision.value} (p={mcnemar.p_value:.3g}), "
+            f"Wilcoxon {wilcoxon.decision.value} (p={wilcoxon.p_value:.3g})")
+
+    if cfg.adversarial is not None:
+        a = cfg.adversarial
+        log(f"[cotter] adversarial: eps={a.epsilon} over {a.n_episodes} episodes")
+        random_result = run_adversarial_test(
+            policy, env, success_fn, epsilon=a.epsilon, n_episodes=a.n_episodes,
+            min_success_rate=a.min_success_rate, base_seed=cfg.base_seed + _ADV_SEED,
+            notes="uniform random baseline",
+            n_workers=a.n_workers, env_factory=factory,
+        )
+        report.add_adversarial(random_result, name="random_baseline")
+        log(f"[cotter]   random baseline: {random_result.clean_success_rate:.0%} clean -> "
+            f"{random_result.adversarial_success_rate:.0%} perturbed")
+        if a.train:
+            adversary, notes = _obtain_adversary(policy_path, policy, env, cfg, a, log)
+            learned_result = run_adversarial_test(
+                policy, env, success_fn, epsilon=a.epsilon, n_episodes=a.n_episodes,
+                adversary=adversary, min_success_rate=a.min_success_rate,
+                base_seed=cfg.base_seed + _ADV_SEED,
+                notes=notes or "trained PPO adversary",
+                n_workers=a.n_workers, env_factory=factory,
+            )
+            report.add_adversarial(learned_result, name=f"learned_{adversary.name}")
+            log(f"[cotter]   {adversary.name} adversary: "
+                f"{learned_result.clean_success_rate:.0%} clean -> "
+                f"{learned_result.adversarial_success_rate:.0%} perturbed")
+
+    if cfg.report is not None:
+        path = report.to_json(cfg.report)
+        log(f"[cotter] JSON report written to {path}")
+
+    return report