alberta-framework 0.3.0__tar.gz → 0.4.0__tar.gz

{alberta_framework-0.3.0 → alberta_framework-0.4.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, NormalizerTrackingConfig, NormalizerHistory, BatchedLearningResult, BatchedNormalizedResult
-│   ├── optimizers.py   # LMS, IDBD, Autostep optimizers
+│   ├── types.py        # TimeStep, LearnerState, optimizer states, TDTimeStep, TDLearnerState, TDIDBDState, AutoTDIDBDState
+│   ├── optimizers.py   # LMS, IDBD, Autostep, TDIDBD, AutoTDIDBD optimizers
 │   ├── normalizers.py  # OnlineNormalizer, NormalizerState
-│   └── learners.py     # LinearLearner, NormalizedLinearLearner, run_learning_loop, run_learning_loop_batched, run_normalized_learning_loop, run_normalized_learning_loop_batched, metrics_to_dicts
+│   └── learners.py     # LinearLearner, TDLinearLearner, run_learning_loop, run_td_learning_loop
 ├── streams/
 │   ├── base.py         # ScanStream protocol (pure function interface for jax.lax.scan)
 │   ├── synthetic.py    # RandomWalkStream, AbruptChangeStream, CyclicStream, PeriodicChangeStream, ScaledStreamWrapper, DynamicScaleShiftStream, ScaleDriftStream
@@ -470,6 +470,20 @@ The publish workflow uses OpenID Connect (no API tokens). Configure on PyPI:
 
 ## Changelog
 
+### v0.4.0 (2026-02-04)
+- **FEATURE**: Implemented TD-IDBD optimizer for temporal-difference learning with per-weight adaptive step-sizes and eligibility traces (Kearney et al., 2019)
+- **FEATURE**: Implemented AutoTDIDBD optimizer with AutoStep-style normalization for improved stability
+- **FEATURE**: Added `TDLinearLearner` class for linear value function approximation in TD learning
+- **FEATURE**: Added `run_td_learning_loop()` for JIT-compiled TD learning via `jax.lax.scan`
+- **FEATURE**: Added TD state types: `TDIDBDState`, `AutoTDIDBDState`, `TDLearnerState`, `TDTimeStep`
+- **FEATURE**: Added `TDStream` protocol for TD experience streams
+- **DOCS**: Updated README with TD learning documentation and Kearney et al. 2019 reference
+
+### v0.3.2 (2026-02-03)
+- **FIX**: Relaxed test tolerance in batched vs sequential comparison tests (`rtol=1e-5`) to account for floating-point differences between vmap and sequential execution paths
+- **FIX**: Added `ignore = ["F722"]` to ruff config for jaxtyping shape annotation syntax that ruff doesn't understand
+- **FIX**: Removed unused `PRNGKeyArray` import from `core/types.py`
+
 ### v0.3.0 (2026-02-03)
 - **FEATURE**: Migrated all state types from NamedTuple to `@chex.dataclass(frozen=True)` for DeepMind-style JAX compatibility
 - **FEATURE**: Added jaxtyping shape annotations for compile-time type safety (`Float[Array, " feature_dim"]`, `PRNGKeyArray`, etc.)

{alberta_framework-0.3.0 → alberta_framework-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alberta-framework
-Version: 0.3.0
+Version: 0.4.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
@@ -113,10 +113,15 @@ state, metrics = run_learning_loop(learner, stream, num_steps=10000, key=jr.key(
 
 ### Optimizers
 
+**Supervised Learning:**
 - **LMS**: Fixed step-size baseline
 - **IDBD**: Per-weight adaptive step-sizes via gradient correlation (Sutton, 1992)
 - **Autostep**: Tuning-free adaptation with gradient normalization (Mahmood et al., 2012)
 
+**TD Learning:**
+- **TDIDBD**: TD learning with per-weight adaptive step-sizes and eligibility traces (Kearney et al., 2019)
+- **AutoTDIDBD**: TD learning with AutoStep-style normalization for improved stability
+
 ### Streams
 
 Non-stationary experience generators implementing the `ScanStream` protocol:
@@ -126,6 +131,17 @@ Non-stationary experience generators implementing the `ScanStream` protocol:
 - `PeriodicChangeStream`: Sinusoidal oscillation
 - `DynamicScaleShiftStream`: Time-varying feature scales
 
+### TD Learning
+
+For temporal-difference learning with value function approximation:
+
+```python
+from alberta_framework import TDLinearLearner, TDIDBD, run_td_learning_loop
+
+learner = TDLinearLearner(optimizer=TDIDBD(trace_decay=0.9))
+state, metrics = run_td_learning_loop(learner, td_stream, num_steps=10000, key=jr.key(42))
+```
+
 ### Gymnasium Integration
 
 ```python
@@ -202,6 +218,13 @@ If you use this framework in your research, please cite:
   booktitle = {IEEE International Conference on Acoustics, Speech and Signal Processing},
   year = {2012}
 }
+
+@inproceedings{kearney2019tidbd,
+  title = {Learning Feature Relevance Through Step Size Adaptation in Temporal-Difference Learning},
+  author = {Kearney, Alex and Veeriah, Vivek and Travnik, Jaden and Sutton, Richard S. and Pilarski, Patrick M.},
+  booktitle = {International Conference on Machine Learning},
+  year = {2019}
+}
 ```
 
 ## License

{alberta_framework-0.3.0 → alberta_framework-0.4.0}/README.md

@@ -66,10 +66,15 @@ state, metrics = run_learning_loop(learner, stream, num_steps=10000, key=jr.key(
 
 ### Optimizers
 
+**Supervised Learning:**
 - **LMS**: Fixed step-size baseline
 - **IDBD**: Per-weight adaptive step-sizes via gradient correlation (Sutton, 1992)
 - **Autostep**: Tuning-free adaptation with gradient normalization (Mahmood et al., 2012)
 
+**TD Learning:**
+- **TDIDBD**: TD learning with per-weight adaptive step-sizes and eligibility traces (Kearney et al., 2019)
+- **AutoTDIDBD**: TD learning with AutoStep-style normalization for improved stability
+
 ### Streams
 
 Non-stationary experience generators implementing the `ScanStream` protocol:
@@ -79,6 +84,17 @@ Non-stationary experience generators implementing the `ScanStream` protocol:
 - `PeriodicChangeStream`: Sinusoidal oscillation
 - `DynamicScaleShiftStream`: Time-varying feature scales
 
+### TD Learning
+
+For temporal-difference learning with value function approximation:
+
+```python
+from alberta_framework import TDLinearLearner, TDIDBD, run_td_learning_loop
+
+learner = TDLinearLearner(optimizer=TDIDBD(trace_decay=0.9))
+state, metrics = run_td_learning_loop(learner, td_stream, num_steps=10000, key=jr.key(42))
+```
+
 ### Gymnasium Integration
 
 ```python
@@ -155,6 +171,13 @@ If you use this framework in your research, please cite:
   booktitle = {IEEE International Conference on Acoustics, Speech and Signal Processing},
   year = {2012}
 }
+
+@inproceedings{kearney2019tidbd,
+  title = {Learning Feature Relevance Through Step Size Adaptation in Temporal-Difference Learning},
+  author = {Kearney, Alex and Veeriah, Vivek and Travnik, Jaden and Sutton, Richard S. and Pilarski, Patrick M.},
+  booktitle = {International Conference on Machine Learning},
+  year = {2019}
+}
 ```
 
 ## License

{alberta_framework-0.3.0 → alberta_framework-0.4.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "alberta-framework"
-version = "0.3.0"
+version = "0.4.0"
 description = "Implementation of the Alberta Plan for AI Research - continual learning with meta-learned step-sizes"
 readme = "README.md"
 license = "Apache-2.0"
@@ -71,6 +71,8 @@ target-version = "py313"
 
 [tool.ruff.lint]
 select = ["E", "F", "I", "N", "W", "UP"]
+# F722: Syntax error in forward annotation - ruff doesn't understand jaxtyping shape annotations
+ignore = ["F722"]
 
 [tool.mypy]
 python_version = "3.13"

{alberta_framework-0.3.0 → alberta_framework-0.4.0}/src/alberta_framework/__init__.py

@@ -39,7 +39,7 @@ References
 - Tuning-free Step-size Adaptation (Mahmood et al., 2012)
 """
 
-__version__ = "0.2.0"
+__version__ = "0.4.0"
 
 # Core types
 # Learners
@@ -47,12 +47,15 @@ from alberta_framework.core.learners import (
     LinearLearner,
     NormalizedLearnerState,
     NormalizedLinearLearner,
+    TDLinearLearner,
+    TDUpdateResult,
     UpdateResult,
     metrics_to_dicts,
     run_learning_loop,
     run_learning_loop_batched,
     run_normalized_learning_loop,
     run_normalized_learning_loop_batched,
+    run_td_learning_loop,
 )
 
 # Normalizers
@@ -63,9 +66,19 @@ from alberta_framework.core.normalizers import (
 )
 
 # Optimizers
-from alberta_framework.core.optimizers import IDBD, LMS, Autostep, Optimizer
+from alberta_framework.core.optimizers import (
+    IDBD,
+    LMS,
+    TDIDBD,
+    Autostep,
+    AutoTDIDBD,
+    Optimizer,
+    TDOptimizer,
+    TDOptimizerUpdate,
+)
 from alberta_framework.core.types import (
     AutostepState,
+    AutoTDIDBDState,
     BatchedLearningResult,
     BatchedNormalizedResult,
     IDBDState,
@@ -78,7 +91,12 @@ from alberta_framework.core.types import (
     StepSizeHistory,
     StepSizeTrackingConfig,
     Target,
+    TDIDBDState,
+    TDLearnerState,
+    TDTimeStep,
     TimeStep,
+    create_autotdidbd_state,
+    create_tdidbd_state,
 )
 
 # Streams - base
@@ -140,7 +158,7 @@ except ImportError:
 __all__ = [
     # Version
     "__version__",
-    # Types
+    # Types - Supervised Learning
     "AutostepState",
     "BatchedLearningResult",
     "BatchedNormalizedResult",
@@ -157,15 +175,28 @@ __all__ = [
     "Target",
     "TimeStep",
     "UpdateResult",
-    # Optimizers
+    # Types - TD Learning
+    "AutoTDIDBDState",
+    "TDIDBDState",
+    "TDLearnerState",
+    "TDTimeStep",
+    "TDUpdateResult",
+    "create_tdidbd_state",
+    "create_autotdidbd_state",
+    # Optimizers - Supervised Learning
     "Autostep",
     "IDBD",
     "LMS",
     "Optimizer",
+    # Optimizers - TD Learning
+    "AutoTDIDBD",
+    "TDIDBD",
+    "TDOptimizer",
+    "TDOptimizerUpdate",
     # Normalizers
     "OnlineNormalizer",
     "create_normalizer_state",
-    # Learners
+    # Learners - Supervised Learning
     "LinearLearner",
     "NormalizedLearnerState",
     "NormalizedLinearLearner",
@@ -174,6 +205,9 @@ __all__ = [
     "run_normalized_learning_loop",
     "run_normalized_learning_loop_batched",
     "metrics_to_dicts",
+    # Learners - TD Learning
+    "TDLinearLearner",
+    "run_td_learning_loop",
     # Streams - protocol
     "ScanStream",
     # Streams - synthetic

alberta_framework-0.4.0/src/alberta_framework/core/__init__.py

@@ -0,0 +1,51 @@
+"""Core components for the Alberta Framework."""
+
+from alberta_framework.core.learners import LinearLearner, TDLinearLearner, TDUpdateResult
+from alberta_framework.core.optimizers import (
+    IDBD,
+    LMS,
+    TDIDBD,
+    AutoTDIDBD,
+    Optimizer,
+    TDOptimizer,
+    TDOptimizerUpdate,
+)
+from alberta_framework.core.types import (
+    AutoTDIDBDState,
+    IDBDState,
+    LearnerState,
+    LMSState,
+    Observation,
+    Prediction,
+    Target,
+    TDIDBDState,
+    TDLearnerState,
+    TDTimeStep,
+    TimeStep,
+)
+
+__all__ = [
+    # Supervised learning
+    "IDBD",
+    "IDBDState",
+    "LMS",
+    "LMSState",
+    "LearnerState",
+    "LinearLearner",
+    "Observation",
+    "Optimizer",
+    "Prediction",
+    "Target",
+    "TimeStep",
+    # TD learning
+    "AutoTDIDBD",
+    "AutoTDIDBDState",
+    "TDIDBD",
+    "TDIDBDState",
+    "TDLearnerState",
+    "TDLinearLearner",
+    "TDOptimizer",
+    "TDOptimizerUpdate",
+    "TDTimeStep",
+    "TDUpdateResult",
+]