Add pluggable trajectory aggregation with learned GBDT reference implementation

.gitignore

@@ -185,4 +185,10 @@ cython_debug/
 
 # local dev
 notebooks/dev.ipynb
-results/
+results/
+
+# Gas City
+.beads/*
+!.beads/config.yaml
+!.beads/metadata.json
+.claude/

its_hub/aggregators/__init__.py

@@ -0,0 +1,4 @@
+from .hardcoded import HardcodedAggregator
+from .learned import LearnedGBDTAggregator, LearnedMLPAggregator
+
+__all__ = ["HardcodedAggregator", "LearnedGBDTAggregator", "LearnedMLPAggregator"]

its_hub/aggregators/hardcoded.py

@@ -0,0 +1,31 @@
+import math
+from typing import Literal
+
+from its_hub.base import AbstractTrajectoryAggregator
+
+
+class HardcodedAggregator(AbstractTrajectoryAggregator):
+    """Reduces per-step PRM scores with a fixed aggregation rule.
+
+    Wraps the three reductions previously hardcoded inside PRM implementations,
+    making the choice explicit and pluggable at the algorithm level.
+
+    prod: product of step probabilities (log-space: sum of log-probs)
+    min:  worst-case step score
+    mean: length-invariant average
+    """
+
+    def __init__(self, reduction: Literal["prod", "min", "mean"] = "prod"):
+        self.reduction = reduction
+
+    def aggregate(self, step_scores: list[float]) -> float:
+        if not step_scores:
+            return 0.0
+        if self.reduction == "prod":
+            return math.prod(step_scores)
+        elif self.reduction == "min":
+            return min(step_scores)
+        elif self.reduction == "mean":
+            return sum(step_scores) / len(step_scores)
+        else:
+            raise ValueError(f"Unknown reduction: {self.reduction!r}")

its_hub/aggregators/learned.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+import numpy as np
+
+from its_hub.base import AbstractTrajectoryAggregator
+
+
+def _extract_features(step_scores: list[float]) -> np.ndarray:
+    """Convert variable-length step scores to a fixed 10-dim feature vector.
+
+    Features (all position indices normalised by length):
+      0  mean
+      1  min
+      2  max
+      3  last step score
+      4  trajectory length (raw count)
+      5  variance
+      6  normalised position of min (argmin / (len-1), or 0 if len==1)
+      7  normalised position of max
+      8  last minus first score
+      9  score gap at min position (score_before_min - score_at_min), or 0 at boundary
+    """
+    n = len(step_scores)
+    if n == 0:
+        return np.zeros(10, dtype=np.float32)
+
+    arr = np.array(step_scores, dtype=np.float64)
+    mean = float(arr.mean())
+    mn = float(arr.min())
+    mx = float(arr.max())
+    last = float(arr[-1])
+    length = float(n)
+    var = float(arr.var())
+
+    if n > 1:
+        pos_min = float(arr.argmin()) / (n - 1)
+        pos_max = float(arr.argmax()) / (n - 1)
+    else:
+        pos_min = 0.0
+        pos_max = 0.0
+
+    delta_last_first = last - float(arr[0])
+
+    amin = int(arr.argmin())
+    gap_at_min = float(arr[amin - 1] - arr[amin]) if amin > 0 else 0.0
+
+    return np.array(
+        [mean, mn, mx, last, length, var, pos_min, pos_max, delta_last_first, gap_at_min],
+        dtype=np.float32,
+    )
+
+
+class LearnedMLPAggregator(AbstractTrajectoryAggregator):
+    """Trajectory aggregator backed by a trained MLP checkpoint.
+
+    The MLP maps a 10-dim feature vector (derived from per-step scores) to a
+    scalar trajectory score via a sigmoid output.  Checkpoint format is a plain
+    PyTorch state-dict saved alongside the architecture hyper-parameters:
+
+        torch.save({"state_dict": model.state_dict(), "hidden_width": 16}, path)
+
+    Requires torch; raises ImportError with a clear message if absent.
+    """
+
+    def __init__(self, checkpoint_path: str):
+        try:
+            import torch
+        except ImportError as exc:
+            raise ImportError(
+                "LearnedMLPAggregator requires PyTorch. "
+                "Install it with: pip install torch"
+            ) from exc
+
+        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+        hidden_width = checkpoint.get("hidden_width", 16)
+
+        self._model = _TrajectoryMLP(input_dim=10, hidden_width=hidden_width)
+        self._model.load_state_dict(checkpoint["state_dict"])
+        self._model.eval()
+        self._torch = torch
+
+    def aggregate(self, step_scores: list[float]) -> float:
+        features = _extract_features(step_scores)
+        x = self._torch.tensor(features, dtype=self._torch.float32).unsqueeze(0)
+        with self._torch.no_grad():
+            score = self._model(x).item()
+        return score
+
+
+class _TrajectoryMLP:
+    """Minimal 2-layer MLP; kept here so the aggregator is self-contained."""
+
+    def __init__(self, input_dim: int = 10, hidden_width: int = 16):
+        try:
+            import torch.nn as nn
+        except ImportError as exc:
+            raise ImportError("torch is required") from exc
+
+        import torch.nn as nn
+
+        self._net = nn.Sequential(
+            nn.Linear(input_dim, hidden_width),
+            nn.ReLU(),
+            nn.Linear(hidden_width, 1),
+            nn.Sigmoid(),
+        )
+
+    def load_state_dict(self, state_dict):
+        self._net.load_state_dict(state_dict)
+
+    def eval(self):
+        self._net.eval()
+
+    def __call__(self, x):
+        return self._net(x).squeeze(-1)
+
+
+class LearnedGBDTAggregator(AbstractTrajectoryAggregator):
+    """Trajectory aggregator backed by a trained GBDT checkpoint.
+
+    Maps a 10-dim feature vector (derived from per-step scores) to a trajectory
+    score via sklearn's GradientBoostingClassifier.predict_proba.  Checkpoint is
+    a joblib-serialised sklearn estimator:
+
+        import joblib
+        joblib.dump(clf, path)
+
+    Requires scikit-learn; raises ImportError with a clear message if absent.
+    """
+
+    def __init__(self, checkpoint_path: str):
+        try:
+            import joblib
+        except ImportError as exc:
+            raise ImportError(
+                "LearnedGBDTAggregator requires scikit-learn. "
+                "Install it with: pip install scikit-learn"
+            ) from exc
+
+        self._clf = joblib.load(checkpoint_path)
+
+    def aggregate(self, step_scores: list[float]) -> float:
+        features = _extract_features(step_scores)
+        return float(self._clf.predict_proba([features])[0, 1])

its_hub/algorithms/beam_search.py

@@ -1,13 +1,16 @@
 import copy
+from dataclasses import field
 
 import numpy as np
 from pydantic.dataclasses import dataclass
 
+from its_hub.aggregators import HardcodedAggregator
 from its_hub.base import (
     AbstractLanguageModel,
     AbstractProcessRewardModel,
     AbstractScalingAlgorithm,
     AbstractScalingResult,
+    AbstractTrajectoryAggregator,
 )
 from its_hub.lms import StepGeneration
 from its_hub.types import ChatMessage, ChatMessages
@@ -30,13 +33,15 @@ class Path:
     steps: list[str]
     is_stopped: bool
     score: float
+    step_scores: list[float] = field(default_factory=list)  # Raw PRM scores per step for trajectory aggregation
 
     def deepcopy(self):
         # create a deep copy of the path object
         return Path(
             steps=copy.deepcopy(self.steps),
             is_stopped=self.is_stopped,
             score=self.score,
+            step_scores=copy.deepcopy(self.step_scores),
         )
 
 
@@ -46,10 +51,12 @@ def __init__(
         sg: StepGeneration,
         prm: AbstractProcessRewardModel,
         beam_width: int,
+        aggregator: AbstractTrajectoryAggregator | None = None,
     ):
         self.sg = sg
         self.prm = prm
         self.beam_width = beam_width
+        self.aggregator: AbstractTrajectoryAggregator = aggregator or HardcodedAggregator("prod")
 
     async def _asearch_one_level(
         self,
@@ -107,6 +114,7 @@ async def _asearch_one_level(
             if is_stopped:
                 continue
             c.score = scores[i]
+            c.step_scores.append(scores[i])
             i += 1
 
         return candidates
@@ -169,7 +177,7 @@ async def ainfer(
                     new_candidates.append(c.deepcopy())
             candidates = new_candidates
 
-        scores = [c.score for c in candidates]
+        scores = [self.aggregator.aggregate(c.step_scores) for c in candidates]
         steps_used = [len(c.steps) for c in candidates]
         result = BeamSearchResult(
             responses=[

its_hub/algorithms/particle_gibbs.py

@@ -1,15 +1,18 @@
 import copy
 import random
+from dataclasses import field
 from enum import Enum
 
 import numpy as np
 from pydantic.dataclasses import dataclass
 
+from its_hub.aggregators import HardcodedAggregator
 from its_hub.base import (
     AbstractLanguageModel,
     AbstractProcessRewardModel,
     AbstractScalingAlgorithm,
     AbstractScalingResult,
+    AbstractTrajectoryAggregator,
 )
 from its_hub.lms import StepGeneration
 from its_hub.types import ChatMessage, ChatMessages
@@ -45,6 +48,7 @@ class Particle:
     steps: list[str]
     is_stopped: bool
     partial_log_weights: list[float]  # Store aggregated log weights until each step
+    step_scores: list[float] = field(default_factory=list)  # Raw PRM scores per step for trajectory aggregation
 
     @property
     def log_weight(self) -> float:
@@ -59,6 +63,7 @@ def deepcopy(self):
             steps=copy.deepcopy(self.steps),
             is_stopped=self.is_stopped,
             partial_log_weights=copy.deepcopy(self.partial_log_weights),
+            step_scores=copy.deepcopy(self.step_scores),
         )
 
 
@@ -115,6 +120,7 @@ def __init__(
         early_phase: float = 0.5,
         resampling_method: str | ResamplingMethod = ResamplingMethod.MULTINOMIAL,
         temperature_method: str | TemperatureMethod = TemperatureMethod.ESS,
+        aggregator: AbstractTrajectoryAggregator | None = None,
     ):
         if isinstance(final_response_selection, str):
             final_response_selection = SelectionMethod(final_response_selection)
@@ -138,6 +144,7 @@ def __init__(
         self.early_phase = early_phase
         self.resampling_method = resampling_method
         self.temperature_method = temperature_method
+        self.aggregator: AbstractTrajectoryAggregator = aggregator or HardcodedAggregator("prod")
 
     async def _apropagate(
         self,
@@ -194,6 +201,7 @@ async def _apropagate(
         for p, is_stopped in zip(particles, is_stopped_in_the_beginning):
             if is_stopped:
                 continue
+            p.step_scores.append(scores[i])
             p.partial_log_weights.append(_inv_sigmoid(scores[i]))
             i += 1
 
@@ -391,7 +399,7 @@ async def ainfer(
             num_free_particles = num_particles - len(ref_particles)
 
             particles = [
-                Particle(steps=[], is_stopped=False, partial_log_weights=[])
+                Particle(steps=[], is_stopped=False, partial_log_weights=[], step_scores=[])
                 for _ in range(num_free_particles)
             ] + ref_particles
 
@@ -478,15 +486,15 @@ async def ainfer(
             ref_indices_lst.append(ref_indices)
             steps_used_lst.append([len(p.steps) for p in particles])
 
-        # select the chosen particle based on final response selection method
-        # log_weights and probabilities are from the last iteration
+        # select the chosen particle using the trajectory aggregator
+        agg_scores = [self.aggregator.aggregate(p.step_scores) for p in particles]
         match self.final_response_selection:
             case SelectionMethod.SAMPLE:
                 selected_index = random.choices(
-                    range(len(particles)), weights=probabilities, k=1
+                    range(len(particles)), weights=_softmax(agg_scores), k=1
                 )[0]
             case SelectionMethod.ARGMAX:
-                selected_index = np.argmax(log_weights).item()
+                selected_index = int(np.argmax(agg_scores))
 
         result = ParticleGibbsResult(
             responses_lst=responses_lst,
@@ -510,6 +518,7 @@ def __init__(
         prm: AbstractProcessRewardModel,
         final_response_selection: str | SelectionMethod = SelectionMethod.ARGMAX,
         resampling_method: str | ResamplingMethod = ResamplingMethod.MULTINOMIAL,
+        aggregator: AbstractTrajectoryAggregator | None = None,
     ):
         # initialize with num_iterations=1
         super().__init__(
@@ -521,6 +530,7 @@ def __init__(
             does_ancestor_sampling=False,
             does_entropic_annealing=False,
             does_lookahead_modulation=False,
+            aggregator=aggregator,
         )
 
     async def ainfer(

its_hub/base.py

@@ -100,7 +100,6 @@ def score(
         pass
 
 
-# TODO(GX) deal with aggregation of PRM scores somehow in a common place, e.g. here
 class AbstractProcessRewardModel(ABC):
     """abstract base class for process reward models"""
 
@@ -121,3 +120,20 @@ def score(
     ) -> list[float]:
         """score steps synchronously"""
         pass
+
+
+class AbstractTrajectoryAggregator(ABC):
+    """abstract base class for trajectory aggregators
+
+    Reduces a sequence of per-step PRM scores to a single trajectory score.
+    Pluggable into ParticleFiltering and BeamSearch via the aggregator parameter.
+    """
+
+    @abstractmethod
+    def aggregate(self, step_scores: list[float]) -> float:
+        """reduce per-step scores to a single trajectory score"""
+        pass
+
+    async def aaggregate(self, step_scores: list[float]) -> float:
+        """async variant; delegates to sync by default"""
+        return self.aggregate(step_scores)