"""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 = []