alberta-framework 0.1.1__tar.gz → 0.2.1__tar.gz

{alberta_framework-0.1.1 → alberta_framework-0.2.1}/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
+│   ├── 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, run_normalized_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
@@ -187,6 +187,59 @@ Return value depends on tracking options:
 - 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
 
 Wrap Gymnasium RL environments as experience streams for the framework.

{alberta_framework-0.1.1 → alberta_framework-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alberta-framework
-Version: 0.1.1
+Version: 0.2.1
 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

{alberta_framework-0.1.1 → alberta_framework-0.2.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "alberta-framework"
-version = "0.1.1"
+version = "0.2.1"
 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.1 → alberta_framework-0.2.1}/src/alberta_framework/__init__.py

@@ -39,7 +39,7 @@ References
 - Tuning-free Step-size Adaptation (Mahmood et al., 2012)
 """
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 
 # Core types
 # Learners
@@ -50,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
@@ -64,6 +66,8 @@ 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,
@@ -138,6 +142,8 @@ __all__ = [
     "__version__",
     # Types
     "AutostepState",
+    "BatchedLearningResult",
+    "BatchedNormalizedResult",
     "IDBDState",
     "LMSState",
     "LearnerState",
@@ -164,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",

{alberta_framework-0.1.1 → alberta_framework-0.2.1}/src/alberta_framework/core/learners.py

@@ -15,6 +15,8 @@ from alberta_framework.core.normalizers import NormalizerState, OnlineNormalizer
 from alberta_framework.core.optimizers import LMS, Optimizer
 from alberta_framework.core.types import (
     AutostepState,
+    BatchedLearningResult,
+    BatchedNormalizedResult,
     IDBDState,
     LearnerState,
     LMSState,
@@ -846,6 +848,196 @@ def run_normalized_learning_loop[StreamStateT](
         return final_learner, metrics
 
 
+def run_learning_loop_batched[StreamStateT](
+    learner: LinearLearner,
+    stream: ScanStream[StreamStateT],
+    num_steps: int,
+    keys: Array,
+    learner_state: LearnerState | None = None,
+    step_size_tracking: StepSizeTrackingConfig | None = None,
+) -> BatchedLearningResult:
+    """Run learning loop across multiple seeds in parallel using jax.vmap.
+
+    This function provides GPU parallelization for multi-seed experiments,
+    typically achieving 2-5x speedup over sequential execution.
+
+    Args:
+        learner: The learner to train
+        stream: Experience stream providing (observation, target) pairs
+        num_steps: Number of learning steps to run per seed
+        keys: JAX random keys with shape (num_seeds,) or (num_seeds, 2)
+        learner_state: Initial state (if None, will be initialized from stream).
+            The same initial state is used for all seeds.
+        step_size_tracking: Optional config for recording per-weight step-sizes.
+            When provided, history arrays have shape (num_seeds, num_recordings, ...)
+
+    Returns:
+        BatchedLearningResult containing:
+            - states: Batched final states with shape (num_seeds, ...) for each array
+            - metrics: Array of shape (num_seeds, num_steps, 3)
+            - step_size_history: Batched history or None if tracking disabled
+
+    Examples:
+    ```python
+    import jax.random as jr
+    from alberta_framework import LinearLearner, IDBD, RandomWalkStream
+    from alberta_framework import run_learning_loop_batched
+
+    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)
+    mean_error = result.metrics[:, :, 0].mean(axis=0)  # Average over seeds
+    ```
+    """
+    # Define single-seed function that returns consistent structure
+    def single_seed_run(key: Array) -> tuple[LearnerState, Array, StepSizeHistory | None]:
+        result = run_learning_loop(
+            learner, stream, num_steps, key, learner_state, step_size_tracking
+        )
+        if step_size_tracking is not None:
+            state, metrics, history = result
+            return state, metrics, history
+        else:
+            state, metrics = result
+            # Return None for history to maintain consistent output structure
+            return state, metrics, None
+
+    # vmap over the keys dimension
+    batched_states, batched_metrics, batched_history = jax.vmap(single_seed_run)(keys)
+
+    # Reconstruct batched history if tracking was enabled
+    if step_size_tracking is not None:
+        batched_step_size_history = StepSizeHistory(
+            step_sizes=batched_history.step_sizes,
+            bias_step_sizes=batched_history.bias_step_sizes,
+            recording_indices=batched_history.recording_indices,
+            normalizers=batched_history.normalizers,
+        )
+    else:
+        batched_step_size_history = None
+
+    return BatchedLearningResult(
+        states=batched_states,
+        metrics=batched_metrics,
+        step_size_history=batched_step_size_history,
+    )
+
+
+def run_normalized_learning_loop_batched[StreamStateT](
+    learner: NormalizedLinearLearner,
+    stream: ScanStream[StreamStateT],
+    num_steps: int,
+    keys: Array,
+    learner_state: NormalizedLearnerState | None = None,
+    step_size_tracking: StepSizeTrackingConfig | None = None,
+    normalizer_tracking: NormalizerTrackingConfig | None = None,
+) -> BatchedNormalizedResult:
+    """Run normalized learning loop across multiple seeds in parallel using jax.vmap.
+
+    This function provides GPU parallelization for multi-seed experiments with
+    normalized learners, typically achieving 2-5x speedup over sequential execution.
+
+    Args:
+        learner: The normalized learner to train
+        stream: Experience stream providing (observation, target) pairs
+        num_steps: Number of learning steps to run per seed
+        keys: JAX random keys with shape (num_seeds,) or (num_seeds, 2)
+        learner_state: Initial state (if None, will be initialized from stream).
+            The same initial state is used for all seeds.
+        step_size_tracking: Optional config for recording per-weight step-sizes.
+            When provided, history arrays have shape (num_seeds, num_recordings, ...)
+        normalizer_tracking: Optional config for recording normalizer state.
+            When provided, history arrays have shape (num_seeds, num_recordings, ...)
+
+    Returns:
+        BatchedNormalizedResult containing:
+            - states: Batched final states with shape (num_seeds, ...) for each array
+            - metrics: Array of shape (num_seeds, num_steps, 4)
+            - step_size_history: Batched history or None if tracking disabled
+            - normalizer_history: Batched history or None if tracking disabled
+
+    Examples:
+    ```python
+    import jax.random as jr
+    from alberta_framework import NormalizedLinearLearner, IDBD, RandomWalkStream
+    from alberta_framework import run_normalized_learning_loop_batched
+
+    stream = RandomWalkStream(feature_dim=10)
+    learner = NormalizedLinearLearner(optimizer=IDBD())
+
+    # Run 30 seeds in parallel
+    keys = jr.split(jr.key(42), 30)
+    result = run_normalized_learning_loop_batched(
+        learner, stream, num_steps=10000, keys=keys
+    )
+
+    # result.metrics has shape (30, 10000, 4)
+    mean_error = result.metrics[:, :, 0].mean(axis=0)  # Average over seeds
+    ```
+    """
+    # Define single-seed function that returns consistent structure
+    def single_seed_run(
+        key: Array,
+    ) -> tuple[
+        NormalizedLearnerState, Array, StepSizeHistory | None, NormalizerHistory | None
+    ]:
+        result = run_normalized_learning_loop(
+            learner, stream, num_steps, key, learner_state,
+            step_size_tracking, normalizer_tracking
+        )
+
+        # Unpack based on what tracking was enabled
+        if step_size_tracking is not None and normalizer_tracking is not None:
+            state, metrics, ss_history, norm_history = result
+            return state, metrics, ss_history, norm_history
+        elif step_size_tracking is not None:
+            state, metrics, ss_history = result
+            return state, metrics, ss_history, None
+        elif normalizer_tracking is not None:
+            state, metrics, norm_history = result
+            return state, metrics, None, norm_history
+        else:
+            state, metrics = result
+            return state, metrics, None, None
+
+    # vmap over the keys dimension
+    batched_states, batched_metrics, batched_ss_history, batched_norm_history = (
+        jax.vmap(single_seed_run)(keys)
+    )
+
+    # Reconstruct batched histories if tracking was enabled
+    if step_size_tracking is not None:
+        batched_step_size_history = StepSizeHistory(
+            step_sizes=batched_ss_history.step_sizes,
+            bias_step_sizes=batched_ss_history.bias_step_sizes,
+            recording_indices=batched_ss_history.recording_indices,
+            normalizers=batched_ss_history.normalizers,
+        )
+    else:
+        batched_step_size_history = None
+
+    if normalizer_tracking is not None:
+        batched_normalizer_history = NormalizerHistory(
+            means=batched_norm_history.means,
+            variances=batched_norm_history.variances,
+            recording_indices=batched_norm_history.recording_indices,
+        )
+    else:
+        batched_normalizer_history = None
+
+    return BatchedNormalizedResult(
+        states=batched_states,
+        metrics=batched_metrics,
+        step_size_history=batched_step_size_history,
+        normalizer_history=batched_normalizer_history,
+    )
+
+
 def metrics_to_dicts(metrics: Array, normalized: bool = False) -> list[dict[str, float]]:
     """Convert metrics array to list of dicts for backward compatibility.
 

{alberta_framework-0.1.1 → alberta_framework-0.2.1}/src/alberta_framework/core/types.py

@@ -4,11 +4,14 @@ This module defines the core data types used throughout the framework,
 following JAX conventions with immutable NamedTuples for state management.
 """
 
-from typing import NamedTuple
+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
@@ -164,6 +167,45 @@ class NormalizerHistory(NamedTuple):
     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.
 

{alberta_framework-0.1.1 → alberta_framework-0.2.1}/tests/test_learners.py

@@ -6,6 +6,8 @@ import pytest
 
 from alberta_framework import (
     Autostep,
+    BatchedLearningResult,
+    BatchedNormalizedResult,
     IDBD,
     LMS,
     LinearLearner,
@@ -17,7 +19,9 @@ from alberta_framework import (
     StepSizeTrackingConfig,
     metrics_to_dicts,
     run_learning_loop,
+    run_learning_loop_batched,
     run_normalized_learning_loop,
+    run_normalized_learning_loop_batched,
 )
 
 
@@ -526,3 +530,248 @@ class TestNormalizedLearningLoopTracking:
         # Should record at steps 0, 25, 50, 75
         expected_indices = jnp.array([0, 25, 50, 75])
         assert jnp.allclose(norm_history.recording_indices, expected_indices)
+
+
+class TestBatchedLearningLoop:
+    """Tests for run_learning_loop_batched."""
+
+    def test_batched_returns_correct_shapes(self, rng_key):
+        """Batched loop should return metrics with shape (num_seeds, num_steps, 3)."""
+        num_seeds = 5
+        num_steps = 100
+        feature_dim = 10
+
+        stream = RandomWalkStream(feature_dim=feature_dim)
+        learner = LinearLearner(optimizer=IDBD())
+        keys = jr.split(rng_key, num_seeds)
+
+        result = run_learning_loop_batched(learner, stream, num_steps, keys)
+
+        assert isinstance(result, BatchedLearningResult)
+        assert result.metrics.shape == (num_seeds, num_steps, 3)
+        assert result.states.weights.shape == (num_seeds, feature_dim)
+        assert result.states.bias.shape == (num_seeds,)
+        assert result.step_size_history is None
+
+    def test_batched_matches_sequential(self, rng_key):
+        """Batched results should match sequential execution."""
+        num_seeds = 3
+        num_steps = 50
+        feature_dim = 5
+
+        stream = RandomWalkStream(feature_dim=feature_dim)
+        learner = LinearLearner(optimizer=IDBD())
+        keys = jr.split(rng_key, num_seeds)
+
+        # Run batched
+        batched_result = run_learning_loop_batched(learner, stream, num_steps, keys)
+
+        # Run sequential
+        sequential_metrics = []
+        for i in range(num_seeds):
+            _, metrics = run_learning_loop(learner, stream, num_steps, keys[i])
+            sequential_metrics.append(metrics)
+        sequential_metrics = jnp.stack(sequential_metrics)
+
+        # Should match
+        assert jnp.allclose(batched_result.metrics, sequential_metrics)
+
+    def test_batched_with_step_size_tracking(self, rng_key):
+        """Batched loop should support step-size tracking."""
+        num_seeds = 4
+        num_steps = 100
+        feature_dim = 8
+        interval = 10
+        expected_recordings = num_steps // interval
+
+        stream = RandomWalkStream(feature_dim=feature_dim)
+        learner = LinearLearner(optimizer=Autostep())
+        keys = jr.split(rng_key, num_seeds)
+        config = StepSizeTrackingConfig(interval=interval)
+
+        result = run_learning_loop_batched(
+            learner, stream, num_steps, keys, step_size_tracking=config
+        )
+
+        assert result.step_size_history is not None
+        assert result.step_size_history.step_sizes.shape == (
+            num_seeds, expected_recordings, feature_dim
+        )
+        assert result.step_size_history.bias_step_sizes.shape == (
+            num_seeds, expected_recordings
+        )
+        assert result.step_size_history.recording_indices.shape == (
+            num_seeds, expected_recordings
+        )
+        # Autostep should have normalizers tracked
+        assert result.step_size_history.normalizers is not None
+        assert result.step_size_history.normalizers.shape == (
+            num_seeds, expected_recordings, feature_dim
+        )
+
+    def test_batched_without_tracking_has_none_history(self, rng_key):
+        """When tracking disabled, step_size_history should be None."""
+        stream = RandomWalkStream(feature_dim=5)
+        learner = LinearLearner(optimizer=IDBD())
+        keys = jr.split(rng_key, 3)
+
+        result = run_learning_loop_batched(learner, stream, num_steps=50, keys=keys)
+
+        assert result.step_size_history is None
+
+    def test_batched_deterministic_with_same_keys(self, rng_key):
+        """Same keys should produce same results."""
+        stream = RandomWalkStream(feature_dim=5)
+        learner = LinearLearner(optimizer=IDBD())
+        keys = jr.split(rng_key, 4)
+
+        result1 = run_learning_loop_batched(learner, stream, num_steps=50, keys=keys)
+        result2 = run_learning_loop_batched(learner, stream, num_steps=50, keys=keys)
+
+        assert jnp.allclose(result1.metrics, result2.metrics)
+        assert jnp.allclose(result1.states.weights, result2.states.weights)
+
+    def test_batched_different_keys_different_results(self, rng_key):
+        """Different keys should produce different results."""
+        stream = RandomWalkStream(feature_dim=5)
+        learner = LinearLearner(optimizer=IDBD())
+
+        keys1 = jr.split(jr.key(42), 3)
+        keys2 = jr.split(jr.key(123), 3)
+
+        result1 = run_learning_loop_batched(learner, stream, num_steps=50, keys=keys1)
+        result2 = run_learning_loop_batched(learner, stream, num_steps=50, keys=keys2)
+
+        assert not jnp.allclose(result1.metrics, result2.metrics)
+
+    def test_batched_with_lms_optimizer(self, rng_key):
+        """Batched loop should work with LMS optimizer."""
+        stream = RandomWalkStream(feature_dim=5)
+        learner = LinearLearner(optimizer=LMS(step_size=0.01))
+        keys = jr.split(rng_key, 3)
+
+        result = run_learning_loop_batched(learner, stream, num_steps=50, keys=keys)
+
+        assert result.metrics.shape == (3, 50, 3)
+        # LMS doesn't report mean_step_size in metrics (defaults to 0.0)
+        assert jnp.allclose(result.metrics[:, :, 2], 0.0)
+
+
+class TestBatchedNormalizedLearningLoop:
+    """Tests for run_normalized_learning_loop_batched."""
+
+    def test_normalized_batched_returns_correct_shapes(self, rng_key):
+        """Batched normalized loop should return metrics with shape (num_seeds, num_steps, 4)."""
+        num_seeds = 5
+        num_steps = 100
+        feature_dim = 10
+
+        stream = RandomWalkStream(feature_dim=feature_dim)
+        learner = NormalizedLinearLearner(optimizer=IDBD())
+        keys = jr.split(rng_key, num_seeds)
+
+        result = run_normalized_learning_loop_batched(learner, stream, num_steps, keys)
+
+        assert isinstance(result, BatchedNormalizedResult)
+        assert result.metrics.shape == (num_seeds, num_steps, 4)
+        assert result.states.learner_state.weights.shape == (num_seeds, feature_dim)
+        assert result.states.normalizer_state.mean.shape == (num_seeds, feature_dim)
+        assert result.step_size_history is None
+        assert result.normalizer_history is None
+
+    def test_normalized_batched_matches_sequential(self, rng_key):
+        """Batched normalized results should match sequential execution."""
+        num_seeds = 3
+        num_steps = 50
+        feature_dim = 5
+
+        stream = RandomWalkStream(feature_dim=feature_dim)
+        learner = NormalizedLinearLearner(optimizer=IDBD())
+        keys = jr.split(rng_key, num_seeds)
+
+        # Run batched
+        batched_result = run_normalized_learning_loop_batched(
+            learner, stream, num_steps, keys
+        )
+
+        # Run sequential
+        sequential_metrics = []
+        for i in range(num_seeds):
+            _, metrics = run_normalized_learning_loop(learner, stream, num_steps, keys[i])
+            sequential_metrics.append(metrics)
+        sequential_metrics = jnp.stack(sequential_metrics)
+
+        # Should match
+        assert jnp.allclose(batched_result.metrics, sequential_metrics)
+
+    def test_normalized_batched_with_both_tracking(self, rng_key):
+        """Batched normalized loop should support both tracking options."""
+        num_seeds = 4
+        num_steps = 100
+        feature_dim = 8
+        ss_interval = 10
+        norm_interval = 20
+        ss_recordings = num_steps // ss_interval
+        norm_recordings = num_steps // norm_interval
+
+        stream = RandomWalkStream(feature_dim=feature_dim)
+        learner = NormalizedLinearLearner(optimizer=Autostep())
+        keys = jr.split(rng_key, num_seeds)
+        ss_config = StepSizeTrackingConfig(interval=ss_interval)
+        norm_config = NormalizerTrackingConfig(interval=norm_interval)
+
+        result = run_normalized_learning_loop_batched(
+            learner, stream, num_steps, keys,
+            step_size_tracking=ss_config, normalizer_tracking=norm_config
+        )
+
+        # Step-size history
+        assert result.step_size_history is not None
+        assert result.step_size_history.step_sizes.shape == (
+            num_seeds, ss_recordings, feature_dim
+        )
+        # Autostep normalizers
+        assert result.step_size_history.normalizers is not None
+
+        # Normalizer history
+        assert result.normalizer_history is not None
+        assert result.normalizer_history.means.shape == (
+            num_seeds, norm_recordings, feature_dim
+        )
+        assert result.normalizer_history.variances.shape == (
+            num_seeds, norm_recordings, feature_dim
+        )
+
+    def test_normalized_batched_step_size_only(self, rng_key):
+        """Batched normalized loop with only step-size tracking."""
+        num_seeds = 3
+        num_steps = 50
+
+        stream = RandomWalkStream(feature_dim=5)
+        learner = NormalizedLinearLearner(optimizer=IDBD())
+        keys = jr.split(rng_key, num_seeds)
+        ss_config = StepSizeTrackingConfig(interval=10)
+
+        result = run_normalized_learning_loop_batched(
+            learner, stream, num_steps, keys, step_size_tracking=ss_config
+        )
+
+        assert result.step_size_history is not None
+        assert result.normalizer_history is None
+
+    def test_normalized_batched_normalizer_only(self, rng_key):
+        """Batched normalized loop with only normalizer tracking."""
+        num_seeds = 3
+        num_steps = 50
+
+        stream = RandomWalkStream(feature_dim=5)
+        learner = NormalizedLinearLearner(optimizer=IDBD())
+        keys = jr.split(rng_key, num_seeds)
+        norm_config = NormalizerTrackingConfig(interval=10)
+
+        result = run_normalized_learning_loop_batched(
+            learner, stream, num_steps, keys, normalizer_tracking=norm_config
+        )
+
+        assert result.step_size_history is None
+        assert result.normalizer_history is not None