Active-Federated Learning

Federated Learning with active learning: This project explores combining active learning strategies with federated learning. Instead of naive Federated Averaging (FedAvg), the central server evaluates and selects updates using Active Weight and Active Data methods to improve the global model.

The framework supports running experiments both locally (via parallel subprocesses) and remotely on Kubernetes — Kubeflow Pipelines owns the round-level DAG, and Temporal launches and watches one Kubernetes Job per worker inside each round.

The Problem: Why Active Learning in Federated Learning?

In many Federated Learning scenarios, standard FedAvg works well when client datasets are relatively static and similar. However, it often struggles due to non-stationarity and catastrophic forgetting:

High Variance in Updates: Workers explore different parts of the environment or train on diverse local data. A client stuck in a bad local optimum will send a gradient update that disrupts the global model's progress.
Sample and Resource Inefficiency: Training locally requires interacting with environments or crunching local data. Blindly averaging bad weights means the global model regresses, wasting the compute and samples workers just collected.

The Solution: Instead of blindly averaging client weights, the central aggregator performs a Target-Environment Probe: evaluating each worker's proposed weight update (ΔW) directly in the target environment before accepting it.

Depending on how well the worker's update performs during this probe, the aggregator applies two active learning strategies:

Active Mode: Weight: We use a 4-Factor Scoring mechanism (which heavily penalizes updates that regress target-environment performance) to compute a weighted FedAvg. Updates that perform poorly are rejected entirely.
Active Mode: Data: For workers whose updates do significantly improve performance in the target environment, we capture the successful trajectories from their evaluation rollouts. We then use Behavioral Cloning (BC) to explicitly fine-tune the aggregated global model on these high-quality, provably successful trajectories.

Experiment Results

Running all active learning combinations locally (make run-experiments) generates these comparison plots automatically:

1. Global Learning Curves

Behavioral Cloning (Data Only): Learns extremely fast early on, but is limited by the pure BC approach and fails to adapt or generalize to the environment, preventing it from achieving high final rewards.
FedAvg: Suffers from catastrophic forgetting as local worker updates diverge, leading to a massive drop in performance midway through training.
Active FL: Active methods learn steadily across all rounds. Combining Active Weight + Active Data (BC) dramatically accelerates learning with a much steeper curve and achieves the highest peak performance.

2. Worker Own-Environment Performance

Active Stability & Independence: Active methods perform better globally and yield higher improvements, even if individual worker rewards in their own environments aren't as extremely high. The global model is stabilized without being strictly bound to the immediate whims of any one worker.
FedAvg Entanglement: In contrast, the FedAvg baseline global reward (solid line) stays very close to the individual worker rewards, struggling with high variance and catastrophic forgetting as divergent local updates constantly overwrite each other.

Architecture & Active Methods

graph TD
    classDef worker fill:#e1f5fe,stroke:#01579b
    classDef agg fill:#f3e5f5,stroke:#4a148c
    classDef probe fill:#fff8e1,stroke:#f57f17
    
    W1[Worker 1: PPO Train locally]:::worker
    W2[Worker 2: PPO Train locally]:::worker
    WN[Worker N: PPO Train locally]:::worker
    
    subgraph Aggregator Node
        Probe1[Eval W_global + ΔW_1]:::probe
        Probe2[Eval W_global + ΔW_2]:::probe
        ProbeN[Eval W_global + ΔW_N]:::probe
        
        Score[4-Factor Scoring]:::agg
        Target[Target-Env Tracking]:::agg
        
        subgraph Active Mode: Weight
            AW[Weighted FedAvg]:::agg
        end
        
        subgraph Active Mode: Data
            AD[BC Fine-Tuning]:::agg
        end
    end
    
    W1 -->|ΔW_1| Probe1
    W2 -->|ΔW_2| Probe2
    WN -->|ΔW_N| ProbeN
    
    Probe1 -->|Rewards| Target
    Probe2 -->|Rewards| Target
    ProbeN -->|Rewards| Target

    Probe1 -->|Trajectories & Gradients| Score
    Probe2 -->|Trajectories & Gradients| Score
    ProbeN -->|Trajectories & Gradients| Score
    
    Score -->|Accepted Weights & Scores| AW
    Target -->|High-value Trajectories| AD
    AW -->|Aggregated Model| AD
    AD -->|New Global Model| NextRound[Next FL Round]

Configure modes via config/local.yaml or config/k8s.yaml:

`weight_mode`	`active_data_mode`	Description
`fedavg`	`none`	Baseline — vanilla equal-weight FedAvg, no active learning
`active`	`none`	Active Weight only — scored, importance-weighted FedAvg
`data_only`	`bc`	Active Data only — skip weight avg, BC fine-tune on probe trajectories
`active`	`bc`	Both active paths — scored FedAvg + BC fine-tune

Method 1: Active Weight (4-Factor Scoring)

Instead of equal-weight averaging, each client's ΔW is scored and softmax-normalized. A client is rejected entirely if its score falls below score_threshold. The score combines four signals:

Target-env improvement (α): How much did this ΔW increase the target environment reward?
Gradient norm (β): Penalize clients that barely moved their weights.
Weight diversity (γ): Cosine distance from the mean update (penalize redundant updates).
TD-error penalty (δ): Negative penalty for training instability.

Method 2: Active Data (Fine-Tuning)

The trajectories captured during the Evaluation Probe are extremely valuable: they are guaranteed to be on-task and come from provably improved policies. We collect these trajectories and explicitly fine-tune the aggregated weights:

bc (Behavioral Cloning): Maximize the log-likelihood of the collected actions. Fast and stable.

Use Case 1: Local Experiments (fast iteration, no infrastructure)

All 4 mode combinations run in parallel subprocesses via ProcessPoolExecutor. No Docker, no K8s, no MinIO — just Python. After all runs finish, comparison plots are generated automatically.

When to use: tuning hyperparameters, comparing modes, validating methodology, CI.

# Install
make install-dev

# Run tests
make test

# Run all combinations from config/local.yaml (parallel, auto-plots)
make run-experiments

# Smoke test — override rounds/workers/episodes via ARGS
make run-experiments ARGS="--rounds 2 --workers 2 --episodes 30"

# Skip auto-plotting
make run-experiments ARGS="--no-viz"

# Cap parallel experiments (e.g. 2 at a time)
make run-experiments ARGS="--jobs 2"

# Run a single combination
make run-single WEIGHT_MODE=data_only ACTIVE_DATA_MODE=bc

# View results in MLflow UI (local ./mlruns store)
make mlflow-ui
# Open http://localhost:5000

# Regenerate comparison plots from existing results
make compare

Outputs saved to results/:

<run_name>.json — per-round metrics (reward, acceptance rate, active data usage)
results/plots/learning_curves.png
results/plots/heatmap.png
results/plots/final_reward_bar.png
results/plots/acceptance_rate.png
results/plots/active_data_usage.png
results/plots/client_improvements.png
results/plots/worker_own_env_curves.png
results/plots/worker_target_env_curves.png

Configure in config/local.yaml — edit combinations:, training rounds, workers, etc.

Use Case 2: Kubernetes / Kubeflow Pipeline (real distributed FL)

Workers run as isolated pods (true process separation), weights and artifacts flow through MinIO, and metrics are tracked in a shared MLflow server. Kubeflow Pipelines orchestrates the round-level DAG; inside each round, Temporal fans out one durable WorkerWorkflow per worker, each launching and watching a plain Kubernetes Job. (There is no PyTorchJob in this path — the pytorchjob launcher was removed in Phase P2.)

When to use: real federated scenario (workers on different machines/data), GPU training, production-scale runs, or when you need the full MLflow + artifact tracking pipeline.

What runs in K8s

graph TD
    classDef job fill:#e3f2fd,stroke:#1565c0
    classDef pod fill:#f3e5f5,stroke:#6a1b9a
    classDef store fill:#fff3e0,stroke:#e65100
    classDef temporal fill:#e8f5e9,stroke:#2e7d32

    Train[train_workers<br><i>thin Temporal client: starts one TrainRoundWorkflow</i>]:::job
    Round[TrainRoundWorkflow<br><i>fans out one WorkerWorkflow per worker</i>]:::temporal
    Jobs[Worker Kubernetes Jobs<br><i>one pod per worker trains PPO</i>]:::pod
    Agg[score_and_aggregate<br><i>Aggregator pod: eval probes, scoring, FedAvg + active data</i>]:::pod
    Eval[evaluate_global<br><i>Evaluation pod: global model eval</i>]:::pod

    MinIO[(MinIO Storage)]:::store
    MLflow[(MLflow Server)]:::store

    Train --> Round
    Round -->|launch_and_watch_pod| Jobs
    Jobs -->|Δw_i| MinIO
    Train --> Agg
    MinIO -->|Fetch Δw_i| Agg
    Agg -->|Write W_new| MinIO
    Agg --> Eval
    MinIO -->|Fetch W_new| Eval
    Eval -->|Log Metrics & Checkpoints| MLflow

KFP still owns the round-level DAG and artifact lineage; Temporal owns the worker fleet within a round (retries, heartbeats, per-worker failure reasons — see Observability below).

What to see in Kubeflow UI

Pipeline graph: per-round DAG with train → aggregate → evaluate chain
Pod logs: per-worker training progress, client scores, improvement values
Artifacts: aggregation report JSON (accepted/rejected clients with scores)

Infrastructure package

Cluster bootstrap is handled by fed-infra, a reusable Bash library vendored as a git submodule at vendor/fed-infra and driven by infra.env (single cluster) / infra.env.multi (multi cluster). It brings up the kind cluster(s) with Kubeflow Pipelines, the Kubeflow Training Operator, Temporal, MinIO, MLflow, and the Kubernetes Dashboard (plus Karmada and its dashboard in multi mode). See the fed-infra README and variable reference for details.

# Prerequisites: kind, kubectl, helm, docker

# 0. Initialize the vendor/fed-infra submodule (first checkout, or after a pull
#    that bumped it) — `make local-setup` also does this, but it's cheap to run
#    explicitly and is required before invoking vendor/fed-infra scripts directly.
git submodule update --init --recursive

# 1. Bootstrap local kind cluster (MinIO, MLflow, Kubeflow)
make local-setup

# 2. Run pipeline (uses config/k8s.yaml)
make run-pipeline
# Flags (via ARGS): --config PATH, --kfp-host URL, --wait (block until runs
# finish), --auto-download (fetch MLflow results after; implies --wait)
make run-pipeline ARGS="--wait --auto-download"

# Open UIs (after port-forward):
#   Kubeflow:             http://localhost:8080
#   MLflow:               http://localhost:5050
#   MinIO:                http://localhost:9001
#   Kubernetes Dashboard: https://localhost:8443

# 3. Fetch K8s MLflow results + generate plots
make compare-k8s

# 4. Teardown
make local-teardown

Multi-cluster topology (Karmada)

Setting orchestration.topology: multi propagates each worker Job via Karmada onto a separate member cluster instead of running the whole fleet on one cluster. Two file pairs drive the two topologies — config/k8s.yaml ↔ infra.env (single) and config/k8s-multi.yaml ↔ infra.env.multi (multi):

infra.env.multi tells vendor/fed-infra to create a host cluster + 2 member clusters (FED_MEMBER_COUNT/FED_MEMBER_PREFIX) and adds the karmada + karmada-dashboard components.
config/k8s-multi.yaml sets topology: multi, members: 2, and runs a single mode combination — it exists to exercise cross-cluster propagation, not to compare strategies.
Keep the pairs in sync: members must equal FED_MEMBER_COUNT and member_prefix must equal FED_MEMBER_PREFIX. A mismatch makes the PropagationPolicy name a cluster that doesn't exist, which Karmada matches silently to nothing rather than erroring.

# Bootstrap host + member clusters (uses infra.env.multi)
make multi-setup

# Run the pipeline against the multi topology
make run-pipeline ARGS="--config config/k8s-multi.yaml"

# Propagation state & member cluster health
# Karmada Dashboard: http://localhost:32000

# Teardown
make multi-teardown

Observability: Four Surfaces

Running the K8s pipeline (make run-pipeline) gives you four UIs, each answering a different question. They never overlap in scope — KFP sequences rounds, Temporal manages the worker fleet inside a round, MLflow tracks ML metrics, and the Kubernetes Dashboard shows pod-level health — so there's no ambiguity about which one to open for a given question. In multi-cluster mode a fifth surface, the Karmada Dashboard, adds propagation state and member cluster health.

Surface	URL	Answers
Kubeflow Pipelines	http://localhost:8080	Round DAG, node logs, artifact lineage
Temporal	http://localhost:8233	Which worker failed and why; retry counts; live progress
MLflow	http://localhost:5050	Reward curves, client scores, acceptance rate, active-data usage
Kubernetes Dashboard	http://localhost:8443	Pod phase, restarts, events, exec
Karmada Dashboard (multi only)	http://localhost:32000	Propagation state, member cluster health

Every round's MLflow run also carries five cross-link tags — kfp_run_id, kfp_run_url, temporal_workflow_id, temporal_workflow_url, topology — set by log_run_context() (src/tracking/mlflow_logger.py) inside evaluate_global. They turn the first three surfaces from independent islands into one connected trail.

Start from a bad reward curve

A concrete walkthrough of the trail above, for the question that comes up most often — "this reward curve looks wrong, which worker failed?":

Open the suspicious round's run in MLflow and follow its temporal_workflow_url tag.
That opens the round's TrainRoundWorkflow in Temporal, fanned out into one WorkerWorkflow per worker — find the one that failed.
Read that WorkerWorkflow's failure reason directly (see Reading per-worker progress in Temporal below for what that looks like in the UI).
Follow the same MLflow run's kfp_run_url tag back to that round's KFP DAG for node logs and artifact lineage.

Reach for the Kubernetes Dashboard (https://localhost:8443) instead when the question is about the pod rather than the workflow — phase, restarts, events, or an interactive exec. Get a login token with fed_dashboard_token active-fed dashboard-admin (from vendor/fed-infra/lib/dashboard.sh; see that repo's README for the full function contract).

Reading per-worker progress in Temporal

Each FL round starts one TrainRoundWorkflow (workflow ID train-<8hex>-r<round>, where <8hex> is the first 8 characters of run_uid — an id run_pipeline.py generates itself per submitted run, not a value read from KFP) that fans out one WorkerWorkflow child per worker (train-<8hex>-r<round>-w<worker>). In the Temporal UI:

Open the TrainRoundWorkflow for the round you care about — it lists N WorkerWorkflow children (one per worker).
Open a child WorkerWorkflow. While its worker pod is training, the Pending Activities panel shows the launch_and_watch_pod activity with heartbeat details — {"worker_id": ..., "active": ..., "waited_s": ...} — updated every 5 seconds. This live heartbeat is what replaces the old PyTorchJob path's opaque 20-minute ceiling: you can see exactly which worker is still running and for how long, instead of waiting on one fleet-wide timer.
If the activity code itself fails or times out (e.g. an unrecoverable Kubernetes API error, or the pod-watch timeout), that WorkerWorkflow shows the retry attempt and the root-cause failure message for that worker specifically, not a generic whole-fleet error. Temporal owns retry entirely: each worker Job runs with backoffLimit: 0 and restartPolicy: Never, so a lost pod (killed, evicted, node drain) fails the Job immediately instead of being silently replaced under Kubernetes' own Job-level self-healing. A retried activity finds that failed Job (Job names are deterministic), deletes it, waits for the deletion to complete, and creates a fresh one — so the WorkerWorkflow's retry attempt and root-cause message reflect that specific pod's failure, not a mid-round replacement Kubernetes made on its own. (This replaces an earlier gate finding where a force-killed pod was absorbed by Kubernetes' Job controller before the 5-second poll ever saw it — see .superpowers/sdd/2026-08-12-p1-temporal-orchestration/gate-fixes-report.md, F5.)
The underlying Kubernetes Job for each worker is named deterministically: aflw-<8hex>-r<round>-w<worker> — one Job per (round, worker), safe to re-attach to on retry instead of racing a second Job onto the same MinIO keys.

# Inspect worker Jobs/Pods for the current pipeline run
kubectl get jobs -n active-fed -l app=active-fl-worker
kubectl get pods -n active-fed -l app=active-fl-worker

What to see in MLflow (Local & K8s)

Because both the local runner (Use Case 1) and the pipeline (Use Case 2) use the exact same MLflow tracking logic, you will see the exact same metrics and artifacts tracked over time in the MLflow UI (regardless of whether it's running locally in ./mlruns or remotely in Kubernetes).

Each round logs:

global_eval_reward_mean / _std / solved
clients_accepted / clients_rejected
effective_weight_norm
active_data_applied / active_data_n_steps / num_active_data_sources
client_<id>_score / client_<id>_improvement / client_<id>_accepted
client_<id>_own_env_reward_mean
client_<id>_target_env_reward_mean / _std
Global model checkpoint artifact in global_models/round_N/
Aggregation JSON report artifact in reports/round_N/

CI/CD

Two GitHub Actions workflows cover validation and publishing:

CI (.github/workflows/ci.yml) — runs on every push and PR to main, three parallel jobs:

Lint, test, compile — make ci (ruff + pytest + pipeline compilation), i.e. exactly what you can run locally before pushing.
Contract dry-runs — make contracts: renders both consumer contracts (infra.env, infra.env.multi) against the pinned vendor/fed-infra submodule, catching a submodule bump that breaks our contracts before it ever reaches a cluster.
Build images — builds both Docker images (no push) with GHA layer caching.

Release (.github/workflows/release.yml) — publishes the worker and aggregator images to GHCR:

On a push to main: waits for CI to finish and only publishes off a green run, tagging main and sha-<short>. A red main is never published.
On a v* tag push: publishes a semver-tagged image (e.g. 1.2.3) and creates a GitHub Release with generated notes. No floating latest tag is ever produced.

# Run the same checks CI runs, locally
make ci          # lint + test + compile-pipeline
make contracts   # dry-run infra.env + infra.env.multi against pinned vendor/fed-infra

Project Structure

src/
  agent/          PPO ActorCritic model, PPO agent, worker entrypoint
  aggregator/     Evaluator, scorer, aggregator (Active Weight + Data), MinIO collector
  experiment/     Local FL runner (no K8s)
  orchestration/  Temporal workflows/activities/worker entrypoint (per-round worker fleet)
  pipelines/      Kubeflow DSL pipeline definition (drives the fleet through Temporal)
  tracking/       MLflow helpers
config/
  local.yaml      Hyperparameters + combinations for local experiments
  k8s.yaml        Hyperparameters for Kubeflow pipelines (single topology)
  k8s-multi.yaml  Multi-cluster variant (topology: multi, pairs with infra.env.multi)
experiments/
  run_experiments.py  Local experiment orchestrator (parallel, auto-viz)
analysis/
  compare_runs.py     Comparison visualization (6 plots)
  fetch_k8s_runs.py   Download results from remote MLflow (K8s runs)
k8s/              Manifests: RBAC (worker Job + Temporal-worker access), Temporal worker Deployment
docker/           Worker + aggregator Dockerfiles
setup/            kind cluster bootstrap + teardown scripts (delegates to vendor/fed-infra)
vendor/fed-infra  Git submodule: reusable kind/KFP/Temporal/MinIO/MLflow/Karmada infra
infra.env         Consumer contract for vendor/fed-infra (single topology)
infra.env.multi   Consumer contract for the multi topology (host + Karmada members)
.github/workflows CI (lint/test/compile, contract dry-runs, image builds) + Release (GHCR)
docs/superpowers/ Design specs, implementation plans, and review records
run_pipeline.sh   Kubeflow pipeline trigger script
tests/            Unit tests across all components
results/          Experiment outputs (JSON + plots)
mlruns/           Local MLflow tracking store (created on first local run, gitignored)

Makefile Quick Reference

make install          install prod deps (uv sync --no-dev)
make install-dev      install all deps including dev (uv sync --extra dev)
make test             run unit tests
make test-fast        run tests, skipping slow training tests
make lint             ruff check src/ tests/
make fmt              ruff format src/ tests/
make type-check       mypy src/
make ci               exactly what CI runs per PR: lint + test + compile-pipeline
make contracts        dry-run infra.env + infra.env.multi against pinned vendor/fed-infra
make run-experiments  run all combinations from config/local.yaml (parallel)
make run-single       WEIGHT_MODE=X ACTIVE_DATA_MODE=Y  (single combo)
make dry-run-worker   smoke-test worker entrypoint locally (no K8s)
make mlflow-ui        launch MLflow UI against ./mlruns (http://localhost:5000)
make temporal-ui      print the Temporal UI URL (http://localhost:8233)
make run-temporal-worker  run the orchestration worker locally (no K8s Deployment)
make compare          regenerate plots from results/
make compare-k8s      fetch remote MLflow results + regenerate plots
make compile-pipeline compile Kubeflow pipeline → /tmp/active_fl_pipeline.yaml
make run-pipeline     trigger the compiled pipeline locally
make build-images     build Docker images
make load-images      load into kind cluster
make local-setup      bootstrap kind cluster (single topology, infra.env)
make local-teardown   destroy kind cluster
make multi-setup      bootstrap host + Karmada member clusters (infra.env.multi)
make multi-teardown   destroy host + member clusters
make clean            remove __pycache__
make clean-results    remove results/ directory

Disclaimer

This is a personal exploration project, not a formal research paper or production framework.

License

This project is licensed under the MIT License - see the LICENSE file for details.

Navigation Menu

Repository files navigation

Active-Federated Learning

The Problem: Why Active Learning in Federated Learning?

Experiment Results

1. Global Learning Curves

2. Worker Own-Environment Performance

Architecture & Active Methods

Method 1: Active Weight (4-Factor Scoring)

Method 2: Active Data (Fine-Tuning)

Use Case 1: Local Experiments (fast iteration, no infrastructure)

Use Case 2: Kubernetes / Kubeflow Pipeline (real distributed FL)

What runs in K8s

What to see in Kubeflow UI

Infrastructure package

Multi-cluster topology (Karmada)

Observability: Four Surfaces

Start from a bad reward curve

Reading per-worker progress in Temporal

What to see in MLflow (Local & K8s)

CI/CD

Project Structure

Makefile Quick Reference

Disclaimer

License

About

Contributors

Languages

Name		Name	Last commit message	Last commit date
.github/workflows		.github/workflows
analysis		analysis
config		config
docker		docker
docs/superpowers		docs/superpowers
experiments		experiments
k8s		k8s
results/plots		results/plots
setup		setup
src		src
tests		tests
vendor		vendor
.dockerignore		.dockerignore
.gitignore		.gitignore
.gitmodules		.gitmodules
LICENSE		LICENSE
Makefile		Makefile
README.md		README.md
infra.env		infra.env
infra.env.multi		infra.env.multi
pyproject.toml		pyproject.toml
run_pipeline.sh		run_pipeline.sh
uv.lock		uv.lock

Latest commit

History

Folders and files

Repository files navigation

Active-Federated Learning

The Problem: Why Active Learning in Federated Learning?

Experiment Results

1. Global Learning Curves

2. Worker Own-Environment Performance

Architecture & Active Methods

Method 1: Active Weight (4-Factor Scoring)

Method 2: Active Data (Fine-Tuning)

Use Case 1: Local Experiments (fast iteration, no infrastructure)

Use Case 2: Kubernetes / Kubeflow Pipeline (real distributed FL)

What runs in K8s

What to see in Kubeflow UI

Infrastructure package

Multi-cluster topology (Karmada)

Observability: Four Surfaces

Start from a bad reward curve

Reading per-worker progress in Temporal

What to see in MLflow (Local & K8s)

CI/CD

Project Structure

Makefile Quick Reference

Disclaimer

License

About

Topics

Resources

Stars

Watchers

Forks

Contributors

Languages