Make the rollout fit about 3 times faster without changing its results

predicators/agent_sdk/parallel_rollouts.py

@@ -16,10 +16,11 @@
 ``scripts/parallel_rollout_bench.py``), both load-bearing:
 
 * Children exit with ``os._exit`` right after a SYNCHRONOUS
-  ``SimpleQueue.put``. Normal interpreter teardown runs PyBullet's
-  atexit handler against the forked copy of the parent's live client
-  and lags exit by seconds; ``mp.Queue`` would instead hand the result
-  to a feeder thread that ``os._exit`` could kill mid-write.
+  ``SimpleQueue.put`` (or :class:`_ResultChannel`'s, the same locked
+  pipe). Normal interpreter teardown runs PyBullet's atexit handler
+  against the forked copy of the parent's live client and lags exit by
+  seconds; ``mp.Queue`` would instead hand the result to a feeder
+  thread that ``os._exit`` could kill mid-write.
 * The parent's window is bookkept by OUTSTANDING RESULTS, never by
   process liveness: a child that delivered its result frees its slot
   even while its exit lags, and a liveness-based window deadlocks.
@@ -46,9 +47,26 @@ def parallel_rollouts_available() -> bool:
     return sys.platform.startswith("linux") and hasattr(os, "fork")
 
 
-def prefetch_parallel(jobs: Sequence[Callable[[], Any]],
-                      label: str,
-                      quiet: bool = False) -> List[Optional[Any]]:
+def parallel_workers() -> int:
+    """How many forked children a wave may run at once.
+
+    One or fewer means waves run nothing in parallel: the setting
+    ``CFG.agent_validation_parallel_workers`` is at most 1, or this
+    platform cannot fork.
+    """
+    # Deferred: settings must stay import-cycle-free from tool modules.
+    # pylint: disable-next=import-outside-toplevel
+    from predicators.settings import CFG
+    if not parallel_rollouts_available():
+        return 0
+    return int(CFG.agent_validation_parallel_workers)
+
+
+def prefetch_parallel(
+        jobs: Sequence[Callable[[], Any]],
+        label: str,
+        quiet: bool = False,
+        costs: Optional[Sequence[float]] = None) -> List[Optional[Any]]:
     """Pre-run independent rollout jobs as forked children when enabled.
 
     Returns an index-aligned result list, or all-``None`` when parallel
@@ -62,17 +80,43 @@ def prefetch_parallel(jobs: Sequence[Callable[[], Any]],
     ``quiet`` demotes the per-call INFO line to DEBUG: the sysID
     objective calls this once per candidate theta, hundreds of times
     per fit, and one log line each would drown the run's info.log.
+    ``costs`` orders the children's starts (see
+    :func:`run_forked_rollouts`).
     """
-    # Deferred: settings must stay import-cycle-free from tool modules.
-    # pylint: disable-next=import-outside-toplevel
-    from predicators.settings import CFG
-    workers = min(int(CFG.agent_validation_parallel_workers), len(jobs))
-    if workers <= 1 or len(jobs) <= 1 or not parallel_rollouts_available():
+    workers = min(parallel_workers(), len(jobs))
+    if workers <= 1:
         return [None] * len(jobs)
     logger.log(logging.DEBUG if quiet else logging.INFO,
                "[%s] prefetching %d rollouts across %d forked children.",
                label, len(jobs), workers)
-    return run_forked_rollouts(jobs, workers, label)
+    return run_forked_rollouts(jobs, workers, label, costs=costs)
+
+
+class _ResultChannel:
+    """The locked pipe of ``SimpleQueue``, which the parent can also wait on
+    with a timeout."""
+
+    def __init__(self, ctx: Any) -> None:
+        self._reader, self._writer = ctx.Pipe(duplex=False)
+        self._lock = ctx.Lock()
+
+    def put(self, obj: Any) -> None:
+        """Send ``obj`` whole (a child's only write)."""
+        with self._lock:
+            self._writer.send(obj)
+
+    def wait(self, timeout: float) -> bool:
+        """Whether a result arrived within ``timeout`` seconds."""
+        return bool(self._reader.poll(timeout))
+
+    def get(self) -> Any:
+        """The next result (call after :meth:`wait` said one arrived)."""
+        return self._reader.recv()
+
+    def close(self) -> None:
+        """Close both ends in this process."""
+        self._reader.close()
+        self._writer.close()
 
 
 def _child_main(idx: int, job: Callable[[], Any], q: Any,
@@ -105,6 +149,7 @@ def run_forked_rollouts(
     label: str,
     per_job_timeout: float = 900.0,
     quiet_child_logging: bool = True,
+    costs: Optional[Sequence[float]] = None,
 ) -> List[Optional[Any]]:
     """Run ``jobs`` in forked children, at most ``max_workers`` at once.
 
@@ -114,37 +159,42 @@ def run_forked_rollouts(
     be picklable; jobs must not depend on each other or on shared
     mutable state surviving into the parent (child-side mutations stay
     in the child).
+
+    ``costs`` (an estimate per job, such as its rollout's length) starts
+    the costliest jobs first, so a wave's longest job never waits for a
+    worker behind short ones; ties keep the job order. The parent takes
+    each result the moment it arrives and starts the next job before
+    the finished child has exited.
     """
     n = len(jobs)
     if n == 0:
         return []
     assert max_workers >= 1
+    order = list(range(n))
+    if costs is not None:
+        assert len(costs) == n
+        order.sort(key=lambda k: -float(costs[k]))
     ctx = mp.get_context("fork")
-    q = ctx.SimpleQueue()
+    q = _ResultChannel(ctx)
     procs: dict = {}
+    exiting: list = []
     out: List[Optional[Any]] = [None] * n
     got = 0
-    next_idx = 0
+    started = 0
     outstanding = 0
     deadline = time.monotonic() + per_job_timeout * max(
         1, (n + max_workers - 1) // max_workers)
     try:
         while got < n:
-            while next_idx < n and outstanding < max_workers:
+            while started < n and outstanding < max_workers:
+                idx = order[started]
                 p = ctx.Process(target=_child_main,
-                                args=(next_idx, jobs[next_idx], q,
-                                      quiet_child_logging))
+                                args=(idx, jobs[idx], q, quiet_child_logging))
                 p.start()
-                procs[next_idx] = p
+                procs[idx] = p
                 outstanding += 1
-                next_idx += 1
-            timed_out = False
-            while q.empty():
-                if time.monotonic() > deadline:
-                    timed_out = True
-                    break
-                time.sleep(0.05)
-            if timed_out:
+                started += 1
+            if not q.wait(max(0.0, deadline - time.monotonic())):
                 logger.warning(
                     "[%s] parallel rollouts deadline exceeded with %d/%d "
                     "results; abandoning %d outstanding children.", label, got,
@@ -158,12 +208,18 @@ def run_forked_rollouts(
             else:
                 logger.warning("[%s] parallel rollout %d failed in child:\n%s",
                                label, idx, payload)
+            # The child exits right after its send: reap it once it has
+            # (is_alive reaps), never waiting for it here.
             done = procs.pop(idx, None)
             if done is not None:
-                done.join(timeout=30)
+                exiting.append(done)
+            exiting = [p for p in exiting if p.is_alive()]
     finally:
         for p in procs.values():
             if p.is_alive():
                 p.terminate()
             p.join(timeout=10)
+        for p in exiting:
+            p.join(timeout=30)
+        q.close()
     return out

predicators/code_sim_learning/grid_seed.py

@@ -20,7 +20,7 @@
     scalar_from_fit_space, scalar_to_fit_space
 from predicators.code_sim_learning.rollout_env import RolloutTrajectory
 from predicators.code_sim_learning.rollout_objective import \
-    compute_rollout_sse, per_trajectory_rms
+    compute_rollout_sse, rollout_sse_by_point, trajectory_terms_by_point
 from predicators.code_sim_learning.trajectory_prep import ResidualScaling
 
 logger = logging.getLogger(__name__)
@@ -286,6 +286,22 @@ def _sse_for(spec: ParamSpec, value: float) -> float:
                                             rules, latent_init, scaling)
         return memo[key]
 
+    def _score_pool(spec: ParamSpec, values: Sequence[float]) -> None:
+        """Score a pool's new candidates together (rollout_sse_by_point), with
+        the SSEs _sse_for would compute one by one."""
+        trials: Dict[Tuple[float, ...], Dict[str, float]] = {}
+        for value in values:
+            trial = dict(current)
+            trial[spec.name] = float(value)
+            key = tuple(trial[n] for n in names)
+            if key not in memo:
+                trials[key] = trial
+        sses = rollout_sse_by_point(base_env, trajectories,
+                                    list(trials.values()), residual_features,
+                                    physical_names, rules, latent_init,
+                                    scaling)
+        memo.update(zip(trials, sses))
+
     pools: Dict[str, List[Tuple[float, float]]] = {}
     for pass_idx in range(num_passes):
         moved_any = False
@@ -297,6 +313,7 @@ def _sse_for(spec: ParamSpec, value: float) -> float:
                 float(spec.init_value), anchor_of[spec.name],
                 current[spec.name]
             ]
+            _score_pool(spec, list(dict.fromkeys(values)))
             pool = [(v, _sse_for(spec, v)) for v in dict.fromkeys(values)]
             pools[spec.name] = pool
             flat, _best, tol = _flat_candidates(pool, noise_floor, flat_frac,
@@ -457,11 +474,15 @@ def min_explainable_fits(
         n: float(base[n])
         for n in physical_names
     } for _ in trajectories]
-    for params in candidates:
-        rms = per_trajectory_rms(base_env, trajectories, params,
-                                 residual_features, physical_names, rules,
-                                 latent_init, scaling)
-        for i, r in enumerate(rms):
+    # Every candidate's trajectories in one batch; the same per-trajectory
+    # terms per_trajectory_rms scores, each candidate's RMS computed alike.
+    per_candidate = trajectory_terms_by_point(base_env, trajectories,
+                                              candidates, residual_features,
+                                              physical_names, rules,
+                                              latent_init, scaling)
+    for params, per_trajectory in zip(candidates, per_candidate):
+        for i, res in enumerate(per_trajectory):
+            r = float(np.sqrt(np.mean(res**2))) if res.size else float("inf")
             if r < best[i]:
                 best[i] = r
                 best_params[i] = {n: float(params[n]) for n in physical_names}

predicators/code_sim_learning/lm.py

@@ -11,7 +11,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import numpy as np
 
@@ -328,6 +328,63 @@ def _eval_at(g: float, col: int = j) -> Tuple[float, float]:
     return z, sse, notes
 
 
+def _batched_jacobian(
+    residuals_z: Callable[[np.ndarray], np.ndarray],
+    batch_z: Callable[[List[np.ndarray]], List[np.ndarray]],
+    lo: np.ndarray,
+    hi: np.ndarray,
+    diff_step: Optional[float],
+) -> Tuple[Callable[[np.ndarray], np.ndarray], Callable[..., np.ndarray]]:
+    """least_squares' own 2-point Jacobian, its column points scored in one
+    batch; returns ``(residuals, jacobian)`` to pass to least_squares.
+
+    At each point, scipy's ``approx_derivative`` (the helper
+    least_squares calls for ``jac='2-point'``) runs twice: first with a
+    stand-in that only records where it would evaluate, then, once
+    ``batch_z`` has scored all those points together, with the cached
+    residuals. The Jacobian is therefore exactly the one least_squares
+    would compute, and least_squares counts no Jacobian evaluation in
+    ``nfev`` either way. ``residuals`` caches by point, so the Jacobian
+    reuses the residuals least_squares already took there.
+    """
+    # pylint: disable-next=import-outside-toplevel
+    from scipy.optimize._numdiff import approx_derivative
+    cache: Dict[bytes, np.ndarray] = {}
+
+    def residuals(z: np.ndarray) -> np.ndarray:
+        key = np.asarray(z, dtype=float).tobytes()
+        if key not in cache:
+            cache[key] = residuals_z(z)
+        return cache[key]
+
+    def jacobian(z: np.ndarray, *_args: Any, **_kwargs: Any) -> np.ndarray:
+        z = np.asarray(z, dtype=float)
+        f0 = residuals(z)
+        wanted: List[np.ndarray] = []
+
+        def record(x: np.ndarray) -> np.ndarray:
+            wanted.append(np.array(x, dtype=float))
+            return f0
+
+        approx_derivative(record,
+                          z,
+                          method="2-point",
+                          rel_step=diff_step,
+                          f0=f0,
+                          bounds=(lo, hi))
+        missing = [x for x in wanted if x.tobytes() not in cache]
+        for x, res in zip(missing, batch_z(missing)):
+            cache[x.tobytes()] = np.asarray(res, dtype=float)
+        return approx_derivative(residuals,
+                                 z,
+                                 method="2-point",
+                                 rel_step=diff_step,
+                                 f0=f0,
+                                 bounds=(lo, hi))
+
+    return residuals, jacobian
+
+
 def solve_lm(
     residuals_fn: Callable[[np.ndarray], np.ndarray],
     param_specs: List[ParamSpec],
@@ -337,6 +394,8 @@ def solve_lm(
     notes_out: Optional[List[str]] = None,
     n_prior_rows: int = 0,
     flat_params_out: Optional[List[str]] = None,
+    batch_residuals_fn: Optional[Callable[[List[np.ndarray]],
+                                          List[np.ndarray]]] = None,
 ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
     """Shared Levenberg-Marquardt core for the per-transition, recurrent, and
     rollout (``physical_sysid``) MAP fits.
@@ -382,6 +441,12 @@ def solve_lm(
     collects the names of parameters the search measured flat across
     their whole box (box-wide insensitivity evidence for the
     identifiability report).
+
+    ``batch_residuals_fn`` scores a list of (external) thetas together,
+    each as ``residuals_fn`` would. When given, the finite-difference
+    Jacobian scores its column points in one batch
+    (:func:`_batched_jacobian`), with exactly the values of the
+    ``diff_step`` 2-point Jacobian least_squares computes itself.
     """
     from scipy.optimize import \
         least_squares  # pylint: disable=import-outside-toplevel
@@ -395,9 +460,21 @@ def solve_lm(
     safe_hi = np.where(np.isfinite(hi), hi - 1e-9, np.inf)
     init = np.minimum(init, safe_hi)
 
-    def internal_residuals(z: np.ndarray) -> np.ndarray:
+    def residuals_at(z: np.ndarray) -> np.ndarray:
         return residuals_fn(from_fit_space(param_specs, z))
 
+    internal_residuals: Callable[[np.ndarray], np.ndarray] = residuals_at
+    # How least_squares differentiates; ``jac`` below is the result's.
+    jac_method: Any = "2-point"
+    if batch_residuals_fn is not None:
+        batch = batch_residuals_fn
+
+        def internal_batch(zs: List[np.ndarray]) -> List[np.ndarray]:
+            return batch([from_fit_space(param_specs, z) for z in zs])
+
+        internal_residuals, jac_method = _batched_jacobian(
+            residuals_at, internal_batch, lo, hi, diff_step)
+
     init_residuals = internal_residuals(init)
     if init_residuals.size == 0:
         logger.warning(
@@ -410,6 +487,7 @@ def internal_residuals(z: np.ndarray) -> np.ndarray:
     try:
         result = least_squares(internal_residuals,
                                init,
+                               jac=jac_method,
                                method='trf',
                                bounds=(lo, hi),
                                diff_step=diff_step,
@@ -450,6 +528,7 @@ def internal_residuals(z: np.ndarray) -> np.ndarray:
             try:
                 polished = least_squares(internal_residuals,
                                          z_new,
+                                         jac=jac_method,
                                          method='trf',
                                          bounds=(lo, hi),
                                          diff_step=diff_step,