Remove trend_phase context mode (v0.2.0)

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.                                                    |
 

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",

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):

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,18 +167,15 @@ 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."""
 
         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

src/pulseopt/scheduler.py

@@ -3,15 +3,13 @@
 from __future__ import annotations
 
 import math
-import warnings
 from typing import Any
 
 import torch
 from torch import Tensor
 
 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,

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)

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."""
 

tests/test_scheduler.py

@@ -148,23 +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:
@@ -240,17 +224,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(
@@ -281,18 +254,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):