alberta-framework 0.2.2__py3-none-any.whl

alberta_framework/core/types.py

@@ -0,0 +1,271 @@
+"""Type definitions for the Alberta Framework.
+
+This module defines the core data types used throughout the framework,
+following JAX conventions with immutable NamedTuples for state management.
+"""
+
+from typing import TYPE_CHECKING, NamedTuple
+
+import jax.numpy as jnp
+from jax import Array
+
+if TYPE_CHECKING:
+    from alberta_framework.core.learners import NormalizedLearnerState
+
+# Type aliases for clarity
+Observation = Array  # x_t: feature vector
+Target = Array  # y*_t: desired output
+Prediction = Array  # y_t: model output
+Reward = float  # r_t: scalar reward
+
+
+class TimeStep(NamedTuple):
+    """Single experience from an experience stream.
+
+    Attributes:
+        observation: Feature vector x_t
+        target: Desired output y*_t (for supervised learning)
+    """
+
+    observation: Observation
+    target: Target
+
+
+class LearnerState(NamedTuple):
+    """State for a linear learner.
+
+    Attributes:
+        weights: Weight vector for linear prediction
+        bias: Bias term
+        optimizer_state: State maintained by the optimizer
+    """
+
+    weights: Array
+    bias: Array
+    optimizer_state: "LMSState | IDBDState | AutostepState"
+
+
+class LMSState(NamedTuple):
+    """State for the LMS (Least Mean Square) optimizer.
+
+    LMS uses a fixed step-size, so state only tracks the step-size parameter.
+
+    Attributes:
+        step_size: Fixed learning rate alpha
+    """
+
+    step_size: Array
+
+
+class IDBDState(NamedTuple):
+    """State for the IDBD (Incremental Delta-Bar-Delta) optimizer.
+
+    IDBD maintains per-weight adaptive step-sizes that are meta-learned
+    based on the correlation of successive gradients.
+
+    Reference: Sutton 1992, "Adapting Bias by Gradient Descent"
+
+    Attributes:
+        log_step_sizes: Log of per-weight step-sizes (log alpha_i)
+        traces: Per-weight traces h_i for gradient correlation
+        meta_step_size: Meta learning rate beta for adapting step-sizes
+        bias_step_size: Step-size for the bias term
+        bias_trace: Trace for the bias term
+    """
+
+    log_step_sizes: Array  # log(alpha_i) for numerical stability
+    traces: Array  # h_i: trace of weight-feature products
+    meta_step_size: Array  # beta: step-size for the step-sizes
+    bias_step_size: Array  # Step-size for bias
+    bias_trace: Array  # Trace for bias
+
+
+class AutostepState(NamedTuple):
+    """State for the Autostep optimizer.
+
+    Autostep is a tuning-free step-size adaptation algorithm that normalizes
+    gradients to prevent large updates and adapts step-sizes based on
+    gradient correlation.
+
+    Reference: Mahmood et al. 2012, "Tuning-free step-size adaptation"
+
+    Attributes:
+        step_sizes: Per-weight step-sizes (alpha_i)
+        traces: Per-weight traces for gradient correlation (h_i)
+        normalizers: Running max absolute gradient per weight (v_i)
+        meta_step_size: Meta learning rate mu for adapting step-sizes
+        normalizer_decay: Decay factor for the normalizer (tau)
+        bias_step_size: Step-size for the bias term
+        bias_trace: Trace for the bias term
+        bias_normalizer: Normalizer for the bias gradient
+    """
+
+    step_sizes: Array  # alpha_i
+    traces: Array  # h_i
+    normalizers: Array  # v_i: running max of |gradient|
+    meta_step_size: Array  # mu
+    normalizer_decay: Array  # tau
+    bias_step_size: Array
+    bias_trace: Array
+    bias_normalizer: Array
+
+
+class StepSizeTrackingConfig(NamedTuple):
+    """Configuration for recording per-weight step-sizes during training.
+
+    Attributes:
+        interval: Record step-sizes every N steps
+        include_bias: Whether to also record the bias step-size
+    """
+
+    interval: int
+    include_bias: bool = True
+
+
+class StepSizeHistory(NamedTuple):
+    """History of per-weight step-sizes recorded during training.
+
+    Attributes:
+        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,)
+
+
+class BatchedLearningResult(NamedTuple):
+    """Result from batched learning loop across multiple seeds.
+
+    Used with `run_learning_loop_batched` for vmap-based GPU parallelization.
+
+    Attributes:
+        states: Batched learner states - each array has shape (num_seeds, ...)
+        metrics: Metrics array with shape (num_seeds, num_steps, 3)
+            where columns are [squared_error, error, mean_step_size]
+        step_size_history: Optional step-size history with batched shapes,
+            or None if tracking was disabled
+    """
+
+    states: "LearnerState"  # Batched: each array has shape (num_seeds, ...)
+    metrics: Array  # Shape: (num_seeds, num_steps, 3)
+    step_size_history: StepSizeHistory | None
+
+
+class BatchedNormalizedResult(NamedTuple):
+    """Result from batched normalized learning loop across multiple seeds.
+
+    Used with `run_normalized_learning_loop_batched` for vmap-based GPU parallelization.
+
+    Attributes:
+        states: Batched normalized learner states - each array has shape (num_seeds, ...)
+        metrics: Metrics array with shape (num_seeds, num_steps, 4)
+            where columns are [squared_error, error, mean_step_size, normalizer_mean_var]
+        step_size_history: Optional step-size history with batched shapes,
+            or None if tracking was disabled
+        normalizer_history: Optional normalizer history with batched shapes,
+            or None if tracking was disabled
+    """
+
+    states: "NormalizedLearnerState"  # Batched: each array has shape (num_seeds, ...)
+    metrics: Array  # Shape: (num_seeds, num_steps, 4)
+    step_size_history: StepSizeHistory | None
+    normalizer_history: NormalizerHistory | None
+
+
+def create_lms_state(step_size: float = 0.01) -> LMSState:
+    """Create initial LMS optimizer state.
+
+    Args:
+        step_size: Fixed learning rate
+
+    Returns:
+        Initial LMS state
+    """
+    return LMSState(step_size=jnp.array(step_size, dtype=jnp.float32))
+
+
+def create_idbd_state(
+    feature_dim: int,
+    initial_step_size: float = 0.01,
+    meta_step_size: float = 0.01,
+) -> IDBDState:
+    """Create initial IDBD optimizer state.
+
+    Args:
+        feature_dim: Dimension of the feature vector
+        initial_step_size: Initial per-weight step-size
+        meta_step_size: Meta learning rate for adapting step-sizes
+
+    Returns:
+        Initial IDBD state
+    """
+    return IDBDState(
+        log_step_sizes=jnp.full(feature_dim, jnp.log(initial_step_size), dtype=jnp.float32),
+        traces=jnp.zeros(feature_dim, dtype=jnp.float32),
+        meta_step_size=jnp.array(meta_step_size, dtype=jnp.float32),
+        bias_step_size=jnp.array(initial_step_size, dtype=jnp.float32),
+        bias_trace=jnp.array(0.0, dtype=jnp.float32),
+    )
+
+
+def create_autostep_state(
+    feature_dim: int,
+    initial_step_size: float = 0.01,
+    meta_step_size: float = 0.01,
+    normalizer_decay: float = 0.99,
+) -> AutostepState:
+    """Create initial Autostep optimizer state.
+
+    Args:
+        feature_dim: Dimension of the feature vector
+        initial_step_size: Initial per-weight step-size
+        meta_step_size: Meta learning rate for adapting step-sizes
+        normalizer_decay: Decay factor for gradient normalizers
+
+    Returns:
+        Initial Autostep state
+    """
+    return AutostepState(
+        step_sizes=jnp.full(feature_dim, initial_step_size, dtype=jnp.float32),
+        traces=jnp.zeros(feature_dim, dtype=jnp.float32),
+        normalizers=jnp.ones(feature_dim, dtype=jnp.float32),
+        meta_step_size=jnp.array(meta_step_size, dtype=jnp.float32),
+        normalizer_decay=jnp.array(normalizer_decay, dtype=jnp.float32),
+        bias_step_size=jnp.array(initial_step_size, dtype=jnp.float32),
+        bias_trace=jnp.array(0.0, dtype=jnp.float32),
+        bias_normalizer=jnp.array(1.0, dtype=jnp.float32),
+    )

alberta_framework/streams/__init__.py

@@ -0,0 +1,83 @@
+"""Experience streams for continual learning."""
+
+from alberta_framework.streams.base import ScanStream
+from alberta_framework.streams.synthetic import (
+    AbruptChangeState,
+    AbruptChangeStream,
+    AbruptChangeTarget,
+    CyclicState,
+    CyclicStream,
+    CyclicTarget,
+    DynamicScaleShiftState,
+    DynamicScaleShiftStream,
+    PeriodicChangeState,
+    PeriodicChangeStream,
+    PeriodicChangeTarget,
+    RandomWalkState,
+    RandomWalkStream,
+    RandomWalkTarget,
+    ScaleDriftState,
+    ScaleDriftStream,
+    ScaledStreamState,
+    ScaledStreamWrapper,
+    SuttonExperiment1State,
+    SuttonExperiment1Stream,
+    make_scale_range,
+)
+
+__all__ = [
+    # Protocol
+    "ScanStream",
+    # Stream classes
+    "AbruptChangeState",
+    "AbruptChangeStream",
+    "AbruptChangeTarget",
+    "CyclicState",
+    "CyclicStream",
+    "CyclicTarget",
+    "DynamicScaleShiftState",
+    "DynamicScaleShiftStream",
+    "PeriodicChangeState",
+    "PeriodicChangeStream",
+    "PeriodicChangeTarget",
+    "RandomWalkState",
+    "RandomWalkStream",
+    "RandomWalkTarget",
+    "ScaleDriftState",
+    "ScaleDriftStream",
+    "ScaledStreamState",
+    "ScaledStreamWrapper",
+    "SuttonExperiment1State",
+    "SuttonExperiment1Stream",
+    # Utilities
+    "make_scale_range",
+]
+
+# Gymnasium streams are optional - only export if gymnasium is installed
+try:
+    from alberta_framework.streams.gymnasium import (
+        GymnasiumStream,
+        PredictionMode,
+        TDStream,
+        collect_trajectory,
+        learn_from_trajectory,
+        learn_from_trajectory_normalized,
+        make_epsilon_greedy_policy,
+        make_gymnasium_stream,
+        make_random_policy,
+    )
+
+    __all__ += [
+        "GymnasiumStream",
+        "PredictionMode",
+        "TDStream",
+        "collect_trajectory",
+        "learn_from_trajectory",
+        "learn_from_trajectory_normalized",
+        "make_epsilon_greedy_policy",
+        "make_gymnasium_stream",
+        "make_random_policy",
+    ]
+except ImportError:
+    # gymnasium not installed
+    pass

alberta_framework/streams/base.py

@@ -0,0 +1,73 @@
+"""Base protocol for experience streams.
+
+Experience streams generate temporally-uniform experience for continual learning.
+Every time step produces a new observation-target pair.
+
+This module defines the ScanStream protocol for JAX scan-compatible streams.
+All streams implement pure functions that can be JIT-compiled.
+"""
+
+from typing import Protocol, TypeVar
+
+from jax import Array
+
+from alberta_framework.core.types import TimeStep
+
+# Type variable for stream state
+StateT = TypeVar("StateT")
+
+
+class ScanStream(Protocol[StateT]):
+    """Protocol for JAX scan-compatible experience streams.
+
+    Streams generate temporally-uniform experience for continual learning.
+    Unlike iterator-based streams, ScanStream uses pure functions that
+    can be compiled with JAX's JIT and used with jax.lax.scan.
+
+    The stream should be non-stationary to test continual learning
+    capabilities - the underlying target function changes over time.
+
+    Type Parameters:
+        StateT: The state type maintained by this stream
+
+    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
+    def feature_dim(self) -> int:
+        """Return the dimension of observation vectors."""
+        ...
+
+    def init(self, key: Array) -> StateT:
+        """Initialize stream state.
+
+        Args:
+            key: JAX random key for initialization
+
+        Returns:
+            Initial stream state
+        """
+        ...
+
+    def step(self, state: StateT, idx: Array) -> tuple[TimeStep, StateT]:
+        """Generate one time step. Must be JIT-compatible.
+
+        This is a pure function that takes the current state and step index,
+        and returns a TimeStep along with the updated state. The step index
+        can be used for time-dependent behavior but is often ignored.
+
+        Args:
+            state: Current stream state
+            idx: Current step index (can be ignored for most streams)
+
+        Returns:
+            Tuple of (timestep, new_state)
+        """
+        ...