FEDORA 논문 리뷰

Federated Ensemble-Directed Offline Reinforcement Learning

---

이번 포스팅에서는 NeurIPS 2024에 게재된 FEDORA 논문을 리뷰한다. 이 논문은 Federated Offline Reinforcement Learning 문제를 다루며, 이질적인(Heterogeneous) 데이터 분포를 가진 클라이언트 환경에서 어떻게 효율적으로 하나의 최적 정책(Optimal Policy)을 학습할 수 있을지를 제안한다.

1. Preliminaries & Problem Settings

기본 설정 (Settings)

Federated Offline RL의 목표는 여러 클라이언트에 분산된 오프라인 데이터셋을 이용해 하나의 글로벌 최적 정책을 학습하는 것이다

• Clients ($N$): $N$개의 클라이언트가 존재하며, 각자 로컬 오프라인 데이터셋 $\mathcal{D}_i$를 가진다.
• Data Constraints: 데이터는 프라이버시 등의 이유로 공유될 수 없으며(Server 접근 불가), 각 클라이언트는 자신의 데이터 퀄리티가 얼마나 좋은지 스스로 알 수 없다고 가정한다.
• Environment: 모든 클라이언트는 동일한 MDP와 Reward Function을 공유한다.

기존 방법론의 한계 (Issues of Federated Offline RL)

논문에서는 기존의 연합학습(Federated Learning) 방법론을 Offline RL에 그대로 적용했을 때 발생하는 세 가지 문제점을 지적한다.

1. Ensemble Heterogeneity: 단순히 클라이언트들의 정책을 평균 내는 것(Simple Averaging)은 RL에서 최적이 아니다6. Supervised Learning에서 사용하는 FedAvg 방식(데이터셋 크기에 비례한 가중치, $w_i = |\mathcal{D}_i| / \sum |\mathcal{D}_j|$)은 데이터의 '질(Quality)'이 중요한 RL에서는 적합하지 않다.
2. Pessimistic Value Computation: Offline RL은 OOD(Out-of-Distribution) 행동에 대해 보수적(Conservative)이어야 한다. 하지만 로컬에서는 OOD일지라도 글로벌 관점에서는 ID(In-Distribution)일 수 있다. 즉, "Unseen"이 항상 "Unsafe"한 것은 아니다. 클라이언트들이 지나치게 보수적으로 학습하면, 실제로는 동료(Peer)들이 경험해본 좋은 행동조차 회피하게 되는 Over-Conservatism 문제가 발생한다.
3. Communication Efficiency vs. Local Drift: 통신 효율을 위해 로컬 업데이트 횟수($K$)를 늘리면, 이질적인 데이터 분포로 인해 각 클라이언트 모델이 서로 다른 방향으로 수렴하는 Client Drift가 발생하고, 이는 글로벌 모델의 성능 저하(Parameter Averaging Collapse)로 이어진다.

2. FEDORA Algorithm

FEDORA는 위 문제들을 해결하기 위해 서버(Server)와 클라이언트(Client) 측면에서 새로운 메커니즘을 제안한다.

2.1. Server Side: Performance-based Aggregation

서버는 클라이언트들의 정책을 통합할 때, 단순히 데이터 양만 보는 것이 아니라 엔트로피 규제된 목적 함수(Entropy-regularized Objective)를 최대화하는 방향으로 가중치 $w_i$를 결정한다.

$$w_i = \frac{e^{\beta J_i^t} |\mathcal{D}_i|}{\sum_{j=1}^{|N|} e^{\beta J_j^t} |\mathcal{D}_j|}$$

 

여기서 $J_i^t = \mathbb{E}_{s \sim \mathcal{D}_i} [Q_i^t(s, \pi_i^t(s))]$는 클라이언트 $i$의 로컬 데이터셋에서 평가된 기대 수익(Q-value)이다. 즉, 더 높은 기대 보상을 주는 정책을 가진 클라이언트에게 더 높은 가중치(Softmax 형태)를 부여하되, 데이터셋 크기도 함께 고려한다. 이는 특정 클라이언트에게만 가중치가 쏠리는 것을 방지하면서도 데이터 품질을 반영하기 위함이다.

이러한 $w$를 얻은 후, global policy와 value function은 다음과 같이 결합된다.

2.2. Client Side: Optimistic Critic & Proximal Update

각 클라이언트는 TD3-BC 알고리즘을 베이스로 학습을 진행한다. 여기에 FEDORA만의 두 가지 핵심 요소가 추가된다.

1) Optimistic Critic Update

로컬 Critic을 업데이트할 때, 타겟 값 계산에 로컬 Critic $Q_i$와 글로벌 Critic $Q_{fed}$ 중 더 큰 값을 사용한다.

$$r + \gamma \max(Q_i^{(t,k)}(s', \pi(s')), Q_{fed}^t(s', \pi(s')))$$

이는 로컬에서의 불확실성을 글로벌 정보로 보완하여, "시스템 전체적으로 봤을 때 유망한 행동"에 대해서는 낙관적으로 평가하게 만든다.

2) Proximal Policy Update

Actor(정책) 업데이트 시, 기존 손실함수(RL Loss + BC Loss)에 Proximal Term을 추가한다.

$$\mathcal{L}_{prox}(\pi) = \mathbb{E}_{s \sim \mathcal{D}_i} [(\pi(s) - \pi_{fed}^t(s))^2]$$

이는 로컬 정책이 글로벌 정책($\pi_{fed}$)에서 너무 멀어지는 것(Drift)을 방지한다.

3) Decay Mechanism

만약 로컬 업데이트 후의 성능이 글로벌 정책보다 떨어진다면($J_i^{fed, t} \ge J_i^t$), 해당 클라이언트의 로컬 학습 영향력을 감소(Decay)시킨다. 이는 퀄리티가 낮은 데이터를 가진 클라이언트가 글로벌 모델을 오염시키는 것을 막아준다. 아래의 실험 결과도 확인해보자:

3. Experiments

우선 다음의 실험 비교군에 대해서 살펴보자:

Algorithm	Aggregation Scope	Key Characteristic
Fed-A	Actor Only	Critics are trained locally; only policy is shared.
Fed-AC	Actor + Critic	Standard FedAvg applied to both networks.
Fed-AC-Prox	Actor + Critic	Includes a Proximal Term to handle data heterogeneity.
HDAFL	Generic Part Only	Actor split into generic (shared) vs. specific (local) modules.
Centralized	N/A (Pooled)	Trained on combined data from all clients

Centralized 세팅보다도 더 좋은 성능이 나오는 것을 확인할 수 있다. 나머지는 논문을 참고하면 될 것이다.

4. Limitations

논문을 읽으면서 들었던 몇 가지 의문점과 한계점이다.

1) Objective의 타당성

논문에서 제안한 엔트로피 기반의 가중치 $w_i$ 산정 방식(Softmax)이 수학적으로 최적성을 보장하는지에 대한 이론적 증명이 부족하다. "이렇게 하면 잘 될 것이다"라는 직관에 의존하고 있다. 즉 RL의 assumption과, 그 optimality에 대한 수학적 보장이 부족하며 하이퍼패러미터가 많기 떄문에 잘 되지 않는 상황이 많을 것 같다.

또한, 논문에서 대부분의 offline RL 알고리즘에서 좋을 것이라고 하는데, 결국 TD3-BC로만 실험한것으로 보아, 아키텍처 자체에 대한 제한이 없다고 설명하는 것에 관한 타당성이 부족하다.

 

2) 실험 설정의 현실성

실험에서 50명의 클라이언트 중 절반은 Expert, 절반은 Medium 데이터만 가지고 있다고 가정했다. 하지만 실제 상황에서는 각 클라이언트가 Expert와 Non-Expert 데이터를 섞어서 가지고 있을 확률이 높다(Mixed Quality). "Only Expert" 클라이언트가 존재하는 상황은 다소 작위적일 수 있어, 더 일반적인 상황에서의 검증이 필요해 보인다.

 

사실 내 연구 주제를 위한 것이었고, 나는 완전히 수학적으로 학습을 도출하여, 이론적 optimality를 달성할 수 있음을 보일 것이다. 만약 내 논문이 NeurIPS에 나온다면... 리뷰해보자 ㅎㅎ