alberta-framework 0.3.2__py3-none-any.whl → 0.4.0__py3-none-any.whl

alberta_framework/streams/gymnasium.py

@@ -117,9 +117,7 @@ def _flatten_action(action: Any, space: gymnasium.spaces.Space[Any]) -> Array:
         raise ValueError(f"Unsupported space type: {type(space).__name__}")
 
 
-def make_random_policy(
-    env: gymnasium.Env[Any, Any], seed: int = 0
-) -> Callable[[Array], Any]:
+def make_random_policy(env: gymnasium.Env[Any, Any], seed: int = 0) -> Callable[[Array], Any]:
     """Create a random action policy for an environment.
 
     Args:
@@ -147,10 +145,7 @@ def make_random_policy(
             return jr.uniform(key, action_space.shape, minval=low, maxval=high)
         elif isinstance(action_space, gymnasium.spaces.MultiDiscrete):
             nvec = action_space.nvec
-            return [
-                int(jr.randint(jr.fold_in(key, i), (), 0, n))
-                for i, n in enumerate(nvec)
-            ]
+            return [int(jr.randint(jr.fold_in(key, i), (), 0, n)) for i, n in enumerate(nvec)]
         else:
             raise ValueError(f"Unsupported action space: {type(action_space).__name__}")
 
@@ -284,9 +279,7 @@ def learn_from_trajectory(
     if learner_state is None:
         learner_state = learner.init(observations.shape[1])
 
-    def step_fn(
-        state: LearnerState, inputs: tuple[Array, Array]
-    ) -> tuple[LearnerState, Array]:
+    def step_fn(state: LearnerState, inputs: tuple[Array, Array]) -> tuple[LearnerState, Array]:
         obs, target = inputs
         result = learner.update(state, obs, target)
         return result.state, result.metrics

alberta_framework/streams/synthetic.py

@@ -540,9 +540,7 @@ class PeriodicChangeStream:
             step_count=jnp.array(0, dtype=jnp.int32),
         )
 
-    def step(
-        self, state: PeriodicChangeState, idx: Array
-    ) -> tuple[TimeStep, PeriodicChangeState]:
+    def step(self, state: PeriodicChangeState, idx: Array) -> tuple[TimeStep, PeriodicChangeState]:
         """Generate one time step.
 
         Args:
@@ -557,9 +555,7 @@ class PeriodicChangeStream:
 
         # Compute oscillating weights: w(t) = base + amplitude * sin(2π * t / period + phase)
         t = state.step_count.astype(jnp.float32)
-        oscillation = self._amplitude * jnp.sin(
-            2.0 * jnp.pi * t / self._period + state.phases
-        )
+        oscillation = self._amplitude * jnp.sin(2.0 * jnp.pi * t / self._period + state.phases)
         true_weights = state.base_weights + oscillation
 
         # Generate observation
@@ -955,9 +951,7 @@ class ScaleDriftStream:
             step_count=jnp.array(0, dtype=jnp.int32),
         )
 
-    def step(
-        self, state: ScaleDriftState, idx: Array
-    ) -> tuple[TimeStep, ScaleDriftState]:
+    def step(self, state: ScaleDriftState, idx: Array) -> tuple[TimeStep, ScaleDriftState]:
         """Generate one time step.
 
         Args:

alberta_framework/utils/experiments.py

@@ -110,9 +110,7 @@ def run_single_experiment(
 
     final_state: LearnerState | NormalizedLearnerState
     if isinstance(learner, NormalizedLinearLearner):
-        norm_result = run_normalized_learning_loop(
-            learner, stream, config.num_steps, key
-        )
+        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:

alberta_framework/utils/export.py

@@ -51,14 +51,16 @@ def _export_summary_csv(
 
         for name, agg in results.items():
             summary = agg.summary[metric]
-            writer.writerow([
-                name,
-                f"{summary.mean:.6f}",
-                f"{summary.std:.6f}",
-                f"{summary.min:.6f}",
-                f"{summary.max:.6f}",
-                summary.n_seeds,
-            ])
+            writer.writerow(
+                [
+                    name,
+                    f"{summary.mean:.6f}",
+                    f"{summary.std:.6f}",
+                    f"{summary.min:.6f}",
+                    f"{summary.max:.6f}",
+                    summary.n_seeds,
+                ]
+            )
 
 
 def _export_timeseries_csv(
@@ -497,13 +499,15 @@ def results_to_dataframe(
     rows = []
     for name, agg in results.items():
         summary = agg.summary[metric]
-        rows.append({
-            "method": name,
-            "mean": summary.mean,
-            "std": summary.std,
-            "min": summary.min,
-            "max": summary.max,
-            "n_seeds": summary.n_seeds,
-        })
+        rows.append(
+            {
+                "method": name,
+                "mean": summary.mean,
+                "std": summary.std,
+                "min": summary.min,
+                "max": summary.max,
+                "n_seeds": summary.n_seeds,
+            }
+        )
 
     return pd.DataFrame(rows)

alberta_framework/utils/statistics.py

@@ -313,9 +313,7 @@ def wilcoxon_comparison(
         stat_val = float(result[0])
         p_val = float(result[1])
     except ImportError:
-        raise ImportError(
-            "scipy is required for Wilcoxon test. Install with: pip install scipy"
-        )
+        raise ImportError("scipy is required for Wilcoxon test. Install with: pip install scipy")
 
     effect = cohens_d(a, b)
 
@@ -443,18 +441,28 @@ def pairwise_comparisons(
 
             if test == "ttest":
                 result = ttest_comparison(
-                    values_a, values_b, paired=True, alpha=alpha,
-                    method_a=name_a, method_b=name_b,
+                    values_a,
+                    values_b,
+                    paired=True,
+                    alpha=alpha,
+                    method_a=name_a,
+                    method_b=name_b,
                 )
             elif test == "mann_whitney":
                 result = mann_whitney_comparison(
-                    values_a, values_b, alpha=alpha,
-                    method_a=name_a, method_b=name_b,
+                    values_a,
+                    values_b,
+                    alpha=alpha,
+                    method_a=name_a,
+                    method_b=name_b,
                 )
             else:  # wilcoxon
                 result = wilcoxon_comparison(
-                    values_a, values_b, alpha=alpha,
-                    method_a=name_a, method_b=name_b,
+                    values_a,
+                    values_b,
+                    alpha=alpha,
+                    method_a=name_a,
+                    method_b=name_b,
                 )
 
             comparisons[(name_a, name_b)] = result

alberta_framework/utils/visualization.py

@@ -69,29 +69,33 @@ def set_publication_style(
         pass
 
     # Configure matplotlib
-    plt.rcParams.update({
-        "font.size": font_size,
-        "axes.labelsize": font_size,
-        "axes.titlesize": font_size + 1,
-        "xtick.labelsize": font_size - 1,
-        "ytick.labelsize": font_size - 1,
-        "legend.fontsize": font_size - 1,
-        "figure.figsize": (_current_style["figure_width"], _current_style["figure_height"]),
-        "figure.dpi": _current_style["dpi"],
-        "savefig.dpi": _current_style["dpi"],
-        "lines.linewidth": _current_style["line_width"],
-        "lines.markersize": _current_style["marker_size"],
-        "axes.linewidth": 0.8,
-        "grid.linewidth": 0.5,
-        "grid.alpha": 0.3,
-    })
+    plt.rcParams.update(
+        {
+            "font.size": font_size,
+            "axes.labelsize": font_size,
+            "axes.titlesize": font_size + 1,
+            "xtick.labelsize": font_size - 1,
+            "ytick.labelsize": font_size - 1,
+            "legend.fontsize": font_size - 1,
+            "figure.figsize": (_current_style["figure_width"], _current_style["figure_height"]),
+            "figure.dpi": _current_style["dpi"],
+            "savefig.dpi": _current_style["dpi"],
+            "lines.linewidth": _current_style["line_width"],
+            "lines.markersize": _current_style["marker_size"],
+            "axes.linewidth": 0.8,
+            "grid.linewidth": 0.5,
+            "grid.alpha": 0.3,
+        }
+    )
 
     if use_latex:
-        plt.rcParams.update({
-            "text.usetex": True,
-            "font.family": "serif",
-            "font.serif": ["Computer Modern Roman"],
-        })
+        plt.rcParams.update(
+            {
+                "text.usetex": True,
+                "font.family": "serif",
+                "font.serif": ["Computer Modern Roman"],
+            }
+        )
 
 
 def plot_learning_curves(
@@ -142,10 +146,12 @@ def plot_learning_curves(
         metric_array = agg.metric_arrays[metric]
 
         # Smooth each seed individually, then compute statistics
-        smoothed = np.array([
-            compute_running_mean(metric_array[seed_idx], window_size)
-            for seed_idx in range(metric_array.shape[0])
-        ])
+        smoothed = np.array(
+            [
+                compute_running_mean(metric_array[seed_idx], window_size)
+                for seed_idx in range(metric_array.shape[0])
+            ]
+        )
 
         mean, ci_lower, ci_upper = compute_timeseries_statistics(smoothed)
 

{alberta_framework-0.3.2.dist-info → alberta_framework-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alberta-framework
-Version: 0.3.2
+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.4.0.dist-info/RECORD

@@ -0,0 +1,22 @@
+alberta_framework/__init__.py,sha256=RB8-ly8UK6IGnDX8Qw3mW_uSJc8iEJT57CCXg6cxj4c,6451
+alberta_framework/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+alberta_framework/core/__init__.py,sha256=wr7VZrbXP0SWw-_MEFn0rfn1wQrUaDkIm1tXFYhjz1k,961
+alberta_framework/core/learners.py,sha256=lePBbeDReGAg73TvpLPMQNI2k5H6OQfECnvi2qTYT-I,45277
+alberta_framework/core/normalizers.py,sha256=GZkmFbRI3lk7HvNMSX9ByOvXCj_3QX_5h_k6-Y35IqY,5893
+alberta_framework/core/optimizers.py,sha256=9fdic6h-vxBm7BXox8QSLhtzEEkVc7VSHn3sufuIMZY,34588
+alberta_framework/core/types.py,sha256=SZr16VQAemJCXIlBoOlp91rqDeWM5rd831Vya7G2ths,16430
+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=3Kg8qORzvNqRkTceQ7THfm3kp3_Skbva1XbtCDBTsT4,21914
+alberta_framework/streams/synthetic.py,sha256=8e5EY3rtiJhdQbLlWyalNE3nRHhn_5T2Z_aHRS4BpG4,33457
+alberta_framework/utils/__init__.py,sha256=zfKfnbikhLp0J6UgVa8HeRo59gZHwqOc8jf03s7AaT4,2845
+alberta_framework/utils/experiments.py,sha256=vxcbCxUloWu2J2mKHjdkM6cLeY9EYIq1JvpR-gyfPwQ,10622
+alberta_framework/utils/export.py,sha256=vGsBTFcr84Ga8Ka0IZFMVqhMUNOCQstVcfyU468V3Cs,15940
+alberta_framework/utils/metrics.py,sha256=1cryNJoboO67vvRhausaucbYZFgdL_06vaf08UXbojg,3349
+alberta_framework/utils/statistics.py,sha256=QZsDVVNqhiY6chGFLzYmtiUVBIHeBfr_LDTuSyBQROY,15594
+alberta_framework/utils/timing.py,sha256=JOLq8CpCAV7LWOWkftxefduSFjaXnVwal1MFBKEMdJI,4049
+alberta_framework/utils/visualization.py,sha256=aQc4PsWGFCycm0uPvaChFjaoWgBOsD7UHOydWF0WKFo,18070
+alberta_framework-0.4.0.dist-info/METADATA,sha256=7jnUWrT9qFNC6UFZ6r9gxpXUsIO1MyeSsJQFNdBgc1g,8769
+alberta_framework-0.4.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+alberta_framework-0.4.0.dist-info/licenses/LICENSE,sha256=TI1avodt5mvxz7sunyxIa0HlNgLQcmKNLeRjCVcgKmE,10754
+alberta_framework-0.4.0.dist-info/RECORD,,

alberta_framework-0.3.2.dist-info/RECORD

@@ -1,22 +0,0 @@
-alberta_framework/__init__.py,sha256=gAafDDmkivDdfnvDVff9zbVY9ilzqqfJ9KvpbRegKqs,5726
-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=gUhX7caXBfpWYgnvYTp5YKXfP6wbzB2T2gkSMMtrHDQ,38042
-alberta_framework/core/normalizers.py,sha256=QmKmha-mFgKi1KD-f8xuB2U175yQL6Ll0D4c8OONIl0,5927
-alberta_framework/core/optimizers.py,sha256=a4gYac5DyXReir9ycudRg8uQ9b53uLWTIldZ1A3Ae5c,14646
-alberta_framework/core/types.py,sha256=XBmT689nRKEBwwfUbpohi4IfT-d-eJRIFH_L2swYW2E,9793
-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=8njzQCFRi_iVgdPA3slyn46vFIHHkIwaZsABZyPwqnU,33507
-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.3.2.dist-info/METADATA,sha256=aD7q4wh1xm0pQiARtRnUrgLU83JQ8JBidzK-bXmn5_s,7872
-alberta_framework-0.3.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-alberta_framework-0.3.2.dist-info/licenses/LICENSE,sha256=TI1avodt5mvxz7sunyxIa0HlNgLQcmKNLeRjCVcgKmE,10754
-alberta_framework-0.3.2.dist-info/RECORD,,