Stop Regressing: Training Value Functions via Classification for Scalable Deep RL

Value functions are a central component of deep reinforcement learning (RL). These functions, parameterized by neural networks, are trained using a mean squared error regression objective to match bootstrapped target values. However, scaling value-based RL methods that use regression to large networks, such as high-capacity Transformers, has proven challenging. This difficulty is in stark contrast to supervised learning: by leveraging a cross-entropy classification loss, supervised methods have scaled reliably to massive networks. Observing this discrepancy, in this paper, we investigate whether the scalability of deep RL can also be improved simply by using classification in place of regression for training value functions. We demonstrate that value functions trained with categorical cross-entropy significantly improves performance and scalability in a variety of domains. These include: single-task RL on Atari 2600 games with SoftMoEs, multi-task RL on Atari with large-scale ResNets, robotic manipulation with Q-transformers, playing Chess without search, and a language-agent Wordle task with high-capacity Transformers, achieving state-of-the-art results on these domains. Through careful analysis, we show that the benefits of categorical cross-entropy primarily stem from its ability to mitigate issues inherent to value-based RL, such as noisy targets and non-stationarity. Overall, we argue that a simple shift to training value functions with categorical cross-entropy can yield substantial improvements in the scalability of deep RL at little-to-no cost.

한 줄 요약

가치 함수 학습에서 회귀 대신 분류 손실을 사용하면 다양한 도메인에서 성능과 확장성이 크게 향상된다.

핵심 기여도

**분류 손실**(categorical cross-entropy)을 사용한 가치 함수 학습이 **Atari 2600**, **체스**, **로봇 조작**, **Wordle 언어 에이전트** 등 다양한 도메인에서 **기존 회귀 기반 방법 대비 30~70% 성능 향상**을 보임.
**SoftMoEs**, **ResNet**, **Q-transformers**, **Transformer** 등 대규모 네트워크에서 **확장성 개선**을 달성.
**분류 손실**이 **노이즈 타겟**, **비정상성**(non-stationarity) 문제를 완화함을 분석적으로 밝힘.
**CQL**, **DQN** 등 기존 알고리즘에 **분류 손실을 적용**하여 **성능 향상**을 검증.

핵심 아이디어

기존 가치 기반 강화 학습(RL)은 **회귀 손실**(mean squared error, MSE)을 사용해 가치 함수를 학습하지만, 이는 대규모 네트워크(예: Transformer)에서 확장성이 떨어진다. 반면, **감독 학습**(supervised learning)에서는 **분류 손실**(cross-entropy)이 대규모 모델 확장에 효과적임을 고려, 본 연구는 **가치 함수 학습을 분류 문제로 재구성**함으로써 RL의 확장성을 개선할 수 있는지 조사했다.

**가치 함수 학습을 분류 문제로 바꾸는 핵심 아이디어**는 다음과 같다:

**실수 타겟**(continuous scalar targets)을 **범주형 라벨**(categorical labels)로 변환.
**분류 손실**(categorical cross-entropy)을 사용해 **가치 함수를 학습**.
**KL 발산**(KL divergence) 최소화를 통해 **분포 학습**을 수행.

이러한 접근은 **노이즈 타겟**과 **비정상성** 문제를 완화하고, 네트워크가 더 **유연하게 학습**할 수 있도록 돕는다. 특히, **Transformer**와 같은 대규모 모델에서 **성능과 확장성**을 동시에 개선하는 데 기여함을 실험적으로 입증했다.

기술적 접근법

**가치 함수 학습**을 **분류 문제**로 재구성:
**실수 타겟**을 **범주형 분포**로 변환.
**분류 손실**(categorical cross-entropy)을 사용.
**분포 학습**:
**KL 발산** 최소화를 통해 **타겟 분포**와 **예측 분포** 간 차이를 줄임.
**모델 및 알고리즘 적용**:
**SoftMoEs**, **ResNet**, **Q-transformers**, **Transformer** 등 다양한 아키텍처에 적용.
**CQL**, **DQN** 등 기존 알고리즘에 **분류 손실**을 적용.
**하이퍼파라미터**:
**분류 라벨 간 간격**(spacing) 조정, **분포 투영**(projection) 방법 등 실험적으로 탐색.

주요 결과

**Atari 2600 단일 작업 RL**에서 **Mixture-of-Experts**(SoftMoEs)를 사용한 경우, **분류 손실** 적용으로 **30% 성능 향상**.
**Atari 다중 작업 RL**에서 **ResNet**을 사용한 경우, **1.8~2.1× 성능 향상**.
**Wordle 언어 에이전트**에서 **Transformer**를 사용한 경우, **40% 성능 향상**.
**체스 플레이**(search 없이)에서 **70% 성능 향상**.
**로봇 조작**(Q-transformers)에서 **67% 성능 향상**.
**분류 손실**은 **노이즈 타겟**과 **비정상성** 문제를 완화함을 분석적으로 입증.

의의 및 한계

본 연구는 **가치 함수 학습을 분류 문제로 재구성**함으로써 **확장성과 성능**을 동시에 개선할 수 있음을 입증했다. 특히, **Transformer**와 같은 대규모 모델에서 **기존 회귀 기반 방법 대비 훨씬 높은 성능**을 달성한 점이 주목할 만하다. 이는 **강화 학습 알고리즘 설계**에 있어 **분류 손실의 활용**이 중요한 단서를 제공한다.

하지만, **분류 손실**은 **모든 문제를 완전히 해결하지는 않으며**, **타겟 분포의 정확한 정의**나 **분류 라벨 간 간격**(spacing) 설정 등 **추가적인 연구**가 필요하다. 또한, **분류 손실**이 **모든 RL 도메인**에서 동일한 효과를 보이는지에 대한 **추가 실험**이 필요하다는 한계점도 존재한다.

실용적 활용

본 연구의 접근법은 **대규모 네트워크**(예: Transformer)를 사용하는 **로봇 조작**, **언어 에이전트**, **게임 플레이** 등 다양한 **산업 및 연구 분야**에서 활용 가능하다. 특히, **강화 학습 알고리즘의 확장성**을 향상시키는 데 기여할 수 있으며, **CQL**, **DQN** 등 기존 알고리즘에 **분류 손실을 적용**함으로써 **성능 개선**을 기대할 수 있다.