fed-playground

A modular federated learning simulation framework for research and education.

Table of Contents

• Overview
• Project Structure
• Requirements
• Installation
• Configuration
• Usage
• Running Tests
• Contributing
• License

Overview

fed-playground is a pure-Python federated learning playground designed to make it easy to prototype, analyse, and teach the core ideas behind federated learning (FL) — without requiring a cluster or real network infrastructure.

The library implements a strategy-pattern architecture: every key concern (encryption, aggregation, model type, data loading) is expressed as an abstract base class that can be swapped out independently. This makes it straightforward to answer questions such as "what happens to model divergence as I increase the number of parties?" or "how does closed-form averaging compare to gradient descent in a federated setting?"

Out of the box the framework ships five swappable component types, each an ABC with multiple pure-NumPy implementations:

• Models — linear regression (gradient-descent / closed-form / ridge), logistic regression, MLP regressor & classifier, plus SVM (Pegasos), Lasso, ElasticNet, Poisson GLM and Huber regression.
• Aggregation strategies — FedAvg (mean), plus Byzantine-robust median, trimmed mean, Krum / Multi-Krum, Bulyan, geometric median (RFA), median-of-means and centered clipping.
• Encryption schemes — passthrough baseline, Gaussian & Laplace local DP, additive secret sharing and pairwise-masking secure aggregation; FHE-ready via the optional [fhe] extra (Concrete ML).
• Attacks — sign-flip, Gaussian, IPM and A-Little-Is-Enough Byzantine poisoning, for stress-testing robust aggregation.
• Visualizers — training-history curves, model comparison, divergence analysis and DP privacy/utility frontiers.

Research-grounded components

Most of the classes above implement published algorithms; each cites its source in its docstring:

Robust aggregators and additive-masking schemes carry an is_linear_only flag: masking schemes (secret sharing, pairwise masks) hide individual updates, so order/distance-based aggregators (Krum, median, Bulyan, …) correctly refuse them — only linear aggregation (MeanAggregation) is sound over masked shares.

Runnable demos for all of the above live in examples/ (e.g. example_byzantine_robust.py, example_svm_lasso.py, example_dp_and_secure_agg.py, example_elasticnet_poisson.py, example_bulyan_mom.py, example_privacy_utility_curve.py).

Benchmarking (fedbench)

Describe an experiment in TOML and run the whole sweep with one command:

fedbench run benchmarks/robustness.toml      # attack × defense matrix
fedbench run benchmarks/impossibility.toml   # privacy × robustness frontier
fedbench run benchmarks/privacy.toml         # DP utility cost

Each run writes a results CSV and a Markdown leaderboard (NaN cells = incompatible combinations, e.g. masking secure-aggregation × order-statistic defense). The committed benchmarks/STUDY.md interprets the three — reproducible byte-for-byte under the configs' fixed seeds.

Optional: Fully Homomorphic Encryption (FHE)

The [fhe] optional extra installs Concrete ML (Zama's FHE library). This requires Python 3.11 or 3.12 (concrete-ml does not yet support 3.13+). When installed, custom EncryptionScheme subclasses can perform genuine ciphertext-level operations during aggregation. Without it the framework operates entirely in plaintext — which is correct for educational purposes.

Project Structure

fed-playground/
├── fed_playground/             # Installable Python package (import from here)
│   ├── __init__.py             # Public API + name registry (__all__)
│   └── src/
│       ├── models.py           # Model ABC + 11 regressors/classifiers
│       ├── aggregation.py      # AggregationStrategy ABC + FedAvg & 7 robust rules
│       ├── encryption.py       # EncryptionScheme ABC + DP / secret-sharing / masking
│       ├── attacks.py          # Attack ABC + Byzantine poisoning strategies
│       ├── environment.py      # Environment: high-level simulation driver
│       ├── orchestrator.py     # Orchestrator: broadcast + aggregate (+ attack hook)
│       ├── party.py            # Party: local training + encrypted updates
│       ├── dataloader.py       # DataLoader + load_dataset (synthetic/sklearn/openml/csv)
│       ├── utils_data.py       # IID + Dirichlet non-IID partitioners
│       ├── benchmark.py        # run_benchmark grid sweep + Markdown leaderboard
│       ├── cli.py              # `fedbench` CLI (run / list-components)
│       └── visualization.py    # Training / Comparison / Divergence / PrivacyUtility plots
├── examples/                   # Runnable demos (example_*.py) + bundled toy CSV
├── benchmarks/                 # TOML configs + generated leaderboards + STUDY.md
├── docs/                       # MkDocs site (api / extending / study) + design specs
├── tests/                      # pytest suite (mirrors src/ layout)
├── .github/workflows/          # CI, docs deploy, PyPI release
├── pyproject.toml              # Metadata, dependencies, tool config
├── mkdocs.yml                  # Docs site config
└── LICENSE                     # MIT

Requirements

• Python ≥ 3.11, < 3.13
• System: no compiled extensions required for the core library
• Optional FHE: concrete-ml ≥ 1.9.0 (Python 3.11 or 3.12 only, macOS/Linux)

Core Python dependencies (installed automatically):

Optional extras: [examples] (scikit-learn, tqdm), [dev] (pytest, ruff, black, mypy), [docs] (mkdocs-material, mkdocstrings), [fhe] (concrete-ml).

Installation

Using uv (recommended)

uv is the fastest way to get started — it manages the Python version and virtual environment automatically.

# 1. Clone the repository
git clone https://github.com/thecola13/fed-playground.git
cd fed-playground

# 2. Install dependencies and create the virtual environment
uv sync

# 3. (Optional) Include development tools
uv sync --extra dev

# 4. (Optional) Include FHE support (Python 3.11 or 3.12 required)
uv sync --extra fhe

Run scripts directly without activating the environment:

uv run python examples/basic_simulation.py --parties 4 --rounds 10
uv run pytest

Or activate the environment first:

source .venv/bin/activate          # Windows: .venv\Scripts\activate
python examples/basic_simulation.py
pytest

Using pip (alternative)

# 1. Clone the repository
git clone https://github.com/thecola13/fed-playground.git
cd fed-playground

# 2. Create and activate a virtual environment
python3.11 -m venv .venv
source .venv/bin/activate          # Windows: .venv\Scripts\activate

# 3. Install the library in editable mode
pip install -e .

# 4. (Optional) Install development tools
pip install -e ".[dev]"

# 5. (Optional) FHE support
pip install -e ".[fhe]"

Verify the installation:

python -c "import fed_playground; print('fed-playground installed successfully')"

Configuration

The library reads a small number of optional environment variables. Copy .env.example to .env and edit as needed — the file is git-ignored.

The library uses Python's standard logging module. To see round-by-round progress, set FED_LOG_LEVEL=INFO or configure logging in your script:

import logging
logging.basicConfig(level=logging.INFO)

Usage

Basic simulation

from fed_playground import (
    Environment,
    NoEncryption,
    MeanAggregation,
    ClosedFormLinearRegressionModel,
    DataLoader,
)

# Load data from a CSV file
loader = DataLoader(file_path="examples/test_data.csv", target_column="target")

# Build and run a 3-party federated simulation for 10 rounds
env = Environment(
    n_parties=3,
    encryption_scheme=NoEncryption(),
    aggregation_strategy=MeanAggregation(),
    model_class=ClosedFormLinearRegressionModel,
    data_loader=loader,
)

history = env.run_simulation(rounds=10)

print(history["global_loss"])   # per-round MSE on the held-out test set
print(history["party_loss"])    # per-round average local training MSE

Using synthetic data

env = Environment(
    n_parties=5,
    n_features=10,
    n_samples=500,
    encryption_scheme=NoEncryption(),
    aggregation_strategy=MeanAggregation(),
    model_class=ClosedFormLinearRegressionModel,
)
history = env.run_simulation(rounds=20)

Visualizing results

from fed_playground import TrainingHistoryVisualizer, ComparisonVisualizer

# Plot training curves (displayed interactively)
viz = TrainingHistoryVisualizer()
viz.plot(
    data={"Global Loss": history["global_loss"], "Party Loss": history["party_loss"]},
    title="Federated Training",
    xlabel="Round",
    ylabel="MSE",
)

# Save to a file instead
viz_save = TrainingHistoryVisualizer(save_dir="./plots")
viz_save.plot(data=..., filename="training.png")

Extending with a custom model

import numpy as np
from fed_playground import Model, Environment, NoEncryption, MeanAggregation

class MyRidgeModel(Model):  # the framework already ships RidgeRegressionModel
    def __init__(self, input_dim: int, alpha: float = 0.01) -> None:
        self.input_dim = input_dim
        self.alpha = alpha
        self.params = np.zeros(input_dim + 1)

    def train(self, X, y):
        n = X.shape[0]
        X_b = np.hstack([X, np.ones((n, 1))])
        I = np.eye(X_b.shape[1])
        I[-1, -1] = 0  # do not regularize bias term
        self.params = np.linalg.solve(X_b.T @ X_b + self.alpha * I, X_b.T @ y)

    def get_parameters(self): return self.params
    def set_parameters(self, p): self.params = p
    def predict(self, X): return np.hstack([X, np.ones((X.shape[0], 1))]) @ self.params
    def evaluate(self, X, y): return float(np.mean((y - self.predict(X)) ** 2))

env = Environment(
    n_parties=4,
    n_features=5,
    n_samples=200,
    model_class=MyRidgeModel,
    model_params={"alpha": 0.1},
)
env.run_simulation(rounds=5)

CLI example scripts

# Basic simulation with a table of per-round metrics
python examples/basic_simulation.py --parties 4 --rounds 10

# Divergence analysis: vary party count, display plots interactively
python examples/divergence_analysis.py \
    --data-path examples/test_data.csv \
    --features feature_1 feature_2 feature_3 feature_4 feature_5 \
    --target target \
    --instances-diff --min-instances 2 --max-instances 8

# Save plots to disk instead of showing them
python examples/divergence_analysis.py \
    --data-path examples/test_data.csv \
    --features feature_1 feature_2 feature_3 feature_4 feature_5 \
    --target target \
    --instances-diff \
    --save-dir ./results

Running Tests

# Run the full suite
uv run pytest

# Run a specific module
uv run pytest tests/test_environment.py

# Run with coverage report
uv run pytest --cov=fed_playground --cov-report=term-missing

A passing run covers the whole suite (≈160 tests across the src/ modules):

tests/test_aggregation.py ...  tests/test_attacks.py ...  tests/test_benchmark.py ...
tests/test_dataloader.py ...   tests/test_encryption.py ...  tests/test_environment.py ...
tests/test_models.py ...       tests/test_orchestrator.py ...  tests/test_party.py ...
tests/test_phase2.py ...       tests/test_utils_data.py ...  tests/test_visualization.py ...

Contributing

Setting up a dev environment

git clone https://github.com/thecola13/fed-playground.git
cd fed-playground
uv sync --extra dev

Branch naming

Commit style

Follow Conventional Commits:

feat(aggregation): add weighted FedAvg strategy
fix(dataloader): correct elif branch for DataFrame input
docs(readme): add custom model example

Pull request checklist

• pytest passes with zero failures
• ruff check . reports zero violations
• black --check . reports no formatting issues
• New public functions/classes have Google-style docstrings
• New behaviour is covered by at least one test

License

This project is licensed under the MIT License — see LICENSE for the full text.