Skip to content

sunblaze-ucb/FedAgent

Folders and files

NameName
Last commit message
Last commit date

Latest commit

Β 

History

157 Commits
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

FedAgent logo

FedAgent: A Library for Decentralized Agent Learning

Train LLM agents collaboratively across decentralized clients, without sharing local data.

Homepage Paper (PDF) License: Apache 2.0

Best Paper Award, AAAI 2026 Workshop on Trust and Control in Agentic AI


Updates


Overview

FedAgent is a library for federated RL training of LLM agents. It implements a federated training loop with FedAvg aggregation (plus optional client-side FedProx), a two-level heterogeneity suite (task vs environment partitioning), and federated PPO/GRPO trainers, built as a thin overlay on stock verl 0.8 (no trainer fork: verl is imported as a library and driven through its public extension points). The paper's full experiment matrix is configured (176 configs; the GRPO arms are GPU-verified end-to-end, while PPO/decentralized/larger backbones are code-audited pending GPU smoke β€” see the verification status), and the framework is extensible with your own datasets, environments, and algorithms.

FedAgent is the reference implementation for the paper, which formalizes agent heterogeneity at two structurally distinct levels (task vs environment) and derives an asymmetric robustness result: federated training is robust to task-level heterogeneity but worst-case non-robust to environment-level heterogeneity. See fedagent/docs/heterogeneity.md for the full construction.

The maintained code lives in fedagent/; this README and the fedagent/docs/ suite document it. The original verl-agent-0.3.1 implementation is preserved on the paper-reproduce-verl-agent branch; what changed and why is in fedagent/docs/migration.md.


Key Features

  • Federated GRPO and PPO on stock verl 0.8: GRPO is the default (group size G=8 via rollout.n=8, no critic); PPO (adv_estimator=gae) additionally federates the value model alongside the actor each round.
  • Two-level heterogeneity suite: task-level (Preference / Coverage / Hardness) and environment-level (Catalog-Split + 4 WebShop transition variants: BM25 field-subset, BM25 reweight, lookalike, rank-wrapper), the first systematic decomposition for agent FL.
  • FedAvg aggregation over FSDP-sharded checkpoints, plus optional client-side FedProx (a proximal term added to local training, injected non-fork via the repo-root sitecustomize.py, not a server rule).
  • Baselines built in: federated (default), centralized (one client on pooled data), and local (one pinned client, no federation), selectable from the same config.
  • Fully configurable protocol: clients N, clients/round M, local epochs E, rounds T, tasks/client |Xα΅’|, with a ready-made 176-config paper matrix.
  • Verified acceleration, results unchanged: optional speed levers (keep one trainer+vLLM process alive across rounds instead of restarting it per client, evaluate on the already-loaded engine, keep env services warm, run ALFWorld's env service as parallel replicas) make paper runs Γ—2.5–×3.5 faster; each lever was checked to leave the final model unchanged (≀ 9.3e-5, below GPU run-to-run noise). All 176 paper configs come pre-accelerated in config/paper_accelerated/, see fedagent/docs/acceleration.md.
  • Any HuggingFace backbone (paper: Qwen2.5-1.5B/3B/7B-Instruct, Llama-3.2-3B-Instruct); WebShop and ALFWorld benchmarks out of the box, each behind a per-client HTTP env service so their conflicting dependencies stay isolated from the trainer.
  • FSDP sharding (single-GPU to 4-GPU), W&B-free (metrics go to JSON / console).

Within a round, clients train sequentially by default (one subprocess per client, then FedAvg) or concurrently with parallel_clients: P (each client on its own GPU slice, numerically identical: FedAvg is order-free and data seeds are client-indexed); the loop is verl-agnostic and resumable. Extension points for new datasets, environments, heterogeneity strategies, and aggregation rules are documented in fedagent/docs/extending.md; the capability→config→source map is in fedagent/docs/features.md.


Repository layout

fedagent/                      ← the maintained verl-0.8 overlay (start here)
β”œβ”€β”€ fed/                       federated round loop (run_fed.py) + JSON metrics logger
β”œβ”€β”€ envs/                      BaseTextEnv contract + registry; tiny_guess + per-env packages:
β”‚   └── {webshop,alfworld}/    └── <env>_env.py (client) + service/ (HTTP backend) + engine/ (vendored sim)
β”œβ”€β”€ agent_loops/               GymTextAgentLoop, multi-turn rollout (verl AgentLoopBase)
β”œβ”€β”€ hetero/                    two-level heterogeneity constructions (task + environment)
β”œβ”€β”€ data/                      AgenticDataset (verl custom_cls) + per-client partitioning
β”œβ”€β”€ config/                    Hydra base, agent registry, env specs, + the 176-config paper matrix (and its accelerated twins)
β”œβ”€β”€ docs/                      full documentation suite (architecture … migration)
β”œβ”€β”€ fedprox.py                 client-side FedProx proximal term
└── main_ppo_fed.py            per-client entry: stock verl run_ppo + FedAgent hooks

sitecustomize.py               repo-root FedProx hook (auto-imported on PYTHONPATH)
data/                          shipped data assets: env_heterogeneity (holdout/lookalike) + hardness reference labels
tools/                         operator toolbox: paper-config generator, hardness-label generator, eval/log/plot helpers (+ setup/: flash-attn build, verl patch)
LICENSE Β· NOTICE Β· CITATION.cff

Per-subpackage READMEs live alongside the code (e.g. fedagent/fed/, fedagent/envs/, fedagent/hetero/); the design and the file→role map are in fedagent/docs/architecture.md.


Installation

FedAgent uses three conda envs because the trainer, WebShop, and ALFWorld have mutually conflicting dependencies; they communicate over HTTP, so each stays isolated:

Env Role Key deps
fedagent-verl08 the trainer / federated runner Python 3.12, verl 0.8, vLLM, FSDP
verl-agent-webshop the WebShop env service Python 3.10, gym 0.24, pyserini (JDK/Lucene), torch 2.6
verl-agent-alfworld the ALFWorld env service Python 3.10, alfworld, textworld, gymnasium

Full step-by-step setup (env creation, the JDK for WebShop, ALFWorld game files via alfworld-download) is in fedagent/docs/installation.md. W&B logging is removed: no tracking account or key needed.


Data

The default WebShop configs run out of the box: the small WebShop catalog files are already shipped inside the vendored env package (fedagent/envs/webshop/engine/webshop/data/), and the env-level heterogeneity holdout/lookalike sets are tracked under data/env_heterogeneity/. Two things are fetched/generated separately:

  • ALFWorld game files: one-time alfworld-download -f β†’ ~/.cache/alfworld/ (see fedagent/docs/installation.md).
  • Hardness arm trajectories: the Hardness heterogeneity configs require per-backbone task-difficulty labels at data/hardness/*.json; generate them before any hardness run with python tools/gen_hardness_trajectories.py (see fedagent/docs/reproducing.md).

Models

Backbones are HuggingFace model ids (default Qwen/Qwen2.5-1.5B-Instruct) and auto-download on first run to ~/.cache/huggingface (set HF_HOME to relocate). The main table's Llama-3.2-3B-Instruct is gated (accept its license + huggingface-cli login); on offline / air-gapped clusters, pre-fetch on a login node and pass --model-path <local snapshot>. See fedagent/docs/installation.md.


Quick Start

Run a FedAgent experiment directly with the federated runner, from the repo root inside the fedagent-verl08 env (WebShop/ALFWorld runs also need their service env available):

# 0) In-process smoke, verifies the federated loop end-to-end, no remote service
python -m fedagent.fed.run_fed --config fedagent/config/examples/tinyguess_2cl_2rd.yaml

# 1) WebShop, homogeneous, GRPO
python -m fedagent.fed.run_fed --config fedagent/config/examples/webshop/homog_long.yaml

# 2) A paper cell (WebShop, Qwen2.5-1.5B, main, GRPO) from the 176-config matrix
python -m fedagent.fed.run_fed \
  --config fedagent/config/paper/uniform/Qwen2.5-1.5B-Instruct/main/grpo/fed_webshop_grpo_total-100_cl-per-rd-2_rd-70_ep-per-cl-3_min-goals-per-cl-100_p-uniform.yaml

# 3) The same cell ACCELERATED: identical science, a fraction of the wall-clock
#    (swap paper/ -> paper_accelerated/ for any cell; see fedagent/docs/gpu_recipes.md)
python -m fedagent.fed.run_fed \
  --config fedagent/config/paper_accelerated/uniform/Qwen2.5-1.5B-Instruct/main/grpo/fed_webshop_grpo_total-100_cl-per-rd-2_rd-70_ep-per-cl-3_min-goals-per-cl-100_p-uniform.yaml

CLI flags override the YAML: --rounds N Β· --clients N Β· --n-gpus 4 Β· --base-seed S Β· --fedprox-mu 0.1 Β· --local-client-id K. Every config key is documented in fedagent/fed/README.md; hardware/run modes are in fedagent/docs/running.md.


Reproducing the paper

The paper's experiments are the 176-config matrix under fedagent/config/paper/, mirroring the original tree 1:1 (uniform/ main table across 4 backbones Γ— WebShop + ALFWorld; env_heterogeneity/, task_heterogeneity/, decentralized/ on Qwen2.5-1.5B). Each cell is one command:

python -m fedagent.fed.run_fed --config fedagent/config/paper/<family>/<...>.yaml

Every cell also has an accelerated twin at the same relative path under fedagent/config/paper_accelerated/; swap paper/ β†’ paper_accelerated/ for the same science at roughly Γ—2.5–×3.5 less wall-clock (fedagent/docs/gpu_recipes.md).

Per-table recipes, seeds, and compute estimates (~1,800 H100 GPU-hours total) are in fedagent/docs/reproducing.md, covering the main table (Local / Centralized / FedAgent Γ— four backbones Γ— WebShop + ALFWorld), the task- and environment-level heterogeneity studies, and the decentralized ablations.


Documentation

Doc Contents
fedagent/docs/architecture.md The overlay design: the round loop + hooks on stock verl 0.8, and the file→role map.
fedagent/docs/installation.md The three-conda-env setup (trainer + WebShop + ALFWorld), JDK / game-file notes.
fedagent/docs/running.md Running run_fed.py: modes, GPUs, baselines, FedProx, eval, worked examples.
fedagent/docs/gpu_recipes.md Best-practice recipes per GPU count (1 / 2 / 4) and the accelerated paper matrix.
fedagent/docs/reproducing.md Per-experiment reproduction recipes, the 176-config matrix, compute, seeds.
fedagent/docs/heterogeneity.md The two-level taxonomy and how to construct/select each arm.
fedagent/docs/configuration.md Config-file decoder and the federated-runner key reference.
fedagent/docs/features.md Each capability β†’ its config key β†’ its source file.
fedagent/docs/extending.md Extension points: new dataset/env, heterogeneity strategy, RL algorithm, aggregation rule.
fedagent/docs/migration.md What changed from the verl-agent-0.3.1 fork to stock verl 0.8, and the equivalence checks.
fedagent/docs/acceleration.md How to accelerate & why: the final recipe, why each lever works, the equivalence bar.

Citation

If you use FedAgent in your research, please cite:

@article{fedagent2026,
  title   = {Is Decentralized LLM Agent RL Robust to Heterogeneity? An Asymmetric Tale},
  author  = {Chen, Canyu and Zhu, Kangyu and Chen, Zhaorun and Zhou, Zhanhui and Diao, Shizhe and Lu, Yiping and Li, Tian and Li, Manling and Song, Dawn},
  journal = {arXiv preprint arXiv:},
  year    = {2026}
}

License

This project is released under the Apache License 2.0, see LICENSE.

Acknowledgements

FedAgent is a thin overlay on stock verl 0.8 and reuses the verl-agent environment packages. We gratefully acknowledge:

  • veRL: Β© ByteDance / the veRL authors (Apache-2.0): the base RL training framework that FedAgent imports as a library. https://github.com/volcengine/verl
  • verl-agent / GiGPO: Feng et al., Group-in-Group Policy Optimization for LLM Agent Reinforcement Learning (arXiv:2505.10978): the agent-RL environment integrations FedAgent builds on. https://github.com/langfengQ/verl-agent
  • WebShop: Yao et al., Princeton NLP (MIT License): the e-commerce agent benchmark.
  • ALFWorld: Shridhar et al., Microsoft Research (MIT License): the embodied household agent benchmark.

Full per-component attributions and license texts are aggregated in NOTICE.

About

FedAgent: A Library for Decentralized Agent Learning. πŸ† Best Paper Award at the AAAI'26 TrustAgent Workshop, πŸ† Outstanding Paper Award in the AAAI'26 PerFM Workshop.

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages