Approximated Behavioral Metric-based State Projection for Federated Reinforcement Learning

1 Introduction

In recent years, federated learning has emerged as a new approach to enable data owners to collaboratively train each one's improved local model with the help of the privacy preserved information from others [Yang et al., 2019a; Yang et al., 2019b; Li et al., 2020a; Wei et al., 2020; Lyu et al., 2020]. Federated reinforcement learning (FRL) applies federated learning principles to reinforcement learning [Zhuo et al., 2019]. In FRL, multiple clients, each with their own local environments, collaborate to learn a collective optimal policy [Qi et al., 2021].

Aggregating knowledge from clients in non-identical environments allows FRL to explore a huge state-action space, enhance sample efficiency and accelerate the learning process [Wang et al., 2020]. However, FRL faces unique challenges primarily due to the different local environments and diverse data distributions among clients. In FRL, clients may experience very different states and rewards in their own environment, resulting in diverse data distribution. This diversity may lead to significant differences in the learning model, making it difficult for clients to converge to a robust common policy [Zhao et al., 2018]. Additionally, FRL must ensure that sensitive information remains protected from exposure to other clients or the central server [Zhu et al., 2019; Anwar and Raychowdhury, 2021].

∗Corresponding author.

Previous researches found that learning representation based behavioral metric can significantly accelerate the reinforcement learning process and enhance the generality of policy [Zhang et al., 2020a; Agarwal et al., 2021; Kemertas and Aumentado-Armstrong, 2021]. This method involves learning a state projection function by evaluating the behavioral similarities between states, which are measured in terms of rewards and state transition probabilities. The state projection function is valuable to the learning process, yet it does not reveal any sensitive task-specific information. In the FRL settings, clients would not directly share the rewards and state information because of the privacy issues. Therefore, sharing the parameters of the state projection function could be a promising research direction for FRL.

In this work, we propose the Federated Reinforcement Learning with Reducing Approximation Gap (FedRAG), a novel FRL framework to share parameters of state projection functions and to learn a local behavioral metric-based state projection function for each client. We detail FedRAG's network architecture in Figure 1, emphasizing how client collaboration is achieved through shared state projection functions. The global state projection function is formed by aggregating local state projection functions, each trained with behavioral metrics to capture the unique transition dynamics and rewards of its respective environment. By integrating these locally learned features, the global state projection function reflects the diverse dynamics and rewards across different environments. Periodically, each client's local state projection function is replaced with the global state projection function, while the L2 regularization is continuously applied to maintain alignment throughout the learning process. Together, these mechanisms improve local state projection function and strategies that are robust and adaptable across varied environments.

The main contributions are as follows:

• We propose FedRAG, a novel federated reinforcement learning framework to share the projection function of states, instead of traditionally sharing the encrypted states information. Subsequent analysis show that our method is beneficial to privacy-preserving as a sideeffect.
• Under the FedRAG framework, we introduce a behavioral metric-based state projection function and develop its practical approximation algorithm in Federated Learning settings. Empirical results demonstrate our method is effective.

2 RelatedWork

Federated Learning. Federated Learning (FL) was first introduced in FedAvg [McMahan et al., 2017], where training data remains distributed across mobile devices, and a shared model is learned by aggregating locally computed updates through iterative model averaging. Subsequently, FedProx [Li et al., 2020b] addresses system heterogeneity and statistical variability in federated networks. It incorporates a proximal term into local optimizations, allowing for variable computational efforts across devices, which helps stabilize diverse local updates. To accommodate the inherent heterogeneity in FL, Per-FedAvg [Fallah et al., 2020] was developed as a personalized approach. This method adapts Model-Agnostic Meta-Learning (MAML) to provide a suitable initial model that quickly adapts to each user's local data after training. Another innovation, pFedMe [T Dinh et al., 2020] tackles the statistical diversity among clients by utilizing Moreau envelopes as client-specific regularized loss functions, effectively decoupling personalized model optimization from global model learning.

Federated Representation Learning. Recently, federated representation learning, which focuses on training models to extract effective feature representations directly from raw data, has become increasingly popular. LG-FedAvg [Liang et al., 2020] optimizes for compact local representations on each device alongside a global model spanning all devices. FedRep [Collins et al., 2021] learns a shared data representation among clients while maintaining unique local heads to enhance each client's model quality. Model Contrastive Learning (MOON) [Li et al., 2021a] improves local update consistency by maximizing alignment between representations learned from local and global models. Additionally, the Federated Prototype-wise Contrastive Learning (FedPCL) approach [Tan et al., 2022] was introduced, leveraging pre-trained neural networks as backbones to enable knowledge sharing through class prototypes while constructing client-specific representations using prototype-wise contrastive learning. FedCA [Zhang et al., 2023] aggregates representations from each client, aligning them with a base model trained on public data to mitigate inconsistencies and misalignment in the representation space across clients. TurboSVM-FL [Wang et al., 2024] accelerates convergence in federated classification tasks by employing support vector machines for selective aggregation and applying max-margin spread-out regularization on class embeddings. Despite these advancements, research in federated representation learning specific to reinforcement learning remains limited.

Federated Reinforcement Learning. Federated Reinforcement Learning enables clients to collaboratively learn a unified policy while preserving privacy by avoiding the exchange of raw trajectories. FedPG-BR [Fan et al., 2021] addresses convergence and fault tolerance against adversarial attacks or random failures in homogeneous environments using variance-reduced policy gradients. However, it does not tackle the challenges of heterogeneous environments, which is the focus of our work. To address environmental heterogeneity, Jin et al. [2022] introduced QAvg and PAvg algorithms, employing value function-based and policy gradient methods. They further proposed personalized policies that embed environment-specific state transitions into lowdimensional vectors, improving generalization and efficiency. Similarly, Tang et al. [2022] developed FeSAC, based on the soft actor-critic framework, which isolates local policies from global integration and employs trend models to adapt to regional disparities. Building on these advancements, our work focuses on learning a federated behavioral metric-based state projection function to effectively generalize across diverse environments.

BehavioralMetrics-basedRepresentationLearning. Behavioral metric-based representation learning aims to create an embedding space that preserves behavioral similarities based on transitions and immediate rewards. Bisimulation metrics [Ferns et al., 2011] measure state behavioral similarities in probabilistic transition systems for continuous statespace Markov Decision Processes (MDPs). On-policy bisimulation metrics [Castro, 2020] focus on behaviors specific to a given policy π, incorporating a reward difference term and the Wasserstein distance between dynamics models. To address the computational challenges associated with the Wasserstein distance, the MICo distance [Castro et al., 2021] was developed to compare dynamics model distributions by measuring the distance between sampled subsequent states. The Conservative State-Action Discrepancy [Liao et al., 2023] separates the learning of the RL policy from the metric itself, focusing on the most divergent reward outcomes between states taking the same actions to define similarity in the embedding space. Chen and Pan [2022] propose the Reducing Approximation Gap distance to recursively measure expected states over dynamics models, focusing on sampling from the policy π rather than the dynamics models. This approach reduces approximation errors and is particularly effective for representation learning. In our work, we apply approximation behavior metric-based representation learning to develop local state projection functions, capturing task-relevant behavioral similarities within each client's environment. Federated Learning then allows for sharing the parameters of these local projection function, enabling clients to benefit from generalized state representations across diverse environments.

3 Preliminaries

This section highlights the Federated Soft Actor-Critic (FeSAC) variant central to our research. Soft Actor-Critic (SAC) is an off-policy actor-critic algorithm based on the maximum entropy reinforcement learning framework [Haarnoja et al., 2018a]. It aims to maximize cumulative future rewards and entropy to enhance robustness and exploration while preventing convergence to suboptimal policies. FeSAC extends SAC to a federated setting, enabling collaborative training among clients operating in diverse environments while ensuring data privacy. The global environment E = {E1 ,E2 , . . . ,EN } is composed of N distinct local environments, and each client k operates within its own unique local environment Ek . The transition probabilities differ across local environments, i.e., P(sit+1 |sit , a) ̸= P(sjt+1 |sjt , a), i ̸= j.

As our study focuses on applying approximated behavioral metric-based representation learning to FRL, we introduce the state projection function when discussing FeSAC. In the scope of representation learning for deep RL, a state projection function ϕωk maps a high-dimensional state to lowdimensional vector, from which the policy πψk(a|ϕωk(s)) is learned. We configure all critic networks, target critic networks, and action networks to take the state representation ϕωk(s) as input instead of the raw state s.

Unlike traditional FRL, the objective of FeSAC is to derive a set of maximum entropy policies that are specifically optimized for their respective local environments. The target policy ˜ πk for client k in its local environment Ek is

where skt and akt represent the state and action made by client k in its local environment Ek at time t; τπk refers to the trajectory generated by the policy πk of client k, which encompasses the sequence of states and actions over time; γ k is the discount rate; αk is the entropy regularization coefficient used to control the importance of entropy; H(πk (·|ϕωk(skt ))) = E[-logπk (·|ϕωk(skt ))] represents the entropy of the policy.

To evaluate the impact of the policy on local environments, the soft state value is defined as:

where Qθk denote the local critic Q network for client k.

Each client adjusts its local Q-network to approximate the global Q-network, thus leveraging global knowledge while retaining its own characteristics:

where V¯ θ denotes use the target critic Q networks to calculate the soft state value.

In FeSAC, the target critic Q network refers to the global critic Q network, which is broadcasted by the server to all clients. The global critic Q network Q¯ θ is formed by aggregating the local critic Q networks of each client through soft updates, considering the reward differences of state-action pairs in each client's environment to obtain a value estimation in a global context:

where ϵ is the aggregation factor.

The updated local Q-network then guides the update of the local policy, which keeps the local variability as well as learning the implicit trend of the global environment:

The temperature parameter αk is adapted to balance exploration and exploitation by controlling the relative importance of the entropy term in the policy's objective. The update objective for αk in client k is as follows [Haarnoja et al., 2018b]:

where ¯H is a target entropy level to tune the degree of exploration and ¯H = -|A|.

4 Methodology

In this section, we present the problem formulation for federated reinforcement learning with heterogeneous environments, introduce the approximated behavioral metric-based state projection function, propose the FedRAG framework and provide a theoretical analysis of its privacy preserving.

4.1 ProblemFormulation

In federated reinforcement learning with heterogeneous environments, N clients each interact with their own unique local environment Ek , each modeled as a unique Markov Decision Process (MDP): {Sk ,A,Rk , P k , γ}. Each client has a unique state space Sk , reward function Rk (s, a), and state transition dynamics P k (s′ |s, a), reflecting the diversity of their environments, while sharing a common action space A and discount factor γ. A central server facilitates collaboration by periodically aggregating and distributing shared model parameters, specifically the state projection function ϕω in FedRAG. This function maps local states to a shared embedding space, enabling clients to benefit from collective learning while preserving privacy. FedRAG optimizes local policies πk (a|ϕωk (s)) by sharing the parameters of ϕω, aiming to maximize cumulative reward and entropy:

where akt ∼ πk (· |ϕωk(skt )), skt+1 ∼ P k (· |skt , akt ). To preserve data privacy, only the parameters of the state projection function ω are shared between clients and the server, while raw states, rewards, and transition dynamics remain local to each client, ensuring sensitive information is not exchanged while enabling effective federated learning.

4.2 Client RAGDistance

In FeSAC, clients across different environments share knowledge by aligning their local Q networks with the global Q network, enabling optimal local policies while adapting to changes. However, in complex environments, clients may struggle to capture task-relevant information (see Section 5.2), leading to unclear global perceptions and hindering adaptation to environmental changes.

To enhance generalization in complex environments, we introduce behavior metric-based representation learning. This approach learns robust state representations that filter out task-irrelevant background information, speeding up the learning process and improving policy generalization across diverse environments.

For each client k, behavioral metric-based representation learning is to learn a local state encoding network ϕωk : Sk → Rn with parameters ωk , which can be cast as a minimization problem of the loss between the distance on the embedding space, ˆ d(ϕωk (ski ) , ϕωk (skj )), and the corresponding behavior metric, dπ (ski , skj ) , between any pair of states ski and skj :

The Reducing Approximation Gap (RAG) distance is a behavioral metric that measures the absolute difference between the reward expectations of two states and the distance between the next state expectations of dynamics models. And it is defined as follows:

where Eaki ∼πk raki ski represents the expected reward obtained by taking action aki in state ski under the policy πk of client k,

state over the dynamics model P(ski , aki ).

Then the approximation of RAG relax the computationally intractable reward difference term without introducing any approximate gap, as shown below:

Since the reward variance Var[rski ] is computationally intractable, we can learn a neural network approximator to estimate it by assuming that the reward rsk on state sk is Gaussian distributed. Let ˆRξk(sk ) = {ˆµ(rsk), ˆσ(rsk)} be the learned reward function approximation parameterized by ξ k , which outputs a Gaussian distribution. The loss function is:

where ˆµ and ˆσ are the mean and the standard deviation, respectively.

Similarly, in order to estimate the expected next states E[ski+1 ], we learn a dynamics model ˆ Pηk(ϕωk(s), a) = {ˆµ( ˆ P a ϕωk (s) ), ˆσ( ˆ P a ϕωk(s) )} for each client, which outputs a Gaussian distribution over the next state embedding:

Based on the above approximation, the RAG loss for each client can be defined as:

where Dk represents the replay buffer or the set of data collected from environment Ek by the RL algorithm, e.g. SAC.

Considering that the behavior metric has non-zero selfdistance, the distance on the embedding space adopts the approximate form proposed in MICo [Castro et al., 2021], which produces a non-zero self-distance and helps in maintaining proximity between similar states rather than pushing them apart:

while φ is absolute angle distance and K is a hyperparameter.

Algorithm 1 FedRAG algorithm

• 1: Initialize local networks ϕωk, ϕ¯ωk, Qθk, Q¯ θk, πψk, ˆRξk, ˆ Pηk for each client k ∈ {1, 2, . . . ,N}, and global network ϕωG at the server.
• 2: Synchronize local and global parameters: ωk , ¯ωk ← ωG for each client k.
• 3: Initialize empty replay memory Dk for each client k.
• 4: while running do
• 5: for each client k do
• 6: Observe state st from local environment Ek , sample action at ∼ π(· |ϕωk (st)) and execute.
• 7: Receive reward rt ← R(st, at) and transition to next state st+1 ∼ P(· |st, at).
• 8: Store transition (st, at, rt, st+1) in Dk .
• 9: Update local networks Qθk, πψk, αk , ˆ Pηk, ˆRξk, ϕωk via gradient descent using Eq. 3,5,6,12,11,15
• 10: Softly update target networks: ¯ θk ← τQθk + (1 - τQ)¯ θk , ¯ωk ← τϕωk + (1 - τϕ)¯ωk .
• 11: if running n iterations then
• 12: Upload ωk to federated center node.
• 13: end if
• 14: end for
• 15: if in federated center node then
• 16: Aggregate global parameters: ωG ← 1N ∑Nk=1 ωk .
• 17: Broadcast updated global parameters: ωk , ¯ωk ← ωG for all clients.
• 18: end if
• 19: end while

4.3 FedRAG Framework

Under the federated learning framework, we share the parameter ω of the state projection function ϕω. The FedRAG framework operates with multiple clients and a federated central node. Each client k generates local parameters ωk for the state projection function and updates policy networks based on their local environment. The federated central node collects these local parameters ωk from all clients, aggregates them into a global distribution, and then distributes the updated global parameters back to the clients. Each client uses the state projection ϕωk(s) as input for both the actor and critic networks. We assume that global ω follows a Gaussian distribution,with each client learning only a portion of the overall distribution. Therefore, we add a Gaussian regularization term after the RAG regression function Eq. 13, leading to the new loss formulation:

-

where ωG represents the expectation of the global Gaussian distribution. The regularization term helps reduce environmental heterogeneity, thereby enhancing collaborative learning, as demonstrated in Section 5.6.

The proposed FedRAG is detailed in Algorithm 1. During the FL process, we upload ωk to the server periodically. According to the central limit theorem, we approximate the global Gaussian distribution by aggregating the mean of all local ωk at the server. The server then distributes the results to each client, aligning local learning with the global distribution. By averaging local state projection function parameters, FedRAG integrates the specialized features learned in each client's environment. Each client can maintain its own local training advantages while incorporating the global nature, and perform better when dealing with data outside of its own.

4.4 Effectiveness of Anti-attack

Note that the data we aim to protect is not directly uploaded, potentially reducing the need for additional privacy techniques such as differential privacy or homomorphic encryption. Below, we analyze the privacy-preserving properties of our approach.

One of the major issues in federated learning is preserving privacy. In our analysis, we consider the existence of semi-honest adversaries. The adversaries may launch privacy attacks to snoop on the training data of other participants by analyzing periodic updates (e.g., gradients) of the joint model during training [Zhu et al., 2019]. Such kind of attacks is referred to as Bayesian inference attack [Zhang et al., 2022].

A Bayesian inference attack is an optimization process that aims to infer the private variable Dk to best fit client k protected exposed information W S k as

where fDk|W S k (d|w) is the posterior of Dk given the protected variable W S k . According to Bayes's theorem, maximizing the log-posterior fDk|W S k (d|w) on Dk involves maximizing summation of log (fW S k |Dk (d|w)) and log (fDk (d)). The former one aims to find Dk to best match W S k , and the latter one aims to make the prior of Dk more significant. The learned conditional distribution fDk|W S k from the Bayesian inference attack reflects the dependency between W S k and Dk, which determines the amount of information that adversaries may infer about Dk after observing W S k . However, in our approach, the parameter ω that we participate in federated learning is related to the representation function ϕ of the state. From the loss LFedRAG(ϕω) in Equation 15, we can also see that ω is only related to the mapped state and reward, and has nothing to do with our private data state. Therefore, our proposed FedRAG protects the privacy of local state information to a certain extent.

5 Experiment

5.1 Experimental Settings

In this section, we evaluate the effectiveness and generalization of FedRAG using DeepMind Control Suite (DMC).

The DMC is a benchmark for control tasks in continuous action spaces with visual input [Tassa et al., 2018]. We simulated different environments by modifying key physical parameters for several tasks: pole length (cartpole-swing), torso length (cheetah-run), finger distal length (finger-spin), and torso length (walker-walk).

As described in the previous section, each client projects state observation to the embedding space by using the approximated behavioral metric-based local state projection network, and updates local SAC network for policy evaluation and improvement. We perform experiments on 2 settings: 1) Local: clients can only interact and update local network in their own different environments without information sharing; 2) Federated: clients interact with their respective environments, update local network with information sharing according to federated methods.

In our study, we render 84×84 pixels and stack 3 frames as observation at each time step. We set an episode to consist of 125 environment steps, training over a total of 4000 episodes, which equates to 500,000 steps. For each setting, we evaluate the performance of each clients in both the same and other environments every 16 local update episodes. In the federated learning scenario, every 4 episodes, clients upload their local parameters, which the server then aggregates and redistributes as global parameters.

5.2 FedRAG vs. Baseline Performance Comparison

As illustrated in Figure 2, we compared our proposed FedRAG method (λ = 0.001) with FedAvg (equivalent to FedRAG with λ = 0), FeSAC, and Local methods in the CartPole task with varying pole lengths. We assessed the average episode reward and standard deviation achieved by the clients in other environments. The results show that clients in the Local group, trained exclusively in their own environment without federated learning, struggled to adapt to other environments, resulting in the lowest performance. FeSAC had limited effectiveness in capturing task-relevant information in complex states, leading to only modest performance improvements. In contrast, FedRAG outperformed FedAvg by effectively integrating the global state projection function during local updates, resulting in significant performance gains in other environments.Unlike traditional FRL methods, FedRAG transmits only lightweight state projection parameters, significantly reducing communication overhead. While maintaining the state projection function incurs minor computational cost, it accelerates the RL process and enhances model generalization, yielding an overall benefit.

Figure 3: The results of varying lambda. In the above experiment, the training data and testing data are from environments with same setting, while in the below experiment, they are come from environments with different settings.

5.3 the Hyper-parameter of FedRAG

To investigate the impact of the regularization term in FedRAG (Eq.15), we evaluated the Local baseline and FedRAG with varying λ values in both same and other environments, as shown in Figure 3. In other environments, increasing λ strengthens the alignment between local and global state projection functions, promoting parameter sharing and enhancing cross-environment generalization. However, overly large λ constrains local updates, slowing convergence by limiting adaptation to local dynamics. The best performance was achieved at λ = 0.001. In the same environment, performance remained stable with minor fluctuations, demonstrating the robustness of FedRAG. This reveals a trade-off: larger λ promotes generalization across clients via global knowledge sharing, but may hinder specialization for local dynamics due to gradient conflicts from heterogeneous environments. Overall, while performance remained stable in the same environment across all λ values, notable improvements were observed in other environments, confirming the effectiveness of our federated approach.

5.4 Performance Improvement for Federated Learning

In Figure 4, we compare the performance of the FedRAG method (λ = 0.1/0.001) with the Local approach by evaluating average episode rewards in both the same and other environments. The Local approach limits clients to their own environments, resulting in local optimal policies that poorly generalize. In contrast, FedRAG aggregates local state projection functions on a central server to create a global state projection function. By sharing this global function during local updates, clients benefit from cross-environment knowledge sharing while maintaining data privacy. With λ = 0.1, FedRAG enhances local performance by leveraging shared knowledge to overcome local optima, while also improving performance in other environments. At λ = 0.001, FedRAG achieves the best results in other environments with minimal loss in the same environment, demonstrating strong generalization and robustness across diverse settings.

5.5 FedRAG Performance on Various DeepMind Control Tasks

To evaluate the robustness and effectiveness of our method, we conducted experiments on several tasks from DMC and compared the average episode rewards of clients using our FedRAG method with λ = 0.001 and the non-federated Local method in both same and other environments, as illustrated in Figure 5. In cartpole-swing and finger-spin tasks, FedRAG significantly outperformed the Local method in other environments while maintaining near-optimal performance in the same environment. This success stems from its federated approach, which integrates global knowledge while preserving local training advantages. In cheetah-run task, Local clients trained only on their own environments exhibited declining performance in other environments over time. In contrast, FedRAG maintained stable performance in other environments, benefiting from global knowledge. By the end of training, FedRAG outperformed the Local method in cross-environment evaluations. In walker-walk task, FedRAG demonstrated faster convergence and higher episode rewards across all environments, benefiting from federated state projection functions that enhanced task-relevant feature extraction and generalization. These results confirm the robustness and generalization of FedRAG across diverse tasks and environments.

5.6 FedRAG Performance with Increasing Clients and Environmental Heterogeneity

We evaluated FedRAG's performance in heterogeneous environments with increasing client numbers N, comparing it to FedAvg. Using N pole lengths sampled uniformly from [0.9, 1], we created diverse CartPole environments, with N ranging from 1 to 4. As N increased, greater environmental heterogeneity hindered policy convergence, while additional clients provided more learning information. As shown in Figure 6, FedAvg's performance deteriorated significantly in other environments, while FedRAG remained stable across both same and other settings, demonstrating its robustness to heterogeneity. This stability is attributed to FedRAG's behavior metric-based state projection, which captures taskrelevant features and filters out environment-specific noise, enabling effective generalization. The L2 regularization ensures local updates align with the global model, capturing shared dynamics across clients and effectively adapting the projection function as the client count increases.

6 Conclusion

Sharing the parameters of the approximated behavior metricbased state projection function enhances the performance of FRL and protects sensitive local information. In this work, we propose FedRAG, a FRL framework that shares the parameters of the state projections among clients. Under the FedRAG framework, we introduce a behavioral metric-based state projection function and develop its practical approximation algorithm in Federated Learning settings. We conduct empirical studies on several reinforcement learning tasks to verify the effectiveness of our proposed method.

Acknowledgements

This work is supported by the Science Foundation of China University of Petroleum, Beijing (Grant No. 2462023YJRC024) and the Frontier Interdisciplinary Exploration Research Program of China University of Petroleum, Beijing (Grant No. 2462024XKQY003). Zhongqi Lu is the corresponding author.