An Asymmetric Tale: when does federated training of AI agents break?

AI agents are escaping the demo. They book your travel, shop for your groceries, file your tickets, and poke around your filesystem. To get good, they learn from interaction — from the tasks people actually give them and the messes those tasks create. Which raises an awkward problem: the data you’d most want to learn from is the data nobody is willing to hand over.

A hospital won’t upload its patient workflows. A household won’t stream its daily routines to a server. A retailer won’t share its customers’ purchase histories. Each of them has exactly the experience an agent needs, and a very good reason to keep it local.

There’s a well-known answer to this shape of problem: federated learning. Don’t move the data — move the model. Every client trains on its own private data, and only the model weights (never the raw interactions) are sent to a server, averaged together, and sent back. Repeat. It worked for predicting the next word on phone keyboards. The question we set out to answer was simple to ask:

1The naive hope

Train locally, average the weights. Could it really be that easy?

The simplest thing you could try is also the classic one: each client trains its own copy of the agent by trial and error on its own tasks (using a standard RL algorithm — GRPO or PPO; the choice won’t matter until the very end), then ships only the trained numbers to a server. The server averages everyone’s numbers into a new shared model and sends it back, and the cycle repeats. No data ever leaves the building. We call the whole setup FedAgent.

Start in the friendliest possible world: every client is drawing from the same pool of tasks in the same environment. Does an agent stitched together by averaging keep up with one trained the normal way, on all the data pooled together? Drag through the training and see for yourself — the line is the real success rate (how often the agent actually finishes the task) as training proceeds, on two test-beds: WebShop, a shopping site, and ALFWorld, a simulated household robot.

The two lines land on top of each other. Average a handful of independently-trained agents and you get back something as good as the centrally-trained one — sometimes a touch better, since the averaging smooths the noise. No data was shared to get there.

The catch is hiding in two words: identical clients. Real clients are never identical. The interesting question — the one this whole post is about — is what happens when they differ.

2But clients aren’t twins

Real clients differ — and that difference comes in two flavors that behave nothing alike.

Standard federated learning pictures “clients differ” as one fixed dataset sliced unevenly — your photos are mostly cats, mine mostly dogs. But an agent doesn’t have a fixed dataset. It has a task (what it’s asked to do) and a world it acts in. So “clients differ” — what researchers call heterogeneity — splits into two very different things:

Task-level heterogeneity — clients act in the same environment but ask for different things. The power user runs 500 varied tasks; the newcomer runs 15. One household mostly cooks; another mostly tidies. The world responds the same way to everyone; only the requests differ.

Environment-level heterogeneity — clients chase the same goals but the world itself responds differently. Same instruction, same action, different consequences. Picture three shopping clients all told “buy a wireless mouse under $30,” all issuing the same search:

These feel like two flavors of the same “clients differ” problem. They are not. One of them federation shrugs off. The other breaks it. Let’s push on each.

3The asymmetric tale

Crank up each kind of difference and watch what the training curves do.

First, task heterogeneity — turned up to the extreme

We can dial task differences with a single knob, all the way to a deliberately pathological split: each client does almost only one kind of task, or holds a tiny sliver of the task space, or sees only easy (or only impossible) problems. Surely that breaks the averaging trick. Toggle the extreme split on and compare it to the near-uniform baseline:

It doesn’t collapse. The extreme-split curve tracks the near-uniform one the whole way — on some axes it even finishes higher. Whether clients differ in what they prefer, how much they’ve seen, or how hard their tasks are, the averaged agent learns fine. This is the good-news half, and the paper proves it’s not luck: when clients differ only in their tasks, training the federated way and pooling all the data are the same problem in disguise. Robustness here isn’t earned — it’s the default.

Now, environment heterogeneity — the same knob, a very different result

Same federated recipe. This time we keep the tasks identical across clients and instead change how the world itself behaves — five escalating ways to make the WebShop store respond differently to the same actions (splitting the catalog, hiding product fields, re-ordering search results, slipping in near-duplicate look-alikes, scrambling the ranking). Step the severity up and watch the ceiling fall:

Under GRPO the lines fan out and slide down a clean ladder — from the 60% ceiling to 44%, 34%, 15%, and finally ~11%, where the agent has essentially stopped learning. Same algorithm that breezed through extreme task heterogeneity, now coming apart. One kind of difference was free. The other is, in the worst case, fatal.

4The one idea: Input–Dynamics Asymmetry

It comes down to where each kind of information enters the agent.

Look at how an LLM (large-language-model) agent actually receives the two things. Its task arrives in the prompt — spelled out in words it can read: “buy a wireless mouse under $30.” The agent can see it and act differently depending on what it says — type a different request, get different behavior. That’s exactly what these models are built to do. This is the explicit channel.

Its environment arrives nowhere. How the world behaves — how this particular store ranks its results, how this app responds to a click — is never written into the prompt. The agent only ever feels it indirectly, through what happens after it acts. That’s the implicit channel.

Now the asymmetry falls out almost for free. When clients differ in their tasks, their prompts differ too — so one shared model can serve all of them, the way a single chatbot already answers different questions differently. Each client’s practice teaches a different slice of the same “read the request, act on it” skill. The lessons stack up instead of clashing.

When clients differ in their environments, the prompts look identical — same goal, same screen, same words — yet the right move is different in each world. The agent has nothing in its input telling it which world it’s in, so one model is forced to commit to a single behavior for a situation that has two different correct answers. Each client’s practice now pushes the shared numbers in the opposite direction, and averaging them gives a mushy compromise that’s wrong in both.

For the curious: the two bounds, in one breath

Task side (robust). The worst per-client gap is bounded by a term that looks like \(\sqrt{(1+\chi^2)\,R_{\max} H\,(\varepsilon_{\text{approx}}+\varepsilon_{\text{opt}})}\). The only “heterogeneity” quantity in it is a \(\chi^2\) divergence between client task mixtures — and it multiplies errors (\(\varepsilon_{\text{approx}}\), \(\varepsilon_{\text{opt}}\)) that shrink toward zero as the model gets bigger and training proceeds. There is no irreducible term. Robustness is the default state.

Environment side (fragile). One can construct environments — the canonical one is a two-client bandit where client 1’s best action is client 2’s worst — with a worst-client gap bounded below by \(\Omega(R_{\max} H\,\delta)\), where \(\delta\) is how far the clients’ dynamics diverge. That floor doesn’t depend on capacity, samples, training time, or optimizer. You cannot scale your way out of it.

The empirical curves you just dragged through are these two theorems, made visible.

5Escape hatches (and why LLMs are special)

A worst-case wall isn’t an always-case wall. Three ways robustness comes back.

“Worst-case fragile” is a statement about the nastiest possible environment split. Most real ones aren’t that nasty, and the theory says exactly when you’re safe. Robustness returns whenever any of three things holds:

1. The disagreement never comes up. If the optimal way to solve the task simply avoids the states where the worlds disagree, the conflict is off the agent’s path and never bites.

2. The rankings still agree. The worlds can differ in their details yet still prefer the same action at each step. Different scenery, same best move — one policy serves both.

3. The world reveals itself. If the agent can tell which environment it’s in from what it observes, it can act differently in each — effectively moving the environment out of the implicit channel and into the explicit one.

That third one is where LLM agents do something almost sneaky. Having read a huge slice of the world during training, they figure out which world am I in on the fly, from clues sitting right in front of them — the site’s branding, the wording of an error, a Windows-style file path — the same way you can tell a British site from a US one by the spelling and the £ signs. The environment was never handed to them in the prompt, but they reconstruct it from the cues and act on it. (The technical name is in-context inference.)

And there’s a free, practical lever. Flip Widget 3 from GRPO to PPO — a more cautious training algorithm — and the collapsed runs climb back up: the 15% case to ~45%, the 11% case to ~40%. PPO refuses to take the big, lurching steps that a tug-of-war between worlds produces, so it bends where GRPO snaps. The worst case is still real — Rank Wrapper stays stubbornly low — but in practice you usually have room to maneuver.

6So what

What to take away if you build, deploy, or study these systems.

The whole thing reduces to a question you can ask about any agent you’re about to train across clients: is the thing that differs between them something the agent can read, or only something it has to feel? If it’s in the prompt, federation is on your side. If it’s in the dynamics, make it readable — or expect the wall.

A FedAgent explainer · built with the paper’s real training data · fed-agent.github.io