alberta-framework 0.1.0__tar.gz → 0.2.0__tar.gz

{alberta_framework-0.1.0 → alberta_framework-0.2.0}/CLAUDE.md

@@ -14,10 +14,10 @@ This framework implements Step 1 of the Alberta Plan: demonstrating that IDBD (I
 ```
 src/alberta_framework/
 ├── core/
-│   ├── types.py        # TimeStep, LearnerState, LMSState, IDBDState, AutostepState, StepSizeTrackingConfig, StepSizeHistory
+│   ├── types.py        # TimeStep, LearnerState, LMSState, IDBDState, AutostepState, StepSizeTrackingConfig, StepSizeHistory, NormalizerTrackingConfig, NormalizerHistory, BatchedLearningResult, BatchedNormalizedResult
 │   ├── optimizers.py   # LMS, IDBD, Autostep optimizers
 │   ├── normalizers.py  # OnlineNormalizer, NormalizerState
-│   └── learners.py     # LinearLearner, NormalizedLinearLearner, run_learning_loop, metrics_to_dicts
+│   └── learners.py     # LinearLearner, NormalizedLinearLearner, run_learning_loop, run_learning_loop_batched, run_normalized_learning_loop, run_normalized_learning_loop_batched, metrics_to_dicts
 ├── streams/
 │   ├── base.py         # ScanStream protocol (pure function interface for jax.lax.scan)
 │   ├── synthetic.py    # RandomWalkStream, AbruptChangeStream, CyclicStream, PeriodicChangeStream, ScaledStreamWrapper, DynamicScaleShiftStream, ScaleDriftStream
@@ -125,15 +125,15 @@ IDBD/Autostep should beat LMS when starting from the same step-size (demonstrate
 With optimal parameters, adaptive methods should match best grid-searched LMS.
 
 ### Step-Size Tracking for Meta-Adaptation Analysis
-The `run_learning_loop` function supports optional per-weight step-size tracking for analyzing how adaptive optimizers evolve their step-sizes during training:
+The `run_learning_loop` and `run_normalized_learning_loop` functions support optional per-weight step-size tracking for analyzing how adaptive optimizers evolve their step-sizes during training:
 
 ```python
-from alberta_framework import LinearLearner, IDBD, StepSizeTrackingConfig, run_learning_loop
+from alberta_framework import LinearLearner, IDBD, Autostep, StepSizeTrackingConfig, run_learning_loop
 from alberta_framework.streams import RandomWalkStream
 import jax.random as jr
 
 stream = RandomWalkStream(feature_dim=10)
-learner = LinearLearner(optimizer=IDBD())
+learner = LinearLearner(optimizer=Autostep())
 config = StepSizeTrackingConfig(interval=100)  # Record every 100 steps
 
 state, metrics, history = run_learning_loop(
@@ -143,6 +143,7 @@ state, metrics, history = run_learning_loop(
 # history.step_sizes: shape (100, 10) - per-weight step-sizes at each recording
 # history.bias_step_sizes: shape (100,) - bias step-size at each recording
 # history.recording_indices: shape (100,) - step indices where recordings were made
+# history.normalizers: shape (100, 10) - Autostep's v_i normalizers (None for IDBD/LMS)
 ```
 
 Key features:
@@ -150,6 +151,94 @@ Key features:
 - Configurable interval to control memory usage
 - Optional `include_bias=False` to skip bias tracking
 - Works with LMS (constant), IDBD, and Autostep optimizers
+- **Autostep's normalizers (v_i)** are tracked automatically when using Autostep
+
+### Normalizer State Tracking for Reactive Lag Analysis
+The `run_normalized_learning_loop` function supports tracking the OnlineNormalizer's per-feature mean and variance estimates over time. This is essential for analyzing reactive lag — how quickly the normalizer adapts to distribution shifts:
+
+```python
+from alberta_framework import (
+    NormalizedLinearLearner, IDBD,
+    StepSizeTrackingConfig, NormalizerTrackingConfig,
+    run_normalized_learning_loop
+)
+from alberta_framework.streams import RandomWalkStream
+import jax.random as jr
+
+stream = RandomWalkStream(feature_dim=10)
+learner = NormalizedLinearLearner(optimizer=IDBD())
+ss_config = StepSizeTrackingConfig(interval=100)
+norm_config = NormalizerTrackingConfig(interval=100)
+
+# Track both step-sizes and normalizer state
+state, metrics, ss_history, norm_history = run_normalized_learning_loop(
+    learner, stream, num_steps=10000, key=jr.key(42),
+    step_size_tracking=ss_config, normalizer_tracking=norm_config
+)
+
+# norm_history.means: shape (100, 10) - per-feature mean estimates at each recording
+# norm_history.variances: shape (100, 10) - per-feature variance estimates at each recording
+# norm_history.recording_indices: shape (100,) - step indices where recordings were made
+```
+
+Return value depends on tracking options:
+- No tracking: `(state, metrics)` — 2-tuple
+- step_size_tracking only: `(state, metrics, ss_history)` — 3-tuple
+- normalizer_tracking only: `(state, metrics, norm_history)` — 3-tuple
+- Both: `(state, metrics, ss_history, norm_history)` — 4-tuple
+
+### Batched Learning Loops (vmap-based GPU Parallelization)
+The `run_learning_loop_batched` and `run_normalized_learning_loop_batched` functions use `jax.vmap` to run multiple seeds in parallel, typically achieving 2-5x speedup over sequential execution:
+
+```python
+import jax.random as jr
+from alberta_framework import (
+    LinearLearner, IDBD, RandomWalkStream,
+    run_learning_loop_batched, StepSizeTrackingConfig
+)
+
+stream = RandomWalkStream(feature_dim=10)
+learner = LinearLearner(optimizer=IDBD())
+
+# Run 30 seeds in parallel
+keys = jr.split(jr.key(42), 30)
+result = run_learning_loop_batched(learner, stream, num_steps=10000, keys=keys)
+
+# result.metrics has shape (30, 10000, 3)
+# result.states.weights has shape (30, 10)
+mean_error = result.metrics[:, :, 0].mean(axis=0)  # Average squared error over seeds
+
+# With step-size tracking
+config = StepSizeTrackingConfig(interval=100)
+result = run_learning_loop_batched(
+    learner, stream, num_steps=10000, keys=keys, step_size_tracking=config
+)
+# result.step_size_history.step_sizes has shape (30, 100, 10)
+```
+
+Key features:
+- `jax.vmap` parallelizes over seeds, not steps — memory scales with num_seeds
+- `jax.lax.scan` processes steps sequentially within each seed
+- Returns `BatchedLearningResult` or `BatchedNormalizedResult` NamedTuples
+- Tracking histories get batched shapes: `(num_seeds, num_recordings, ...)`
+- Same initial state used for all seeds (controlled variation via different keys)
+
+For normalized learners:
+```python
+from alberta_framework import (
+    NormalizedLinearLearner, run_normalized_learning_loop_batched,
+    NormalizerTrackingConfig
+)
+
+learner = NormalizedLinearLearner(optimizer=IDBD())
+result = run_normalized_learning_loop_batched(
+    learner, stream, num_steps=10000, keys=keys,
+    step_size_tracking=StepSizeTrackingConfig(interval=100),
+    normalizer_tracking=NormalizerTrackingConfig(interval=100)
+)
+# result.metrics has shape (30, 10000, 4)
+# result.step_size_history and result.normalizer_history both batched
+```
 
 ## Gymnasium Integration
 
@@ -267,6 +356,10 @@ The API Reference section is auto-generated from docstrings in the source code.
 ### Docstring Style
 Use NumPy-style docstrings for all public functions and classes. See `core/optimizers.py` for examples.
 
+**Code examples**: Use fenced markdown code blocks (triple backticks with `python`) inside an `Examples` section, not doctest `>>>` format. This ensures proper syntax highlighting in mkdocstrings. See `streams/base.py` or `utils/timing.py` for examples.
+
+**Math formulas**: Wrap inline math expressions in backticks for monospace rendering, e.g., `` `y = w @ x + b` `` or `` `alpha_i = exp(log_alpha_i)` ``. See `core/optimizers.py` for examples.
+
 ## Streams for Factorial Studies
 
 The framework supports factorial experiment designs with multiple non-stationarity types and scale ranges:

{alberta_framework-0.1.0 → alberta_framework-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alberta-framework
-Version: 0.1.0
+Version: 0.2.0
 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.0 → alberta_framework-0.2.0}/README.md

@@ -5,7 +5,7 @@
 [![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*
 
@@ -13,6 +13,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.0 → alberta_framework-0.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "alberta-framework"
-version = "0.1.0"
+version = "0.2.0"
 description = "Implementation of the Alberta Plan for AI Research - continual learning with meta-learned step-sizes"
 readme = "README.md"
 license = "Apache-2.0"

{alberta_framework-0.1.0 → alberta_framework-0.2.0}/src/alberta_framework/__init__.py

@@ -1,28 +1,45 @@
-"""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"
+__version__ = "0.2.0"
 
 # Core types
 # Learners
@@ -33,7 +50,9 @@ from alberta_framework.core.learners import (
     UpdateResult,
     metrics_to_dicts,
     run_learning_loop,
+    run_learning_loop_batched,
     run_normalized_learning_loop,
+    run_normalized_learning_loop_batched,
 )
 
 # Normalizers
@@ -47,9 +66,13 @@ from alberta_framework.core.normalizers import (
 from alberta_framework.core.optimizers import IDBD, LMS, Autostep, Optimizer
 from alberta_framework.core.types import (
     AutostepState,
+    BatchedLearningResult,
+    BatchedNormalizedResult,
     IDBDState,
     LearnerState,
     LMSState,
+    NormalizerHistory,
+    NormalizerTrackingConfig,
     Observation,
     Prediction,
     StepSizeHistory,
@@ -119,10 +142,14 @@ __all__ = [
     "__version__",
     # Types
     "AutostepState",
+    "BatchedLearningResult",
+    "BatchedNormalizedResult",
     "IDBDState",
     "LMSState",
     "LearnerState",
+    "NormalizerHistory",
     "NormalizerState",
+    "NormalizerTrackingConfig",
     "Observation",
     "Prediction",
     "StepSizeHistory",
@@ -143,7 +170,9 @@ __all__ = [
     "NormalizedLearnerState",
     "NormalizedLinearLearner",
     "run_learning_loop",
+    "run_learning_loop_batched",
     "run_normalized_learning_loop",
+    "run_normalized_learning_loop_batched",
     "metrics_to_dicts",
     # Streams - protocol
     "ScanStream",