From a663d6b241872a98b6bdcd3224fce57132f25c7b Mon Sep 17 00:00:00 2001 From: David Kneringer Foss Date: Thu, 7 May 2026 16:17:32 +0200 Subject: [PATCH 1/2] Remove trend_phase context mode and bump to 0.2.0 --- README.md | 4 +-- src/pulseopt/__init__.py | 4 +-- src/pulseopt/controller.py | 12 -------- src/pulseopt/episode.py | 60 ++++++++------------------------------ src/pulseopt/scheduler.py | 22 ++------------ src/pulseopt/types.py | 5 +--- tests/test_episode.py | 33 +-------------------- tests/test_scheduler.py | 44 +--------------------------- 8 files changed, 20 insertions(+), 164 deletions(-) diff --git a/README.md b/README.md index 8285a15..f83db97 100644 --- a/README.md +++ b/README.md @@ -55,7 +55,7 @@ The wrapper owns `optimizer.step()` and `lr_scheduler.step()`; you keep `zero_gr - **Episode**: a fixed-length window of training steps with one frozen candidate: LR multiplier and/or noise std. - **Reward**: log-EMA-loss improvement over the episode, minus an optional instability penalty proportional to within-episode loss variance, clipped to `[-1, 1]`. -- **Controller**: discounted-UCB by default; an optional bucketed-contextual variant uses a coarse loss-trend and optional training-phase bucket to share information across similar regimes. +- **Controller**: discounted-UCB by default; an optional bucketed-contextual variant uses a coarse loss-trend bucket to share information across similar regimes. Axes with a single candidate are treated as fixed constants and get no controller. Passing `lr_candidates=[1.0]` keeps the LR multiplier disabled, and `noise_candidates=[0.0]` keeps gradient noise off. @@ -68,7 +68,7 @@ Axes with a single candidate are treated as fixed constants and get no controlle | `episode_length` | Steps per episode; reward is computed at episode end. | | `lr_scheduler` | Optional `torch.optim.lr_scheduler.*` instance; `step()` is called for you. | | `structured_control_mode` | `"independent"` (default) or `"conditional"` (one noise controller per LR arm). | -| `context_mode` | `"none"` (default), `"trend"`, or `"trend_phase"`; `"trend_phase"` requires `total_training_steps`. | +| `context_mode` | `"none"` (default) or `"trend"`. | | `reward_instability_lambda` | Weight on the variance penalty in the reward. | | `seed` | Seeds controllers and gradient-noise generators. | diff --git a/src/pulseopt/__init__.py b/src/pulseopt/__init__.py index bdb5ef2..f748162 100644 --- a/src/pulseopt/__init__.py +++ b/src/pulseopt/__init__.py @@ -1,10 +1,9 @@ """Adaptive Episodic Exploration Scheduling package.""" -__version__ = "0.1.6rc1" +__version__ = "0.2.0" from pulseopt.controller import ( TREND_CONTEXT_BUCKETS, - TREND_PHASE_CONTEXT_BUCKETS, BaseController, BucketedContextualController, DiscountedUCBController, @@ -35,7 +34,6 @@ "DiscountedUCBController", "RandomController", "TREND_CONTEXT_BUCKETS", - "TREND_PHASE_CONTEXT_BUCKETS", # Reward "BaseReward", "NormalizedLossImprovementReward", diff --git a/src/pulseopt/controller.py b/src/pulseopt/controller.py index 80b2817..b8ae456 100644 --- a/src/pulseopt/controller.py +++ b/src/pulseopt/controller.py @@ -7,18 +7,6 @@ from typing import Protocol TREND_CONTEXT_BUCKETS = ["improving", "stable", "worsening"] -TREND_PHASE_CONTEXT_BUCKETS = [ - "early_improving", - "early_stable", - "early_worsening", - "middle_improving", - "middle_stable", - "middle_worsening", - "late_improving", - "late_stable", - "late_worsening", -] -PHASE_TREND_CONTEXT_BUCKETS = list(TREND_PHASE_CONTEXT_BUCKETS) class BaseController(Protocol): diff --git a/src/pulseopt/episode.py b/src/pulseopt/episode.py index af51840..df322d2 100644 --- a/src/pulseopt/episode.py +++ b/src/pulseopt/episode.py @@ -82,7 +82,6 @@ def __init__( episode_length: int, structured_control_mode: str = "independent", context_mode: str = "none", - total_training_steps: int | None = None, context_trend_window: int = 3, context_trend_epsilon: float = 1e-3, ema_alpha: float = 0.1, @@ -96,16 +95,12 @@ def __init__( raise ValueError("episode_length must be a positive integer.") if structured_control_mode not in {"independent", "conditional"}: raise ValueError("structured_control_mode must be 'independent' or 'conditional'.") - if context_mode not in {"none", "trend", "trend_phase"}: - raise ValueError("context_mode must be 'none', 'trend', or 'trend_phase'.") + if context_mode not in {"none", "trend"}: + raise ValueError("context_mode must be 'none' or 'trend'.") if context_trend_window <= 0: raise ValueError("context_trend_window must be positive.") if context_trend_epsilon < 0.0: raise ValueError("context_trend_epsilon must be non-negative.") - if context_mode == "trend_phase" and ( - total_training_steps is None or total_training_steps <= 0 - ): - raise ValueError("trend_phase context requires a positive total_training_steps.") if len(lr_candidates) == 1: if lr_controller is not None: raise TypeError("Single-candidate LR axes must not create a controller.") @@ -142,7 +137,6 @@ def __init__( self._episode_length = episode_length self._structured_control_mode = structured_control_mode self._context_mode = context_mode - self._total_training_steps = total_training_steps self._context_trend_window = context_trend_window self._context_trend_epsilon = context_trend_epsilon self._ema_alpha = ema_alpha @@ -173,10 +167,7 @@ def __init__( "mean_update_norms": [], } if context_mode != "none": - self._logs["context_bucket_ids"] = [] - self._logs["context_bucket_names"] = [] - self._logs["context_trends"] = [] - self._logs["context_phases"] = [] + self._logs["context_buckets"] = [] def on_step_start(self, global_step: int) -> CandidateConfig: """Return the active structured config, selecting it only at episode start.""" @@ -184,7 +175,7 @@ def on_step_start(self, global_step: int) -> CandidateConfig: if global_step < 0: raise ValueError("global_step must be non-negative.") if self._active_episode is None: - selection = self._select_config(global_step) + selection = self._select_config() self._active_episode = _ActiveStructuredEpisode( episode_index=self._next_episode_index, selection=selection, @@ -287,24 +278,21 @@ def _append_structured_logs( self._logs["reward_final_clipped"].append(summary.reward_final_clipped) self._logs["mean_update_norms"].append(summary.mean_update_norm) if self._context_mode != "none": - self._logs["context_bucket_ids"].append(selection.context_bucket_id) - self._logs["context_bucket_names"].append(selection.context_bucket_name) - self._logs["context_trends"].append(selection.context_trend) - self._logs["context_phases"].append(selection.context_phase) + self._logs["context_buckets"].append(selection.context_bucket) - def _select_config(self, global_step: int) -> StructuredSelection: - bucket_id, bucket_name, trend, phase = self._resolve_context(global_step) + def _select_config(self) -> StructuredSelection: + bucket = self._resolve_context() if self._lr_controller is None: lr_index = 0 else: - _set_controller_context(self._lr_controller, bucket_id) + _set_controller_context(self._lr_controller, bucket) lr_index = self._lr_controller.select_mode() self._validate_index(lr_index, len(self._lr_candidates), "lr_controller") noise_controller = self._select_noise_controller(lr_index) if noise_controller is None: noise_index = 0 else: - _set_controller_context(noise_controller, bucket_id) + _set_controller_context(noise_controller, bucket) noise_index = noise_controller.select_mode() self._validate_index(noise_index, len(self._noise_candidates), "noise_controller") lr_value = self._lr_candidates[lr_index] @@ -315,25 +303,12 @@ def _select_config(self, global_step: int) -> StructuredSelection: lr_value=lr_value, noise_value=noise_value, combined_name=build_generated_candidate_name(lr_value, noise_value), - context_bucket_id=bucket_id, - context_bucket_name=bucket_name, - context_trend=trend, - context_phase=phase, + context_bucket=bucket, ) - def _resolve_context( - self, global_step: int - ) -> tuple[str | None, str | None, str | None, str | None]: + def _resolve_context(self) -> str | None: if self._context_mode == "none": - return None, None, None, None - trend = self._resolve_context_trend() - if self._context_mode == "trend": - return trend, trend, trend, None - phase = self._resolve_context_phase(global_step) - bucket_name = f"{phase}_{trend}" - return bucket_name, bucket_name, trend, phase - - def _resolve_context_trend(self) -> str: + return None window_losses = self._episode_end_loss_history[-self._context_trend_window :] if len(window_losses) < 2: return "stable" @@ -344,17 +319,6 @@ def _resolve_context_trend(self) -> str: return "worsening" return "stable" - def _resolve_context_phase(self, global_step: int) -> str: - total_training_steps = self._total_training_steps - if total_training_steps is None or total_training_steps <= 0: - return "early" - progress = min(max(global_step / total_training_steps, 0.0), 1.0) - if progress < (1.0 / 3.0): - return "early" - if progress < (2.0 / 3.0): - return "middle" - return "late" - def _select_noise_controller(self, lr_index: int) -> BaseController | None: if self._noise_controller is None: return None diff --git a/src/pulseopt/scheduler.py b/src/pulseopt/scheduler.py index db727a2..bbe3ce9 100644 --- a/src/pulseopt/scheduler.py +++ b/src/pulseopt/scheduler.py @@ -3,7 +3,6 @@ from __future__ import annotations import math -import warnings from typing import Any import torch @@ -11,7 +10,6 @@ from pulseopt.controller import ( TREND_CONTEXT_BUCKETS, - TREND_PHASE_CONTEXT_BUCKETS, BucketedContextualController, DiscountedUCBController, RandomController, @@ -21,7 +19,7 @@ from pulseopt.types import CandidateConfig _VALID_CONTROL_MODES = frozenset({"adaptive", "random"}) -_VALID_CONTEXT_MODES = frozenset({"none", "trend", "trend_phase"}) +_VALID_CONTEXT_MODES = frozenset({"none", "trend"}) _VALID_STRUCTURED_CONTROL_MODES = frozenset({"independent", "conditional"}) @@ -76,7 +74,6 @@ def __init__( context_mode: str = "none", context_trend_window: int = 3, context_trend_epsilon: float = 1e-3, - total_training_steps: int | None = None, reward_instability_lambda: float = 0.0, reward_epsilon: float = 1e-8, reward_clip_min: float = -1.0, @@ -112,12 +109,6 @@ def __init__( raise ValueError("noise_candidates values must all be non-negative.") if episode_length < 1: raise ValueError("episode_length must be a positive integer.") - if total_training_steps is not None and context_mode != "trend_phase": - warnings.warn( - "total_training_steps is only used when context_mode='trend_phase'; " - "the value is ignored for the current context_mode.", - stacklevel=2, - ) reward_fn = NormalizedLossImprovementReward( reward_epsilon=reward_epsilon, @@ -126,7 +117,7 @@ def __init__( reward_clip_max=reward_clip_max, ) - bucket_names = _resolve_bucket_names(context_mode) + bucket_names = list(TREND_CONTEXT_BUCKETS) if context_mode == "trend" else None lr_controller = _build_controller( n_arms=len(lr_list), @@ -152,7 +143,6 @@ def __init__( episode_length=episode_length, structured_control_mode=structured_control_mode, context_mode=context_mode, - total_training_steps=total_training_steps if context_mode == "trend_phase" else None, context_trend_window=context_trend_window, context_trend_epsilon=context_trend_epsilon, ema_alpha=ema_alpha, @@ -327,14 +317,6 @@ def _get_noise_generator( # ------------------------------------------------------------------ -def _resolve_bucket_names(context_mode: str) -> list[str] | None: - if context_mode == "trend": - return list(TREND_CONTEXT_BUCKETS) - if context_mode == "trend_phase": - return list(TREND_PHASE_CONTEXT_BUCKETS) - return None - - def _build_controller( n_arms: int, control_mode: str, diff --git a/src/pulseopt/types.py b/src/pulseopt/types.py index b7ce49b..9a6f774 100644 --- a/src/pulseopt/types.py +++ b/src/pulseopt/types.py @@ -31,10 +31,7 @@ class StructuredSelection: lr_value: float noise_value: float combined_name: str - context_bucket_id: str | None = None - context_bucket_name: str | None = None - context_trend: str | None = None - context_phase: str | None = None + context_bucket: str | None = None @dataclass(frozen=True) diff --git a/tests/test_episode.py b/tests/test_episode.py index c49c561..8a807df 100644 --- a/tests/test_episode.py +++ b/tests/test_episode.py @@ -224,42 +224,11 @@ def test_trend_context_is_shared_across_independent_controllers() -> None: manager.on_step_end(loss=loss) logs = manager.get_logs() - assert logs["context_trends"] == ["stable", "stable", "improving"] - assert logs["context_phases"] == [None, None, None] + assert logs["context_buckets"] == ["stable", "stable", "improving"] assert lr_controller.contexts == ["stable", "stable", "improving"] assert noise_controller.contexts == ["stable", "stable", "improving"] -def test_trend_phase_context_logs_phase_buckets() -> None: - """Trend-phase context should combine progress phase with the shared trend bucket.""" - - manager = StructuredEpisodeManager( - lr_candidates=[1.0, 1.2], - noise_candidates=[0.0, 0.05], - lr_controller=ContextAwareStubController([0, 0, 0]), - noise_controller=ContextAwareStubController([0, 0, 0]), - reward_fn=build_reward(), - episode_length=1, - structured_control_mode="independent", - context_mode="trend_phase", - total_training_steps=3, - context_trend_window=2, - context_trend_epsilon=0.1, - ) - - for step, loss in enumerate([5.0, 4.0, 4.5]): - manager.on_step_start(global_step=step) - manager.on_step_end(loss=loss) - - logs = manager.get_logs() - assert logs["context_phases"] == ["early", "middle", "late"] - assert logs["context_bucket_names"] == [ - "early_stable", - "middle_stable", - "late_improving", - ] - - def test_finalize_closes_partial_structured_episode_safely() -> None: """A partial structured episode with completed steps should be finalized cleanly.""" diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 4b96843..e07c807 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -148,24 +148,7 @@ def test_trend_context_mode(self): _run_steps(aees, model, optimizer, 20) aees.finalize() logs = aees.get_logs() - assert "context_bucket_ids" in logs - - def test_trend_phase_context_mode(self): - model, optimizer = _make_model_and_optimizer() - aees = AEES( - optimizer, - lr_candidates=[0.5, 1.0], - noise_candidates=[0.0], - episode_length=4, - context_mode="trend_phase", - total_training_steps=40, - seed=0, - ) - _run_steps(aees, model, optimizer, 20) - aees.finalize() - logs = aees.get_logs() - assert "context_bucket_ids" in logs - + assert "context_buckets" in logs class TestAEESConditionalMode: def test_conditional_structured_control(self): @@ -240,18 +223,6 @@ def test_invalid_context_mode_raises(self): with pytest.raises(ValueError, match="context_mode"): AEES(optimizer, lr_candidates=[1.0], noise_candidates=[0.0], context_mode="invalid") - def test_trend_phase_without_total_steps_raises(self): - _, optimizer = _make_model_and_optimizer() - with pytest.raises(ValueError): - AEES( - optimizer, - lr_candidates=[0.5, 1.0], - noise_candidates=[0.0], - context_mode="trend_phase", - total_training_steps=None, - ) - - class TestAEESInputValidation: @pytest.mark.parametrize( "lr_candidates", @@ -281,19 +252,6 @@ def test_zero_episode_length_raises(self): episode_length=0, ) - def test_total_training_steps_outside_trend_phase_warns(self): - _, optimizer = _make_model_and_optimizer() - with pytest.warns(UserWarning, match="total_training_steps"): - AEES( - optimizer, - lr_candidates=[1.0], - noise_candidates=[0.0], - episode_length=5, - context_mode="trend", - total_training_steps=100, - ) - - class TestAEESCurrentCandidate: def test_current_candidate_lifecycle(self): model, optimizer = _make_model_and_optimizer() From c93a0fa461de4ee69bd511f848b58935ce2fcf4c Mon Sep 17 00:00:00 2001 From: David Kneringer Foss Date: Thu, 7 May 2026 16:24:53 +0200 Subject: [PATCH 2/2] reformat --- tests/test_scheduler.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index e07c807..7c451fc 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -150,6 +150,7 @@ def test_trend_context_mode(self): logs = aees.get_logs() assert "context_buckets" in logs + class TestAEESConditionalMode: def test_conditional_structured_control(self): model, optimizer = _make_model_and_optimizer() @@ -223,6 +224,7 @@ def test_invalid_context_mode_raises(self): with pytest.raises(ValueError, match="context_mode"): AEES(optimizer, lr_candidates=[1.0], noise_candidates=[0.0], context_mode="invalid") + class TestAEESInputValidation: @pytest.mark.parametrize( "lr_candidates", @@ -252,6 +254,7 @@ def test_zero_episode_length_raises(self): episode_length=0, ) + class TestAEESCurrentCandidate: def test_current_candidate_lifecycle(self): model, optimizer = _make_model_and_optimizer()