ALE-Agent
Domain-guided tabu best-first search with 30 sibling branches and three-turn solution refinement.
"""ALE-Agent's tabu best-first expansion with K sibling branches per parent."""
from __future__ import annotations
from typing import Any
from ...components.selection import SelectionPolicy
from ...records import Genome, LoopContext, Selection
from .population import ale_priority
class ALEAgentSelectionPolicy(SelectionPolicy):
"""Expand the best unexpanded frontier node into ``K`` multi-turn branches.
A branch is refined for ``T`` evaluated implementation turns. Only the best version
produced within those turns becomes a new frontier node; all versions remain in the
archive as search evidence. Once a parent starts expansion it is marked tabu and can
never be selected as an outer parent again.
Galapagos executes iterations sequentially for reproducible leaderboard runs. The
policy therefore schedules the paper's K parallel siblings deterministically one branch
at a time while preserving the same frontier and parent-expansion semantics.
"""
def __init__(self, seed: int = 0, *, children_per_parent: int = 30,
refinement_turns: int = 3):
super().__init__(seed)
self.children_per_parent = max(1, int(children_per_parent))
self.refinement_turns = max(1, int(refinement_turns))
self.current_outer_parent_id: str | None = None
self.outer_historical_best_id: str | None = None
self.branches_started = 0
self.active: dict[str, Any] | None = None
self._active_candidates: list[Genome] = []
@staticmethod
def _by_id(population) -> dict[str, Genome]:
return {g.id: g for g in population.all()}
def _resolve_active(self, population) -> None:
if self.active is None:
self._active_candidates = []
return
by_id = self._by_id(population)
ids = self.active.get("candidate_ids") or []
self._active_candidates = [by_id[i] for i in ids if i in by_id]
def _choose_outer_parent(self, population) -> Genome:
frontier = population.query({"frontier": True, "unexpanded": True})
if not frontier:
raise RuntimeError("ALE-Agent search exhausted: no unexpanded frontier state remains")
parent = max(frontier, key=ale_priority)
# All K children are concurrent in the paper and therefore start from the same view of the
# historical best. Freeze that state for the cohort so sequential Galapagos scheduling does
# not let a completed earlier sibling leak into a later sibling's prompt.
historical_best = population.best()
self.outer_historical_best_id = historical_best.id if historical_best is not None else parent.id
parent.metadata["ale_frontier"] = False
parent.metadata["ale_expanded"] = True
self.current_outer_parent_id = parent.id
self.branches_started = 0
return parent
def _start_branch(self) -> None:
branch_index = self.branches_started + 1
self.branches_started = branch_index
self.active = {
"branch_index": branch_index,
"completed_turns": 0,
"candidate_ids": [],
# The paper samples one of four domain prompts stochastically for each refinement.
"guidance_index": self.rng.randrange(4),
}
self._active_candidates = []
def _outer_parent(self, population) -> Genome:
by_id = self._by_id(population)
parent = by_id.get(self.current_outer_parent_id or "")
if parent is None:
return self._choose_outer_parent(population)
return parent
def select(self, population, ctx: LoopContext | None = None) -> Selection:
self._resolve_active(population)
expanded_outer = False
if self.active is None:
if self.current_outer_parent_id is None or self.branches_started >= self.children_per_parent:
outer = self._choose_outer_parent(population)
expanded_outer = True
else:
outer = self._outer_parent(population)
self._start_branch()
else:
outer = self._outer_parent(population)
assert self.active is not None
parent = self._active_candidates[-1] if self._active_candidates else outer
turn = int(self.active["completed_turns"]) + 1
by_id = self._by_id(population)
historical_best = by_id.get(self.outer_historical_best_id or "") or population.best() or outer
# Keep the frozen best explicit even when it is also the parent. The prompt builder must not
# fall back to ctx.best, which can change as earlier sequentially scheduled siblings finish.
inspirations = [historical_best]
if ctx is not None:
ctx.blackboard["ale_agent"] = {
"outer_parent_id": outer.id,
"branch_index": int(self.active["branch_index"]),
"children_per_parent": self.children_per_parent,
"refinement_turn": turn,
"refinement_turns": self.refinement_turns,
"guidance_index": int(self.active["guidance_index"]),
"historical_best_id": historical_best.id,
}
details = {
"selection_strategy": "tabu_best_first",
"selection_mode": "refine_branch" if self._active_candidates else "start_branch",
"outer_parent_candidate_id": outer.id,
"branch_index": int(self.active["branch_index"]),
"branch_count_limit": self.children_per_parent,
"refinement_step": turn,
"refinement_step_limit": self.refinement_turns,
"historical_best_candidate_id": historical_best.id,
}
if expanded_outer:
details["population_updates"] = [{
"collection_type": "ale_frontier",
"collection_id": "unexpanded_frontier",
"added_candidate_ids": [],
"removed_candidate_ids": [outer.id],
"current_candidate_ids": [
genome.id for genome in population.query({"frontier": True})
],
"reason": "outer_parent_expanded",
"location": {"outer_parent_candidate_id": outer.id},
}]
return Selection(
parent=parent,
inspirations=inspirations,
pool=population.all(),
details=details,
)
def _finish_active_branch(self, ctx: LoopContext | None) -> None:
if self._active_candidates:
winner = max(self._active_candidates, key=ale_priority)
for candidate in self._active_candidates:
candidate.metadata["ale_frontier"] = False
candidate.metadata["ale_branch_complete"] = True
candidate.metadata["ale_branch_winner"] = candidate.id == winner.id
winner.metadata["ale_frontier"] = True
winner.metadata["ale_expanded"] = False
if ctx is not None:
ctx.blackboard.setdefault("ale_agent", {})["branch_winner_id"] = winner.id
self.active = None
self._active_candidates = []
def observe(self, genome: Genome, ctx: LoopContext | None = None) -> None:
"""Advance one evaluated implementation turn and finalize at turn T."""
if self.active is None:
return
genome.metadata["ale_frontier"] = False
genome.metadata["ale_expanded"] = False
self._active_candidates.append(genome)
self.active["candidate_ids"] = [g.id for g in self._active_candidates]
self.active["completed_turns"] = int(self.active["completed_turns"]) + 1
if int(self.active["completed_turns"]) >= self.refinement_turns:
self._finish_active_branch(ctx)
def observe_rejected(self, genome: Genome | None, ctx: LoopContext | None = None) -> None:
"""Spend a turn that produced no admissible code (unevaluable or evaluator-invalid)."""
if self.active is None:
return
self.active["completed_turns"] = int(self.active["completed_turns"]) + 1
if int(self.active["completed_turns"]) >= self.refinement_turns:
self._finish_active_branch(ctx)
def can_select(self, population) -> bool:
"""Whether another turn can start without violating the no-reexpansion tabu rule."""
if self.active is not None:
return True
if self.current_outer_parent_id is not None and self.branches_started < self.children_per_parent:
return True
return bool(population.query({"frontier": True, "unexpanded": True, "top": 1}))
def state_dict(self) -> dict:
active = None if self.active is None else {
"branch_index": int(self.active["branch_index"]),
"completed_turns": int(self.active["completed_turns"]),
"candidate_ids": list(self.active.get("candidate_ids") or []),
"guidance_index": int(self.active["guidance_index"]),
}
return {
**super().state_dict(),
"current_outer_parent_id": self.current_outer_parent_id,
"outer_historical_best_id": self.outer_historical_best_id,
"branches_started": int(self.branches_started),
"active": active,
}
def load_state_dict(self, state: dict) -> None:
super().load_state_dict(state)
if not isinstance(state, dict):
return
self.current_outer_parent_id = state.get("current_outer_parent_id") or None
self.outer_historical_best_id = state.get("outer_historical_best_id") or None
self.branches_started = max(0, int(state.get("branches_started", 0) or 0))
active = state.get("active")
if isinstance(active, dict):
self.active = {
"branch_index": int(active.get("branch_index", 1)),
"completed_turns": int(active.get("completed_turns", 0)),
"candidate_ids": [str(i) for i in active.get("candidate_ids") or []],
"guidance_index": int(active.get("guidance_index", 0)) % 4,
}
else:
self.active = None
self._active_candidates = []