조합형 행동공간을 위한 셀프‑어텐션 정책 네트워크 SAINT

초록

SAINT은 전역 상태에 조건화된 셀프‑어텐션을 이용해, 순서에 무관한 서브액션 집합을 표현하고 상호작용을 모델링한다. 이를 통해 수천에서 10^18에 달하는 조합형 이산 행동공간에서도 샘플 효율성을 유지하며, 기존의 독립 가정이나 고정 순서 기반 정책보다 일관된 성능 향상을 보인다.

상세 분석

본 논문은 조합형 이산 행동공간에서 발생하는 지수적 규모의 폭발 문제를 해결하기 위해, 서브액션을 “집합”으로 취급하고 순열 불변성을 보장하는 새로운 정책 구조인 SAINT을 제안한다. 핵심 아이디어는 세 단계로 구성된다. 첫째, 각 서브액션에 고정 차원의 임베딩을 할당하고, 전역 상태 s를 FiLM 방식으로 조절하여 상태‑조건화된 표현 ˜e_i를 만든다. FiLM은 파라미터 효율성이 높으며, 상태 정보를 균등하게 모든 서브액션에 주입한다는 점에서 기존의 단순 연결(concatenation)보다 학습 안정성이 좋다. 둘째, 이러한 상태‑조건화된 임베딩 행렬 ˜E∈ℝ^{A×d}를 위치 인코딩 없이 다중 헤드 셀프‑어텐션 블록에 통과시킨다. 위치 인코딩을 배제함으로써 입력 순서에 의존하지 않는 순열‑동등성을 유지한다. 어텐션 메커니즘은 각 서브액션이 다른 서브액션과 얼마나 관련 있는지를 학습적으로 판단하고, 복잡한 고차 상호작용을 효율적으로 캡처한다. 셋째, 최종 컨텍스트‑aware 표현 x_i를 각각의 Decision MLP f_i에 입력해 로그잇을 출력하고, 소프트맥스를 통해 서브액션별 카테고리 분포 π_i(a_i|s)를 얻는다. 중요한 점은 모든 서브액션이 병렬적으로 디코딩되므로, 전체 조합을 직접 모델링하는 비용 O(∏|A_i|)을 피하면서도, 각 서브액션이 다른 서브액션의 정보를 반영한 조건부 분포를 학습한다는 것이다.

알고리즘적 호환성 측면에서, SAINT은 로그우도 기반의 정책 업데이트 식을 그대로 사용한다. 따라서 PPO, A2C와 같은 온라인 정책 경사법은 물론, IQL, AWAC 등 오프라인 가치‑가중치 방식에도 바로 적용 가능하다. 다만, 전체 행동공간에 대한 기대값이나 최대값을 직접 계산하는 Q‑max 기반 정책 그라디언트는 여전히 비현실적이므로, 이러한 경우에는 정책을 팩터라이즈하거나 후보 집합을 별도로 생성해야 한다.

실험에서는 교통 신호 제어, 로봇 내비게이션, DM‑Control 이산화 작업 등 18개의 다양한 환경을 사용했으며, 행동공간 규모는 10^3부터 1.35×10^18까지 다양했다. SAINT은 모든 도메인에서 기존의 팩터라이즈드 정책(독립 가정)과 순차적 오토레그레시브 정책(고정 순서)보다 높은 평균 보상을 기록했으며, 특히 서브액션 간 복잡한 상호작용이 중요한 약물 조합 시뮬레이션에서 큰 폭의 성능 차이를 보였다. Ablation 실험을 통해 (1) 상태‑조건화 FiLM 모듈 제거 시 성능 저하, (2) 어텐션 레이어를 빼고 단순 MLP만 사용할 경우 샘플 효율성 감소, (3) 순열‑불변성을 깨는 위치 인코딩을 삽입했을 때 학습 불안정성을 확인했다.

이러한 결과는 “서브액션 인덱싱이 의미 없거나 약하게 의미 있는” 상황에서, 정책이 순열‑동등성을 유지하면서도 복잡한 상호작용을 학습할 수 있음을 입증한다. 또한, Transformer 기반의 셀프‑어텐션이 고차원 이산 행동공간에서도 계산적으로 효율적이며, 기존 RL 파이프라인에 최소한의 수정만으로 통합될 수 있음을 보여준다. 앞으로는 더 큰 스케일의 실세계 문제(예: 대규모 자원 할당, 복합 치료 설계)와, 행동‑가치 함수에 어텐션을 결합한 하이브리드 구조 탐색이 기대된다.