"""Evaluation-aware objective and plateau termination for training phases."""
from __future__ import annotations
from collections import deque
from dataclasses import dataclass, field
from enum import Enum
import hashlib
import json
import logging
import math
from typing import Any, Deque, Dict, Iterable, Mapping, Optional, Tuple
from palaestrai.core.protocol.agent_evaluation_rsp import (
AgentEvaluationResponse,
)
from palaestrai.core.protocol.simulation_controller_termination_req import (
SimulationControllerTerminationRequest,
)
from palaestrai.types import Mode, SimulationFlowControl
from .environment_termination_condition import EnvironmentTerminationCondition
from .max_episodes_termination_condition import MaxEpisodesTerminationCondition
from .termination_condition import TerminationCondition
LOG = logging.getLogger(__name__)
class EvaluationTerminationReason(Enum):
"""Terminal outcome recorded by the evaluation-aware condition."""
SUCCESS = "success"
PLATEAU = "plateau"
MAX_EPISODES = "max_episodes"
CANCELLED = "cancelled"
FAILED = "failed"
@dataclass(frozen=True)
class _SuccessRule:
metric: str
threshold: float
direction: str
consecutive: int
min_evaluations: int
@dataclass(frozen=True)
class _PlateauRule:
metric: str
direction: str
min_delta: float
patience: int
min_evaluations: int
window: int
@dataclass
class _AgentState:
last_evaluation: Optional[int] = None
seen: Dict[str, Dict[str, Any]] = field(default_factory=dict)
success_streak: int = 0
plateau_patience: int = 0
best_metric: Optional[float] = None
best_evaluation_id: Optional[str] = None
history: Deque[Dict[str, Any]] = field(default_factory=deque)
[docs]
class AgentEvaluationObjectiveTerminationCondition(TerminationCondition):
"""Terminate a training phase only from completed evaluation outcomes.
The condition combines normal environment episode handling with an optional
finite ``phase_config.episodes`` cap. It never inspects training rollout
metrics. A Brain must return
:class:`~palaestrai.agent.EvaluationResult` from ``evaluate_rollouts``;
its named finite metrics are forwarded after all evaluation workers finish.
Configuration uses one entry per participating agent::
run_config:
condition:
name: palaestrai.experiment:AgentEvaluationObjectiveTerminationCondition
params:
agents:
SAC:
success:
metric: objective
threshold: 0.8
direction: greater_equal
consecutive: 2
min_evaluations: 3
plateau:
metric: objective
direction: maximize
min_delta: 0.01
patience: 4
min_evaluations: 5
window: 1
``success`` and ``plateau`` are independently optional, but each agent
needs at least one. For a multi-agent phase, each enabled category must
be configured for every agent: that makes a global ``SUCCESS`` or
``PLATEAU`` outcome unambiguous. All configured agents must satisfy the
corresponding rule before the phase ends. A success threshold is checked
before plateau at each checkpoint.
Invalid, incomplete, non-finite, duplicate-with-conflicting-data, or
out-of-order outcomes never count as success or improvement. An invalid
outcome resets only that agent's success streak and preserves its plateau
state; it cannot silently create a plateau termination. Exact duplicate
outcome identities are idempotent. The implementation supports one
deterministic worker (``worker: 1``) and rejects a configuration with more
workers to avoid assigning a mixed asynchronous batch a false checkpoint
identity.
Omitting ``episodes`` is the idiomatic cap-free training form; an explicit
``episodes: null`` is accepted as a compatibility alias. Cap-free mode is
valid only with this condition, a positive ``evaluate_every``, one worker,
and at least one success or plateau rule. It can run indefinitely if
evaluations never arrive; deploy it only with an external wall-time and
resource guard.
"""
STATE_VERSION = 1
_SUCCESS_DIRECTIONS = {"greater_equal", "less_equal"}
_PLATEAU_DIRECTIONS = {"maximize", "minimize"}
def __init__(self, *, agents: Mapping[str, Mapping[str, Any]]):
if not isinstance(agents, Mapping) or not agents:
raise ValueError("'agents' must be a non-empty mapping")
self._success_rules: Dict[str, Optional[_SuccessRule]] = {}
self._plateau_rules: Dict[str, Optional[_PlateauRule]] = {}
for agent_name, rules in agents.items():
if not isinstance(agent_name, str) or not agent_name:
raise ValueError("agent names must be non-empty strings")
if not isinstance(rules, Mapping):
raise ValueError(
f"rules for agent {agent_name!r} must be a mapping"
)
unknown = set(rules) - {"success", "plateau"}
if unknown:
raise ValueError(
f"unknown evaluation termination rules for {agent_name!r}: "
f"{sorted(unknown)!r}"
)
success = (
self._parse_success(agent_name, rules["success"])
if "success" in rules
else None
)
plateau = (
self._parse_plateau(agent_name, rules["plateau"])
if "plateau" in rules
else None
)
if success is None and plateau is None:
raise ValueError(
f"agent {agent_name!r} needs a success and/or plateau rule"
)
self._success_rules[agent_name] = success
self._plateau_rules[agent_name] = plateau
self._require_complete_global_rules("success", self._success_rules)
self._require_complete_global_rules("plateau", self._plateau_rules)
self._states = {agent: _AgentState() for agent in self._success_rules}
self._termination_reason: Optional[EvaluationTerminationReason] = None
self._configuration_fingerprint = (
self._make_configuration_fingerprint()
)
self._environment_tc = EnvironmentTerminationCondition()
self._max_episodes_tc = MaxEpisodesTerminationCondition()
@property
def termination_reason(self) -> Optional[EvaluationTerminationReason]:
"""The terminal reason, or ``None`` while the phase is active."""
return self._termination_reason
@property
def requires_evaluation_outcomes(self) -> bool:
"""Request Brain results only after deterministic evaluation rollouts."""
return True
def validate_phase_configuration(
self, phase_config: Mapping[str, Any]
) -> None:
"""Reject unsupported or unsafe phase configuration combinations."""
episodes = phase_config.get("episodes")
mode = str(phase_config.get("mode", "train")).lower()
evaluate_every = phase_config.get("evaluate_every")
worker = phase_config.get("worker", 1)
if (
isinstance(worker, bool)
or not isinstance(worker, int)
or worker != 1
):
raise ValueError(
f"{self.__class__.__name__} supports only worker: 1; got {worker!r}"
)
if mode == "train" and (
isinstance(evaluate_every, bool)
or not isinstance(evaluate_every, int)
or evaluate_every <= 0
):
raise ValueError(
f"{self.__class__.__name__} needs a positive evaluate_every "
"for train phases"
)
if episodes is None:
if mode != "train":
raise ValueError(
"cap-free evaluation termination is supported only for train phases"
)
if not isinstance(evaluate_every, int) or evaluate_every <= 0:
raise ValueError(
"cap-free evaluation termination requires positive evaluate_every"
)
elif (
isinstance(episodes, bool)
or not isinstance(episodes, int)
or episodes <= 0
):
raise ValueError(
"episodes must be a positive integer when provided"
)
def phase_flow_control(
self,
run_governor: Any,
message: Optional[SimulationControllerTerminationRequest],
) -> Tuple[SimulationFlowControl, Any]:
if self._termination_reason in {
EvaluationTerminationReason.SUCCESS,
EvaluationTerminationReason.PLATEAU,
}:
return SimulationFlowControl.STOP_PHASE, {
"termination_reason": self._termination_reason.value
}
if not isinstance(message, SimulationControllerTerminationRequest):
return SimulationFlowControl.CONTINUE, None
flows = [
self._environment_tc.phase_flow_control(run_governor, message)
]
episodes = run_governor.experiment_run.phase_configuration(
run_governor.current_phase
).get("episodes")
if episodes is not None:
max_flow = self._max_episodes_tc.phase_flow_control(
run_governor, message
)
flows.append(max_flow)
if max_flow[0] == SimulationFlowControl.STOP_PHASE:
self._set_terminal_reason(
EvaluationTerminationReason.MAX_EPISODES
)
# In cap-free mode, pause evaluation workers after exactly one
# evaluation rollout. Normal workers continue until a valid
# evaluation rule requests phase termination.
if message.mode == Mode.EVALUATE and episodes is None:
environment_flow = flows[0][0]
if environment_flow.value > SimulationFlowControl.CONTINUE.value:
return SimulationFlowControl.RESTART, None
return max(flows, key=lambda item: item[0].value)
def evaluation_flow_control(
self,
run_governor: Any,
outcomes: Iterable[AgentEvaluationResponse],
) -> Tuple[SimulationFlowControl, Any]:
"""Consume one completed evaluation checkpoint per configured agent."""
if self._termination_reason is not None:
return SimulationFlowControl.STOP_PHASE, {
"termination_reason": self._termination_reason.value
}
outcome_by_agent: Dict[str, AgentEvaluationResponse] = {}
for outcome in outcomes:
if outcome.agent_name not in self._states:
continue
if outcome.agent_name in outcome_by_agent:
raise ValueError(
"duplicate evaluation outcome for agent "
f"{outcome.agent_name!r} in one checkpoint"
)
outcome_by_agent[outcome.agent_name] = outcome
if set(outcome_by_agent) != set(self._states):
for agent_name in set(self._states) - set(outcome_by_agent):
self._invalid_outcome(agent_name, "missing agent outcome")
LOG.warning(
"Ignoring incomplete evaluation checkpoint; expected agents %s, got %s",
sorted(self._states),
sorted(outcome_by_agent),
)
return SimulationFlowControl.CONTINUE, None
for agent_name in sorted(self._states):
self._record_outcome(agent_name, outcome_by_agent[agent_name])
if self._all_successful():
self._set_terminal_reason(EvaluationTerminationReason.SUCCESS)
elif self._all_plateaued():
self._set_terminal_reason(EvaluationTerminationReason.PLATEAU)
if self._termination_reason is None:
return SimulationFlowControl.CONTINUE, None
return SimulationFlowControl.STOP_PHASE, {
"termination_reason": self._termination_reason.value
}
def mark_cancelled(self) -> None:
"""Record an external cancellation as a terminal non-normal outcome."""
self._set_terminal_reason(EvaluationTerminationReason.CANCELLED)
def mark_failed(self) -> None:
"""Record an external framework failure as a terminal non-normal outcome."""
self._set_terminal_reason(EvaluationTerminationReason.FAILED)
def state_dict(self) -> Dict[str, Any]:
"""Return a complete JSON-serializable state for an exact resume."""
return {
"version": self.STATE_VERSION,
"configuration_fingerprint": self._configuration_fingerprint,
"configuration": self._configuration(),
"termination_reason": (
self._termination_reason.value
if self._termination_reason is not None
else None
),
"agents": {
agent: {
"last_evaluation": state.last_evaluation,
"seen": state.seen,
"success_streak": state.success_streak,
"plateau_patience": state.plateau_patience,
"best_metric": state.best_metric,
"best_evaluation_id": state.best_evaluation_id,
"history": list(state.history),
}
for agent, state in self._states.items()
},
}
def load_state_dict(self, state: Mapping[str, Any]) -> None:
"""Restore state only when it exactly matches this rule configuration."""
if state.get("version") != self.STATE_VERSION:
raise ValueError(
"unsupported evaluation termination state version"
)
if (
state.get("configuration_fingerprint")
!= self._configuration_fingerprint
):
raise ValueError(
"evaluation termination state does not match this configuration"
)
if state.get("configuration") != self._configuration():
raise ValueError(
"evaluation termination state configuration does not match"
)
agent_states = state.get("agents")
if not isinstance(agent_states, Mapping) or set(agent_states) != set(
self._states
):
raise ValueError(
"evaluation termination state has incompatible agents"
)
restored: Dict[str, _AgentState] = {}
for agent, raw in agent_states.items():
if not isinstance(raw, Mapping):
raise ValueError(f"invalid saved state for agent {agent!r}")
seen = raw.get("seen")
history = raw.get("history")
if not isinstance(seen, Mapping) or not isinstance(history, list):
raise ValueError(
f"invalid saved outcome history for agent {agent!r}"
)
if (
not isinstance(raw.get("success_streak"), int)
or isinstance(raw.get("success_streak"), bool)
or raw["success_streak"] < 0
or not isinstance(raw.get("plateau_patience"), int)
or isinstance(raw.get("plateau_patience"), bool)
or raw["plateau_patience"] < 0
):
raise ValueError(f"invalid saved counters for agent {agent!r}")
if raw.get("best_metric") is not None:
self._finite(raw["best_metric"], "best_metric")
if raw.get("best_evaluation_id") is not None and not isinstance(
raw["best_evaluation_id"], str
):
raise ValueError("best_evaluation_id must be a string or null")
restored[agent] = _AgentState(
last_evaluation=(
None
if raw["last_evaluation"] is None
else self._nonnegative_int(
raw["last_evaluation"], "last_evaluation"
)
),
seen=dict(seen),
success_streak=raw["success_streak"],
plateau_patience=raw["plateau_patience"],
best_metric=(
None
if raw["best_metric"] is None
else self._finite(raw["best_metric"], "best_metric")
),
best_evaluation_id=raw["best_evaluation_id"],
history=deque(history),
)
reason = state.get("termination_reason")
self._states = restored
self._termination_reason = (
EvaluationTerminationReason(reason) if reason is not None else None
)
def _record_outcome(
self, agent_name: str, outcome: AgentEvaluationResponse
) -> None:
state = self._states[agent_name]
evaluation_id = outcome.evaluation_id
if not isinstance(evaluation_id, str) or not evaluation_id:
self._invalid_outcome(agent_name, "missing evaluation identity")
return
evaluation = outcome.evaluation_episode
if not isinstance(outcome.evaluation_metrics, Mapping):
raise ValueError("evaluation metrics must be a mapping")
if outcome.checkpoint_id is not None and not isinstance(
outcome.checkpoint_id, str
):
raise ValueError("checkpoint_id must be a string or null")
if outcome.checkpoint_hash is not None and not isinstance(
outcome.checkpoint_hash, str
):
raise ValueError("checkpoint_hash must be a string or null")
canonical = {
"evaluation_id": evaluation_id,
"evaluation_episode": evaluation,
"checkpoint_id": outcome.checkpoint_id,
"checkpoint_hash": outcome.checkpoint_hash,
"metrics": dict(outcome.evaluation_metrics),
}
previous = state.seen.get(evaluation_id)
if previous is not None:
if previous != canonical:
raise ValueError(
f"conflicting duplicate evaluation result {evaluation_id!r}"
)
return
expected_evaluation = (
evaluation
if state.last_evaluation is None
else state.last_evaluation + 1
)
if (
isinstance(evaluation, bool)
or not isinstance(evaluation, int)
or evaluation < 0
or evaluation != expected_evaluation
):
raise ValueError(
f"out-of-order evaluation {evaluation!r} for {agent_name!r}; "
f"expected {expected_evaluation}"
)
state.seen[evaluation_id] = canonical
state.last_evaluation = evaluation
if not self._is_valid_outcome(agent_name, outcome):
self._invalid_outcome(
agent_name, "missing or non-finite configured metric"
)
return
state.history.append(canonical)
LOG.info(
"Accepted deterministic evaluation_id=%s for agent=%s "
"episode=%d metrics=%s checkpoint_id=%r checkpoint_hash=%r",
evaluation_id,
agent_name,
evaluation,
canonical["metrics"],
outcome.checkpoint_id,
outcome.checkpoint_hash,
)
success = self._success_rules[agent_name]
plateau = self._plateau_rules[agent_name]
if success is not None:
value = float(outcome.evaluation_metrics[success.metric])
if self._qualifies_success(value, success):
state.success_streak += 1
else:
state.success_streak = 0
if plateau is not None:
plateau_value = self._smoothed_metric(
state, plateau.metric, plateau.window
)
if plateau_value is not None:
if self._is_improvement(
plateau_value, state.best_metric, plateau
):
state.best_metric = plateau_value
state.best_evaluation_id = evaluation_id
state.plateau_patience = 0
elif len(state.history) >= plateau.min_evaluations:
state.plateau_patience += 1
def _invalid_outcome(self, agent_name: str, reason: str) -> None:
state = self._states[agent_name]
state.success_streak = 0
LOG.warning(
"Ignoring invalid deterministic evaluation for %s: %s",
agent_name,
reason,
)
def _is_valid_outcome(
self, agent_name: str, outcome: AgentEvaluationResponse
) -> bool:
required_metrics = {
rule.metric
for rule in (
self._success_rules[agent_name],
self._plateau_rules[agent_name],
)
if rule is not None
}
return all(
metric in outcome.evaluation_metrics
and isinstance(outcome.evaluation_metrics[metric], (int, float))
and math.isfinite(float(outcome.evaluation_metrics[metric]))
for metric in required_metrics
)
def _all_successful(self) -> bool:
if not self._success_rules:
return False
for agent, rule in self._success_rules.items():
if rule is None:
return False
state = self._states[agent]
if (
len(state.history) < rule.min_evaluations
or state.success_streak < rule.consecutive
):
return False
return True
def _all_plateaued(self) -> bool:
if not self._plateau_rules:
return False
for agent, rule in self._plateau_rules.items():
if (
rule is None
or self._states[agent].plateau_patience < rule.patience
):
return False
return True
@staticmethod
def _qualifies_success(value: float, rule: _SuccessRule) -> bool:
if rule.direction == "greater_equal":
return value >= rule.threshold
return value <= rule.threshold
@staticmethod
def _is_improvement(
value: float, best: Optional[float], rule: _PlateauRule
) -> bool:
if best is None:
return True
if rule.direction == "maximize":
return value >= best + rule.min_delta
return value <= best - rule.min_delta
@staticmethod
def _smoothed_metric(
state: _AgentState, metric: str, window: int
) -> Optional[float]:
if len(state.history) < window:
return None
values = [
float(item["metrics"][metric])
for item in list(state.history)[-window:]
if metric in item["metrics"]
]
return sum(values) / len(values) if len(values) == window else None
@staticmethod
def _parse_success(agent_name: str, raw: Any) -> _SuccessRule:
if not isinstance(raw, Mapping):
raise ValueError(
f"success rule for {agent_name!r} must be a mapping"
)
required = {
"metric",
"threshold",
"direction",
"consecutive",
"min_evaluations",
}
if set(raw) != required:
raise ValueError(
f"success rule for {agent_name!r} must contain exactly "
f"{sorted(required)!r}"
)
direction = raw["direction"]
if (
direction
not in AgentEvaluationObjectiveTerminationCondition._SUCCESS_DIRECTIONS
):
raise ValueError(
"success direction must be greater_equal or less_equal"
)
return _SuccessRule(
metric=AgentEvaluationObjectiveTerminationCondition._parse_metric(
raw["metric"]
),
threshold=AgentEvaluationObjectiveTerminationCondition._finite(
raw["threshold"], "threshold"
),
direction=direction,
consecutive=AgentEvaluationObjectiveTerminationCondition._positive_int(
raw["consecutive"], "consecutive"
),
min_evaluations=AgentEvaluationObjectiveTerminationCondition._positive_int(
raw["min_evaluations"], "min_evaluations"
),
)
@staticmethod
def _parse_plateau(agent_name: str, raw: Any) -> _PlateauRule:
if not isinstance(raw, Mapping):
raise ValueError(
f"plateau rule for {agent_name!r} must be a mapping"
)
required = {
"metric",
"direction",
"min_delta",
"patience",
"min_evaluations",
"window",
}
if set(raw) != required:
raise ValueError(
f"plateau rule for {agent_name!r} must contain exactly "
f"{sorted(required)!r}"
)
direction = raw["direction"]
if (
direction
not in AgentEvaluationObjectiveTerminationCondition._PLATEAU_DIRECTIONS
):
raise ValueError("plateau direction must be maximize or minimize")
return _PlateauRule(
metric=AgentEvaluationObjectiveTerminationCondition._parse_metric(
raw["metric"]
),
direction=direction,
min_delta=AgentEvaluationObjectiveTerminationCondition._nonnegative_finite(
raw["min_delta"], "min_delta"
),
patience=AgentEvaluationObjectiveTerminationCondition._positive_int(
raw["patience"], "patience"
),
min_evaluations=AgentEvaluationObjectiveTerminationCondition._positive_int(
raw["min_evaluations"], "min_evaluations"
),
window=AgentEvaluationObjectiveTerminationCondition._positive_int(
raw["window"], "window"
),
)
@staticmethod
def _parse_metric(value: Any) -> str:
if not isinstance(value, str) or not value:
raise ValueError("metric must be a non-empty string")
return value
@staticmethod
def _finite(value: Any, field_name: str) -> float:
try:
result = float(value)
except (TypeError, ValueError) as exc:
raise ValueError(f"{field_name} must be a finite number") from exc
if not math.isfinite(result):
raise ValueError(f"{field_name} must be a finite number")
return result
@staticmethod
def _nonnegative_finite(value: Any, field_name: str) -> float:
result = AgentEvaluationObjectiveTerminationCondition._finite(
value, field_name
)
if result < 0:
raise ValueError(f"{field_name} must be non-negative")
return result
@staticmethod
def _positive_int(value: Any, field_name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
raise ValueError(f"{field_name} must be a positive integer")
return value
@staticmethod
def _nonnegative_int(value: Any, field_name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError(f"{field_name} must be a non-negative integer")
return value
@staticmethod
def _require_complete_global_rules(
rule_name: str, rules: Mapping[str, Optional[Any]]
) -> None:
"""Avoid a mixed multi-agent rule with no unambiguous global outcome."""
present = {agent for agent, rule in rules.items() if rule is not None}
if present and len(present) != len(rules):
missing = sorted(set(rules) - present)
raise ValueError(
f"{rule_name} must be configured for every agent or no agent; "
f"missing for {missing!r}"
)
def _make_configuration_fingerprint(self) -> str:
return hashlib.sha256(
json.dumps(
self._configuration(), sort_keys=True, separators=(",", ":")
).encode()
).hexdigest()
def _configuration(self) -> Dict[str, Dict[str, Optional[Dict[str, Any]]]]:
return {
"success": {
agent: None if rule is None else rule.__dict__
for agent, rule in self._success_rules.items()
},
"plateau": {
agent: None if rule is None else rule.__dict__
for agent, rule in self._plateau_rules.items()
},
}
def _set_terminal_reason(
self, reason: EvaluationTerminationReason
) -> None:
if self._termination_reason is None:
self._termination_reason = reason
LOG.info("Evaluation-aware termination reason: %s", reason.value)