SakanaAI/ale_agent

"""ALE-Agent candidate archive and best-first priority.

The paper ranks states first by the fraction of locally accepted cases and then by the
problem score.  Galapagos task evaluators may expose that fraction under one of several
descriptive metric names; the bundled ALE-Bench tasks currently expose the coarser
``judge_accepted`` all-cases signal.  ``ale_priority`` supports both without coupling the
scaffold to one task implementation.
"""
from __future__ import annotations

import math

from ...components.population import Population
from ...records import Genome

_ACCEPTANCE_KEYS = (
    "acceptance_ratio",
    "accepted_ratio",
    "acceptance_rate",
    "judge_accepted",
)


def acceptance_ratio(genome: Genome) -> float:
    """Return the paper's primary priority signal, normalized to ``[0, 1]``.

    A generic non-ALE task need not publish an acceptance metric.  In that case an
    evaluator-valid program is treated as accepted and an invalid one as rejected, so the
    scaffold remains usable while preserving the acceptance-before-score ordering.
    """
    for key in _ACCEPTANCE_KEYS:
        value = genome.scores.get(key)
        if isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value):
            return max(0.0, min(1.0, float(value)))
    return 0.0 if genome.metadata.get("valid") is False else 1.0


def ale_priority(genome: Genome) -> tuple[float, float]:
    """Lexicographic best-first priority: acceptance ratio, then task fitness."""
    score = genome.fitness
    if math.isnan(score):
        score = float("-inf")
    return acceptance_ratio(genome), score


class ALEAgentPopulation(Population):
    """Keep every evaluator-valid state; frontier/tabu flags live on each genome's metadata.

    ALE-Agent needs the full archive for historical-best context even after a node has been
    expanded. The selection policy alone decides which valid frontier node to expand. Invalid
    attempts remain in the trajectory rather than this selectable archive.
    """

    def __init__(self):
        self._members: dict[str, Genome] = {}

    def add(self, genome: Genome) -> bool:
        if genome.metadata.get("valid") is False:
            genome.metadata.update(admitted=False, eval_failed=True)
            return False
        self._members[genome.id] = genome
        genome.metadata.pop("eval_failed", None)
        genome.metadata["admitted"] = True
        return True

    def query(self, spec: dict | None = None) -> list[Genome]:
        spec = spec or {}
        members = list(self._members.values())
        if spec.get("frontier"):
            members = [g for g in members if g.metadata.get("ale_frontier") is True]
        if spec.get("unexpanded"):
            members = [g for g in members if g.metadata.get("ale_expanded") is not True]
        members.sort(key=ale_priority, reverse=True)
        top = spec.get("top")
        return members[: int(top)] if top is not None else members

    def all(self) -> list[Genome]:
        return list(self._members.values())

    def best(self) -> Genome | None:
        return max(self._members.values(), key=ale_priority) if self._members else None