diff --git a/predicators/agent_sdk/tools/synthesis.py b/predicators/agent_sdk/tools/synthesis.py index 657ed999c..4536e607a 100644 --- a/predicators/agent_sdk/tools/synthesis.py +++ b/predicators/agent_sdk/tools/synthesis.py @@ -2,6 +2,7 @@ import dataclasses import hashlib import os +import time from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union import numpy as np @@ -336,6 +337,7 @@ def create_synthesis_tools( from predicators.code_sim_learning.base_simulator import \ base_simulator_class + from predicators.code_sim_learning.config import SysIdConfig from predicators.code_sim_learning.evidence import format_evidence_lines from predicators.code_sim_learning.fit_space import ParamSpec from predicators.code_sim_learning.fitting import compute_sse, \ @@ -527,8 +529,16 @@ def _evaluate_rollout_fit(rules: list, for s in list(physical_specs) + list(rule_specs) } anchors = approach.fit_prior_anchors(physical_specs) + # pylint: disable-next=import-outside-toplevel + from predicators.agent_sdk.belief_probe import ProbeBudgetExceeded + fit_timeout = float(CFG.agent_sdk_fit_call_timeout) + # The fit's own search yields to the budget the watchdog enforces. + fit_config = dataclasses.replace( + SysIdConfig.from_cfg(), + deadline=(time.monotonic() + + fit_timeout if fit_timeout > 0 else None)) try: - with suspend_budget_watchdog(CFG.agent_sdk_fit_call_timeout): + with suspend_budget_watchdog(fit_timeout): outcome = run_rollout_sysid( fit_env, rollouts, @@ -548,7 +558,19 @@ def _evaluate_rollout_fit(rules: list, approach, "_explainability_cache", None)), fit_cache=(None if exploratory else getattr( approach, "_sysid_fit_cache", None)), - fit_cache_key=version_tag) + fit_cache_key=version_tag, + config=fit_config) + except ProbeBudgetExceeded: + return ( + f"[{version_tag}] Error: the rollout system-ID fit ran past " + f"its {fit_timeout:.0f} s limit and was stopped; nothing was " + "fitted or applied, and the same fit will stop again. A fit " + "replays every scored recording once per candidate parameter " + "point, so its time grows with the declared parameters and " + "with the steps each replay runs. Recordings are cut where " + "every scored feature rests, so a scored feature that keeps " + "changing (a robot that is always moving, a running counter) " + "keeps a recording whole.") except Exception as e: # pylint: disable=broad-except return ( f"[{version_tag}] Error: rollout system-ID fit failed:\n{e}") @@ -725,6 +747,8 @@ def _evaluate_rollout_fit(rules: list, if outcome.belief is not None: lines.extend(["", *_format_parameter_belief(outcome.belief), ""]) + if outcome.fit_result.search_note: + lines.extend([outcome.fit_result.search_note, ""]) if exploratory: lines.append( "EXPLORATORY subset fit: nothing was applied or recorded " diff --git a/predicators/code_sim_learning/config.py b/predicators/code_sim_learning/config.py index 77a621923..8eb10c943 100644 --- a/predicators/code_sim_learning/config.py +++ b/predicators/code_sim_learning/config.py @@ -114,6 +114,10 @@ class SysIdConfig: # code_sim_learning_rollout_global_design_points / _global_starts). global_design_points: int = 0 global_starts: int = 0 + # The time.monotonic() time by which the fit call must have returned + # (the sim.fit tool's budget), or None for no limit. Not a flag: the + # caller that owns the budget sets it (run_rollout_sysid). + deadline: Optional[float] = None @property def flat_band_frac(self) -> float: diff --git a/predicators/code_sim_learning/fit_space.py b/predicators/code_sim_learning/fit_space.py index f1b51d918..812546c1d 100644 --- a/predicators/code_sim_learning/fit_space.py +++ b/predicators/code_sim_learning/fit_space.py @@ -129,6 +129,10 @@ class FitResult: # error that persists through the recordings. None = no multi-start # ran. misfit: Optional[float] = None + # How the whole-box search stopped short of its full design and starts + # to keep the fit call inside its time budget, for the report; empty + # when it ran in full or did not run. + search_note: str = "" @property def point_estimate(self) -> Dict[str, float]: diff --git a/predicators/code_sim_learning/physical_sysid.py b/predicators/code_sim_learning/physical_sysid.py index a65c588d8..1aaab687e 100644 --- a/predicators/code_sim_learning/physical_sysid.py +++ b/predicators/code_sim_learning/physical_sysid.py @@ -293,6 +293,7 @@ def fit_params_rollout( # info-seeking explorer's calibrated ensemble. lm_notes: List[str] = [] lm_flat: List[str] = [] + lm_t0 = time.monotonic() lm_theta, lm_jac = fit_map_lm_rollout(base_env, trajectories, lm_physical_specs, @@ -306,6 +307,7 @@ def fit_params_rollout( noise_sigma=noise_sigma, notes_out=lm_notes, flat_params_out=lm_flat) + lm_seconds = time.monotonic() - lm_t0 if (config.log_hessian_identifiability and lm_jac is not None and lm_jac.size > 0): log_hessian_identifiability(lm_jac, names, noise_sigma, prior_sigma) @@ -335,13 +337,15 @@ def fit_params_rollout( } basins: Optional[List[Tuple[Dict[str, float], float, float]]] = None misfit: Optional[float] = None + search_note = "" if (config.anchors_are_guesses and config.global_design_points > 0 and config.global_starts > 0 and trajectories): - lm_theta, lm_jac, basins, misfit = _global_multistart( + (lm_theta, lm_jac, basins, misfit, search_note) = _global_multistart( base_env, trajectories, all_specs, physical_specs, rule_specs, residual_features, rules, latent_init, scaling, center_int, prior_sigma, noise_sigma, config, (lm_theta, lm_jac), noise_sse, - sigma_tol) + sigma_tol, + _SearchClock(config.deadline, fit_t0, n_start, lm_seconds)) elif config.anchors_are_guesses and trajectories and noise_sse > 0.0: # The same misfit the multi-start measures, at the local fit. map_sse = compute_rollout_sse( @@ -359,7 +363,8 @@ def fit_params_rollout( sensitivity=sensitivity, lm_notes=lm_notes, basins=basins, - misfit=misfit) + misfit=misfit, + search_note=search_note) n_lm = num_rollouts_run() - n_start - n_grid # The ablation pins moved parameters back to their anchors, which # only calibrated baselines warrant; a guess has no claim beyond the @@ -398,6 +403,42 @@ def fit_params_rollout( # fit-space box width are one basin. _SAME_BASIN_FRACTION = 0.05 +# The share of the time a fit call has left when its fit starts that the +# whole-box search may use; the rest is the report's and the belief's, +# which trace every basin. +_SEARCH_BUDGET_SHARE = 0.5 + + +@dataclasses.dataclass(frozen=True) +class _SearchClock: + """The whole-box search's time budget: ``deadline`` (monotonic, or None for + no limit) is when the fit call must have returned, ``fit_t0`` and + ``rollouts_before`` when and after how many rollouts the fit started, and + ``lm_seconds`` how long the local LM took.""" + + deadline: Optional[float] + fit_t0: float + rollouts_before: int + lm_seconds: float + + @property + def budget(self) -> float: + """The seconds the fit had when it started (inf without a limit).""" + if self.deadline is None: + return float("inf") + return max(0.0, self.deadline - self.fit_t0) + + def allows(self, seconds: float) -> bool: + """Whether a phase projected to take ``seconds`` ends inside the + search's share of the budget.""" + return (time.monotonic() + seconds <= + self.fit_t0 + _SEARCH_BUDGET_SHARE * self.budget) + + def seconds_per_rollout(self) -> float: + """The fit's wall time per rollout so far, waves included.""" + done = max(1, num_rollouts_run() - self.rollouts_before) + return (time.monotonic() - self.fit_t0) / done + def _global_multistart( base_env: Any, @@ -416,8 +457,9 @@ def _global_multistart( incumbent: Tuple[np.ndarray, Optional[np.ndarray]], noise_sse: float, sigma_tol: float, -) -> Tuple[np.ndarray, Optional[np.ndarray], List[Tuple[Dict[ - str, float], float, float]], Optional[float]]: + clock: _SearchClock, +) -> Tuple[np.ndarray, Optional[np.ndarray], Optional[List[Tuple[Dict[ + str, float], float, float]]], Optional[float], str]: """Search the whole box for the fit when the starting values are guesses. A scrambled Sobol design of ``config.global_design_points`` over @@ -438,9 +480,36 @@ def _global_multistart( is returned too: the belief's temperature counts it the same way. Without a declared noise channel (``noise_sse`` 0) the misfit is unknown: the point estimate is the only basin and the misfit None. + + The search improves on a fit that already exists, so it never makes + the fit call overrun its budget (``clock``): the design and each LM + start run only if they are projected to end within the search's + share of it, timed by the fit so far. A search cut short says so in + the returned note (empty when it ran in full); one that cannot + afford the design and a start keeps the incumbent, whose misfit is + still measured, and returns no basins. """ names = [s.name for s in all_specs] physical_names = [s.name for s in physical_specs] + design_seconds = (clock.seconds_per_rollout() * + config.global_design_points * len(trajectories)) + if not clock.allows(design_seconds + clock.lm_seconds): + point = dict(zip(names, (float(v) for v in incumbent[0]))) + sse = rollout_sse_by_point(base_env, trajectories, [point], + residual_features, physical_names, rules, + latent_init, scaling)[0] + incumbent_misfit = (max(sse - noise_sse, sigma_tol) + if noise_sse > 0.0 else None) + note = (f"The whole-box search did not run: the fit had used " + f"{time.monotonic() - clock.fit_t0:.0f} s of its " + f"{clock.budget:.0f} s budget, and the design and one LM " + f"start (about {design_seconds + clock.lm_seconds:.0f} s) " + "would not have ended within half of it. The point estimate " + "is the fit from the starting values; other parameter sets " + "that fit the recordings about as well were not looked for.") + logger.info("Rollout sysID global multi-start skipped: %s", note) + return (np.asarray(incumbent[0], dtype=float), incumbent[1], None, + incumbent_misfit, note) lo, hi = fit_space_bounds(list(all_specs)) lo = np.where(np.isfinite(lo), lo, center_int - 2.0 * prior_sigma) hi = np.where(np.isfinite(hi), hi, center_int + 2.0 * prior_sigma) @@ -471,9 +540,19 @@ def started_at(specs: Sequence[ParamSpec], thetas = [np.asarray(incumbent[0], dtype=float)] jacobians = [incumbent[1]] - for index in np.argsort(np.asarray(sses), - kind="stable")[:config.global_starts]: + starts = np.argsort(np.asarray(sses), kind="stable")[:config.global_starts] + lm_seconds = clock.lm_seconds + note = "" + for ran, index in enumerate(starts): + if not clock.allows(lm_seconds): + note = (f"The whole-box search ran the LM from {ran} of its " + f"{len(starts)} best design points: another start (about " + f"{lm_seconds:.0f} s) would not have ended within half " + f"of the fit's {clock.budget:.0f} s budget.") + logger.info("Rollout sysID global multi-start cut short: %s", note) + break start = design[int(index)] + lm_t0 = time.monotonic() theta, jac = fit_map_lm_rollout(base_env, trajectories, started_at(physical_specs, start), @@ -485,6 +564,7 @@ def started_at(specs: Sequence[ParamSpec], prior_centers=center_int, prior_sigmas=prior_sigma, noise_sigma=noise_sigma) + lm_seconds = max(lm_seconds, time.monotonic() - lm_t0) thetas.append(np.asarray(theta, dtype=float)) jacobians.append(jac) points = [dict(zip(names, (float(v) for v in theta))) for theta in thetas] @@ -525,7 +605,7 @@ def started_at(specs: Sequence[ParamSpec], float(np.min(sses)), len(thetas) - 1, finals[0], [round(float(s), 4) for s in finals[1:]], len(basins), band) - return thetas[best], jacobians[best], basins, misfit + return thetas[best], jacobians[best], basins, misfit, note # A MAP within this fit-space distance of its anchor counts as unmoved diff --git a/tests/approaches/test_agent_continual_approach.py b/tests/approaches/test_agent_continual_approach.py index aafd52f42..584c3f2dd 100644 --- a/tests/approaches/test_agent_continual_approach.py +++ b/tests/approaches/test_agent_continual_approach.py @@ -298,6 +298,59 @@ def fake_query(*_args: Any, **_kwargs: Any) -> List[Dict[str, Any]]: assert card.levels[0].sandbox.get("fits", 0) == int(fit_then_edit) +@pytest.mark.slow +def test_fit_that_overruns_its_budget_says_so(tmp_path: Any, + monkeypatch: Any) -> None: + """A fit stopped by its time limit tells the agent what happened and what + makes fits long, where the watchdog's bare exception used to print an empty + "fit failed:".""" + # pylint: disable=import-outside-toplevel,protected-access + from predicators.agent_sdk.belief_probe import ProbeBudgetExceeded + from predicators.code_sim_learning import orchestrator + + seen: Dict[str, Any] = {} + + def overrun(*_args: Any, **kwargs: Any) -> Any: + seen["deadline"] = kwargs["config"].deadline + raise ProbeBudgetExceeded + + # The fit tool binds run_rollout_sysid when the approach builds it. + monkeypatch.setattr(orchestrator, "run_rollout_sysid", overrun) + _config(tmp_path, agent_sdk_fit_call_timeout=1200.0) + env, approach = _make_approach() + source = ''' +class Counter(BaseSimulator): + AGENT_PARAM_SPECS = [ParamSpec("rate", .25, lo=0.0, hi=1.0)] + MODEL_STATE_INIT = {"charge": 0.0} + RESIDUAL_FEATURES = {} + + @classmethod + def update_model_state(cls, observation, model_state, params, action): + model_state["charge"] += params["rate"] + +RESIDUAL_ENV = Counter +''' + + def fake_query(*_args: Any, **_kwargs: Any) -> List[Dict[str, Any]]: + assert "step applied" in _call(approach, + "env_step", + action=[0.0] * + env.action_space.shape[0]) + path = os.path.join(approach._tool_context.sandbox_dir, "simulator.py") + with open(path, "w", encoding="utf-8") as file: + file.write(source) + seen["output"] = _call(approach, "run_python", code="print(sim.fit())") + assert "Give-up recorded" in _call(approach, "give_up", note="done") + return _result() + + monkeypatch.setattr(approach, "_query_agent_sync", fake_query) + approach.prepare_for_continual(Dataset([])) + ContinualRun(env, approach, create_level_player(env, approach)).run() + assert "ran past its 1200 s limit" in seen["output"] + assert "nothing was fitted or applied" in seen["output"] + assert seen["deadline"] is not None + + @pytest.mark.slow def test_play_loop_with_a_scripted_agent(tmp_path: Any) -> None: """Two rounds of one conversation: act in the first, which also carries the diff --git a/tests/code_sim_learning/test_physical_sysid.py b/tests/code_sim_learning/test_physical_sysid.py index bc8a30cf7..8183db751 100644 --- a/tests/code_sim_learning/test_physical_sysid.py +++ b/tests/code_sim_learning/test_physical_sysid.py @@ -1553,7 +1553,7 @@ def run(noise_sse): None, [_trajectory([0.0, 0.1])], [spec], [spec], [], _RESIDUAL_FEATURES, (), None, None, np.array([0.2]), np.array([1e6]), 0.05, config, (np.array([0.2]), None), noise_sse, - 0.0) + 0.0, physical_sysid._SearchClock(None, 0.0, 0, 0.0))[:4] # The best fit leaves 0.05 above the noise: the 0.2 basin, 0.2 worse, # is distinguishable. @@ -1577,6 +1577,66 @@ def run(noise_sse): assert misfit is None +class _FakeTime: + """A monotonic clock that only the test advances.""" + + def __init__(self): + self.now = 1000.0 + + def monotonic(self): + """The current fake time.""" + return self.now + + +def test_global_multistart_stays_inside_the_fit_budget(monkeypatch): + """The whole-box search runs the LM starts that end within its share of the + fit's time budget and says how far it got; with no room for the design and + one start it keeps the guesses' fit, still measures its misfit, and mixes + no basins.""" + import dataclasses + + from predicators.code_sim_learning.config import SysIdConfig + _stub_multistart(monkeypatch) + fake = _FakeTime() + monkeypatch.setattr(physical_sysid, "time", fake) + stub_lm = physical_sysid.fit_map_lm_rollout + + def ten_second_lm(*args, **kwargs): + fake.now += 10.0 + return stub_lm(*args, **kwargs) + + monkeypatch.setattr(physical_sysid, "fit_map_lm_rollout", ten_second_lm) + spec = ParamSpec("k", 0.2, lo=0.0, hi=1.0) + config = dataclasses.replace(SysIdConfig.from_cfg(), + global_design_points=16, + global_starts=3) + + def run(budget): + deadline = None if budget is None else fake.now + budget + clock = physical_sysid._SearchClock(deadline, fake.now, + physical_sysid.num_rollouts_run(), + 10.0) + return physical_sysid._global_multistart( + None, [_trajectory([0.0, 0.1])], [spec], [spec], + [], _RESIDUAL_FEATURES, (), None, None, np.array([0.2]), + np.array([1e6]), 0.05, config, (np.array([0.2]), None), 10.0, 0.0, + clock) + + # Half of a 50 s budget holds two 10 s starts, not a third. + theta, _jac, basins, _misfit, note = run(50.0) + assert theta[0] == 0.8 and basins + assert "from 2 of its 3 best design points" in note + # Half of a 10 s budget holds no start: the guesses' fit stays, with + # its misfit (SSE 10.5 at 0.2, against 10.0 of noise). + theta, _jac, basins, misfit, note = run(10.0) + assert theta[0] == 0.2 and basins is None + assert misfit == pytest.approx(0.5) + assert "did not run" in note + # Without a limit the whole search runs. + theta, _jac, basins, _misfit, note = run(None) + assert theta[0] == 0.8 and note == "" + + def test_global_multistart_runs_for_guessed_starting_values_only(monkeypatch): """Calibrated starting values sit in the right basin, so only guessed ones get the whole-box search.""" @@ -1596,7 +1656,7 @@ def fake_lm(_env, def fake_multistart(*args, **_kwargs): calls.append(True) - return args[13][0], args[13][1], [({"mu": 0.5}, 1.0, 1.0)], 0.2 + return args[13][0], args[13][1], [({"mu": 0.5}, 1.0, 1.0)], 0.2, "" monkeypatch.setattr(physical_sysid, "fit_map_lm_rollout", fake_lm) monkeypatch.setattr(physical_sysid, "_grid_seed_physical_specs",