From b0e3f7e185e2b09b681e22852f7306fa24af18e3 Mon Sep 17 00:00:00 2001
From: Yichao Liang <ycliang6@gmail.com>
Date: Fri, 9 Oct 2026 03:22:23 -0400
Subject: [PATCH 1/3] Score the fit's independent rollouts together

A rollout fit scores many parameter points whose rollouts do not depend
on each other: the explainability sweep's candidates, each grid pool,
each LM Jacobian's column points and the parameter belief's line
points. The objective scored them one point at a time. Each point forked
one wave of children, one per segment, and waited on its slowest segment
before the next point started, and the explainability sweep scored its
segments serially. One sim.fit of seed 0 of the Domino round fixes_r2
took 39 minutes.

trajectory_terms_by_point scores every (point, trajectory) pair in one
fork wave, with exactly the serial path's terms (per-step scoring with a
factory env; interval scoring and a shared env keep the point-by-point
path). The explainability sweep and the grid pools score their
candidates through it. solve_lm takes a batch scorer and computes
least_squares' own 2-point Jacobian with the column points scored
together: approx_derivative runs once to record the points and once on
their cached residuals. The belief traces its lines in lockstep, one
batch per round. Every value is the one the point-by-point path
computes; the new tests check it bit for bit.
---
 predicators/code_sim_learning/grid_seed.py    |  33 +-
 predicators/code_sim_learning/lm.py           |  83 ++++-
 predicators/code_sim_learning/orchestrator.py |  30 +-
 .../code_sim_learning/parameter_belief.py     | 113 +++++--
 .../code_sim_learning/rollout_objective.py    | 297 +++++++++++++++---
 .../test_parameter_belief.py                  |  44 +++
 .../code_sim_learning/test_physical_sysid.py  | 118 ++++++-
 .../test_zero_gradient_search.py              |  33 ++
 8 files changed, 670 insertions(+), 81 deletions(-)

diff --git a/predicators/code_sim_learning/grid_seed.py b/predicators/code_sim_learning/grid_seed.py
index 45437621f5..d19debe6b9 100644
--- a/predicators/code_sim_learning/grid_seed.py
+++ b/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}
diff --git a/predicators/code_sim_learning/lm.py b/predicators/code_sim_learning/lm.py
index 10a9ba6144..2eb7c7abd1 100644
--- a/predicators/code_sim_learning/lm.py
+++ b/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,
diff --git a/predicators/code_sim_learning/orchestrator.py b/predicators/code_sim_learning/orchestrator.py
index d3400789e3..78fdcdddbc 100644
--- a/predicators/code_sim_learning/orchestrator.py
+++ b/predicators/code_sim_learning/orchestrator.py
@@ -45,7 +45,8 @@
 from predicators.code_sim_learning.rollout_env import RolloutTrajectory, \
     num_rollouts_run
 from predicators.code_sim_learning.rollout_objective import \
-    compute_rollout_residuals, compute_rollout_sse
+    compute_rollout_residuals, compute_rollout_sse, \
+    rollout_residuals_by_point
 from predicators.code_sim_learning.trajectory_prep import \
     compute_residual_scaling
 
@@ -401,9 +402,16 @@ def rollout_residuals_fn(params: Dict[str, float]) -> np.ndarray:
                                              residual_features, physical_names,
                                              rules, latent_init, scaling)
 
+        def rollout_residuals_batch_fn(
+                points: List[Dict[str, float]]) -> List[np.ndarray]:
+            return rollout_residuals_by_point(fit_env, survivors, points,
+                                              residual_features,
+                                              physical_names, rules,
+                                              latent_init, scaling)
+
         belief = _build_belief(result, all_specs, anchors,
                                rollout_residuals_fn, belief_config,
-                               len(survivors))
+                               len(survivors), rollout_residuals_batch_fn)
     return _FitComputation(fit_result=result,
                            report=report,
                            num_survivors=len(survivors),
@@ -452,11 +460,16 @@ def prior_parameter_belief(all_specs: Sequence[ParamSpec],
                         seed=seed)
 
 
-def _build_belief(result: FitResult, all_specs: Sequence[ParamSpec],
-                  anchors: Dict[str, float],
-                  residuals_fn: Callable[[Dict[str, float]], np.ndarray],
-                  belief_config: BeliefConfig,
-                  num_survivors: int) -> ParameterBelief:
+def _build_belief(
+    result: FitResult,
+    all_specs: Sequence[ParamSpec],
+    anchors: Dict[str, float],
+    residuals_fn: Callable[[Dict[str, float]], np.ndarray],
+    belief_config: BeliefConfig,
+    num_survivors: int,
+    batch_residuals_fn: Optional[Callable[[List[Dict[str, float]]],
+                                          List[np.ndarray]]] = None
+) -> ParameterBelief:
     """The parameter factor around ``result``'s MAP, under its own prior.
 
     The prior is the one the fit folded in: centred on each parameter's
@@ -486,4 +499,5 @@ def _build_belief(result: FitResult, all_specs: Sequence[ParamSpec],
             for s, sigma in zip(all_specs, sigmas)
         },
         config=belief_config,
-        seed=stable_seed(CFG.seed, sorted(fitted.items()), num_survivors))
+        seed=stable_seed(CFG.seed, sorted(fitted.items()), num_survivors),
+        batch_residuals=batch_residuals_fn)
diff --git a/predicators/code_sim_learning/parameter_belief.py b/predicators/code_sim_learning/parameter_belief.py
index 67db2745b1..f9e76bb0e5 100644
--- a/predicators/code_sim_learning/parameter_belief.py
+++ b/predicators/code_sim_learning/parameter_belief.py
@@ -34,7 +34,8 @@
 import hashlib
 import logging
 from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+from typing import Any, Callable, Dict, Generator, Iterator, List, Optional, \
+    Sequence, Tuple
 
 import numpy as np
 import numpy.typing as npt
@@ -55,6 +56,8 @@
 
 # One pass of the fit's objective: the residual vector at params.
 ResidualsFn = Callable[[Dict[str, float]], np.ndarray]
+# Scores a list of parameter dicts together, each as a ResidualsFn would.
+BatchResidualsFn = Callable[[List[Dict[str, float]]], List[np.ndarray]]
 
 
 @dataclass(frozen=True)
@@ -411,13 +414,17 @@ def build_parameter_belief(
     prior_sigmas: Dict[str, float],
     config: BeliefConfig,
     seed: int,
+    batch_residuals: Optional[BatchResidualsFn] = None,
 ) -> ParameterBelief:
     """Build ``q(theta)`` around ``map_params`` and draw from it.
 
     ``residuals`` is the fit's objective (scaled residuals over the
     segments it pooled) and ``noise_sigma`` the Gaussian width it scores
     them with; ``prior_centers`` / ``prior_sigmas`` (fit space) are the
-    Gaussian prior it folded in.
+    Gaussian prior it folded in. ``batch_residuals``, when given, scores
+    a list of parameter dicts together, each as ``residuals`` would; the
+    lines are then traced in lockstep (:func:`_trace_lines`), with the
+    same result.
     """
     assert noise_sigma > 0.0, noise_sigma
     at_map = np.asarray(residuals(dict(map_params)), dtype=float)
@@ -432,6 +439,7 @@ def build_parameter_belief(
     lines: Dict[str, LinePosterior] = {}
     discrete: Dict[str, DiscretePosterior] = {}
     held: List[str] = []
+    traced: List[Tuple[ParamSpec, _LineEvaluator, float, float]] = []
     for spec in specs:
         center = prior_centers.get(spec.name)
         sigma_p = prior_sigmas.get(spec.name)
@@ -453,7 +461,14 @@ def build_parameter_belief(
                 values, neg_log)
             continue
         lo_arr, hi_arr = fit_space_bounds([spec])
-        _trace_line(line, float(lo_arr[0]), float(hi_arr[0]), config)
+        traced.append((spec, line, float(lo_arr[0]), float(hi_arr[0])))
+    if batch_residuals is not None:
+        _trace_lines([(line, lo, hi) for _spec, line, lo, hi in traced],
+                     config, batch_residuals)
+    else:
+        for _spec, line, lo, hi in traced:
+            _trace_line(line, lo, hi, config)
+    for spec, line, _lo, _hi in traced:
         evaluations += line.count
         grid = line.points()
         lines[spec.name] = LinePosterior.from_neg_log(grid,
@@ -590,14 +605,23 @@ def _prior(self, z: float) -> float:
         center, sigma = self.prior
         return 0.5 * ((z - center) / sigma)**2
 
+    def params_at(self, z: float) -> Dict[str, float]:
+        """The MAP with only this parameter moved to fit-space ``z``."""
+        params = dict(self.map_params)
+        params[self.spec.name] = scalar_from_fit_space(self.spec, float(z))
+        return params
+
+    def record(self, z: float, residuals: np.ndarray) -> None:
+        """Store the objective's residuals at ``z``, however they were
+        scored."""
+        res = np.asarray(residuals, dtype=float)
+        self.sse[float(z)] = float(np.dot(res, res))
+        self.count += 1
+
     def __call__(self, z: float) -> float:
         z = float(z)
         if z not in self.sse:
-            params = dict(self.map_params)
-            params[self.spec.name] = scalar_from_fit_space(self.spec, z)
-            res = np.asarray(self.residuals(params), dtype=float)
-            self.sse[z] = float(np.dot(res, res))
-            self.count += 1
+            self.record(z, self.residuals(self.params_at(z)))
         return ((self.sse[z] - self.sse_min) * self.k + self._prior(z) -
                 self._prior(self.z0))
 
@@ -606,15 +630,26 @@ def points(self) -> List[float]:
         return sorted(self.sse)
 
 
-def _trace_line(line: _LineEvaluator, lo: float, hi: float,
-                config: BeliefConfig) -> None:
-    """Evaluate a parameter's line until its posterior mass is resolved.
+def _line_points(line: _LineEvaluator, lo: float, hi: float,
+                 config: BeliefConfig) -> Iterator[float]:
+    """The points :func:`_trace_line` evaluates, as a generator.
+
+    It yields each point the line has not scored yet and resumes once the
+    point is scored (``line.record``), so driving it alone
+    (:func:`_trace_line`) or together with other lines
+    (:func:`_trace_lines`) evaluates the same points.
 
     Each side is probed outward from the MAP, first shrinking the step
     until a probe falls inside the cutoff, then doubling it until the
     cutoff or the declared bound is reached. The remaining budget
     bisects the grid interval holding the most mass.
     """
+
+    def value(z: float) -> Generator[float, None, float]:
+        if float(z) not in line.sse:
+            yield float(z)
+        return line(z)
+
     z0 = line.z0
     if line.sigma_p is not None:
         h0 = 0.25 * line.sigma_p
@@ -629,24 +664,27 @@ def _trace_line(line: _LineEvaluator, lo: float, hi: float,
         if (edge - z0) * side <= 0.0:
             continue
 
-        def probe(step: float,
-                  edge: float = edge,
-                  side: float = side) -> Tuple[float, float]:
+        def probe(
+                step: float,
+                edge: float = edge,
+                side: float = side
+        ) -> Generator[float, None, Tuple[float, float]]:
             z = z0 + side * step
             if (z - edge) * side > 0.0:
                 z = edge
-            return z, line(z)
+            g = yield from value(z)
+            return z, g
 
         step = h0
-        z, g = probe(step)
+        z, g = yield from probe(step)
         shrinks = 0
         while g > cutoff and shrinks < _MAX_SHRINKS and line.count < budget:
             step /= 4.0
-            z, g = probe(step)
+            z, g = yield from probe(step)
             shrinks += 1
         while g <= cutoff and z != edge and line.count < budget:
             step *= 2.0
-            z, g = probe(step)
+            z, g = yield from probe(step)
     while line.count < budget:
         grid = np.array(line.points())
         if grid.size < 2:
@@ -663,4 +701,41 @@ def probe(step: float,
         mid = 0.5 * (grid[widest] + grid[widest + 1])
         if mid in (grid[widest], grid[widest + 1]):
             return
-        line(float(mid))
+        yield from value(float(mid))
+
+
+def _trace_line(line: _LineEvaluator, lo: float, hi: float,
+                config: BeliefConfig) -> None:
+    """Evaluate a parameter's line until its posterior mass is resolved
+    (:func:`_line_points`), one point at a time."""
+    for z in _line_points(line, lo, hi, config):
+        line(z)
+
+
+def _trace_lines(traced: Sequence[Tuple[_LineEvaluator, float,
+                                        float]], config: BeliefConfig,
+                 batch_residuals: BatchResidualsFn) -> None:
+    """Trace several lines in lockstep: each round scores every unfinished
+    line's next point in one ``batch_residuals`` call.
+
+    The lines are independent (each moves one parameter off the MAP), so
+    each evaluates exactly the points it would alone, while a round
+    costs one batch where tracing line by line costs a call per point.
+    """
+    generators = [
+        _line_points(line, lo, hi, config) for line, lo, hi in traced
+    ]
+    pending: Dict[int, float] = {}
+    for i, generator in enumerate(generators):
+        z = next(generator, None)
+        if z is not None:
+            pending[i] = z
+    while pending:
+        order = list(pending)
+        results = batch_residuals(
+            [traced[i][0].params_at(pending[i]) for i in order])
+        for i, residuals in zip(order, results):
+            traced[i][0].record(pending.pop(i), residuals)
+            z = next(generators[i], None)
+            if z is not None:
+                pending[i] = z
diff --git a/predicators/code_sim_learning/rollout_objective.py b/predicators/code_sim_learning/rollout_objective.py
index 1b4752155e..6a9e119337 100644
--- a/predicators/code_sim_learning/rollout_objective.py
+++ b/predicators/code_sim_learning/rollout_objective.py
@@ -158,7 +158,26 @@ def _huberize(residual: float, delta: float) -> float:
     return float(np.sign(residual) * np.sqrt(2.0 * delta * a - delta * delta))
 
 
-def _iter_rollout_residual_terms(
+@dataclasses.dataclass(frozen=True)
+class _PointScorer:
+    """How the objective scores one parameter point on a trajectory set.
+
+    ``per_step_terms`` scores one trajectory on its own: its rollout,
+    per-step residuals and endpoint and onset summaries. Under interval
+    scoring (``tracks`` non-empty) the trajectories of one episode are
+    scored together from their ``rollout``s instead.
+    """
+    rollout: Callable[[List[State], List[Action]],
+                      Tuple[List[State], List[Dict[Any, Dict[str, Any]]]]]
+    per_step_terms: Callable[[List[State], List[Action]], List[float]]
+    tracks: List[Any]
+    id_maps: List[Any]
+    paired_tracks: bool
+    summary_w: float
+    config: SysIdConfig
+
+
+def _point_scorer(
     base_env: Any,
     trajectories: List[RolloutTrajectory],
     params: Dict[str, float],
@@ -169,31 +188,9 @@ def _iter_rollout_residual_terms(
     scaling: Optional[ResidualScaling] = None,
     config: Optional[SysIdConfig] = None,
     episode_count: Optional[int] = None,
-) -> Iterator[float]:
-    """Yield per-feature residuals for the joint forward model.
-
-    Shared prediction pipeline behind :func:`compute_rollout_sse` and
-    :func:`compute_rollout_residuals`: free-run the base sim at the
-    physical slice of ``params``, apply rules (latents threaded) on each
-    rolled-out state, and score every in-scope feature against its
-    observation. Deterministic iteration order. Without ``scaling`` the
-    residual is the raw ``pred - obs`` difference (legacy objective).
-    With a factory ``base_env`` and fork-parallel workers enabled
-    (``agent_validation_parallel_workers``), per-step scoring fans
-    whole trajectories out to forked children; the terms and their
-    order are identical to the serial path's (see
-    :func:`_prefetch_trajectory_terms`).
-
-    Each per-step residual is Huber-capped (``huber_delta``), and two
-    kinds of per-trajectory SUMMARY residuals are appended with weight
-    ``summary_weight`` (see the flags in ``settings.py``): the settled
-    ENDPOINT residuals (the trajectory's final scored features - the
-    rest poses a plan actually depends on, re-emphasized because they
-    are 1 of N steps in the per-step sum) and per-object motion ONSET
-    residuals (first step any scored feature moves beyond
-    ``settle_tol``, normalized by the horizon - event timing is smooth
-    in the physical params where mid-flight paths are chaos).
-    """
+) -> _PointScorer:
+    """The scorer of one parameter point; see
+    :func:`_iter_rollout_residual_terms` for what it scores."""
     # pylint: disable=import-outside-toplevel
     from predicators.code_sim_learning.commands import CommandBuffer
     from predicators.code_sim_learning.utils import apply_rules, \
@@ -332,13 +329,61 @@ def _per_step_terms(states: List[State],
                                      summary_w))
         return terms
 
-    if not score_intervals:
+    return _PointScorer(rollout=_rollout_with_rules,
+                        per_step_terms=_per_step_terms,
+                        tracks=tracks,
+                        id_maps=id_maps,
+                        paired_tracks=paired_tracks,
+                        summary_w=float(summary_w),
+                        config=config)
+
+
+def _iter_rollout_residual_terms(
+    base_env: Any,
+    trajectories: List[RolloutTrajectory],
+    params: Dict[str, float],
+    residual_features: Dict[str, List[str]],
+    physical_names: Sequence[str],
+    rules: Sequence[Any],
+    latent_init: Any,
+    scaling: Optional[ResidualScaling] = None,
+    config: Optional[SysIdConfig] = None,
+    episode_count: Optional[int] = None,
+) -> Iterator[float]:
+    """Yield per-feature residuals for the joint forward model.
+
+    Shared prediction pipeline behind :func:`compute_rollout_sse` and
+    :func:`compute_rollout_residuals`: free-run the base sim at the
+    physical slice of ``params``, apply rules (latents threaded) on each
+    rolled-out state, and score every in-scope feature against its
+    observation. Deterministic iteration order. Without ``scaling`` the
+    residual is the raw ``pred - obs`` difference (legacy objective).
+    With a factory ``base_env`` and fork-parallel workers enabled
+    (``agent_validation_parallel_workers``), per-step scoring fans
+    whole trajectories out to forked children; the terms and their
+    order are identical to the serial path's (see
+    :func:`_prefetch_trajectory_terms`).
+
+    Each per-step residual is Huber-capped (``huber_delta``), and two
+    kinds of per-trajectory SUMMARY residuals are appended with weight
+    ``summary_weight`` (see the flags in ``settings.py``): the settled
+    ENDPOINT residuals (the trajectory's final scored features - the
+    rest poses a plan actually depends on, re-emphasized because they
+    are 1 of N steps in the per-step sum) and per-object motion ONSET
+    residuals (first step any scored feature moves beyond
+    ``settle_tol``, normalized by the horizon - event timing is smooth
+    in the physical params where mid-flight paths are chaos).
+    """
+    scorer = _point_scorer(base_env, trajectories, params, residual_features,
+                           physical_names, rules, latent_init, scaling, config,
+                           episode_count)
+    if not scorer.tracks:
         prefetched = _prefetch_trajectory_terms(base_env, trajectories,
-                                                _per_step_terms)
+                                                scorer.per_step_terms)
         for idx, (states, actions) in enumerate(trajectories):
             terms = None if prefetched is None else prefetched[idx]
             if terms is None:
-                terms = _per_step_terms(states, actions)
+                terms = scorer.per_step_terms(states, actions)
             yield from terms
         return
 
@@ -349,23 +394,25 @@ def _per_step_terms(states: List[State],
     # outvote the real evidence by thousands of terms to a handful.
     episode_rollouts: List[List[State]] = []
     for traj_index, (states, actions) in enumerate(trajectories):
-        sim_states, _ = _rollout_with_rules(states, actions)
-        if paired_tracks:
+        sim_states, _ = scorer.rollout(states, actions)
+        if scorer.paired_tracks:
             # One track per trajectory: each trajectory IS an episode, so
             # scoring it on its own already is per-episode.
             yield from _interval_residual_terms(sim_states, states,
-                                                tracks[traj_index],
-                                                id_maps[traj_index], config,
-                                                summary_w)
+                                                scorer.tracks[traj_index],
+                                                scorer.id_maps[traj_index],
+                                                scorer.config,
+                                                scorer.summary_w)
         else:
             # Segments of ONE episode. Held, not scored: see the episode
             # -level yield after this loop.
             episode_rollouts.append(sim_states)
-    if not paired_tracks and episode_rollouts:
+    if not scorer.paired_tracks and episode_rollouts:
         yield from _episode_interval_terms(episode_rollouts,
                                            [s for s, _ in trajectories],
-                                           tracks[-1], id_maps[0], config,
-                                           summary_w)
+                                           scorer.tracks[-1],
+                                           scorer.id_maps[0], scorer.config,
+                                           scorer.summary_w)
 
 
 def _prefetch_trajectory_terms(
@@ -409,6 +456,157 @@ def _prefetch_trajectory_terms(
     return results
 
 
+def trajectory_terms_by_point(
+    base_env: Any,
+    trajectories: List[RolloutTrajectory],
+    points: Sequence[Dict[str, float]],
+    residual_features: Dict[str, List[str]],
+    physical_names: Sequence[str],
+    rules: Sequence[Any] = (),
+    latent_init: Any = None,
+    scaling: Optional[ResidualScaling] = None,
+    config: Optional[SysIdConfig] = None,
+) -> List[List[np.ndarray]]:
+    """Each point's residual terms per trajectory, each trajectory scored
+    alone.
+
+    Entry ``[p][i]`` holds the terms of ``trajectories[i]`` at
+    ``points[p]``, as :func:`compute_rollout_residuals` scores it with
+    ``episode_count=len(trajectories)`` (what :func:`per_trajectory_rms`
+    evaluates). Under per-step scoring with a factory env, every (point,
+    trajectory) pair goes to forked children in one wave; scoring point
+    by point forks a wave per point and waits on its slowest trajectory
+    each time. Each pair's terms are the serial path's, and a pair a
+    child failed on is scored here. Interval scoring and a shared env
+    instance score pair by pair in this process.
+    """
+    n = len(trajectories)
+    if not points or not n:
+        return [[] for _ in points]
+    out: List[List[Optional[np.ndarray]]] = [[None] * n for _ in points]
+    first = _point_scorer(base_env, trajectories, dict(points[0]),
+                          residual_features, physical_names, rules,
+                          latent_init, scaling, config, n)
+    if callable(base_env) and not first.tracks and len(points) * n > 1:
+        # Deferred for the same layering reason as in
+        # _prefetch_trajectory_terms.
+        # pylint: disable-next=import-outside-toplevel
+        from predicators.agent_sdk.parallel_rollouts import prefetch_parallel
+        scorers = [first] + [
+            _point_scorer(base_env, trajectories, dict(point),
+                          residual_features, physical_names, rules,
+                          latent_init, scaling, config, n)
+            for point in points[1:]
+        ]
+        jobs: List[Callable[[], List[float]]] = [
+            functools.partial(scorer.per_step_terms, states, actions)
+            for scorer in scorers for states, actions in trajectories
+        ]
+        results = prefetch_parallel(jobs, "sysid batch", quiet=True)
+        for k, terms in enumerate(results):
+            if terms is not None:
+                out[k // n][k % n] = np.asarray(terms, dtype=float)
+        add_rollouts_run(sum(1 for r in results if r is not None))
+    filled: List[List[np.ndarray]] = []
+    for point, row in zip(points, out):
+        filled.append([
+            terms if terms is not None else compute_rollout_residuals(
+                base_env, [traj],
+                dict(point),
+                residual_features,
+                physical_names,
+                rules,
+                latent_init,
+                scaling,
+                config,
+                episode_count=n) for traj, terms in zip(trajectories, row)
+        ])
+    return filled
+
+
+def rollout_residuals_by_point(
+    base_env: Any,
+    trajectories: List[RolloutTrajectory],
+    points: Sequence[Dict[str, float]],
+    residual_features: Dict[str, List[str]],
+    physical_names: Sequence[str],
+    rules: Sequence[Any] = (),
+    latent_init: Any = None,
+    scaling: Optional[ResidualScaling] = None,
+    config: Optional[SysIdConfig] = None,
+) -> List[np.ndarray]:
+    """:func:`compute_rollout_residuals` at each point, the points scored
+    together.
+
+    Under per-step scoring a point's residual vector is its
+    trajectories' terms in order, so batching every point's
+    trajectories into one fork wave (:func:`trajectory_terms_by_point`)
+    changes no value; the LM's finite-difference Jacobian scores its
+    column points this way. Interval scoring and a shared env instance
+    score point by point.
+    """
+    if not points:
+        return []
+    first = _point_scorer(base_env, trajectories, dict(points[0]),
+                          residual_features, physical_names, rules,
+                          latent_init, scaling, config)
+    if not callable(base_env) or first.tracks or len(trajectories) == 0:
+        return [
+            compute_rollout_residuals(base_env, trajectories, dict(point),
+                                      residual_features, physical_names, rules,
+                                      latent_init, scaling, config)
+            for point in points
+        ]
+    per_point = trajectory_terms_by_point(base_env, trajectories, points,
+                                          residual_features, physical_names,
+                                          rules, latent_init, scaling, config)
+    return [np.concatenate(per_trajectory) for per_trajectory in per_point]
+
+
+def rollout_sse_by_point(
+    base_env: Any,
+    trajectories: List[RolloutTrajectory],
+    points: Sequence[Dict[str, float]],
+    residual_features: Dict[str, List[str]],
+    physical_names: Sequence[str],
+    rules: Sequence[Any] = (),
+    latent_init: Any = None,
+    scaling: Optional[ResidualScaling] = None,
+    config: Optional[SysIdConfig] = None,
+) -> List[float]:
+    """:func:`compute_rollout_sse` at each point, the points scored together.
+
+    Under per-step scoring a point's SSE is its trajectories' terms
+    squared and summed in order, the same float sum
+    :func:`compute_rollout_sse` takes, so batching every point's
+    trajectories into one fork wave (:func:`trajectory_terms_by_point`)
+    changes no value. Interval scoring scores the segments of an episode
+    together, and a shared env instance cannot fork, so both score point
+    by point.
+    """
+    if not points:
+        return []
+    first = _point_scorer(base_env, trajectories, dict(points[0]),
+                          residual_features, physical_names, rules,
+                          latent_init, scaling, config)
+    if not callable(base_env) or first.tracks or len(trajectories) == 0:
+        return [
+            compute_rollout_sse(base_env, trajectories, dict(point),
+                                residual_features, physical_names, rules,
+                                latent_init, scaling, config)
+            for point in points
+        ]
+    per_point = trajectory_terms_by_point(base_env, trajectories, points,
+                                          residual_features, physical_names,
+                                          rules, latent_init, scaling, config)
+    return [
+        float(
+            sum(
+                float(r) * float(r) for terms in per_trajectory
+                for r in terms)) for per_trajectory in per_point
+    ]
+
+
 def _load_scored_track(config: SysIdConfig) -> Optional[Any]:
     """The observation track to score against, or None to fall back.
 
@@ -915,18 +1113,32 @@ def fit_map_lm_rollout(
     physical_names = list(fixed) + [s.name for s in physical_specs]
     use_prior = prior_centers is not None and prior_sigmas is not None
 
-    def residuals_fn(theta: np.ndarray) -> np.ndarray:
+    def params_of(theta: np.ndarray) -> Dict[str, float]:
         params = {n: float(theta[i]) for i, n in enumerate(names)}
         params.update(fixed)
-        res = compute_rollout_residuals(base_env, trajectories, params,
-                                        residual_features, physical_names,
-                                        rules, latent_init, scaling)
+        return params
+
+    def with_prior(res: np.ndarray, theta: np.ndarray) -> np.ndarray:
         if not use_prior:
             return res
         z = to_fit_space(all_specs, theta)
         prior_res = noise_sigma * (z - prior_centers) / prior_sigmas
         return np.concatenate([res, prior_res])
 
+    def residuals_fn(theta: np.ndarray) -> np.ndarray:
+        res = compute_rollout_residuals(base_env, trajectories,
+                                        params_of(theta), residual_features,
+                                        physical_names, rules, latent_init,
+                                        scaling)
+        return with_prior(res, theta)
+
+    def batch_residuals_fn(thetas: List[np.ndarray]) -> List[np.ndarray]:
+        data = rollout_residuals_by_point(base_env, trajectories,
+                                          [params_of(t) for t in thetas],
+                                          residual_features, physical_names,
+                                          rules, latent_init, scaling)
+        return [with_prior(res, t) for res, t in zip(data, thetas)]
+
     theta_map, jac = solve_lm(residuals_fn,
                               all_specs,
                               max_nfev,
@@ -934,7 +1146,8 @@ def residuals_fn(theta: np.ndarray) -> np.ndarray:
                               diff_step=_ROLLOUT_LM_DIFF_STEP,
                               notes_out=notes_out,
                               n_prior_rows=len(all_specs) if use_prior else 0,
-                              flat_params_out=flat_params_out)
+                              flat_params_out=flat_params_out,
+                              batch_residuals_fn=batch_residuals_fn)
     if use_prior and jac is not None and jac.shape[0] > len(all_specs):
         jac = jac[:-len(all_specs)]
     return theta_map, jac
diff --git a/tests/code_sim_learning/test_parameter_belief.py b/tests/code_sim_learning/test_parameter_belief.py
index 8b90210499..7f34d071c5 100644
--- a/tests/code_sim_learning/test_parameter_belief.py
+++ b/tests/code_sim_learning/test_parameter_belief.py
@@ -78,6 +78,50 @@ def test_gaussian_lines_recover_conditional_variances():
                                                 rel=0.06)
 
 
+def test_lines_traced_in_lockstep_are_the_same():
+    """Tracing every line in lockstep through a batch scorer evaluates the same
+    points and gives the same belief as tracing line by line, with one batch
+    per round."""
+    specs = [
+        ParamSpec("a", 0.3, lo=-10, hi=10),
+        ParamSpec("b", 0.7, lo=0.01, hi=10, scale="log"),
+        ParamSpec("c", 0.5, lo=0.0, hi=1.0),
+    ]
+    residuals = _gaussian_residuals(["a", "b", "c"], [0.3, 0.7, 0.5],
+                                    np.diag([20.0, 15.0, 0.001]),
+                                    offset=[0.01])
+    map_params = {"a": 0.3, "b": 0.7, "c": 0.5}
+    batches = []
+
+    def batch(points):
+        batches.append(len(points))
+        return [residuals(params) for params in points]
+
+    alone = _build(specs, map_params, residuals, prior_sigma=2.0)
+    together = build_parameter_belief(
+        specs,
+        map_params,
+        residuals,
+        noise_sigma=_SIGMA_N,
+        prior_centers={s.name: s.init_value
+                       for s in specs},
+        prior_sigmas={s.name: 2.0
+                      for s in specs},
+        config=BeliefConfig(num_draws=64, line_max_evals=40),
+        seed=0,
+        batch_residuals=batch)
+    for name in ("a", "b", "c"):
+        assert np.array_equal(alone.lines[name].grid,
+                              together.lines[name].grid)
+        assert np.array_equal(alone.lines[name].density,
+                              together.lines[name].density)
+    assert np.array_equal(alone.draws, together.draws)
+    assert alone.evaluations == together.evaluations
+    # Rounds score up to one point per line, and fewer rounds than points.
+    assert max(batches) <= len(specs)
+    assert len(batches) < alone.evaluations
+
+
 def test_flat_line_keeps_the_prior():
     """A parameter the data do not touch gets its (bounded) prior."""
     specs = [
diff --git a/tests/code_sim_learning/test_physical_sysid.py b/tests/code_sim_learning/test_physical_sysid.py
index 944d37c9b3..52eef19da5 100644
--- a/tests/code_sim_learning/test_physical_sysid.py
+++ b/tests/code_sim_learning/test_physical_sysid.py
@@ -480,6 +480,20 @@ def test_select_trustworthy_params_weakly_identified_applies():
     assert applied == {"friction": 0.12}
 
 
+def _batched(fake_rms):
+    """A trajectory_terms_by_point stand-in built from a per_trajectory_rms
+    stub: each candidate's terms are one residual per trajectory, of the
+    stubbed RMS."""
+
+    def fake_terms(env, trajectories, points, *args, **kwargs):
+        return [[
+            np.array([r])
+            for r in fake_rms(env, trajectories, point, *args, **kwargs)
+        ] for point in points]
+
+    return fake_terms
+
+
 def _patch_fit_and_rms(monkeypatch, fit_thetas, rms_by_count):
     """Stub the heavy PyBullet fit/residual calls for trimming tests.
 
@@ -502,9 +516,10 @@ def fake_rms(_env, trajectories, *_args, **_kwargs):
 
     monkeypatch.setattr(physical_sysid, "fit_params_rollout", fake_fit)
     monkeypatch.setattr(physical_sysid, "per_trajectory_rms", fake_rms)
-    # min_explainable_rms (called by the trimmed fit) evaluates its
-    # candidate grid through grid_seed's namespace.
-    monkeypatch.setattr(grid_seed, "per_trajectory_rms", fake_rms)
+    # min_explainable_rms (called by the trimmed fit) scores its
+    # candidate grid in one batch through grid_seed's namespace.
+    monkeypatch.setattr(grid_seed, "trajectory_terms_by_point",
+                        _batched(fake_rms))
     # The stub trajectories are plain strings; skip the data-derived
     # residual scaling (exercised by its own tests).
     monkeypatch.setattr(physical_sysid, "compute_residual_scaling",
@@ -580,7 +595,8 @@ def fake_rms(_env, trajectories, *_args, **_kwargs):
 
     monkeypatch.setattr(physical_sysid, "fit_params_rollout", fake_fit)
     monkeypatch.setattr(physical_sysid, "per_trajectory_rms", fake_rms)
-    monkeypatch.setattr(grid_seed, "per_trajectory_rms", fake_rms)
+    monkeypatch.setattr(grid_seed, "trajectory_terms_by_point",
+                        _batched(fake_rms))
     monkeypatch.setattr(physical_sysid, "compute_residual_scaling",
                         lambda *_a, **_k: None)
     result, survivors, _rms, hull = fit_params_rollout_trimmed(
@@ -1169,7 +1185,12 @@ def _run_sweep(monkeypatch, sse_fn, specs, anchors, **flags):
     def fake_sse(_env, _trajs, params, *_args, **_kwargs):
         return sse_fn(params)
 
+    def fake_batch(_env, _trajs, points, *_args, **_kwargs):
+        return [sse_fn(params) for params in points]
+
     monkeypatch.setattr(grid_seed, "compute_rollout_sse", fake_sse)
+    # A pool's new candidates are scored together.
+    monkeypatch.setattr(grid_seed, "rollout_sse_by_point", fake_batch)
     seeded, info = grid_seed._grid_seed_physical_specs(None, ["traj"],
                                                        specs,
                                                        {"domino": ["x"]}, [],
@@ -1657,6 +1678,95 @@ def factory():
     assert all(not env.connected for env in built)
 
 
+def test_batched_points_match_point_by_point(monkeypatch):
+    """Scoring many points in one fork wave returns every (point, trajectory)
+    pair's serial terms, and each point's SSE equals compute_rollout_sse
+    exactly, with workers or without.
+
+    The grid sweep and the explainability sweep score their candidates
+    this way; any value difference would change fits.
+    """
+    from predicators.agent_sdk.parallel_rollouts import \
+        parallel_rollouts_available
+    from predicators.settings import CFG
+    if not parallel_rollouts_available():
+        pytest.skip("fork not available on this platform")
+    trajectories = [
+        _trajectory([0.0, 0.1, 0.2]),
+        _trajectory([0.0, 0.2, 0.4]),
+        _trajectory([0.0, 0.05, 0.1]),
+    ]
+    built = []
+
+    def factory():
+        return _LinearParamEnv(built)
+
+    points = [{"k": k} for k in (0.0, 0.05, 0.1, 0.2)]
+    monkeypatch.setattr(CFG, "agent_validation_parallel_workers", 0)
+    want_terms = [[
+        rollout_objective.compute_rollout_residuals(
+            factory, [traj],
+            point,
+            _RESIDUAL_FEATURES, ["k"],
+            episode_count=len(trajectories)) for traj in trajectories
+    ] for point in points]
+    want_sse = [
+        rollout_objective.compute_rollout_sse(factory, trajectories, point,
+                                              _RESIDUAL_FEATURES, ["k"])
+        for point in points
+    ]
+    for workers in (0, 3):
+        monkeypatch.setattr(CFG, "agent_validation_parallel_workers", workers)
+        n0 = rollout_env.num_rollouts_run()
+        terms = rollout_objective.trajectory_terms_by_point(
+            factory, trajectories, points, _RESIDUAL_FEATURES, ["k"])
+        # Every pair is credited once, in a child or here.
+        assert rollout_env.num_rollouts_run() - n0 == \
+            len(points) * len(trajectories)
+        for got_row, want_row in zip(terms, want_terms):
+            assert len(got_row) == len(want_row)
+            for got, want in zip(got_row, want_row):
+                assert np.array_equal(got, want)
+        assert rollout_objective.rollout_sse_by_point(factory, trajectories,
+                                                      points,
+                                                      _RESIDUAL_FEATURES,
+                                                      ["k"]) == want_sse
+    assert all(not env.connected for env in built)
+
+
+def test_rollout_lm_fit_is_the_same_with_and_without_workers(monkeypatch):
+    """The rollout LM fit returns the same MAP and Jacobian whether its
+    Jacobian columns are scored together in forked children or one by one in
+    this process."""
+    from predicators.agent_sdk.parallel_rollouts import \
+        parallel_rollouts_available
+    from predicators.settings import CFG
+    if not parallel_rollouts_available():
+        pytest.skip("fork not available on this platform")
+    trajectories = [
+        _trajectory([0.0, 0.1, 0.2]),
+        _trajectory([0.0, 0.2, 0.4]),
+        _trajectory([0.0, 0.05, 0.1]),
+    ]
+    built = []
+
+    def factory():
+        return _LinearParamEnv(built)
+
+    specs = [ParamSpec("k", 0.3, lo=0.0, hi=1.0)]
+    fits = []
+    for workers in (0, 3):
+        monkeypatch.setattr(CFG, "agent_validation_parallel_workers", workers)
+        fits.append(
+            rollout_objective.fit_map_lm_rollout(factory, trajectories, specs,
+                                                 _RESIDUAL_FEATURES))
+    theta, jac = fits[0]
+    fast_theta, fast_jac = fits[1]
+    assert np.array_equal(theta, fast_theta)
+    assert jac is not None and np.array_equal(jac, fast_jac)
+    assert all(not env.connected for env in built)
+
+
 def test_rollout_residuals_shared_env_instance_stays_serial(monkeypatch):
     """A caller-owned env instance never takes the parallel path: its rollouts
     mutate the caller's env, which must happen in-process."""
diff --git a/tests/code_sim_learning/test_zero_gradient_search.py b/tests/code_sim_learning/test_zero_gradient_search.py
index e95a888ce6..ac575d49aa 100644
--- a/tests/code_sim_learning/test_zero_gradient_search.py
+++ b/tests/code_sim_learning/test_zero_gradient_search.py
@@ -186,3 +186,36 @@ def residuals(z: np.ndarray) -> np.ndarray:
     assert not flat  # edge screen alone never suppresses the probe
     assert len(notes) == 1 and "NOT fit from data" in notes[0]
     assert "flat at both box edges" in notes[0]
+
+
+def test_solve_lm_batched_jacobian_matches_least_squares() -> None:
+    """With a batch scorer, solve_lm returns exactly the theta and Jacobian of
+    least_squares' own 2-point Jacobian, and scores each Jacobian's column
+    points through the batch."""
+    specs = [
+        ParamSpec("a", 0.5, lo=0.0, hi=2.0),
+        ParamSpec("b", 0.3, lo=0.01, hi=3.0, scale="log"),
+    ]
+    xs = np.linspace(0.0, 1.0, 7)
+
+    def residuals(theta: np.ndarray) -> np.ndarray:
+        a, b = theta
+        return np.sin(a * xs) + b * xs**2 - (np.sin(1.1 * xs) + 0.7 * xs**2)
+
+    batches: List[int] = []
+
+    def batch(thetas: List[np.ndarray]) -> List[np.ndarray]:
+        batches.append(len(thetas))
+        return [residuals(theta) for theta in thetas]
+
+    theta, jac = solve_lm(residuals, specs, 200, "test", diff_step=2e-2)
+    fast_theta, fast_jac = solve_lm(residuals,
+                                    specs,
+                                    200,
+                                    "test",
+                                    diff_step=2e-2,
+                                    batch_residuals_fn=batch)
+    assert np.array_equal(theta, fast_theta)
+    assert jac is not None and fast_jac is not None
+    assert np.array_equal(jac, fast_jac)
+    assert batches and all(0 < n <= len(specs) for n in batches)

From ad68c3256ad105e8465857715a574827d40f3779 Mon Sep 17 00:00:00 2001
From: Yichao Liang <ycliang6@gmail.com>
Date: Fri, 9 Oct 2026 03:41:50 -0400
Subject: [PATCH 2/3] Replay the fit's forked rollouts on copies of one
 pre-built world

Building the agent's world took 0.39 s of a 0.51 s segment replay on the
from-assets Domino scene. A forked child's copy of a world the parent
built and never stepped is a fresh world, so a fit (and any other fork
wave of the rollout objective) builds one template world and each child
replays on its copy instead of building its own. A rollout the parent
runs itself still builds a fresh world, and nothing is built when waves
do not fork.

The new test replays a row of dominoes in real PyBullet contact physics
both ways: the terms are bit-identical, and no child built a world.
---
 predicators/agent_sdk/parallel_rollouts.py    |  22 ++-
 predicators/code_sim_learning/orchestrator.py |  11 +-
 predicators/code_sim_learning/rollout_env.py  |  61 +++++++-
 .../code_sim_learning/rollout_objective.py    |  18 ++-
 .../code_sim_learning/test_physical_sysid.py  | 148 ++++++++++++++++++
 5 files changed, 240 insertions(+), 20 deletions(-)

diff --git a/predicators/agent_sdk/parallel_rollouts.py b/predicators/agent_sdk/parallel_rollouts.py
index 4304854c1f..873d7c8e4c 100644
--- a/predicators/agent_sdk/parallel_rollouts.py
+++ b/predicators/agent_sdk/parallel_rollouts.py
@@ -46,6 +46,21 @@ def parallel_rollouts_available() -> bool:
     return sys.platform.startswith("linux") and hasattr(os, "fork")
 
 
+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) -> List[Optional[Any]]:
@@ -63,11 +78,8 @@ def prefetch_parallel(jobs: Sequence[Callable[[], Any]],
     objective calls this once per candidate theta, hundreds of times
     per fit, and one log line each would drown the run's info.log.
     """
-    # 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.",
diff --git a/predicators/code_sim_learning/orchestrator.py b/predicators/code_sim_learning/orchestrator.py
index 78fdcdddbc..f7245bbac2 100644
--- a/predicators/code_sim_learning/orchestrator.py
+++ b/predicators/code_sim_learning/orchestrator.py
@@ -43,7 +43,7 @@
     _ROLLOUT_PRIOR_SIGMA_SCALE, _explainability_cache_key, \
     fit_params_rollout_trimmed
 from predicators.code_sim_learning.rollout_env import RolloutTrajectory, \
-    num_rollouts_run
+    fork_template, num_rollouts_run
 from predicators.code_sim_learning.rollout_objective import \
     compute_rollout_residuals, compute_rollout_sse, \
     rollout_residuals_by_point
@@ -203,9 +203,12 @@ def run_rollout_sysid(
 
     from_cache = core is not None
     if core is None:
-        core = _compute_fit(fit_env, rollouts, physical_specs,
-                            residual_features, rules, rule_specs, latent_init,
-                            anchors, rms_cache, config, belief_config)
+        # One template world serves every fork wave of the fit.
+        with fork_template(fit_env):
+            core = _compute_fit(fit_env, rollouts, physical_specs,
+                                residual_features, rules, rule_specs,
+                                latent_init, anchors, rms_cache, config,
+                                belief_config)
         if fit_cache is not None and cache_key is not None:
             fit_cache[cache_key] = core
 
diff --git a/predicators/code_sim_learning/rollout_env.py b/predicators/code_sim_learning/rollout_env.py
index bdcad5d3b2..031a29e24c 100644
--- a/predicators/code_sim_learning/rollout_env.py
+++ b/predicators/code_sim_learning/rollout_env.py
@@ -9,8 +9,10 @@
 
 from __future__ import annotations
 
-from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, \
-    Tuple
+import contextlib
+import os
+from typing import Any, Callable, Dict, Iterator, List, Mapping, Optional, \
+    Sequence, Tuple
 
 import pybullet as p
 
@@ -26,6 +28,53 @@
 # around their stages to log where the budget actually went.
 _NUM_ROLLOUTS = 0
 
+# Worlds a process built for the forked children of one parallel wave and
+# never steps itself: id(factory) -> (the building process's pid, the
+# world). A forked child replays on its private copy of the world instead
+# of building one (fork_template).
+_FORK_TEMPLATES: Dict[int, Tuple[int, Any]] = {}
+
+
+@contextlib.contextmanager
+def fork_template(factory: Any) -> Iterator[None]:
+    """Build one world from ``factory`` for the forked children of a wave.
+
+    Building a world dominates a short rollout (0.39 s of a 22-step
+    segment's 0.51 s on the from-assets Domino scene), and a fresh world
+    per rollout is what keeps rollouts deterministic. A forked child's
+    copy of a world this process built and never stepped is such a fresh
+    world, so :func:`rollout_states` in a forked child replays on that
+    copy (once per child) instead of building its own. This process
+    never touches the template, so a rollout it runs itself still builds
+    a fresh world. Nothing is built when waves do not fork, when
+    ``factory`` is an env instance, or inside a template of the same
+    factory.
+    """
+    # Deferred: agent_sdk imports this package.
+    # pylint: disable-next=import-outside-toplevel
+    from predicators.agent_sdk.parallel_rollouts import parallel_workers
+    if not callable(factory) or id(factory) in _FORK_TEMPLATES or \
+            parallel_workers() <= 1:
+        yield
+        return
+    world = factory()
+    _FORK_TEMPLATES[id(factory)] = (os.getpid(), world)
+    try:
+        yield
+    finally:
+        _FORK_TEMPLATES.pop(id(factory), None)
+        dispose_env(world)
+
+
+def _forked_template_copy(factory: Any) -> Optional[Any]:
+    """This forked child's copy of ``factory``'s template world, at most once;
+    None in the process that built it."""
+    entry = _FORK_TEMPLATES.get(id(factory))
+    if entry is None or entry[0] == os.getpid():
+        return None
+    del _FORK_TEMPLATES[id(factory)]
+    return entry[1]
+
 
 def num_rollouts_run() -> int:
     """Total :func:`rollout_states` invocations so far (cost telemetry)."""
@@ -159,7 +208,9 @@ def rollout_states(
 
     ``base_env`` is either an env instance or a zero-arg FACTORY: a
     factory is invoked to build a fresh env for this single rollout and
-    the fresh env's PyBullet client is disconnected before returning.
+    the fresh env's PyBullet client is disconnected before returning. A
+    forked child of a :func:`fork_template` wave uses its copy of the
+    template world instead, which is the same fresh world.
     Fresh per-rollout worlds are what make repeated evaluations of the
     same theta deterministic - state-level resets on a shared env leave
     history-dependent residuals (near-matching bodies skipped by the
@@ -179,7 +230,9 @@ def rollout_states(
     """
     global _NUM_ROLLOUTS  # pylint: disable=global-statement
     _NUM_ROLLOUTS += 1
-    env = base_env() if callable(base_env) else base_env
+    env = base_env
+    if callable(base_env):
+        env = _forked_template_copy(base_env) or base_env()
     try:
         _pin_all_physical_params(env, physical_params)
         env._set_state(init_state)  # pylint: disable=protected-access
diff --git a/predicators/code_sim_learning/rollout_objective.py b/predicators/code_sim_learning/rollout_objective.py
index 6a9e119337..d225560329 100644
--- a/predicators/code_sim_learning/rollout_objective.py
+++ b/predicators/code_sim_learning/rollout_objective.py
@@ -21,7 +21,7 @@
 from predicators.code_sim_learning.fit_space import ParamSpec, to_fit_space
 from predicators.code_sim_learning.lm import solve_lm
 from predicators.code_sim_learning.rollout_env import RolloutTrajectory, \
-    add_rollouts_run, rollout_states
+    add_rollouts_run, fork_template, rollout_states
 from predicators.code_sim_learning.trajectory_prep import ResidualScaling
 from predicators.settings import CFG
 from predicators.structs import Action, State
@@ -422,10 +422,12 @@ def _prefetch_trajectory_terms(
 ) -> Optional[List[Optional[List[float]]]]:
     """Fan per-trajectory scoring out to forked children when enabled.
 
-    Factory envs only: each rollout builds (and disposes) its own fresh
-    world whichever process runs it, so a child's term list is
-    bit-identical to what the serial path would compute - which is what
-    the LM finite-difference Jacobian requires of repeated same-theta
+    Factory envs only: each rollout runs in a fresh world whichever
+    process runs it (a child's copy of the wave's
+    :func:`~predicators.code_sim_learning.rollout_env.fork_template`
+    world, or one it builds), so a child's term list is bit-identical
+    to what the serial path would compute - which is what the LM
+    finite-difference Jacobian requires of repeated same-theta
     evaluations. A shared env instance keeps the serial path untouched
     (its rollouts mutate the caller's env, which must happen in this
     process).
@@ -448,7 +450,8 @@ def _prefetch_trajectory_terms(
         functools.partial(score_fn, states, actions)
         for states, actions in trajectories
     ]
-    results = prefetch_parallel(jobs, "sysid objective", quiet=True)
+    with fork_template(base_env):
+        results = prefetch_parallel(jobs, "sysid objective", quiet=True)
     done = sum(1 for r in results if r is not None)
     if done == 0:
         return None
@@ -502,7 +505,8 @@ def trajectory_terms_by_point(
             functools.partial(scorer.per_step_terms, states, actions)
             for scorer in scorers for states, actions in trajectories
         ]
-        results = prefetch_parallel(jobs, "sysid batch", quiet=True)
+        with fork_template(base_env):
+            results = prefetch_parallel(jobs, "sysid batch", quiet=True)
         for k, terms in enumerate(results):
             if terms is not None:
                 out[k // n][k % n] = np.asarray(terms, dtype=float)
diff --git a/tests/code_sim_learning/test_physical_sysid.py b/tests/code_sim_learning/test_physical_sysid.py
index 52eef19da5..0ab172511a 100644
--- a/tests/code_sim_learning/test_physical_sysid.py
+++ b/tests/code_sim_learning/test_physical_sysid.py
@@ -7,6 +7,8 @@
 """
 # pylint: disable=protected-access,import-outside-toplevel
 
+import os
+
 import numpy as np
 import pybullet as p
 import pytest
@@ -1767,6 +1769,152 @@ def factory():
     assert all(not env.connected for env in built)
 
 
+class _DominoRowEnv:
+    """A real DIRECT PyBullet world with contact: a row of six dominoes on a
+    plane, the first pushed near its top on every physics step.
+
+    ``k`` is every domino's lateral friction; ``r0`` reports the first
+    domino's x and ``d0`` the last one's, so the residual sees the whole
+    cascade. Only ``builder_pid`` may build one: a forked child of a
+    template wave must replay on its copy of the template world.
+    """
+
+    def __init__(self, log, builder_pid):
+        if os.getpid() != builder_pid:
+            raise RuntimeError("a forked child built a world")
+        cid = p.connect(p.DIRECT)
+        self._physics_client_id = cid
+        log.append(self)
+        p.setGravity(0.0, 0.0, -9.81, physicsClientId=cid)
+        p.createMultiBody(0,
+                          p.createCollisionShape(p.GEOM_PLANE,
+                                                 physicsClientId=cid),
+                          physicsClientId=cid)
+        shape = p.createCollisionShape(p.GEOM_BOX,
+                                       halfExtents=[0.005, 0.03, 0.05],
+                                       physicsClientId=cid)
+        self._bodies = [
+            p.createMultiBody(0.05,
+                              shape,
+                              basePosition=[0.06 * i, 0.0, 0.05],
+                              physicsClientId=cid) for i in range(6)
+        ]
+        self._domino = Object("d0", _DOMINO_TYPE)
+        self._robot = Object("r0", _ROBOT_TYPE)
+
+    def apply_physical_param_overrides(self, params):
+        """Set every domino's lateral friction to ``k``."""
+        for body in self._bodies:
+            p.changeDynamics(body,
+                             -1,
+                             lateralFriction=float(params.get("k", 0.5)),
+                             physicsClientId=self._physics_client_id)
+
+    def _set_state(self, state):
+        """Space the row evenly from ``r0``'s x to ``d0``'s x."""
+        first = float(state.get(self._robot, "x"))
+        last = float(state.get(self._domino, "x"))
+        for i, body in enumerate(self._bodies):
+            x = first + i * (last - first) / (len(self._bodies) - 1)
+            p.resetBasePositionAndOrientation(
+                body, [x, 0.0, 0.05], [0.0, 0.0, 0.0, 1.0],
+                physicsClientId=self._physics_client_id)
+
+    def step(self, action):
+        """Push the first domino for 24 physics steps; report the row's
+        ends."""
+        del action
+        cid = self._physics_client_id
+        for _ in range(24):
+            p.applyExternalForce(self._bodies[0],
+                                 -1, [0.05, 0.0, 0.0], [0.0, 0.0, 0.04],
+                                 p.LINK_FRAME,
+                                 physicsClientId=cid)
+            p.stepSimulation(physicsClientId=cid)
+        ends = [
+            p.getBasePositionAndOrientation(body, physicsClientId=cid)[0][0]
+            for body in (self._bodies[0], self._bodies[-1])
+        ]
+        return State({
+            self._robot: np.array([ends[0]], dtype=float),
+            self._domino: np.array([ends[1]], dtype=float),
+        })
+
+    @property
+    def connected(self):
+        """Whether this world's PyBullet client is still connected."""
+        return bool(
+            p.getConnectionInfo(self._physics_client_id)["isConnected"])
+
+
+def test_forked_children_replay_on_copies_of_one_template_world(monkeypatch):
+    """A fork wave builds one template world, and every forked child replays on
+    its copy of it instead of building a world; the terms equal fresh-world
+    rollouts bit for bit on contact-rich physics.
+
+    Building dominates a short rollout of a real scene, and a fresh
+    world per rollout is what keeps the fit deterministic, so a copy
+    that differs from a fresh world in any bit would change fits. A
+    template around several waves (as around a whole fit) serves them
+    all, and a rollout in the building process still gets a fresh world.
+    """
+    from predicators.agent_sdk.parallel_rollouts import \
+        parallel_rollouts_available
+    from predicators.settings import CFG
+    if not parallel_rollouts_available():
+        pytest.skip("fork not available on this platform")
+    features = {"domino": ["x"], "robot": ["x"]}
+    trajectories = [
+        _trajectory([0.30] * 8, robot_xs=[0.0] * 8),
+        _trajectory([0.325] * 8, robot_xs=[0.0] * 8),
+        _trajectory([0.275] * 8, robot_xs=[0.0] * 8),
+    ]
+    builder = os.getpid()
+    built = []
+
+    def factory():
+        return _DominoRowEnv(built, builder)
+
+    def terms_by_point():
+        return rollout_objective.trajectory_terms_by_point(
+            factory, trajectories, points, features, ["k"])
+
+    def assert_same(got, want):
+        for got_row, want_row in zip(got, want):
+            assert len(got_row) == len(want_row)
+            for got_terms, want_terms in zip(got_row, want_row):
+                assert np.array_equal(got_terms, want_terms)
+
+    points = [{"k": k} for k in (0.2, 0.5, 0.9)]
+    pairs = len(points) * len(trajectories)
+    monkeypatch.setattr(CFG, "agent_validation_parallel_workers", 0)
+    fresh = terms_by_point()
+    assert len(built) == pairs  # no template without forks
+    # The cascade reaches the far end and depends on the friction.
+    assert not np.array_equal(fresh[0][0], fresh[-1][0])
+
+    monkeypatch.setattr(CFG, "agent_validation_parallel_workers", 3)
+    assert_same(terms_by_point(), fresh)
+    # Only the wave's template was built here: no child built a world
+    # (one that tried would raise and be rescored here, building more).
+    assert len(built) == pairs + 1
+
+    with rollout_env.fork_template(factory):
+        assert len(built) == pairs + 2
+        for _ in range(2):
+            assert_same(terms_by_point(), fresh)
+        assert len(built) == pairs + 2  # both waves used the outer one
+        serial = rollout_objective.compute_rollout_residuals(
+            factory,
+            trajectories[:1],
+            points[0],
+            features, ["k"],
+            episode_count=len(trajectories))
+        assert len(built) == pairs + 3  # this process built a fresh world
+        assert np.array_equal(serial, fresh[0][0])
+    assert all(not env.connected for env in built)
+
+
 def test_rollout_residuals_shared_env_instance_stays_serial(monkeypatch):
     """A caller-owned env instance never takes the parallel path: its rollouts
     mutate the caller's env, which must happen in-process."""

From 7337fe87b032e20eced6e8044e8921688b5e5bce Mon Sep 17 00:00:00 2001
From: Yichao Liang <ycliang6@gmail.com>
Date: Fri, 9 Oct 2026 04:15:55 -0400
Subject: [PATCH 3/3] Start a wave's longest rollouts first and take each
 result at once

A single-point wave of the fit scores 7 segments of 22 to 99 steps on 6
workers. The dispatcher started jobs in segment order, so the 99-step
replay waited for a worker behind short ones; it polled for results
every 50 ms; and it joined each finished child before starting the next
job. Jobs can now carry costs (the rollout objective passes segment
lengths) and start costliest first. The parent waits on the result pipe,
takes each result the moment it arrives, starts the next job at once
and reaps exited children without blocking. Results stay at their job's
index, so only the wall time changes.
---
 predicators/agent_sdk/parallel_rollouts.py    | 90 ++++++++++++++-----
 .../code_sim_learning/rollout_objective.py    | 14 ++-
 tests/agent_sdk/test_parallel_rollouts.py     | 24 +++++
 3 files changed, 103 insertions(+), 25 deletions(-)

diff --git a/predicators/agent_sdk/parallel_rollouts.py b/predicators/agent_sdk/parallel_rollouts.py
index 873d7c8e4c..42b5548e69 100644
--- a/predicators/agent_sdk/parallel_rollouts.py
+++ b/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.
@@ -61,9 +62,11 @@ def parallel_workers() -> int:
     return int(CFG.agent_validation_parallel_workers)
 
 
-def prefetch_parallel(jobs: Sequence[Callable[[], Any]],
-                      label: str,
-                      quiet: bool = False) -> List[Optional[Any]]:
+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
@@ -77,6 +80,8 @@ 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`).
     """
     workers = min(parallel_workers(), len(jobs))
     if workers <= 1:
@@ -84,7 +89,34 @@ def prefetch_parallel(jobs: Sequence[Callable[[], Any]],
     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,
@@ -117,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.
 
@@ -126,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,
@@ -170,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
diff --git a/predicators/code_sim_learning/rollout_objective.py b/predicators/code_sim_learning/rollout_objective.py
index d225560329..34a5fbb191 100644
--- a/predicators/code_sim_learning/rollout_objective.py
+++ b/predicators/code_sim_learning/rollout_objective.py
@@ -450,8 +450,13 @@ def _prefetch_trajectory_terms(
         functools.partial(score_fn, states, actions)
         for states, actions in trajectories
     ]
+    # The longest segments start first (run_forked_rollouts' costs).
+    lengths = [len(actions) for _states, actions in trajectories]
     with fork_template(base_env):
-        results = prefetch_parallel(jobs, "sysid objective", quiet=True)
+        results = prefetch_parallel(jobs,
+                                    "sysid objective",
+                                    quiet=True,
+                                    costs=lengths)
     done = sum(1 for r in results if r is not None)
     if done == 0:
         return None
@@ -505,8 +510,13 @@ def trajectory_terms_by_point(
             functools.partial(scorer.per_step_terms, states, actions)
             for scorer in scorers for states, actions in trajectories
         ]
+        # The longest segments start first (run_forked_rollouts' costs).
+        lengths = [len(actions) for _states, actions in trajectories]
         with fork_template(base_env):
-            results = prefetch_parallel(jobs, "sysid batch", quiet=True)
+            results = prefetch_parallel(jobs,
+                                        "sysid batch",
+                                        quiet=True,
+                                        costs=lengths * len(scorers))
         for k, terms in enumerate(results):
             if terms is not None:
                 out[k // n][k % n] = np.asarray(terms, dtype=float)
diff --git a/tests/agent_sdk/test_parallel_rollouts.py b/tests/agent_sdk/test_parallel_rollouts.py
index 8013fbea85..5624a81a00 100644
--- a/tests/agent_sdk/test_parallel_rollouts.py
+++ b/tests/agent_sdk/test_parallel_rollouts.py
@@ -1,6 +1,7 @@
 """Tests for the fork-based parallel rollout helper."""
 
 import time
+from typing import Tuple
 
 import pytest
 
@@ -28,6 +29,29 @@ def job() -> int:
     assert out == [0, 10, 20, 30]
 
 
+def test_costliest_jobs_start_first() -> None:
+    """With costs, the costliest jobs start first (ties in job order), and
+    results stay at their job's index."""
+    if not parallel_rollouts_available():
+        pytest.skip("fork not available on this platform")
+
+    def make_job(i: int):
+
+        def job() -> Tuple[int, float]:
+            return i, time.monotonic()
+
+        return job
+
+    out = run_forked_rollouts([make_job(i) for i in range(4)],
+                              max_workers=1,
+                              label="test",
+                              costs=[1.0, 3.0, 2.0, 3.0])
+    assert [result[0] for result in out if result is not None] == [0, 1, 2, 3]
+    starts = [result[1] for result in out if result is not None]
+    # One worker runs the jobs one after another, in start order.
+    assert sorted(range(4), key=lambda k: starts[k]) == [1, 3, 2, 0]
+
+
 def test_child_exception_yields_none_slot() -> None:
     """A raising job costs only its own slot."""
     if not parallel_rollouts_available():