alberta-framework 0.1.0__tar.gz

alberta_framework-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,65 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python 3.13
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.13"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[dev,gymnasium]"
+
+      - name: Run tests
+        run: pytest --cov=alberta_framework --cov-report=term-missing
+
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python 3.13
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.13"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[dev]"
+
+      - name: Run ruff
+        run: ruff check src/
+
+      - name: Run mypy
+        run: mypy src/
+
+  docs:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python 3.13
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.13"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[docs]"
+
+      - name: Build docs
+        run: mkdocs build --strict

alberta_framework-0.1.0/.github/workflows/docs.yml

@@ -0,0 +1,50 @@
+name: Deploy Docs
+
+on:
+  push:
+    branches: [main]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+  pages: write
+  id-token: write
+
+concurrency:
+  group: "pages"
+  cancel-in-progress: false
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python 3.13
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.13"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[docs]"
+
+      - name: Build docs
+        run: mkdocs build --strict
+
+      - name: Upload artifact
+        uses: actions/upload-pages-artifact@v3
+        with:
+          path: site/
+
+  deploy:
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    runs-on: ubuntu-latest
+    needs: build
+    steps:
+      - name: Deploy to GitHub Pages
+        id: deployment
+        uses: actions/deploy-pages@v4

alberta_framework-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,62 @@
+name: Publish to PyPI
+
+on:
+  push:
+    tags:
+      - "v*"
+
+permissions:
+  id-token: write
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python 3.13
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.13"
+
+      - name: Install build dependencies
+        run: python -m pip install --upgrade pip build
+
+      - name: Build package
+        run: python -m build
+
+      - name: Upload dist artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish-testpypi:
+    runs-on: ubuntu-latest
+    needs: build
+    environment: testpypi
+    steps:
+      - name: Download dist artifacts
+        uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+      - name: Publish to TestPyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          repository-url: https://test.pypi.org/legacy/
+
+  publish-pypi:
+    runs-on: ubuntu-latest
+    needs: publish-testpypi
+    environment: pypi
+    steps:
+      - name: Download dist artifacts
+        uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

alberta_framework-0.1.0/.gitignore

@@ -0,0 +1,107 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# IDEs and editors
+.idea/
+.vscode/
+*.swp
+*.swo
+*~
+.project
+.pydevproject
+.settings/
+*.sublime-project
+*.sublime-workspace
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Ruff
+.ruff_cache/
+
+# JAX compilation cache
+__jax_cache__/
+
+# macOS
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# Thumbnails
+._*
+
+# Local development
+*.local
+.python-version
+
+# Documentation builds
+docs/_build/
+site/
+
+# Experiment outputs
+outputs/
+results/
+logs/
+*.png
+*.pdf
+!docs/*.png
+!docs/*.pdf

alberta_framework-0.1.0/ALBERTA_PLAN.md

@@ -0,0 +1,37 @@
+<role>
+You are a Senior Research Engineer specializing in Reinforcement Learning and Functional Programming. Your expertise is in JAX, the Alberta Plan for AI Research (Sutton et al., 2022) found in papers/The Alberta Plan.pdf, and the mathematical implementation of Temporal Difference (TD) learning.
+</role>
+
+<context>
+I am building a research framework called "Alberta Framework" that can be used by researchers to implement and conduct experiments as they work on the 12 steps of the Alberta Plan. The goal is to move from linear TD(λ) with IDBD meta-learning (Step 1) to a full OaK (Option-and-Knowledge) architecture (Step 11).
+The framework should serve both as a tool for implementing new algorithms to work on the Alberta Plan as well as reference implementations to learn from and test with.
+Key Research Principles:
+1. Temporal Uniformity: Every component must update at every time step.
+2. Continual Learning: No experience replay; the agent learns from a non-stationary stream.
+3. Functional JAX: State is immutable; updates are pure functions using jax.jit, vmap, and grad.
+</context>
+
+<current_milestone>
+Step 1: Representation I. 
+Objective: Step 1 is exemplary of the primary strategy of the Alberta Plan: to focus on a particular issue by considering the simplest setting in which it arises and attempting to deal with it there, fully, before generalizing to more complex settings. The issues focussed on in Step 1 are continual learning and the meta-learning of representations. How can learning be rapid, robust, and efficient while continuing over long periods of time? How can long periods of learning be taken advantage of to meta-learn better representations, and thus to learn most efficiently?  The simple setting employed in Step 1 is that of supervised learning and stochastic gradient descent with a linear function approximator with static, given features. In this setting, conventional stochastic-gradient-descent methods such as the Least-Mean-Square learning rule work reasonably well even if the problem is non-stationary. However, these methods can be significantly improved in their efficiency and robustness, and that is the purpose of Step 1. First, these methods often involve a global step-size parameter that must be set by an expert user aided by knowledge of the target outputs, the features, the number of features, and heuristics. All of that user expertise should be replaced by a meta-algorithm for setting the step-size parameter so that the same method can be used on any problem or for any part of a large problem. Second, instead of a global step-size parameter, there should be different step-size parameters for each feature depending on how much generalization should be done along that feature. If this is done, then there will be many step-size parameters to set, and it will be even more important to set them algorithmically.  In this setting, the representations are the features, which are given and fixed, so it may seem surprising to offer the setting up as a way of exploring representation learning. It is true that the setting cannot be used to discover features or to search for new features, but it can be used to assess the utility of given features—an important precursor to full representation discovery. Even without changing the features it is possible to learn which features are relevant and which are not. The relevant features can be given large step-size parameters and the irrelevant features small ones; this itself is a kind of representation learning than can affect learning efficiency even without changing the features themselves.  Finally, there are normalizations of the features (scalings and offsets) that can greatly affect learning efficiency without changing the representational power of the linear function approximator, and we include these in Step 1.  In particular, we consider an infinite sequence of examples of desired behavior, each consisting of a real-valued input vector paired with a real-valued desired output. Let the tth example be a pair denoted (xt, yt∗). The learner seeks to find an affine mapping from each input vector xt to an output yt that closely approximates the desired output yt∗. That affine map is represented as a vector of weights wt and a scalar bias or offset term bt. That is, the output is yt  .= wt>xt + bt. The objective is to minimize the squared error (yt∗ − yt)2 by learning wt and bt. Each example is independent, but the distribution from which it is generated changes over time, making the problem nonstationary. In particular, we can take the desired output as being affine in the input vector and an unknown target weight vector wt∗ that changes slowly over time, plus an  additional, independent mean-zero noise signal: yt∗  =. wt∗>xt + bt∗ + ηt. The problem  is non-stationary if wt∗ or bt∗ change over time or if the distribution of xt changes over time.  In this simple setting, there are still essential questions that have yet to be definitively answered. We are particularly interested in questions of normalization and step-size adaptation. Without changing the expressive power of our linear learning unit or the order of its computational complexity, we can transform the individual inputs xit to  produce normalized signals x ̃it  =. xit−μit  σti  where μit and σti are non-stationary (tracking)  estimates of the ith signal’s mean and standard deviation. Surprisingly, the effect of such online normalization has yet to be definitively established in the literature. We consider learning rules of the form:  wi  t+1  .= wi  t + αi  t (y∗  t − yt) x ̃i  t, ∀i, (1)  where each αit is a meta-learned, per-weight, step-size parameter, and  bt+1  =. bt + αb  t (y∗  t − bt) , (2)  where αtb is another potentially-meta-learned step-size parameter. Our initial studies for Step 1 will focus on algorithms for meta-learning the step-size parameters, building on existing algorithms,11 and on demonstrating their improved robustness.  The overall idea of Step 1 is to design as powerful an algorithm as possible given a fixed feature representation. It should include all the most important issues of nonstationarity in the problem (for a fixed set of linear features), including the tracking of changes in feature relevance. It should include the meta-learning of feature relevance, a challenging issue in representation learning—arguably the most challenging issue—but it does not include actually changing the set of features under consideration; that is explored in Step 2.
+Existing algorithms include NADALINE (Sutton 1988b), IDBD (Sutton 1992a,b), Autostep (Mahmood et al. 2012), Autostep for GTD(λ) (Kearney et al. 2022), Auto (Degris in prep.), Adam (Kingma & Ba 2014), RMSprop (Tieleman & Hinton 2012), and Batch Normalization (e.g., Ioffe & Szegedy 2015). 
+</current_milestone>
+
+<instructions>
+1. Always maintain a 'Project Memory' of our progress through the tools reqiured for the 12 steps. 
+2. When writing code, prioritize JAX-native patterns (NamedTuples for state, pure update functions).
+3. Before implementing a feature, suggest a 'Mathematical Plan' using the notation from the Alberta Plan papers.
+4. Ensure all implementations are modular so that Step 3 (GVFs) and Step 8 (Models) can be integrated without breaking Step 1.
+5. If I suggest an approach that violates 'Temporal Uniformity' or 'Continual Learning' (e.g., using a batch or replay buffer), politely correct me as per the Alberta Plan's 'Massive Retreat' discipline.
+6. All examples and implementations of environments should be gymnasium compatible as well as all of the features of this framework.
+</instructions>
+
+<output_format>
+- Use bullet points for summaries.
+- No sycophancy; provide honest, technical critiques of architecture.
+- Format math using LaTeX.
+</output_format>
+
+<ask_clarification>
+If you understand the vision for the Alberta framework and our current focus on Step 1 Representation I: Continual supervised learning with given features, please summarize the core algorithm update equations we will be using to ensure we are aligned.
+</ask_clarification>

alberta_framework-0.1.0/CHANGELOG.md

@@ -0,0 +1,30 @@
+# Changelog
+
+All notable changes to this project will be documented in this file.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [0.1.0] - 2026-01-19
+
+### Added
+
+- **Core Optimizers**: LMS (baseline), IDBD (Sutton 1992), and Autostep (Mahmood et al. 2012) with per-weight adaptive step-sizes
+- **Linear Learners**: `LinearLearner` and `NormalizedLinearLearner` with pluggable optimizers
+- **Scan-based Learning Loops**: JIT-compiled training with `jax.lax.scan` for efficiency
+- **Online Normalization**: Streaming feature normalization with exponential moving averages
+- **Experience Streams**: `RandomWalkStream`, `AbruptChangeStream`, `CyclicStream`, `SuttonExperiment1Stream`
+- **Gymnasium Integration**: Trajectory collection and learning from Gymnasium RL environments
+- **Step-Size Tracking**: Optional per-weight step-size history recording for meta-adaptation analysis
+- **Multi-Seed Experiments**: `run_multi_seed_experiment` with optional parallelization via joblib
+- **Statistical Analysis**: Pairwise comparisons, confidence intervals, effect sizes (requires scipy)
+- **Publication Visualization**: Learning curves, bar charts, heatmaps with matplotlib
+- **Export Utilities**: CSV, JSON, LaTeX, and Markdown table generation
+- **Documentation**: MkDocs-based documentation with auto-generated API reference
+
+### Notes
+
+- Requires Python 3.13+
+- Implements Step 1 of the Alberta Plan: demonstrating that IDBD/Autostep can match or beat hand-tuned LMS
+- All state uses immutable NamedTuples for JAX compatibility
+- Follows temporal uniformity principle: every component updates at every time step