Andrii Bidochko

• Updated: June 16, 2026
• 7 min read

Personalized Observation Normalization for Federated Reinforcement Learning in Simulation Environments with Heterogeneity

Direct Answer

The paper introduces Personalized Observation Normalization (PON), a technique that lets each federated reinforcement‑learning (FedRL) client locally scale its raw state observations using its own running mean and variance. By keeping normalization statistics private, PON prevents the distortion that occurs when heterogeneous agents with different dynamics share a single global scaling factor, thereby accelerating convergence and boosting final performance.

Background: Why This Problem Is Hard

Federated reinforcement learning promises collaborative policy training without exposing raw trajectories—a crucial advantage for privacy‑sensitive domains such as autonomous robotics, healthcare simulators, and industrial control. In practice, however, the agents that participate in a FedRL round often operate in heterogeneous simulation environments. Different robots may have distinct mass, joint limits, or friction coefficients, leading to divergent state‑transition dynamics.

These differences manifest as non‑identical input distributions. When agents send their local policy gradients to a central server, the aggregation step assumes that the underlying feature scales are comparable. If one client’s observations are naturally larger (e.g., higher joint velocities) while another’s are smaller, the server’s averaging will overweight the larger‑scale updates and under‑represent the smaller ones. The result is:

• Slower convergence because the global optimizer chases a moving target.
• Unstable training, often manifested as oscillating rewards.
• Reduced final policy quality, especially for agents whose data distribution deviates most from the mean.

Traditional remedies—such as global batch normalization or sharing a single set of statistics—fail in this setting because they implicitly force all agents to adopt a common scale that does not reflect any individual’s reality. Consequently, the core challenge is to preserve the benefits of federated learning while respecting each client’s unique observation distribution.

What the Researchers Propose

The authors present Personalized Observation Normalization (PON), a lightweight, client‑side module that computes a running mean µ and variance σ² for the raw state vector observed by each agent. PON then normalizes every incoming observation s as:

ŝ = (s − µ) / √(σ² + ε)

where ε is a small constant for numerical stability. Crucially, these statistics are never transmitted to the server; they remain strictly local. The global training loop proceeds exactly as in standard FedRL: each client runs its own environment, updates its local policy using the normalized observations, and periodically sends the policy parameters (or gradients) to the aggregator.

Key components of the PON framework include:

• Local Statistic Tracker: An online algorithm (e.g., Welford’s method) that updates µ and σ² after each step.
• Normalization Layer: Inserted at the front of the policy network, ensuring that all downstream layers receive consistently scaled inputs.
• Federated Aggregator: Unchanged from baseline FedRL; it simply averages the received model weights.

By decoupling observation scaling from the federated aggregation, PON eliminates the “scale‑mismatch” problem that has plagued prior attempts at heterogeneous FedRL.

How It Works in Practice

Conceptual Workflow

1. Initialization: Each client downloads the initial global policy and creates an empty µ/σ² buffer.
2. Interaction Loop: For every environment step, the client:
   • Receives raw state s from its simulator.
   • Updates µ and σ² using the new sample.
   • Normalizes s to ŝ and feeds ŝ into the policy network.
   • Executes the action, observes reward, and stores transition data.
3. Local Optimization: After a fixed number of steps (or episodes), the client performs standard RL updates (e.g., PPO, SAC) on its policy using the normalized observations.
4. Federated Sync: The client encrypts and sends only the updated policy parameters to the central server.
5. Aggregation: The server averages the received weights, producing a new global model that is broadcast back to all clients.
6. Repeat: The cycle continues for multiple communication rounds.

Interaction Between Components

The only new data flow introduced by PON is the local statistic stream (µ, σ²). Because these values never cross the network boundary, privacy is preserved and communication overhead remains unchanged. The normalization layer is effectively a “plug‑and‑play” pre‑processor that can be attached to any existing policy architecture without redesign.

What Sets PON Apart

• Personalization at the observation level: Prior work focused on personalized model heads or fine‑tuning after aggregation; PON personalizes the input itself.
• Zero communication cost: No extra bits are sent, keeping the federated protocol lightweight.
• Compatibility: Works with any off‑the‑shelf RL algorithm and any neural‑network backbone.

Evaluation & Results

Experimental Setup

The authors benchmarked PON on a suite of heterogeneous MuJoCo tasks, including variations of the classic HalfCheetah and Ant environments. Heterogeneity was introduced by randomly perturbing physical parameters (mass, friction, joint limits) for each client, creating distinct observation distributions while keeping the reward function identical.

Four baselines were compared:

• Standard FedRL with no normalization.
• Global batch normalization (shared µ/σ²).
• Local batch normalization without federated updates (each client trains in isolation).
• Personalized model heads (only the final layer is client‑specific).

Key Findings

• Faster Convergence: PON reached 80 % of the asymptotic reward in roughly half the communication rounds required by the vanilla FedRL baseline.
• Higher Final Performance: Across all tasks, the PON‑augmented policy achieved a 12‑15 % boost in average cumulative reward compared to the best non‑personalized baseline.
• Robustness to Extreme Heterogeneity: When client dynamics diverged by more than 30 % from the mean, global batch normalization actually degraded performance, whereas PON maintained stable learning curves.
• Negligible Overhead: The additional computation for updating µ and σ² was less than 0.5 % of total training time per client.

These results collectively demonstrate that personalizing observation scaling is a simple yet powerful lever for improving federated RL in realistic, non‑uniform environments.

Why This Matters for AI Systems and Agents

For practitioners building privacy‑preserving AI agents, PON offers a pragmatic path to scale federated training without sacrificing performance. The technique directly addresses two pain points:

• Data Heterogeneity: Real‑world deployments rarely involve identical sensors or dynamics. PON lets each device speak its own “language” while still contributing to a shared policy.
• Communication Efficiency: Since no extra statistics travel over the network, existing bandwidth budgets remain intact—a critical factor for edge devices operating on cellular or satellite links.

In concrete terms, developers can integrate PON into existing pipelines on the UBOS platform overview, enabling rapid prototyping of federated agents for robotics, smart manufacturing, or autonomous vehicles. Moreover, the approach aligns with emerging regulatory expectations around data minimization, because raw sensor streams never leave the device.

Beyond pure RL, any multi‑agent system that aggregates learned models—such as AI marketing agents that personalize campaign strategies across regions—can benefit from observation‑level personalization to handle locale‑specific feature distributions.

What Comes Next

While PON marks a significant step forward, several avenues remain open for exploration:

• Adaptive Statistic Sharing: Investigate hybrid schemes where clients occasionally exchange anonymized summary statistics to accelerate convergence without compromising privacy.
• Extension to Multi‑Modal Observations: Apply PON to settings where agents receive image, LiDAR, and proprioceptive data simultaneously, requiring per‑modality normalization.
• Integration with Secure Aggregation: Combine PON with cryptographic protocols (e.g., homomorphic encryption) to further harden the federated pipeline.
• Real‑World Deployments: Validate the method on physical robot fleets or IoT sensor networks, measuring energy consumption and latency impacts.

Developers interested in building end‑to‑end federated workflows can leverage the Workflow automation studio to orchestrate data collection, model updates, and deployment across heterogeneous devices. For teams seeking to experiment quickly, the UBOS templates for quick start provide pre‑configured pipelines that include custom normalization hooks.

Finally, the research community is encouraged to release benchmark suites that capture a broader spectrum of heterogeneity—beyond MuJoCo—to foster reproducibility and accelerate innovation.

References

For the full technical details, see the original arXiv paper.

Andrii Bidochko

CTO UBOS

Andrii Bidochko is an AI entrepreneur and researcher focused on AI agents, reinforcement learning, and autonomous systems. He writes about the technologies shaping the future of machine intelligence, from frontier models and agent architectures to real-world AI applications.