diff --git a/src/cbfkit/benchmarks/sweep.py b/src/cbfkit/benchmarks/sweep.py index 3fbe5513..bf85d20a 100644 --- a/src/cbfkit/benchmarks/sweep.py +++ b/src/cbfkit/benchmarks/sweep.py @@ -210,8 +210,7 @@ def _build_combo_summary( def _format_combo_desc(combo: dict[str, Any], max_len: int = 40) -> str: """Format a parameter combo as a short description string.""" desc = ", ".join( - f"{k}={v:.3g}" if isinstance(v, float) else f"{k}={v}" - for k, v in combo.items() + f"{k}={v:.3g}" if isinstance(v, float) else f"{k}={v}" for k, v in combo.items() ) if len(desc) > max_len: desc = desc[: max_len - 3] + "..." @@ -242,9 +241,15 @@ def _process_combo( progress.update(seed_task_id, description=f" Seeds ({combo_desc})") combo_records, was_falsified = _run_combo( - runner, seeds, combo, combo_idx, records, - falsifier=falsifier, falsifier_metric=falsifier_metric, - progress=progress, seed_task_id=seed_task_id, + runner, + seeds, + combo, + combo_idx, + records, + falsifier=falsifier, + falsifier_metric=falsifier_metric, + progress=progress, + seed_task_id=seed_task_id, batch_runner=batch_runner, ) @@ -293,7 +298,10 @@ def run_sweep( colour-coded results table and scatter plot in the terminal. """ falsifier, falsifier_metric = _resolve_falsifier_kwargs( - falsifier, falsifier_metric, skip_on_failure, failure_metric, + falsifier, + falsifier_metric, + skip_on_failure, + failure_metric, ) records: list[dict[str, Any]] = [] @@ -307,26 +315,38 @@ def _build_live_renderable(): return Group(viz.render_header(), progress, viz.render()) return progress - with Live(_build_live_renderable(), console=_console, refresh_per_second=4, - transient=True, vertical_overflow="visible") as live, \ - _quiet_stdout(): + with Live( + _build_live_renderable(), + console=_console, + refresh_per_second=4, + transient=True, + vertical_overflow="visible", + ) as live, _quiet_stdout(): combo_task = progress.add_task("Combos", total=len(param_combos)) seed_task = progress.add_task(" Seeds", total=len(seeds)) for combo_idx, combo in enumerate(param_combos): remaining = _process_combo( - runner, seeds, combo, combo_idx, records, per_combo_summaries, - falsifier=falsifier, falsifier_metric=falsifier_metric, - progress=progress, seed_task_id=seed_task, combo_task_id=combo_task, - batch_runner=batch_runner, viz=viz, live=live, + runner, + seeds, + combo, + combo_idx, + records, + per_combo_summaries, + falsifier=falsifier, + falsifier_metric=falsifier_metric, + progress=progress, + seed_task_id=seed_task, + combo_task_id=combo_task, + batch_runner=batch_runner, + viz=viz, + live=live, live_renderable_fn=_build_live_renderable, ) skipped += remaining if skipped > 0: - _console.print( - f"[dim]Skipped {skipped} runs (moved to next combo on failure)[/dim]" - ) + _console.print(f"[dim]Skipped {skipped} runs (moved to next combo on failure)[/dim]") # Print final viz so it persists after Live exits if viz is not None: @@ -377,6 +397,7 @@ def run_optuna_sweep( objective_metric: str = "safety_violation_rate", direction: str = "minimize", *, + seed: int | None = 0, falsifier: bool = False, falsifier_metric: str = "safety_violations", safety_constraint: "tuple[str, float] | None" = None, @@ -393,6 +414,11 @@ def run_optuna_sweep( Parameters ---------- + seed : int or None + Seed for Optuna's sampler, so a sweep is reproducible across runs + (matching :func:`sample_param_combos`). Pass *None* to let Optuna + draw a fresh random seed, which makes the trials chosen, and + therefore the results, vary between runs. falsifier : bool When *True*, stop iterating seeds on first failure and move to the next trial. @@ -408,15 +434,17 @@ def run_optuna_sweep( Requires ``pip install cbfkit[optuna]``. """ falsifier, falsifier_metric = _resolve_falsifier_kwargs( - falsifier, falsifier_metric, skip_on_failure, failure_metric, + falsifier, + falsifier_metric, + skip_on_failure, + failure_metric, ) try: import optuna except ImportError as exc: raise ImportError( - "Optuna is required for method='optuna'. " - "Install it with: pip install cbfkit[optuna]" + "Optuna is required for method='optuna'. " "Install it with: pip install cbfkit[optuna]" ) from exc optuna.logging.set_verbosity(optuna.logging.WARNING) @@ -436,15 +464,24 @@ def _build_live_renderable(): return progress def objective(trial) -> float: - combo = {pname: _suggest_param(trial, pname, pspec) - for pname, pspec in parameters.items()} + combo = {pname: _suggest_param(trial, pname, pspec) for pname, pspec in parameters.items()} param_combos.append(combo) _process_combo( - runner, seeds, combo, trial.number, records, per_combo_summaries, - falsifier=falsifier, falsifier_metric=falsifier_metric, - progress=progress, seed_task_id=seed_task, combo_task_id=trial_task, - batch_runner=batch_runner, viz=None, live=None, + runner, + seeds, + combo, + trial.number, + records, + per_combo_summaries, + falsifier=falsifier, + falsifier_metric=falsifier_metric, + progress=progress, + seed_task_id=seed_task, + combo_task_id=trial_task, + batch_runner=batch_runner, + viz=None, + live=None, ) summary = per_combo_summaries[-1] @@ -462,11 +499,20 @@ def objective(trial) -> float: return obj_val - study = optuna.create_study(direction=direction) - - with Live(_build_live_renderable(), console=_console, refresh_per_second=4, - transient=True, vertical_overflow="visible") as live, \ - _quiet_stdout(): + # Seed the sampler so which trials get explored, and therefore the + # results, are reproducible across runs. Left unseeded, the startup + # trials are drawn uniformly with replacement, so a small sweep can + # miss part of the grid entirely from one run to the next. + sampler = optuna.samplers.TPESampler(seed=seed) if seed is not None else None + study = optuna.create_study(direction=direction, sampler=sampler) + + with Live( + _build_live_renderable(), + console=_console, + refresh_per_second=4, + transient=True, + vertical_overflow="visible", + ) as live, _quiet_stdout(): live_instance = live trial_task = progress.add_task("Trials", total=n_trials) seed_task = progress.add_task(" Seeds", total=len(seeds)) diff --git a/tests/benchmarks/test_sweep.py b/tests/benchmarks/test_sweep.py index 3931ea0e..0385f184 100644 --- a/tests/benchmarks/test_sweep.py +++ b/tests/benchmarks/test_sweep.py @@ -66,10 +66,12 @@ def test_single_param(self): assert grid[2] == {"a": 3} def test_cartesian_product(self): - grid = build_param_grid({ - "x": {"values": [1, 2]}, - "y": {"values": [10, 20]}, - }) + grid = build_param_grid( + { + "x": {"values": [1, 2]}, + "y": {"values": [10, 20]}, + } + ) assert len(grid) == 4 assert {"x": 1, "y": 10} in grid assert {"x": 2, "y": 20} in grid @@ -158,8 +160,12 @@ def test_falsifier_skips_seeds_after_failure(self): combos = [{"alpha": 3.0}] # alpha > 2 fails on seed >= 1 seeds = [0, 1, 2, 3] result = run_sweep( - "test_falsify", seeds, combos, _mock_failing_runner, - falsifier=True, falsifier_metric="safety_violations", + "test_falsify", + seeds, + combos, + _mock_failing_runner, + falsifier=True, + falsifier_metric="safety_violations", ) # seed=0 passes, seed=1 fails -> seeds 2,3 skipped assert len(result.records) == 2 @@ -170,7 +176,10 @@ def test_falsifier_no_failure_runs_all_seeds(self): combos = [{"alpha": 1.0}] # alpha <= 2 never fails seeds = [0, 1, 2] result = run_sweep( - "test_falsify_pass", seeds, combos, _mock_failing_runner, + "test_falsify_pass", + seeds, + combos, + _mock_failing_runner, falsifier=True, ) assert len(result.records) == 3 @@ -181,7 +190,10 @@ def test_falsifier_disabled_runs_all(self): combos = [{"alpha": 3.0}] seeds = [0, 1, 2, 3] result = run_sweep( - "test_no_falsify", seeds, combos, _mock_failing_runner, + "test_no_falsify", + seeds, + combos, + _mock_failing_runner, falsifier=False, ) assert len(result.records) == 4 @@ -370,8 +382,12 @@ def test_optuna_int_range(self): params = {"count": {"int_range": [1, 10]}} def runner(seed, p): - return {"success": 1, "safety_violations": 0, "solver_failures": 0, - "avg_step_ms": float(p["count"])} + return { + "success": 1, + "safety_violations": 0, + "solver_failures": 0, + "avg_step_ms": float(p["count"]), + } result = run_optuna_sweep( "test_optuna_int", @@ -416,13 +432,58 @@ def test_optuna_falsifier(self): falsifier=True, falsifier_metric="safety_violations", ) - # alpha=1.0 runs all 4 seeds (no failures) - # alpha=3.0 and alpha=5.0 fail on seed>=1, so only 2 seeds each - # Total records should be less than 3*4=12 + # A combo with alpha > 2.0 fails at seed>=1, so it runs 2 seeds + # instead of 4; alpha=1.0 never fails and runs all 4. The sampler is + # seeded (run_optuna_sweep defaults to seed=0), so the three trials + # are a fixed draw that includes alpha > 2.0 and the total stays + # under 3*4=12. Unseeded, a 1-in-27 all-alpha=1.0 draw made this + # assertion fail intermittently in CI. assert len(result.records) < 3 * len(seeds) # At least one combo should be falsified assert any(s["falsified"] for s in result.per_combo_summaries) + def test_optuna_sweep_is_reproducible(self): + """The same seed explores the same trials. + + Regression: the study was created without a sampler, so Optuna drew a + fresh random seed per run. Which points got explored, and therefore + the sweep's results, changed between identical invocations. + """ + params = {"alpha": {"values": [1.0, 3.0, 5.0]}} + + def run(seed): + return run_optuna_sweep( + "test_optuna_repro", + seeds=[0, 1], + parameters=params, + runner=_mock_runner, + n_trials=4, + objective_metric="avg_step_ms", + seed=seed, + ).param_combos + + assert run(0) == run(0) + assert run(7) == run(7) + + def test_optuna_sweep_seed_none_is_unseeded(self): + """seed=None opts back into Optuna choosing its own seed.""" + params = {"alpha": {"range": [0.1, 5.0]}} + + def run(): + return run_optuna_sweep( + "test_optuna_unseeded", + seeds=[0], + parameters=params, + runner=_mock_runner, + n_trials=6, + objective_metric="avg_step_ms", + seed=None, + ).param_combos + + # Continuous range over 6 trials: two unseeded runs matching exactly + # would be an essentially impossible coincidence. + assert run() != run() + def test_optuna_config_loading(self): from cbfkit.benchmarks.sweep_config import load_sweep_config, resolve_param_combos @@ -555,14 +616,22 @@ def test_batch_runner_used(self): def batch_fn(seeds, params): calls.append(seeds) return [ - {"success": 1, "safety_violations": 0, "solver_failures": 0, - "avg_step_ms": 1.0, "score": params.get("alpha", 1.0) + s} + { + "success": 1, + "safety_violations": 0, + "solver_failures": 0, + "avg_step_ms": 1.0, + "score": params.get("alpha", 1.0) + s, + } for s in seeds ] combos = [{"alpha": 2.0}] result = run_sweep( - "test_batch", [0, 1, 2], combos, _mock_runner, + "test_batch", + [0, 1, 2], + combos, + _mock_runner, batch_runner=batch_fn, ) assert len(calls) == 1 # single batched call @@ -575,13 +644,19 @@ def test_batch_runner_skipped_with_falsifier(self): def batch_fn(seeds, params): calls.append(seeds) - return [{"success": 1, "safety_violations": 0, "solver_failures": 0, - "avg_step_ms": 1.0} for _ in seeds] + return [ + {"success": 1, "safety_violations": 0, "solver_failures": 0, "avg_step_ms": 1.0} + for _ in seeds + ] combos = [{"alpha": 3.0}] run_sweep( - "test_batch_falsifier", [0, 1, 2], combos, _mock_failing_runner, - falsifier=True, batch_runner=batch_fn, + "test_batch_falsifier", + [0, 1, 2], + combos, + _mock_failing_runner, + falsifier=True, + batch_runner=batch_fn, ) assert len(calls) == 0 # batch runner not used @@ -680,9 +755,7 @@ def test_parse_circular_sweepable_radius(self): config = load_sweep_config(f.name) assert config.obstacles is not None - assert config.obstacles.items[0].sweepable == { - "radius": {"linspace": [0.3, 1.5, 5]} - } + assert config.obstacles.items[0].sweepable == {"radius": {"linspace": [0.3, 1.5, 5]}} # Synthetic param merged into parameters assert "obstacle_0_radius" in config.parameters assert config.parameters["obstacle_0_radius"] == {"linspace": [0.3, 1.5, 5]} @@ -777,7 +850,11 @@ def test_no_obstacles_block_returns_none(self): assert config.obstacles is None def test_build_obstacle_fixed_params(self): - from cbfkit.benchmarks.sweep_config import ObstacleItemSpec, ObstaclesSpec, _build_obstacle_fixed_params + from cbfkit.benchmarks.sweep_config import ( + ObstacleItemSpec, + ObstaclesSpec, + _build_obstacle_fixed_params, + ) spec = ObstaclesSpec( type="circular", @@ -802,12 +879,18 @@ def test_build_obstacle_fixed_params(self): assert "_obstacle_1_radius" not in fixed def test_extract_obstacle_sweep_params(self): - from cbfkit.benchmarks.sweep_config import ObstacleItemSpec, ObstaclesSpec, _extract_obstacle_sweep_params + from cbfkit.benchmarks.sweep_config import ( + ObstacleItemSpec, + ObstaclesSpec, + _extract_obstacle_sweep_params, + ) spec = ObstaclesSpec( type="circular", items=[ - ObstacleItemSpec(fixed={"center": [3.0, 4.0]}, sweepable={"radius": {"values": [0.5, 1.0]}}), + ObstacleItemSpec( + fixed={"center": [3.0, 4.0]}, sweepable={"radius": {"values": [0.5, 1.0]}} + ), ObstacleItemSpec(fixed={"center": [6.0, 5.0], "radius": 1.2}, sweepable={}), ], )