diff --git a/predicators/agent_sdk/parallel_rollouts.py b/predicators/agent_sdk/parallel_rollouts.py index 4304854c1..42b5548e6 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. @@ -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 diff --git a/predicators/code_sim_learning/grid_seed.py b/predicators/code_sim_learning/grid_seed.py index 45437621f..d19debe6b 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 10a9ba614..2eb7c7abd 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 d3400789e..f7245bbac 100644 --- a/predicators/code_sim_learning/orchestrator.py +++ b/predicators/code_sim_learning/orchestrator.py @@ -43,9 +43,10 @@ _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 + compute_rollout_residuals, compute_rollout_sse, \ + rollout_residuals_by_point from predicators.code_sim_learning.trajectory_prep import \ compute_residual_scaling @@ -202,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 @@ -401,9 +405,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 +463,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 +502,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 67db2745b..f9e76bb0e 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_env.py b/predicators/code_sim_learning/rollout_env.py index bdcad5d3b..031a29e24 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 1b4752155..34a5fbb19 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 @@ -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( @@ -375,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). @@ -401,7 +450,13 @@ def _prefetch_trajectory_terms( functools.partial(score_fn, states, actions) for states, actions in trajectories ] - results = prefetch_parallel(jobs, "sysid objective", quiet=True) + # 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, + costs=lengths) done = sum(1 for r in results if r is not None) if done == 0: return None @@ -409,6 +464,163 @@ 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 + ] + # 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, + 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) + 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 +1127,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 +1160,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/agent_sdk/test_parallel_rollouts.py b/tests/agent_sdk/test_parallel_rollouts.py index 8013fbea8..5624a81a0 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(): diff --git a/tests/code_sim_learning/test_parameter_belief.py b/tests/code_sim_learning/test_parameter_belief.py index 8b9021049..7f34d071c 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 944d37c9b..0ab172511 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 @@ -480,6 +482,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 +518,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 +597,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 +1187,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 +1680,241 @@ 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) + + +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.""" diff --git a/tests/code_sim_learning/test_zero_gradient_search.py b/tests/code_sim_learning/test_zero_gradient_search.py index e95a888ce..ac575d49a 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)