Reinforcement Learning with Discrete Diffusion Policies for Combinatorial Action Spaces

Motivation

• Known: diffusion model은 continuous control 분야에서 policy로 활발히 활용되어 왔고, discrete diffusion model은 텍스트나 bio-sequence 생성 등에서 강력한 생성 성능을 보여왔다. 또한 policy mirror descent(PMD)는 안정적인 policy improvement를 보장하는 이론적 기반으로 잘 알려져 있다.
• Gap: continuous control과 달리 combinatorial(대규모 이산) action space에서 discrete diffusion policy를 online RL로 학습하는 원칙적이고 효율적인 프레임워크는 아직 탐구되지 않았으며, 기존 autoregressive policy나 reward 기반 fine-tuning 방식은 causal ordering 제약, 높은 추론 비용, 단일 스텝 최적화라는 한계를 가진다.
• Why: macro-action을 이용한 hierarchical RL, multi-agent 협력, slate recommendation, DNA sequence 설계 등 실세계 문제 다수가 거대한 combinatorial discrete action space를 가지므로, 이를 안정적으로 다룰 수 있는 policy 클래스와 학습 알고리즘은 RL의 실용적 적용 범위를 크게 넓힐 수 있다.
• Approach: PMD의 최적화 목표로부터 이상적인 정규화된 target policy 분포를 정의하고, 이를 discrete diffusion model이 forward KL(FKL) 또는 reverse KL(RKL) 최소화를 통해 모방하도록 하는 distributional matching 문제로 policy update를 재구성한다.

Achievement

1. RL-D2 프레임워크 제안: PMD 기반 target distribution을 이용해 policy update를 distributional matching 문제로 재구성함으로써, RL objective 최적화와 representation learning(diffusion model)을 분리해 학습을 안정화하는 새로운 online 학습 프레임워크를 제시했다.
2. FKL/RKL 두 가지 policy improvement 방법 도출: forward KL과 reverse KL 최소화에 기반한 두 실용적 학습 방식을 유도하고 분석하여, FKL은 뛰어난 sample efficiency와 빠른 초기 수렴을, RKL은 안정적 학습과 높은 asymptotic 성능을 제공한다는 trade-off를 규명했다.
3. 다양한 challenging 도메인에서 state-of-the-art 달성: DNA sequence generation, Atari에서의 macro-action RL, Google Research Football의 cooperative multi-agent RL 등 세 가지 이질적 도메인에서 baseline 대비 동등하거나 우수한 성능과 뛰어난 sample efficiency를 입증했다.

How

• MDP (S, AK, P, γ, r, ρ0)를 정의하고 action space AK를 macro-action, multi-agent joint action, slate 등 combinatorial 구조로 일반화
• policy mirror descent(PMD)를 이용해 현재 policy와 q-function으로부터 정규화된 ideal target policy 분포를 유도
• 이 target 분포를 discrete diffusion model이 모방하도록 forward KL(FKL) 또는 reverse KL(RKL) divergence를 최소화하는 distributional matching objective로 policy update 문제를 재정의
• FKL variant는 off-policy 학습을 지원하여 데이터 재사용을 통한 샘플 효율성을 극대화
• diffusion model의 multi-step transition dynamics를 expressive prior로 활용해 대규모 combinatorial space에서 탐색(exploration)을 효율적으로 유도
• DNA sequence generation, Atari macro-action RL, Google Research Football multi-agent RL 세 벤치마크에서 실험적 검증 수행

Originality

• discrete diffusion model을 online RL의 policy parameterization으로 사용하는 원칙적이고 효율적인 프레임워크를 최초로 제시
• policy mirror descent(PMD)를 활용해 policy update를 RL objective 최적화가 아닌 distributional matching(생성 모델 학습) 문제로 재구성한 decoupled 접근이 독창적
• FKL과 RKL 두 방식의 trade-off(샘플 효율성 vs 안정성/asymptotic 성능)를 이론적·실증적으로 분석하여 discrete diffusion policy 학습에 대한 실용적 지침 제공
• autoregressive policy의 causal ordering 제약과 순차적 생성으로 인한 높은 추론 비용을 diffusion 기반 비순차적 생성으로 극복

Limitation & Further Study

• diffusion model 자체의 inference 단계(denoising steps) 수에 따른 계산 비용과 latency trade-off에 대한 심층 분석이 본문 발췌에서는 제한적으로 보임
• PMD 기반 target distribution의 이론적 수렴 보장이 discrete diffusion 근사 오차와 결합했을 때 어떤 성능 손실을 야기하는지에 대한 엄밀한 이론적 분석이 더 필요
• 실험이 DNA sequence, Atari macro-action, multi-agent football 세 도메인에 국한되어 있어, 훨씬 더 크거나 연속적/혼합형(discrete-continuous hybrid) action space로의 확장성 검증이 후속 연구로 필요
• FKL의 off-policy 학습이 제공하는 샘플 효율성과 RKL의 안정성을 결합하는 하이브리드 또는 적응적 전환 전략에 대한 추가 연구 여지 존재

Evaluation

총평: PMD 기반 distributional matching이라는 명확한 이론적 근거를 바탕으로 discrete diffusion policy를 online RL에 안정적으로 통합한 참신하고 실용적인 기여이며, 다양한 combinatorial 벤치마크에서 우수한 성능을 입증한 점이 인상적이다.