alberta-framework 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

alberta_framework/__init__.py

@@ -1,25 +1,42 @@
-"""Alberta Framework: Implementation of the Alberta Plan for AI Research.
-
-This framework implements Step 1 of the Alberta Plan: continual supervised
-learning with meta-learned step-sizes.
-
-Core Philosophy: Temporal uniformity - every component updates at every time step.
-
-Quick Start:
-    >>> import jax.random as jr
-    >>> from alberta_framework import LinearLearner, IDBD, RandomWalkStream, run_learning_loop
-    >>>
-    >>> # Create a non-stationary stream
-    >>> stream = RandomWalkStream(feature_dim=10, drift_rate=0.001)
-    >>>
-    >>> # Create a learner with adaptive step-sizes
-    >>> learner = LinearLearner(optimizer=IDBD())
-    >>>
-    >>> # Run learning loop with scan
-    >>> key = jr.key(42)
-    >>> state, metrics = run_learning_loop(learner, stream, num_steps=10000, key=key)
-
-Reference: The Alberta Plan for AI Research (Sutton et al.)
+"""Alberta Framework: A JAX-based research framework for continual AI.
+
+The Alberta Framework provides foundational components for continual reinforcement
+learning research. Built on JAX for hardware acceleration, the framework emphasizes
+temporal uniformity — every component updates at every time step, with no special
+training phases or batch processing.
+
+Roadmap
+-------
+| Step | Focus | Status |
+|------|-------|--------|
+| 1 | Meta-learned step-sizes (IDBD, Autostep) | **Complete** |
+| 2 | Feature generation and testing | Planned |
+| 3 | GVF predictions, Horde architecture | Planned |
+| 4 | Actor-critic with eligibility traces | Planned |
+| 5-6 | Off-policy learning, average reward | Planned |
+| 7-12 | Hierarchical, multi-agent, world models | Future |
+
+Examples
+--------
+```python
+import jax.random as jr
+from alberta_framework import LinearLearner, IDBD, RandomWalkStream, run_learning_loop
+
+# Non-stationary stream where target weights drift over time
+stream = RandomWalkStream(feature_dim=10, drift_rate=0.001)
+
+# Learner with IDBD meta-learned step-sizes
+learner = LinearLearner(optimizer=IDBD())
+
+# JIT-compiled training via jax.lax.scan
+state, metrics = run_learning_loop(learner, stream, num_steps=10000, key=jr.key(42))
+```
+
+References
+----------
+- The Alberta Plan for AI Research (Sutton et al., 2022): https://arxiv.org/abs/2208.11173
+- Adapting Bias by Gradient Descent (Sutton, 1992)
+- Tuning-free Step-size Adaptation (Mahmood et al., 2012)
 """
 
 __version__ = "0.1.0"
@@ -50,6 +67,8 @@ from alberta_framework.core.types import (
     IDBDState,
     LearnerState,
     LMSState,
+    NormalizerHistory,
+    NormalizerTrackingConfig,
     Observation,
     Prediction,
     StepSizeHistory,
@@ -122,7 +141,9 @@ __all__ = [
     "IDBDState",
     "LMSState",
     "LearnerState",
+    "NormalizerHistory",
     "NormalizerState",
+    "NormalizerTrackingConfig",
     "Observation",
     "Prediction",
     "StepSizeHistory",

alberta_framework/core/learners.py

@@ -18,6 +18,8 @@ from alberta_framework.core.types import (
     IDBDState,
     LearnerState,
     LMSState,
+    NormalizerHistory,
+    NormalizerTrackingConfig,
     Observation,
     Prediction,
     StepSizeHistory,
@@ -48,7 +50,7 @@ class UpdateResult(NamedTuple):
 class LinearLearner:
     """Linear function approximator with pluggable optimizer.
 
-    Computes predictions as: y = w @ x + b
+    Computes predictions as: `y = w @ x + b`
 
     The learner maintains weights and bias, delegating the adaptation
     of learning rates to the optimizer (e.g., LMS or IDBD).
@@ -93,7 +95,7 @@ class LinearLearner:
             observation: Input feature vector
 
         Returns:
-            Scalar prediction y = w @ x + b
+            Scalar prediction `y = w @ x + b`
         """
         return jnp.atleast_1d(jnp.dot(state.weights, observation) + state.bias)
 
@@ -238,10 +240,25 @@ def run_learning_loop[StreamStateT](
         )
         recording_indices = jnp.zeros(num_recordings, dtype=jnp.int32)
 
+        # Check if we need to track Autostep normalizers
+        # We detect this at trace time by checking the initial optimizer state
+        track_normalizers = hasattr(learner_state.optimizer_state, "normalizers")
+        normalizer_history = (
+            jnp.zeros((num_recordings, feature_dim), dtype=jnp.float32)
+            if track_normalizers
+            else None
+        )
+
         def step_fn_with_tracking(
-            carry: tuple[LearnerState, StreamStateT, Array, Array | None, Array], idx: Array
-        ) -> tuple[tuple[LearnerState, StreamStateT, Array, Array | None, Array], Array]:
-            l_state, s_state, ss_history, b_history, rec_indices = carry
+            carry: tuple[
+                LearnerState, StreamStateT, Array, Array | None, Array, Array | None
+            ],
+            idx: Array,
+        ) -> tuple[
+            tuple[LearnerState, StreamStateT, Array, Array | None, Array, Array | None],
+            Array,
+        ]:
+            l_state, s_state, ss_history, b_history, rec_indices, norm_history = carry
 
             # Perform learning step
             timestep, new_s_state = stream.step(s_state, idx)
@@ -291,12 +308,25 @@ def run_learning_loop[StreamStateT](
                 None,
             )
 
+            # Track Autostep normalizers (v_i) if applicable
+            new_norm_history = norm_history
+            if norm_history is not None and hasattr(opt_state, "normalizers"):
+                new_norm_history = jax.lax.cond(
+                    should_record,
+                    lambda _: norm_history.at[recording_idx].set(
+                        opt_state.normalizers  # type: ignore[union-attr]
+                    ),
+                    lambda _: norm_history,
+                    None,
+                )
+
             return (
                 result.state,
                 new_s_state,
                 new_ss_history,
                 new_b_history,
                 new_rec_indices,
+                new_norm_history,
             ), result.metrics
 
         initial_carry = (
@@ -305,16 +335,25 @@ def run_learning_loop[StreamStateT](
             step_size_history,
             bias_history,
             recording_indices,
+            normalizer_history,
         )
 
-        (final_learner, _, final_ss_history, final_b_history, final_rec_indices), metrics = (
-            jax.lax.scan(step_fn_with_tracking, initial_carry, jnp.arange(num_steps))
+        (
+            final_learner,
+            _,
+            final_ss_history,
+            final_b_history,
+            final_rec_indices,
+            final_norm_history,
+        ), metrics = jax.lax.scan(
+            step_fn_with_tracking, initial_carry, jnp.arange(num_steps)
         )
 
         history = StepSizeHistory(
             step_sizes=final_ss_history,
             bias_step_sizes=final_b_history,
             recording_indices=final_rec_indices,
+            normalizers=final_norm_history,
         )
 
         return final_learner, metrics, history
@@ -473,7 +512,14 @@ def run_normalized_learning_loop[StreamStateT](
     num_steps: int,
     key: Array,
     learner_state: NormalizedLearnerState | None = None,
-) -> tuple[NormalizedLearnerState, Array]:
+    step_size_tracking: StepSizeTrackingConfig | None = None,
+    normalizer_tracking: NormalizerTrackingConfig | None = None,
+) -> (
+    tuple[NormalizedLearnerState, Array]
+    | tuple[NormalizedLearnerState, Array, StepSizeHistory]
+    | tuple[NormalizedLearnerState, Array, NormalizerHistory]
+    | tuple[NormalizedLearnerState, Array, StepSizeHistory, NormalizerHistory]
+):
     """Run the learning loop with normalization using jax.lax.scan.
 
     Args:
@@ -482,29 +528,322 @@ def run_normalized_learning_loop[StreamStateT](
         num_steps: Number of learning steps to run
         key: JAX random key for stream initialization
         learner_state: Initial state (if None, will be initialized from stream)
+        step_size_tracking: Optional config for recording per-weight step-sizes.
+            When provided, returns StepSizeHistory including Autostep normalizers if applicable.
+        normalizer_tracking: Optional config for recording per-feature normalizer state.
+            When provided, returns NormalizerHistory with means and variances over time.
 
     Returns:
-        Tuple of (final_state, metrics_array) where metrics_array has shape
-        (num_steps, 4) with columns [squared_error, error, mean_step_size, normalizer_mean_var]
+        If no tracking:
+            Tuple of (final_state, metrics_array) where metrics_array has shape
+            (num_steps, 4) with columns [squared_error, error, mean_step_size, normalizer_mean_var]
+        If step_size_tracking only:
+            Tuple of (final_state, metrics_array, step_size_history)
+        If normalizer_tracking only:
+            Tuple of (final_state, metrics_array, normalizer_history)
+        If both:
+            Tuple of (final_state, metrics_array, step_size_history, normalizer_history)
+
+    Raises:
+        ValueError: If tracking interval is invalid
     """
+    # Validate tracking configs
+    if step_size_tracking is not None:
+        if step_size_tracking.interval < 1:
+            raise ValueError(
+                f"step_size_tracking.interval must be >= 1, got {step_size_tracking.interval}"
+            )
+        if step_size_tracking.interval > num_steps:
+            raise ValueError(
+                f"step_size_tracking.interval ({step_size_tracking.interval}) "
+                f"must be <= num_steps ({num_steps})"
+            )
+
+    if normalizer_tracking is not None:
+        if normalizer_tracking.interval < 1:
+            raise ValueError(
+                f"normalizer_tracking.interval must be >= 1, got {normalizer_tracking.interval}"
+            )
+        if normalizer_tracking.interval > num_steps:
+            raise ValueError(
+                f"normalizer_tracking.interval ({normalizer_tracking.interval}) "
+                f"must be <= num_steps ({num_steps})"
+            )
+
     # Initialize states
     if learner_state is None:
         learner_state = learner.init(stream.feature_dim)
     stream_state = stream.init(key)
 
-    def step_fn(
-        carry: tuple[NormalizedLearnerState, StreamStateT], idx: Array
-    ) -> tuple[tuple[NormalizedLearnerState, StreamStateT], Array]:
-        l_state, s_state = carry
+    feature_dim = stream.feature_dim
+
+    # No tracking - simple case
+    if step_size_tracking is None and normalizer_tracking is None:
+
+        def step_fn(
+            carry: tuple[NormalizedLearnerState, StreamStateT], idx: Array
+        ) -> tuple[tuple[NormalizedLearnerState, StreamStateT], Array]:
+            l_state, s_state = carry
+            timestep, new_s_state = stream.step(s_state, idx)
+            result = learner.update(l_state, timestep.observation, timestep.target)
+            return (result.state, new_s_state), result.metrics
+
+        (final_learner, _), metrics = jax.lax.scan(
+            step_fn, (learner_state, stream_state), jnp.arange(num_steps)
+        )
+
+        return final_learner, metrics
+
+    # Tracking enabled - need to set up history arrays
+    ss_interval = step_size_tracking.interval if step_size_tracking else num_steps + 1
+    norm_interval = (
+        normalizer_tracking.interval if normalizer_tracking else num_steps + 1
+    )
+
+    ss_num_recordings = num_steps // ss_interval if step_size_tracking else 0
+    norm_num_recordings = num_steps // norm_interval if normalizer_tracking else 0
+
+    # Pre-allocate step-size history arrays
+    ss_history = (
+        jnp.zeros((ss_num_recordings, feature_dim), dtype=jnp.float32)
+        if step_size_tracking
+        else None
+    )
+    ss_bias_history = (
+        jnp.zeros(ss_num_recordings, dtype=jnp.float32)
+        if step_size_tracking and step_size_tracking.include_bias
+        else None
+    )
+    ss_rec_indices = (
+        jnp.zeros(ss_num_recordings, dtype=jnp.int32) if step_size_tracking else None
+    )
+
+    # Check if we need to track Autostep normalizers
+    track_autostep_normalizers = hasattr(
+        learner_state.learner_state.optimizer_state, "normalizers"
+    )
+    ss_normalizers = (
+        jnp.zeros((ss_num_recordings, feature_dim), dtype=jnp.float32)
+        if step_size_tracking and track_autostep_normalizers
+        else None
+    )
+
+    # Pre-allocate normalizer state history arrays
+    norm_means = (
+        jnp.zeros((norm_num_recordings, feature_dim), dtype=jnp.float32)
+        if normalizer_tracking
+        else None
+    )
+    norm_vars = (
+        jnp.zeros((norm_num_recordings, feature_dim), dtype=jnp.float32)
+        if normalizer_tracking
+        else None
+    )
+    norm_rec_indices = (
+        jnp.zeros(norm_num_recordings, dtype=jnp.int32) if normalizer_tracking else None
+    )
+
+    def step_fn_with_tracking(
+        carry: tuple[
+            NormalizedLearnerState,
+            StreamStateT,
+            Array | None,
+            Array | None,
+            Array | None,
+            Array | None,
+            Array | None,
+            Array | None,
+            Array | None,
+        ],
+        idx: Array,
+    ) -> tuple[
+        tuple[
+            NormalizedLearnerState,
+            StreamStateT,
+            Array | None,
+            Array | None,
+            Array | None,
+            Array | None,
+            Array | None,
+            Array | None,
+            Array | None,
+        ],
+        Array,
+    ]:
+        (
+            l_state,
+            s_state,
+            ss_hist,
+            ss_bias_hist,
+            ss_rec,
+            ss_norm,
+            n_means,
+            n_vars,
+            n_rec,
+        ) = carry
+
+        # Perform learning step
         timestep, new_s_state = stream.step(s_state, idx)
         result = learner.update(l_state, timestep.observation, timestep.target)
-        return (result.state, new_s_state), result.metrics
 
-    (final_learner, _), metrics = jax.lax.scan(
-        step_fn, (learner_state, stream_state), jnp.arange(num_steps)
+        # Step-size tracking
+        new_ss_hist = ss_hist
+        new_ss_bias_hist = ss_bias_hist
+        new_ss_rec = ss_rec
+        new_ss_norm = ss_norm
+
+        if ss_hist is not None:
+            should_record_ss = (idx % ss_interval) == 0
+            recording_idx = idx // ss_interval
+
+            # Extract current step-sizes from the inner learner state
+            opt_state = result.state.learner_state.optimizer_state
+            if hasattr(opt_state, "log_step_sizes"):
+                # IDBD stores log step-sizes
+                weight_ss = jnp.exp(opt_state.log_step_sizes)  # type: ignore[union-attr]
+                bias_ss = opt_state.bias_step_size  # type: ignore[union-attr]
+            elif hasattr(opt_state, "step_sizes"):
+                # Autostep stores step-sizes directly
+                weight_ss = opt_state.step_sizes  # type: ignore[union-attr]
+                bias_ss = opt_state.bias_step_size  # type: ignore[union-attr]
+            else:
+                # LMS has a single fixed step-size
+                weight_ss = jnp.full(feature_dim, opt_state.step_size)
+                bias_ss = opt_state.step_size
+
+            new_ss_hist = jax.lax.cond(
+                should_record_ss,
+                lambda _: ss_hist.at[recording_idx].set(weight_ss),
+                lambda _: ss_hist,
+                None,
+            )
+
+            if ss_bias_hist is not None:
+                new_ss_bias_hist = jax.lax.cond(
+                    should_record_ss,
+                    lambda _: ss_bias_hist.at[recording_idx].set(bias_ss),
+                    lambda _: ss_bias_hist,
+                    None,
+                )
+
+            if ss_rec is not None:
+                new_ss_rec = jax.lax.cond(
+                    should_record_ss,
+                    lambda _: ss_rec.at[recording_idx].set(idx),
+                    lambda _: ss_rec,
+                    None,
+                )
+
+            # Track Autostep normalizers (v_i) if applicable
+            if ss_norm is not None and hasattr(opt_state, "normalizers"):
+                new_ss_norm = jax.lax.cond(
+                    should_record_ss,
+                    lambda _: ss_norm.at[recording_idx].set(
+                        opt_state.normalizers  # type: ignore[union-attr]
+                    ),
+                    lambda _: ss_norm,
+                    None,
+                )
+
+        # Normalizer state tracking
+        new_n_means = n_means
+        new_n_vars = n_vars
+        new_n_rec = n_rec
+
+        if n_means is not None:
+            should_record_norm = (idx % norm_interval) == 0
+            norm_recording_idx = idx // norm_interval
+
+            norm_state = result.state.normalizer_state
+
+            new_n_means = jax.lax.cond(
+                should_record_norm,
+                lambda _: n_means.at[norm_recording_idx].set(norm_state.mean),
+                lambda _: n_means,
+                None,
+            )
+
+            if n_vars is not None:
+                new_n_vars = jax.lax.cond(
+                    should_record_norm,
+                    lambda _: n_vars.at[norm_recording_idx].set(norm_state.var),
+                    lambda _: n_vars,
+                    None,
+                )
+
+            if n_rec is not None:
+                new_n_rec = jax.lax.cond(
+                    should_record_norm,
+                    lambda _: n_rec.at[norm_recording_idx].set(idx),
+                    lambda _: n_rec,
+                    None,
+                )
+
+        return (
+            result.state,
+            new_s_state,
+            new_ss_hist,
+            new_ss_bias_hist,
+            new_ss_rec,
+            new_ss_norm,
+            new_n_means,
+            new_n_vars,
+            new_n_rec,
+        ), result.metrics
+
+    initial_carry = (
+        learner_state,
+        stream_state,
+        ss_history,
+        ss_bias_history,
+        ss_rec_indices,
+        ss_normalizers,
+        norm_means,
+        norm_vars,
+        norm_rec_indices,
+    )
+
+    (
+        final_learner,
+        _,
+        final_ss_hist,
+        final_ss_bias_hist,
+        final_ss_rec,
+        final_ss_norm,
+        final_n_means,
+        final_n_vars,
+        final_n_rec,
+    ), metrics = jax.lax.scan(
+        step_fn_with_tracking, initial_carry, jnp.arange(num_steps)
     )
 
-    return final_learner, metrics
+    # Build return values based on what was tracked
+    ss_history_result = None
+    if step_size_tracking is not None and final_ss_hist is not None:
+        ss_history_result = StepSizeHistory(
+            step_sizes=final_ss_hist,
+            bias_step_sizes=final_ss_bias_hist,
+            recording_indices=final_ss_rec,  # type: ignore[arg-type]
+            normalizers=final_ss_norm,
+        )
+
+    norm_history_result = None
+    if normalizer_tracking is not None and final_n_means is not None:
+        norm_history_result = NormalizerHistory(
+            means=final_n_means,
+            variances=final_n_vars,  # type: ignore[arg-type]
+            recording_indices=final_n_rec,  # type: ignore[arg-type]
+        )
+
+    # Return appropriate tuple based on what was tracked
+    if ss_history_result is not None and norm_history_result is not None:
+        return final_learner, metrics, ss_history_result, norm_history_result
+    elif ss_history_result is not None:
+        return final_learner, metrics, ss_history_result
+    elif norm_history_result is not None:
+        return final_learner, metrics, norm_history_result
+    else:
+        return final_learner, metrics
 
 
 def metrics_to_dicts(metrics: Array, normalized: bool = False) -> list[dict[str, float]]:

alberta_framework/core/normalizers.py

@@ -35,7 +35,7 @@ class OnlineNormalizer:
     """Online feature normalizer for continual learning.
 
     Normalizes features using running estimates of mean and standard deviation:
-        x_normalized = (x - mean) / (std + epsilon)
+    `x_normalized = (x - mean) / (std + epsilon)`
 
     The normalizer updates its estimates at every time step, following
     temporal uniformity. Uses exponential moving average for non-stationary

alberta_framework/core/optimizers.py

@@ -72,7 +72,7 @@ class Optimizer[StateT: (LMSState, IDBDState, AutostepState)](ABC):
 class LMS(Optimizer[LMSState]):
     """Least Mean Square optimizer with fixed step-size.
 
-    The simplest gradient-based optimizer: w_{t+1} = w_t + alpha * delta * x_t
+    The simplest gradient-based optimizer: `w_{t+1} = w_t + alpha * delta * x_t`
 
     This serves as a baseline. The challenge is that the optimal step-size
     depends on the problem and changes as the task becomes non-stationary.
@@ -108,7 +108,7 @@ class LMS(Optimizer[LMSState]):
     ) -> OptimizerUpdate:
         """Compute LMS weight update.
 
-        Update rule: delta_w = alpha * error * x
+        Update rule: `delta_w = alpha * error * x`
 
         Args:
             state: Current LMS state
@@ -195,10 +195,11 @@ class IDBD(Optimizer[IDBDState]):
         """Compute IDBD weight update with adaptive step-sizes.
 
         The IDBD algorithm:
-        1. Compute step-sizes: alpha_i = exp(log_alpha_i)
-        2. Update weights: w_i += alpha_i * error * x_i
-        3. Update log step-sizes: log_alpha_i += beta * error * x_i * h_i
-        4. Update traces: h_i = h_i * max(0, 1 - alpha_i * x_i^2) + alpha_i * error * x_i
+
+        1. Compute step-sizes: `alpha_i = exp(log_alpha_i)`
+        2. Update weights: `w_i += alpha_i * error * x_i`
+        3. Update log step-sizes: `log_alpha_i += beta * error * x_i * h_i`
+        4. Update traces: `h_i = h_i * max(0, 1 - alpha_i * x_i^2) + alpha_i * error * x_i`
 
         The trace h_i tracks the correlation between current and past gradients.
         When gradients consistently point the same direction, h_i grows,
@@ -335,12 +336,13 @@ class Autostep(Optimizer[AutostepState]):
         """Compute Autostep weight update with normalized gradients.
 
         The Autostep algorithm:
-        1. Compute gradient: g_i = error * x_i
-        2. Normalize gradient: g_i' = g_i / max(|g_i|, v_i)
-        3. Update weights: w_i += alpha_i * g_i'
-        4. Update step-sizes: alpha_i *= exp(mu * g_i' * h_i)
-        5. Update traces: h_i = h_i * (1 - alpha_i) + alpha_i * g_i'
-        6. Update normalizers: v_i = max(|g_i|, v_i * tau)
+
+        1. Compute gradient: `g_i = error * x_i`
+        2. Normalize gradient: `g_i' = g_i / max(|g_i|, v_i)`
+        3. Update weights: `w_i += alpha_i * g_i'`
+        4. Update step-sizes: `alpha_i *= exp(mu * g_i' * h_i)`
+        5. Update traces: `h_i = h_i * (1 - alpha_i) + alpha_i * g_i'`
+        6. Update normalizers: `v_i = max(|g_i|, v_i * tau)`
 
         Args:
             state: Current Autostep state

alberta_framework/core/types.py

@@ -126,11 +126,42 @@ class StepSizeHistory(NamedTuple):
         step_sizes: Per-weight step-sizes at each recording, shape (num_recordings, num_weights)
         bias_step_sizes: Bias step-sizes at each recording, shape (num_recordings,) or None
         recording_indices: Step indices where recordings were made, shape (num_recordings,)
+        normalizers: Autostep's per-weight normalizers (v_i) at each recording,
+            shape (num_recordings, num_weights) or None. Only populated for Autostep optimizer.
     """
 
     step_sizes: Array  # (num_recordings, num_weights)
     bias_step_sizes: Array | None  # (num_recordings,) or None
     recording_indices: Array  # (num_recordings,)
+    normalizers: Array | None = None  # (num_recordings, num_weights) - Autostep v_i
+
+
+class NormalizerTrackingConfig(NamedTuple):
+    """Configuration for recording per-feature normalizer state during training.
+
+    Attributes:
+        interval: Record normalizer state every N steps
+    """
+
+    interval: int
+
+
+class NormalizerHistory(NamedTuple):
+    """History of per-feature normalizer state recorded during training.
+
+    Used for analyzing how the OnlineNormalizer adapts to distribution shifts
+    (reactive lag diagnostic).
+
+    Attributes:
+        means: Per-feature mean estimates at each recording, shape (num_recordings, feature_dim)
+        variances: Per-feature variance estimates at each recording,
+            shape (num_recordings, feature_dim)
+        recording_indices: Step indices where recordings were made, shape (num_recordings,)
+    """
+
+    means: Array  # (num_recordings, feature_dim)
+    variances: Array  # (num_recordings, feature_dim)
+    recording_indices: Array  # (num_recordings,)
 
 
 def create_lms_state(step_size: float = 0.01) -> LMSState:

alberta_framework/streams/base.py

@@ -30,11 +30,14 @@ class ScanStream(Protocol[StateT]):
     Type Parameters:
         StateT: The state type maintained by this stream
 
-    Example:
-        >>> stream = RandomWalkStream(feature_dim=10, drift_rate=0.001)
-        >>> key = jax.random.key(42)
-        >>> state = stream.init(key)
-        >>> timestep, new_state = stream.step(state, jnp.array(0))
+    Examples
+    --------
+    ```python
+    stream = RandomWalkStream(feature_dim=10, drift_rate=0.001)
+    key = jax.random.key(42)
+    state = stream.init(key)
+    timestep, new_state = stream.step(state, jnp.array(0))
+    ```
     """
 
     @property

alberta_framework/streams/synthetic.py

@@ -32,7 +32,7 @@ class RandomWalkState(NamedTuple):
 class RandomWalkStream:
     """Non-stationary stream where target weights drift via random walk.
 
-    The true target function is linear: y* = w_true @ x + noise
+    The true target function is linear: `y* = w_true @ x + noise`
     where w_true evolves via random walk at each time step.
 
     This tests the learner's ability to continuously track a moving target.
@@ -590,12 +590,15 @@ class ScaledStreamWrapper:
     scale factor. Useful for testing how learners handle features at different
     scales, which is important for understanding normalization benefits.
 
-    Example:
-        >>> stream = ScaledStreamWrapper(
-        ...     AbruptChangeStream(feature_dim=10, change_interval=1000),
-        ...     feature_scales=jnp.array([0.001, 0.01, 0.1, 1.0, 10.0,
-        ...                               100.0, 1000.0, 0.001, 0.01, 0.1])
-        ... )
+    Examples
+    --------
+    ```python
+    stream = ScaledStreamWrapper(
+        AbruptChangeStream(feature_dim=10, change_interval=1000),
+        feature_scales=jnp.array([0.001, 0.01, 0.1, 1.0, 10.0,
+                                  100.0, 1000.0, 0.001, 0.01, 0.1])
+    )
+    ```
 
     Attributes:
         inner_stream: The wrapped stream instance
@@ -693,9 +696,12 @@ def make_scale_range(
     Returns:
         Array of shape (feature_dim,) with scale factors
 
-    Example:
-        >>> scales = make_scale_range(10, min_scale=0.01, max_scale=100.0)
-        >>> stream = ScaledStreamWrapper(RandomWalkStream(10), scales)
+    Examples
+    --------
+    ```python
+    scales = make_scale_range(10, min_scale=0.01, max_scale=100.0)
+    stream = ScaledStreamWrapper(RandomWalkStream(10), scales)
+    ```
     """
     if log_spaced:
         return jnp.logspace(

alberta_framework/utils/experiments.py

@@ -110,13 +110,14 @@ def run_single_experiment(
 
     final_state: LearnerState | NormalizedLearnerState
     if isinstance(learner, NormalizedLinearLearner):
-        final_state, metrics = run_normalized_learning_loop(
+        norm_result = run_normalized_learning_loop(
             learner, stream, config.num_steps, key
         )
+        final_state, metrics = cast(tuple[NormalizedLearnerState, Any], norm_result)
         metrics_history = metrics_to_dicts(metrics, normalized=True)
     else:
-        result = run_learning_loop(learner, stream, config.num_steps, key)
-        final_state, metrics = cast(tuple[LearnerState, Any], result)
+        linear_result = run_learning_loop(learner, stream, config.num_steps, key)
+        final_state, metrics = cast(tuple[LearnerState, Any], linear_result)
         metrics_history = metrics_to_dicts(metrics)
 
     return SingleRunResult(

alberta_framework/utils/timing.py

@@ -3,19 +3,22 @@
 This module provides a simple Timer context manager for measuring execution time
 and formatting durations in a human-readable format.
 
-Example:
-    >>> from alberta_framework.utils.timing import Timer
-    >>>
-    >>> with Timer("Training"):
-    ...     # run training code
-    ...     pass
-    Training completed in 1.23s
-    >>>
-    >>> # Or capture the duration:
-    >>> with Timer("Experiment") as t:
-    ...     # run experiment
-    ...     pass
-    >>> print(f"Took {t.duration:.2f} seconds")
+Examples
+--------
+```python
+from alberta_framework.utils.timing import Timer
+
+with Timer("Training"):
+    # run training code
+    pass
+# Output: Training completed in 1.23s
+
+# Or capture the duration:
+with Timer("Experiment") as t:
+    # run experiment
+    pass
+print(f"Took {t.duration:.2f} seconds")
+```
 """
 
 import time
@@ -32,13 +35,13 @@ def format_duration(seconds: float) -> str:
     Returns:
         Formatted string like "1.23s", "2m 30.5s", or "1h 5m 30s"
 
-    Examples:
-        >>> format_duration(0.5)
-        '0.50s'
-        >>> format_duration(90.5)
-        '1m 30.50s'
-        >>> format_duration(3665)
-        '1h 1m 5.00s'
+    Examples
+    --------
+    ```python
+    format_duration(0.5)   # Returns: '0.50s'
+    format_duration(90.5)  # Returns: '1m 30.50s'
+    format_duration(3665)  # Returns: '1h 1m 5.00s'
+    ```
     """
     if seconds < 60:
         return f"{seconds:.2f}s"
@@ -66,22 +69,25 @@ class Timer:
         start_time: Timestamp when timing started
         end_time: Timestamp when timing ended
 
-    Example:
-        >>> with Timer("Training loop"):
-        ...     for i in range(1000):
-        ...         pass
-        Training loop completed in 0.01s
-
-        >>> # Silent timing (no print):
-        >>> with Timer("Silent", verbose=False) as t:
-        ...     time.sleep(0.1)
-        >>> print(f"Elapsed: {t.duration:.2f}s")
-        Elapsed: 0.10s
-
-        >>> # Custom print function:
-        >>> with Timer("Custom", print_fn=lambda msg: print(f">> {msg}")):
-        ...     pass
-        >> Custom completed in 0.00s
+    Examples
+    --------
+    ```python
+    with Timer("Training loop"):
+        for i in range(1000):
+            pass
+    # Output: Training loop completed in 0.01s
+
+    # Silent timing (no print):
+    with Timer("Silent", verbose=False) as t:
+        time.sleep(0.1)
+    print(f"Elapsed: {t.duration:.2f}s")
+    # Output: Elapsed: 0.10s
+
+    # Custom print function:
+    with Timer("Custom", print_fn=lambda msg: print(f">> {msg}")):
+        pass
+    # Output: >> Custom completed in 0.00s
+    ```
     """
 
     def __init__(

{alberta_framework-0.1.0.dist-info → alberta_framework-0.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alberta-framework
-Version: 0.1.0
+Version: 0.1.1
 Summary: Implementation of the Alberta Plan for AI Research - continual learning with meta-learned step-sizes
 Project-URL: Homepage, https://github.com/j-klawson/alberta-framework
 Project-URL: Repository, https://github.com/j-klawson/alberta-framework
@@ -49,7 +49,7 @@ Description-Content-Type: text/markdown
 [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 [![Python 3.13+](https://img.shields.io/badge/python-3.13+-blue.svg)](https://www.python.org/downloads/)
 
-A JAX-based research framework implementing components of [The Alberta Plan](https://arxiv.org/abs/2208.11173) in the pursuit of building the foundations of Continual AI.
+A JAX-based research framework implementing components of [The Alberta Plan for AI Research](https://arxiv.org/abs/2208.11173) in the pursuit of building the foundations of Continual AI.
 
 > "The agents are complex only because they interact with a complex world... their initial design is as simple, general, and scalable as possible." — *Sutton et al., 2022*
 
@@ -57,6 +57,14 @@ A JAX-based research framework implementing components of [The Alberta Plan](htt
 
 The Alberta Framework provides foundational components for continual reinforcement learning research. Built on JAX for hardware acceleration, the framework emphasizes temporal uniformity every component updates at every time step, with no special training phases or batch processing.
 
+## Project Context
+
+This framework is developed as part of my D.Eng. work focusing on the foundations of Continual AI. For more background and context see:
+
+* **Research Blog**: [blog.9600baud.net](https://blog.9600baud.net)
+* **Replicating Sutton '92**: [The Foundation of Step-size Adaptation](https://blog.9600baud.net/sutton92.html)
+* **About the Author**: [Keith Lawson](https://blog.9600baud.net/about.html)
+
 ### Roadmap
 
 Depending on my research trajectory I may or may not implement components required for the plan. The current focus of this framework is the Step 1 Baseline Study, investigating the interaction between adaptive optimizers and online normalization.

alberta_framework-0.1.1.dist-info/RECORD

@@ -0,0 +1,22 @@
+alberta_framework/__init__.py,sha256=LUrsm6WFh5-Mxg78d1G-Qe015nkGgcCDhSw5lf3UkFo,5460
+alberta_framework/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+alberta_framework/core/__init__.py,sha256=PSrC4zSxgm_6YXWEQ80aZaunpbQ58QexxKmDDU-jp6c,522
+alberta_framework/core/learners.py,sha256=dnRQ5B16oGYpamDJIRYzR54ED9bvW0lpa8c_suC6YBA,29879
+alberta_framework/core/normalizers.py,sha256=Z_d3H17qoXh87DE7k41imvWzkVJQ2xQgDUP7GYSNzAY,5903
+alberta_framework/core/optimizers.py,sha256=OefVuDDG1phh1QQIUyVPsQckl41VrpWFG7hY2eqyc64,14585
+alberta_framework/core/types.py,sha256=mtpVEr2qJ0XzZyjOsUdChmS7T7mrXBDMHb-jfkrT9JY,7503
+alberta_framework/streams/__init__.py,sha256=bsDgWjWjotDQHMI2lno3dgk8N14pd-2mYAQpXAtCPx4,2035
+alberta_framework/streams/base.py,sha256=9rJxvUgmzd5u2bRV4vi5PxhUvj39EZTD4bZHo-Ptn-U,2168
+alberta_framework/streams/gymnasium.py,sha256=s733X7aEgy05hcSazjZEhBiJChtEL7uVpxwh0fXBQZA,21980
+alberta_framework/streams/synthetic.py,sha256=4R9GR7Kh0LT7GmGtPhzMJGr8HbhrAMUOjvPwLZk6nDg,32979
+alberta_framework/utils/__init__.py,sha256=zfKfnbikhLp0J6UgVa8HeRo59gZHwqOc8jf03s7AaT4,2845
+alberta_framework/utils/experiments.py,sha256=ekGAzveCRgv9YZ5mfAD5Uf7h_PvQnxsNw2KeZN2eu00,10644
+alberta_framework/utils/export.py,sha256=W9RKfeTiyZcLColOGNjBfZU0N6QMXrfPn4pdYcm-OSk,15832
+alberta_framework/utils/metrics.py,sha256=1cryNJoboO67vvRhausaucbYZFgdL_06vaf08UXbojg,3349
+alberta_framework/utils/statistics.py,sha256=4fbzNlmsdUaM5lLW1BhL5B5MUpnqimQlwJklZ4x0y0U,15416
+alberta_framework/utils/timing.py,sha256=JOLq8CpCAV7LWOWkftxefduSFjaXnVwal1MFBKEMdJI,4049
+alberta_framework/utils/visualization.py,sha256=PmKBD3KGabNhgDizcNiGJEbVCyDL1YMUE5yTwgJHu2o,17924
+alberta_framework-0.1.1.dist-info/METADATA,sha256=Ny-LxHiqZVNXZbu5f8ZyBSLCEZd2KsBhA9iROV7tNiU,7763
+alberta_framework-0.1.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+alberta_framework-0.1.1.dist-info/licenses/LICENSE,sha256=TI1avodt5mvxz7sunyxIa0HlNgLQcmKNLeRjCVcgKmE,10754
+alberta_framework-0.1.1.dist-info/RECORD,,

alberta_framework-0.1.0.dist-info/RECORD

@@ -1,22 +0,0 @@
-alberta_framework/__init__.py,sha256=gPLBA2EiPcElsYp_U_Rs7C6wlrGHr8w5IL6C0F90zec,4739
-alberta_framework/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-alberta_framework/core/__init__.py,sha256=PSrC4zSxgm_6YXWEQ80aZaunpbQ58QexxKmDDU-jp6c,522
-alberta_framework/core/learners.py,sha256=Abq_iOb9CWy9DMWeIX_7PXxo59gXEmlvNnf8zrDKQpo,18157
-alberta_framework/core/normalizers.py,sha256=-OFkdKcfx4VTwm4WLXu1hxrh8DTdwFlIG6CgVYgaZBk,5905
-alberta_framework/core/optimizers.py,sha256=uHEOhE1ThLcXJ18zagW9SQAEiNpsWapNWSMelbhkdNY,14559
-alberta_framework/core/types.py,sha256=Op9EHIIoEZGKbbr3b7xijaOurlQ-mxohBRv7rnVybro,6307
-alberta_framework/streams/__init__.py,sha256=bsDgWjWjotDQHMI2lno3dgk8N14pd-2mYAQpXAtCPx4,2035
-alberta_framework/streams/base.py,sha256=81zqXTF30Orj0N2BXSLYVHF9wUYZSqthqQi1MG5Kzxs,2165
-alberta_framework/streams/gymnasium.py,sha256=s733X7aEgy05hcSazjZEhBiJChtEL7uVpxwh0fXBQZA,21980
-alberta_framework/streams/synthetic.py,sha256=kRQktC4NNlFvoF_FmY_WG9VkiASGOudz8qdI5VjoRq8,32963
-alberta_framework/utils/__init__.py,sha256=zfKfnbikhLp0J6UgVa8HeRo59gZHwqOc8jf03s7AaT4,2845
-alberta_framework/utils/experiments.py,sha256=8N_JrffUa1S_lIZQIqKDuBxyv4UYt9QXzLlo-YnMAEU,10554
-alberta_framework/utils/export.py,sha256=W9RKfeTiyZcLColOGNjBfZU0N6QMXrfPn4pdYcm-OSk,15832
-alberta_framework/utils/metrics.py,sha256=1cryNJoboO67vvRhausaucbYZFgdL_06vaf08UXbojg,3349
-alberta_framework/utils/statistics.py,sha256=4fbzNlmsdUaM5lLW1BhL5B5MUpnqimQlwJklZ4x0y0U,15416
-alberta_framework/utils/timing.py,sha256=05NwXrIc9nS2p2MCHjdOgglPbE1CHZsLLdSB6em7YNY,4110
-alberta_framework/utils/visualization.py,sha256=PmKBD3KGabNhgDizcNiGJEbVCyDL1YMUE5yTwgJHu2o,17924
-alberta_framework-0.1.0.dist-info/METADATA,sha256=G4lIPB7NJlkJGOlmbo7BqeIsHNUl2M-L4f8Gd1T2Ro8,7332
-alberta_framework-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-alberta_framework-0.1.0.dist-info/licenses/LICENSE,sha256=TI1avodt5mvxz7sunyxIa0HlNgLQcmKNLeRjCVcgKmE,10754
-alberta_framework-0.1.0.dist-info/RECORD,,