agent-learning 0.4.1__py3-none-any.whl

agent_learning/policy/contextual_softmax.py

@@ -0,0 +1,243 @@
+"""Contextual softmax bandit policy.
+
+The policy stores a weight matrix ``W`` of shape ``(K, d)`` where ``K``
+is the number of actions and ``d`` is the feature dimension. Action
+logits at inference time are computed as ``logits = W @ phi`` for a
+caller-supplied feature vector ``phi`` of length ``d``. A numerically
+stable softmax then turns the logits into a probability distribution
+the policy samples from.
+
+The weight matrix lives in :attr:`PolicySnapshot.metadata`:
+
+- ``metadata["feature_dim"]`` (int) - the dimension ``d``.
+- ``metadata["context_weights"]`` (dict[action_id, list[float]]) -
+  one length-``d`` row per action.
+
+Using ``metadata`` keeps the persisted snapshot schema unchanged
+(the existing ``logits`` map stays untouched and unused). Storage
+and Cosmos round-trip continue to work with no migration.
+"""
+
+from __future__ import annotations
+
+import math
+import random
+from typing import Any, Dict, List, Optional, Sequence
+
+import numpy as np
+
+from ..types import Action, PolicySnapshot
+from .base import Decision, Policy
+
+
+_METADATA_WEIGHTS_KEY = "context_weights"
+_METADATA_DIM_KEY = "feature_dim"
+
+
+class ContextualSoftmaxPolicy(Policy):
+    """Linear contextual softmax over a fixed action set.
+
+    Construction:
+        policy = ContextualSoftmaxPolicy.from_actions(
+            actions, feature_dim=25, agent_id="dq"
+        )
+
+    Continuation from a snapshot:
+        policy = ContextualSoftmaxPolicy.from_snapshot(snapshot)
+    """
+
+    def __init__(
+        self,
+        snapshot: PolicySnapshot,
+        *,
+        rng: Optional[random.Random] = None,
+        max_weight_abs: float = 10.0,
+    ) -> None:
+        if not snapshot.actions:
+            raise ValueError("ContextualSoftmaxPolicy requires at least one action.")
+        feature_dim = int(snapshot.metadata.get(_METADATA_DIM_KEY, 0))
+        if feature_dim <= 0:
+            raise ValueError(
+                "ContextualSoftmaxPolicy requires snapshot.metadata['feature_dim'] > 0."
+            )
+        weights = snapshot.metadata.get(_METADATA_WEIGHTS_KEY)
+        if not isinstance(weights, dict):
+            weights = {}
+        # Materialise a row of zeros for any action missing a weight row.
+        for action in snapshot.actions:
+            row = weights.get(action.id)
+            if row is None or len(row) != feature_dim:
+                weights[action.id] = [0.0] * feature_dim
+            else:
+                weights[action.id] = [float(x) for x in row]
+        snapshot.metadata[_METADATA_WEIGHTS_KEY] = weights
+        snapshot.metadata[_METADATA_DIM_KEY] = feature_dim
+
+        self._snapshot = snapshot
+        self._rng = rng or random.Random()
+        self._max_weight_abs = float(max_weight_abs)
+        self._feature_dim = feature_dim
+
+    # ------------------------------------------------------------------
+    # Constructors
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def from_actions(
+        cls,
+        actions: Sequence[Action],
+        *,
+        feature_dim: int,
+        agent_id: str = "default",
+        task_id: str = "default",
+        rng: Optional[random.Random] = None,
+        max_weight_abs: float = 10.0,
+        initial_weights: Optional[Dict[str, Sequence[float]]] = None,
+    ) -> "ContextualSoftmaxPolicy":
+        if feature_dim <= 0:
+            raise ValueError("feature_dim must be a positive integer.")
+        weights: Dict[str, List[float]] = {
+            a.id: [0.0] * feature_dim for a in actions
+        }
+        if initial_weights:
+            for action_id, row in initial_weights.items():
+                if action_id in weights:
+                    row_list = [float(x) for x in row]
+                    if len(row_list) != feature_dim:
+                        raise ValueError(
+                            f"initial_weights['{action_id}'] has length "
+                            f"{len(row_list)}, expected {feature_dim}."
+                        )
+                    weights[action_id] = row_list
+        snapshot = PolicySnapshot(
+            agent_id=agent_id,
+            task_id=task_id,
+            version=0,
+            actions=list(actions),
+            logits={a.id: 0.0 for a in actions},  # kept for back-compat, unused
+            metadata={
+                _METADATA_DIM_KEY: int(feature_dim),
+                _METADATA_WEIGHTS_KEY: weights,
+                "policy_kind": "contextual_softmax",
+            },
+        )
+        return cls(snapshot, rng=rng, max_weight_abs=max_weight_abs)
+
+    @classmethod
+    def from_snapshot(
+        cls,
+        snapshot: PolicySnapshot,
+        *,
+        rng: Optional[random.Random] = None,
+        max_weight_abs: float = 10.0,
+    ) -> "ContextualSoftmaxPolicy":
+        return cls(snapshot, rng=rng, max_weight_abs=max_weight_abs)
+
+    # ------------------------------------------------------------------
+    # Policy interface
+    # ------------------------------------------------------------------
+
+    def actions(self) -> Sequence[Action]:
+        return tuple(self._snapshot.actions)
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def _weights_matrix(self) -> np.ndarray:
+        rows = [
+            self._snapshot.metadata[_METADATA_WEIGHTS_KEY][a.id]
+            for a in self._snapshot.actions
+        ]
+        return np.asarray(rows, dtype=np.float64)
+
+    def probabilities(self, state: Optional[Any] = None) -> List[float]:
+        """Softmax over ``W @ phi`` in deterministic action order.
+
+        If ``state`` is None the policy returns the uniform distribution
+        over actions (equivalent to ``phi == 0``).
+        """
+        n_actions = len(self._snapshot.actions)
+        if state is None:
+            uniform = 1.0 / n_actions
+            return [uniform] * n_actions
+        phi = np.asarray(state, dtype=np.float64).reshape(-1)
+        if phi.shape[0] != self._feature_dim:
+            raise ValueError(
+                f"State vector has length {phi.shape[0]}, expected "
+                f"{self._feature_dim}."
+            )
+        logits = self._weights_matrix() @ phi
+        logits -= logits.max()  # numerical stability
+        exp = np.exp(logits)
+        denom = float(exp.sum())
+        if denom <= 0.0 or not np.isfinite(denom):  # pragma: no cover
+            uniform = 1.0 / n_actions
+            return [uniform] * n_actions
+        return (exp / denom).tolist()
+
+    def choose(self, state: Optional[Any] = None) -> Decision:
+        probs = self.probabilities(state)
+        idx = self._weighted_sample(probs)
+        chosen = self._snapshot.actions[idx]
+        prob = max(probs[idx], 1e-12)
+        return Decision(action=chosen, logprob=math.log(prob), probabilities=probs)
+
+    def snapshot(self) -> PolicySnapshot:
+        return PolicySnapshot.from_dict(self._snapshot.to_dict())
+
+    def apply_update(
+        self,
+        deltas: Dict[str, Any],
+        *,
+        baseline: float,
+        episodes_seen: int,
+    ) -> None:
+        """Apply additive updates to the context weights, clip, and bump version.
+
+        ``deltas`` accepts two shapes for backward compatibility with
+        the marginal :class:`SoftmaxPolicy` API:
+
+        - ``{action_id: float}`` - apply ``delta`` to every dimension of
+          that action's weight row (rare; useful only as a smoke test).
+        - ``{action_id: Sequence[float]}`` of length ``feature_dim`` -
+          element-wise add to the weight row.
+        """
+        weights: Dict[str, List[float]] = self._snapshot.metadata[_METADATA_WEIGHTS_KEY]
+        for action_id, delta in deltas.items():
+            if action_id not in weights:
+                continue
+            current = np.asarray(weights[action_id], dtype=np.float64)
+            if isinstance(delta, (int, float)):
+                current = current + float(delta)
+            else:
+                delta_arr = np.asarray(delta, dtype=np.float64).reshape(-1)
+                if delta_arr.shape[0] != self._feature_dim:
+                    raise ValueError(
+                        f"Delta for '{action_id}' has length "
+                        f"{delta_arr.shape[0]}, expected {self._feature_dim}."
+                    )
+                current = current + delta_arr
+            current = np.clip(current, -self._max_weight_abs, self._max_weight_abs)
+            weights[action_id] = current.tolist()
+        self._snapshot.metadata[_METADATA_WEIGHTS_KEY] = weights
+        self._snapshot.baseline = float(baseline)
+        self._snapshot.episodes_seen = int(episodes_seen)
+        self._snapshot.updates_applied += 1
+        self._snapshot.advance_version()
+
+    # ------------------------------------------------------------------
+    # Internals
+    # ------------------------------------------------------------------
+
+    def _weighted_sample(self, probs: Sequence[float]) -> int:
+        roll = self._rng.random()
+        cumulative = 0.0
+        for idx, p in enumerate(probs):
+            cumulative += p
+            if roll <= cumulative:
+                return idx
+        return len(probs) - 1  # pragma: no cover - rounding tail
+
+
+__all__ = ["ContextualSoftmaxPolicy"]

agent_learning/policy/softmax_bandit.py

@@ -0,0 +1,157 @@
+"""Softmax bandit policy over a fixed set of discrete actions.
+
+The policy stores one logit per action. Action probabilities are
+computed via a numerically stable softmax. The action space is
+fixed at construction time; the learner mutates only the logits and
+the EMA baseline.
+"""
+
+from __future__ import annotations
+
+import math
+import random
+from typing import Any, Dict, List, Optional, Sequence
+
+import numpy as np
+
+from ..types import Action, PolicySnapshot
+from .base import Decision, Policy
+
+
+class SoftmaxPolicy(Policy):
+    """Discrete softmax policy with one logit per action.
+
+    Construction:
+        policy = SoftmaxPolicy.from_actions(actions, agent_id="dq")
+
+    Continuation from a snapshot (e.g. after process restart):
+        policy = SoftmaxPolicy.from_snapshot(snapshot)
+    """
+
+    def __init__(
+        self,
+        snapshot: PolicySnapshot,
+        *,
+        rng: Optional[random.Random] = None,
+        max_logit_abs: float = 10.0,
+    ) -> None:
+        if not snapshot.actions:
+            raise ValueError("SoftmaxPolicy requires at least one action.")
+        self._snapshot = snapshot
+        self._rng = rng or random.Random()
+        self._max_logit_abs = max_logit_abs
+        # Make sure every action has a logit
+        for action in snapshot.actions:
+            snapshot.logits.setdefault(action.id, 0.0)
+
+    # ------------------------------------------------------------------
+    # Constructors
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def from_actions(
+        cls,
+        actions: Sequence[Action],
+        *,
+        agent_id: str = "default",
+        task_id: str = "default",
+        rng: Optional[random.Random] = None,
+        max_logit_abs: float = 10.0,
+        initial_logits: Optional[Dict[str, float]] = None,
+    ) -> "SoftmaxPolicy":
+        logits = {a.id: 0.0 for a in actions}
+        if initial_logits:
+            for action_id, value in initial_logits.items():
+                if action_id in logits:
+                    logits[action_id] = float(value)
+        snapshot = PolicySnapshot(
+            agent_id=agent_id,
+            task_id=task_id,
+            version=0,
+            actions=list(actions),
+            logits=logits,
+        )
+        return cls(snapshot, rng=rng, max_logit_abs=max_logit_abs)
+
+    @classmethod
+    def from_snapshot(
+        cls,
+        snapshot: PolicySnapshot,
+        *,
+        rng: Optional[random.Random] = None,
+        max_logit_abs: float = 10.0,
+    ) -> "SoftmaxPolicy":
+        return cls(snapshot, rng=rng, max_logit_abs=max_logit_abs)
+
+    # ------------------------------------------------------------------
+    # Policy interface
+    # ------------------------------------------------------------------
+
+    def actions(self) -> Sequence[Action]:
+        return tuple(self._snapshot.actions)
+
+    def probabilities(self) -> List[float]:
+        """Softmax over the action logits in deterministic order."""
+        logits = np.array(
+            [self._snapshot.logits[a.id] for a in self._snapshot.actions],
+            dtype=np.float64,
+        )
+        logits -= logits.max()  # numerical stability
+        exp = np.exp(logits)
+        denom = exp.sum()
+        if denom <= 0.0 or not np.isfinite(denom):  # pragma: no cover
+            uniform = 1.0 / len(self._snapshot.actions)
+            return [uniform] * len(self._snapshot.actions)
+        probs = exp / denom
+        return probs.tolist()
+
+    def choose(self, state: Optional[Any] = None) -> Decision:
+        # ``state`` is accepted for interface compatibility with
+        # contextual policies and is intentionally ignored: this is a
+        # marginal softmax that conditions only on its stored logits.
+        del state
+        probs = self.probabilities()
+        idx = self._weighted_sample(probs)
+        chosen = self._snapshot.actions[idx]
+        # log probability is clamped to avoid -inf for tiny probabilities
+        prob = max(probs[idx], 1e-12)
+        return Decision(action=chosen, logprob=math.log(prob), probabilities=probs)
+
+    def snapshot(self) -> PolicySnapshot:
+        return PolicySnapshot.from_dict(self._snapshot.to_dict())
+
+    def apply_update(
+        self,
+        deltas: Dict[str, float],
+        *,
+        baseline: float,
+        episodes_seen: int,
+    ) -> None:
+        """Apply per-action logit deltas, clip, and bump version."""
+        for action_id, delta in deltas.items():
+            if action_id not in self._snapshot.logits:
+                continue
+            new_value = self._snapshot.logits[action_id] + float(delta)
+            new_value = max(-self._max_logit_abs, min(self._max_logit_abs, new_value))
+            self._snapshot.logits[action_id] = new_value
+        self._snapshot.baseline = float(baseline)
+        self._snapshot.episodes_seen = int(episodes_seen)
+        self._snapshot.updates_applied += 1
+        self._snapshot.advance_version()
+
+    # ------------------------------------------------------------------
+    # Internals
+    # ------------------------------------------------------------------
+
+    def _weighted_sample(self, probs: Sequence[float]) -> int:
+        """Sample an index from ``probs`` using the policy's RNG."""
+        roll = self._rng.random()
+        cumulative = 0.0
+        for idx, p in enumerate(probs):
+            cumulative += p
+            if roll <= cumulative:
+                return idx
+        return len(probs) - 1  # pragma: no cover - rounding tail
+
+
+__all__ = ["SoftmaxPolicy"]

agent_learning/py.typed

@@ -0,0 +1 @@
+"""Marker file used by PEP 561 type checkers."""

agent_learning/rewards/__init__.py

@@ -0,0 +1,6 @@
+"""Reward shaping (metrics → scalar reward) and reward persistence."""
+
+from .shaping import RewardShaper, shape_episode_reward
+from .writer import RewardWriter
+
+__all__ = ["RewardShaper", "RewardWriter", "shape_episode_reward"]

agent_learning/rewards/shaping.py

@@ -0,0 +1,121 @@
+"""Reward shaping: combine multiple metric results into one scalar reward.
+
+The shaper turns each metric's normalised score (in ``[0, 1]``) into
+a signed contribution (in ``[-1, 1]``) and sums them with the
+configured weights. Additional behavioural penalties (latency, cost)
+can be applied after the metric sum. The final reward is clamped to
+``[-1, 1]`` so it is always safe to feed into a policy-gradient
+update.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable, List, Optional
+
+from ..config import ShapingConfig
+from ..types import Episode, MetricName, MetricResult
+
+
+@dataclass
+class ShapedReward:
+    """Result of shaping a metric collection into a scalar reward.
+
+    ``metric_contributions`` exposes the signed contribution of every
+    metric individually so callers can persist per-metric rewards
+    alongside the aggregate.
+    """
+
+    value: float
+    metric_contributions: List[tuple]
+    penalties: List[tuple]
+
+
+class RewardShaper:
+    """Combine multiple :class:`MetricResult` into a scalar in ``[-1, 1]``."""
+
+    def __init__(self, config: Optional[ShapingConfig] = None) -> None:
+        self._config = config or ShapingConfig()
+
+    @property
+    def config(self) -> ShapingConfig:
+        return self._config
+
+    def shape(
+        self,
+        results: Iterable[MetricResult],
+        *,
+        latency_ms: Optional[int] = None,
+        routing_correct: Optional[bool] = None,
+        hallucinated_class: bool = False,
+    ) -> ShapedReward:
+        """Combine scores + behavioural signals into a single scalar.
+
+        Args:
+            results: Per-score metric outputs.
+            latency_ms: Episode latency in milliseconds; if above the
+                configured threshold the latency penalty is applied.
+            routing_correct: ``True`` if the calling system routed the
+                request to a class id in the allowed set
+                (→ add ``route_correct_reward``); ``False`` if it
+                picked a class outside the allowed set (→ add
+                ``route_wrong_penalty``, which is negative). ``None``
+                disables the term entirely (no signal either way).
+            hallucinated_class: ``True`` if the rendered output
+                references an entity/class id not in the allowed set
+                (→ add ``hallucinated_class_penalty``).
+        """
+        weights = {
+            MetricName.INTENT_RESOLUTION: self._config.intent_resolution_weight,
+            MetricName.TASK_ADHERENCE: self._config.task_adherence_weight,
+            MetricName.TASK_COMPLETION: self._config.task_completion_weight,
+        }
+
+        contributions: List[tuple] = []
+        total = 0.0
+        for r in results:
+            if r.normalized is None or r.status != "completed":
+                continue
+            weight = weights.get(r.metric, 0.0)
+            # Map [0, 1] -> [-1, 1]
+            signed = 2.0 * r.normalized - 1.0
+            contribution = weight * signed
+            total += contribution
+            contributions.append((r.metric, weight, signed, contribution))
+
+        penalties: List[tuple] = []
+        if (
+            latency_ms is not None
+            and latency_ms > self._config.latency_penalty_threshold_ms
+        ):
+            total += self._config.latency_penalty_value
+            penalties.append(("latency", latency_ms, self._config.latency_penalty_value))
+
+        if routing_correct is True:
+            total += self._config.route_correct_reward
+            penalties.append(("route_correct", True, self._config.route_correct_reward))
+        elif routing_correct is False:
+            total += self._config.route_wrong_penalty
+            penalties.append(("route_wrong", False, self._config.route_wrong_penalty))
+
+        if hallucinated_class:
+            total += self._config.hallucinated_class_penalty
+            penalties.append(
+                ("hallucinated_class", True, self._config.hallucinated_class_penalty)
+            )
+
+        clamped = max(-1.0, min(1.0, total))
+        return ShapedReward(value=clamped, metric_contributions=contributions, penalties=penalties)
+
+
+def shape_episode_reward(
+    episode: Episode,
+    results: Iterable[MetricResult],
+    config: Optional[ShapingConfig] = None,
+) -> ShapedReward:
+    """Convenience wrapper that pulls latency from the episode automatically."""
+    shaper = RewardShaper(config)
+    return shaper.shape(results, latency_ms=episode.request_latency_ms)
+
+
+__all__ = ["RewardShaper", "ShapedReward", "shape_episode_reward"]

agent_learning/rewards/writer.py

@@ -0,0 +1,130 @@
+"""Persist metric results and aggregate rewards for an episode."""
+
+from __future__ import annotations
+
+import logging
+from typing import Iterable, List, Optional
+
+from ..storage.base import LearningStore
+from ..storage.cosmos import get_default_store
+from ..types import Episode, MetricName, MetricResult, Reward, RewardSource
+from .shaping import ShapedReward
+
+logger = logging.getLogger(__name__)
+
+
+class RewardWriter:
+    """Translate metric results into stored :class:`Reward` records.
+
+    The writer stores one ``Reward`` per metric (``source=METRIC``)
+    plus a single aggregate reward (``source=AGGREGATE``) that the
+    learner consumes. Storing per-metric values keeps the learning history
+    complete and lets downstream analytics decompose the aggregate.
+    """
+
+    def __init__(self, store: Optional[LearningStore] = None) -> None:
+        self._store = store
+
+    @property
+    def store(self) -> LearningStore:
+        if self._store is None:
+            self._store = get_default_store()
+        return self._store
+
+    def write(
+        self,
+        episode: Episode,
+        results: Iterable[MetricResult],
+        shaped: ShapedReward,
+        *,
+        rubric: Optional[str] = None,
+    ) -> List[Reward]:
+        """Persist metric and aggregate rewards. Returns the stored rows."""
+        stored: List[Reward] = []
+        results_list = list(results)
+
+        # 1) Persist the raw metric results themselves (richer than rewards)
+        try:
+            self.store.store_metric_results(episode.id, episode.agent_id, results_list)
+        except Exception as exc:  # pragma: no cover - persistence error path
+            logger.warning("Failed to persist metric results for %s: %s", episode.id, exc)
+
+        # 2) Per-metric reward records (signed contributions)
+        contributions = {m: signed for m, _w, signed, _c in shaped.metric_contributions}
+        for result in results_list:
+            if result.status != "completed" or result.normalized is None:
+                continue
+            signed_value = contributions.get(
+                result.metric,
+                2.0 * result.normalized - 1.0,
+            )
+            reward = Reward(
+                episode_id=episode.id,
+                agent_id=episode.agent_id,
+                source=RewardSource.METRIC,
+                value=max(-1.0, min(1.0, signed_value)),
+                raw_value=result.score,
+                metric=result.metric,
+                rubric=rubric or result.metric.value,
+                evaluator=result.evaluator,
+                metadata={"reason": result.reason} if result.reason else {},
+            )
+            try:
+                self.store.store_reward(reward)
+                stored.append(reward)
+            except Exception as exc:  # pragma: no cover
+                logger.warning(
+                    "Failed to persist %s reward for %s: %s",
+                    result.metric.value,
+                    episode.id,
+                    exc,
+                )
+
+        # 3) Latency / cost penalty rewards (kept separate for analytics)
+        for kind, raw_value, value in shaped.penalties:
+            reward = Reward(
+                episode_id=episode.id,
+                agent_id=episode.agent_id,
+                source=RewardSource.LATENCY_PENALTY
+                if kind == "latency"
+                else RewardSource.COST_PENALTY,
+                value=max(-1.0, min(1.0, value)),
+                raw_value=raw_value,
+                rubric=kind,
+            )
+            try:
+                self.store.store_reward(reward)
+                stored.append(reward)
+            except Exception as exc:  # pragma: no cover
+                logger.warning("Failed to persist %s penalty for %s: %s", kind, episode.id, exc)
+
+        # 4) Aggregate reward consumed by the learner
+        aggregate = Reward(
+            episode_id=episode.id,
+            agent_id=episode.agent_id,
+            source=RewardSource.AGGREGATE,
+            value=shaped.value,
+            rubric=rubric or "aggregate",
+            metadata={
+                "metric_contributions": [
+                    {
+                        "metric": m.value if isinstance(m, MetricName) else str(m),
+                        "weight": w,
+                        "signed": signed,
+                        "contribution": c,
+                    }
+                    for m, w, signed, c in shaped.metric_contributions
+                ],
+                "penalties": [{"kind": k, "raw": r, "value": v} for k, r, v in shaped.penalties],
+            },
+        )
+        try:
+            self.store.store_reward(aggregate)
+            stored.append(aggregate)
+        except Exception as exc:  # pragma: no cover
+            logger.warning("Failed to persist aggregate reward for %s: %s", episode.id, exc)
+
+        return stored
+
+
+__all__ = ["RewardWriter"]