[논문 정리] Reasoning with Latent Diffusion in Offline Reinforcement Learning

Why This Paper?

앞에서 정리했던 GFlowNet에 이어 이번에는 디퓨전 모델을 강화학습에 더한 논문이다. 이유는 앞의 논문과 동일하다.

---

Introduction

오프라인 강화학습에서는 다양한 행동 정책으로 수집된 행동 시퀀스를 활용하여 학습이 이루어진다. 하지만 이러한 방식에서는 Q-함수의 지원 범위를 벗어난 샘플(out-of-support samples)에 대한 외삽 오류(extrapolation error) 가 발생할 수 있다. 예를 들어, Q-learning 기법에서는 최적의 정책을 추출하기 위해 argmax 연산을 사용하지만, 이 과정에서 불완전한 Q-함수를 기반으로 잘못된 정책이 학습될 가능성이 존재한다.

 

최근 Diffusion Model이 조건부 이미지 생성 분야에서 높은 성능을 보이며 주목받고 있으며, 이를 오프라인 강화학습에 적용하려는 연구가 이루어지고 있다. 그러나 기존 연구들은 원시 상태-행동(state-action) 공간에서 직접 확산(diffusion)을 수행하여 최적 행동을 선택하는 데 어려움을 겪었으며, 낮은 수준의 trajectory 표현만을 사용하여 학습 효율이 저하되는 문제가 있었다.

 

이에 본 연구에서는 강력한 Diffusion 생성 모델을 활용하여 시간적 추상화(temporal abstraction)가 가능한 RL 정책을 학습하는 방법을 제안한다. Latent Diffusion Model(LDM)의 성과를 기반으로, 잠재 공간(latent space)에서 고수준 행동을 학습하고, 이를 바탕으로 정책 디코더(policy decoder)를 통해 저수준 행동 시퀀스를 복원하는 방식을 도입한다. 이를 통해 더 효과적인 행동 학습과 신뢰성 높은 정책 생성이 가능해진다.

---

Related work

Offline RL

앞에서 말했던 것처럼 오프라인 강화학습은 suboptimal(최적이 아닌) 동선을 연결하는 과정에서 분포 이동(distributional shift, 학습 데이터의 분포와 실제 환경의 데이터 분포가 다른 경우)문제를 가진다. 이를 해결하기 위해 여러가지 기법이 제안되었다.

• Conservative Q-Learning(CQL)은 보수적인 Q-함수를 학습하여 외삽 오류를 줄이지만, 최적 정책을 학습하기 어렵다.
• Implicit Q-Learning(IQL)은 SARSA와 DQN을 trade-off하는 방식으로 외삽 오류를 줄이고자 한다. 그러나 과정에서 최적화가 어려워지는 단점이 있다.
• 반면 본 연구에서는 보수적인 Q-함수를 도입하지 않고, trade-off 없이 직접 batch-constrained Q-함수를 학습하는 방식을 제안한다.

최근 생성 모델은 텍스트 생성, 음성 합성, 이미지 생성 등 다양한 분야에서 성과를 내고 있다. Chen et al. 의 연구에서는 Q-learning이나 bootstrapping 없이 return-conditioning(보상에 따라 행동을 조절하는 방식)을 활용한 오프라인 강화학습 방식을 제안했지만 특정 상태에서 올바른 return을 추정하는 것이 어렵다는 단점이 있다.

 

본 연구에서는 return-conditioning을 사용하지 않고 batch-constraining(학습된 정책이 배치에서 벗어나지 않도록 제한)을 활용하여 데이터 분포를 모델링하는 생성 모델을 사용하며, 이를 활용해 후보 행동을 생성하고 외삽 오류 없이 Q-함수를 학습하는 방식을 제안한다. 샘플링 과정에서 지원 범위를 벗어난 샘플이 포함될 위험이 있었으나, 본 연구에서는 잠재 확산 모델(latent diffusion model)을 도입하여 이를 해결한다.

---

Diffusion Probabilistic Models

Diffusion 모델은 이미지 생성, 초해상도, 이미지 복원 등의 분야에서 강력한 생성 모델로 자리 잡았으며, 기존의 VAEs 및 GANs보다 더 나은 안정성과 표현력을 제공한다.

 

최근 연구에서는 Diffusion을 오프라인 강화학습에 적용하여 suboptimal 행동 시퀀스를 연결하는 방법이 제안되었으나, 기존 연구들은 외삽 오류(extrapolation error) 문제나 분포 이동(distributional shift) 문제를 해결하지 못했다.

 

본 연구에서는 Latent Diffusion을 활용하여 Q-함수를 잠재 공간에서 학습하는 방법을 제안하며, 긴 시퀀스(stitching) 연결과 신용 할당(credit assignment) 을 효과적으로 수행하는 방법을 탐구한다.

---

Preliminaries

Diffusion Probabilistic Models

Diffusion 모델은 잠재 변수 기반 생성 모델로, 확률 분포 $p(x)$에서 샘플을 생성하는 모델이다. 이 모델은 확률 분포에서 샘플을 생성하는 과정을 가우시안 노이즈를점진적으로 타깃 분포로 변환하는 방식으로 수행한다. 확률 밀도 함수는 아래와 같이 표현된다.

$$p_\psi (x_0) := \int p_\psi (x_{0:T})dx_{1:T}$$

 

여기서 $x_0, ..., x_T$는 잠재 변수이고, 모델은 고정된 마르코프 체인(fixed Markov chain)을 통해 가우시안 노이즈를 주입하여 근사 사후 확률 $q(x_{1:T}|x_0)$을 정의한다. 이 과정을 전방 확산 과정(forward diffusion process)라고 한다.

 

시간 $t = 1$부터 $T$까지 반복해서 노이즈를 추가하는 과정이며, 수식은 아래와 같다.

$$q(x_{1:T}|x_0) := \prod_{t=1}^{T}q(x_t|x_{t-1}), \quad q(x_t|x_{t-1}) := \mathcal{N}(x_t; \sqrt{1-B_t}x_{t-1}, B_tI)$$

• $B_t$ : 작은 노이즈를 추가하는 정도
• $x_{t-1}$ : 이전 단계의 데이터
• $x_t$ : 노이즈가 더 추가된 데이터
• $I$ : 단위 행렬

그리고 Diffusion 모델은 목표 확률 분포 $p(x)$에서 샘플을 생성하기 위해 역방향 확산 과정(reverse diffusion process)을 수행한다. 먼저 가우시안 노이즈에서 시작한 후 점진적으로 노이즈를 제거하는 방식으로 타깃 분포 내의 샘플을 생성한다.

$$p_\psi(x_{0:T}) := p(x_T)\prod_{t=1}^{T}p_\psi(x_{t-1}|x_t), \quad p_\psi(x_{t-1}|x_t) := \mathcal{N}(x_{t-1};\mu_\psi(x_t,t), \Sigma_\psi(x_t,t))$$

 

여기서 $\mu_\psi(x_t,t)$와 $\Sigma_\psi(x_t,t)$는 학습 가능한 매개변수이며 각각 다음을 의미한다.

• $\mu_\psi(x_t,t)$  : 신경망이 예측한 $x_{t-1}$의 평균값
• $\Sigma_\psi(x_t,t)$  : 예측의 불확실성을 나타내는 공분산 행렬

Diffusion 모델의 목표는 $\mu_\psi$를 학습하는 것이다. 즉 노이즈를 제거하고 원래 이미지를 복원하는 방법을 배우는 것이 모델의 핵심이다.

 

이러한 Reverse Diffusion Process가 올바르게 작동하기 위해서는 손실함수가 필요하다. Diffusion 모델에서는 아래와 같은 손실 함수를 사용한다.

$$\mathcal{L}(\psi) = \mathbb{E}_{t \sim [1,T], x_0 \sim q(x_0), \epsilon \sim \mathcal{N}(0, I)}||\epsilon - \epsilon_\psi(x_t,t)||^2$$

• $x_t$ : 원래 데이터 $x_0$에 노이즈를 추가한 데이터
• $\epsilon$ : 실제 노이즈 값
• $\epsilon_\psi(x_t,t)$ : 모델이 예측한 노이즈
• $||\epsilon - \epsilon_{\psi}(x_t,t)||^2$ : 실제 노이즈와 모델이 예측한 노이즈의 차이 (MSE Loss)

기존 Diffusion 모델은 높은 차원의 데이터를 직접 다루기 때문에 계산량이 매우 크다. 이를 해결하기 위해 Latent Diffusion 방법을 사용한다.

• 데이터 $x$를 낮은 차원의 잠재 공간 $z$에 인코딩
• 잠재 공간에서 Diffusion을 수행한 후 다시 복원(디코딩)

---

Offline Reinforcement Learning

강화학습(RL) 문제는 마르코프 결정 프로세스(MDP, Markov Decision Process)로 정의될 수 있다. MDP는 아래와 같은 5가지로 구성되며 튜플 $(\rho_0, S, A, r, P, \gamma)$로 표현된다.

• $\rho_0$ : 초기 상태 분포
• $S$ : 상태(state)들의 집합
• $A$ : 행동(action)들의 집합
• $r$ : 보상 함수(reward function), $S \times A \to \mathbb{R}$
• $P$ : 상태 전이 확률(transition function), 특정 상태에서 특정 행동을 했을 때 다음 상태로 이동할 확률, $S \times A \times S \to [0,1]$
• $\gamma$ : 할인 계수, 미래 보상의 중요도를 조절하는 값(0과 1 사이의 값)

강화학습의 목표는 정책(policy) $\pi$를 학습하는 것이며, 각 상태에서 최적의 행동을 선택하는 매핑 $\pi : S \to A$를 학습하여 누적 보상을 최대화하는것이다.

 

오프라인 강화학습은 에이전트가 환경과 직접 상호작용하지 않고, 주어진 데이터셋만을 사용하여 새로운 정책을 학습한다. 오프라인 강화학습은 행동 복제(Behavioral Cloning)와는 다르게 주어진 데이터를 바탕으로 행동 정책을 개선하는 것을 목표로 한다. 

 

그러나 행동 정책이 사용되지 않았던 지원 범위 밖의 행동이 학습 과정에서 발생할 수 있으며 이러한 경우 Q-함수를 잘못 추정하여 외삽 오류(extrapolation error)가 발생할 수 있다.

---

Latent Diffusion Reinforcement Learning

Two-Stage LDM training

Latent Representation and Low-Level Policy. 첫 번째 단계는 잠재 경로 표현(latent trajectory representation)을 학습하는 것이다. 길이 $H$의 경로 집합 $\tau_H$가 상태 및 행동 시퀀스로 구성된 경우, 저수준 정책 $\pi_\theta(a|s,z)$를 학습하여 $z$가 경로에서 고수준 행동(high-level behaviors)을 표현하도록 하는 것이 목표이다. 이를 위해 $\beta$-Variantional Autoencoder(VAE)를 사용하며 손실함수로 Evidence Lower Bound(ELBO)를 사용한다.

$$L(\theta, \phi, w) = \mathbb{E}_{\tau_H \sim D}\left [ \mathbb{E}_{q_\phi (z|\tau_H)}\sum_{t=0}^{H-1} \log \pi_\theta (a_t|s_t,z) - \beta D_{KL}(q_\phi (z|\tau_H)||p_\omega (z|s_0)) \right ]$$

• $q_\phi (z|\tau_H)$ : 주어진 경로 $\tau_H$에 대한 잠재 변수 $z$의 근사 사후 확률
• $p_\omega (z|s_0))$ : 초기 상태 $s_0$에서 조건부 가우시안 사전 확률

BCQ(Behavior Cloning with Q-learning)와 달리, 이 방법은 VAE의 조건부 가우시안 사전 확률을 정책 생성 모델로 사용하지 않고, 잠개 공간에서만 확산을 수행하도록 설계되었다.

 

본 연구에서는 잠재 확산(latent diffusion)을 사용하여 더욱 표현력이 높은 잠재 공간을 학습하고, 다양한 행동 변화를 포착할 수 있도록 개선했다.

 

Latent Diffusion Prior. Diffusion 모델을 학습하기 위해, 상태-잠재 변수 쌍 $(s_0, z)$을 포함하는 데이터셋을 수집한다. 여기서 $\tau_H$는 길이 $H$의 경로이고 $z \sim q_\phi (z | \tau_H)$는 이전 단계에서 학습된 VAE 인코더를 통해 샘플링된 잠재 변수이다. 또한 $s_0$는 $\tau_H$의 첫 번째 상태이다. 우리는 $p(z|s_0)$의 사전 확률을 모델링하고자 한다.

 

이어서 조건부 잠재 확산 모델 $p_\psi (z|s_0)$를 학습하는데 여기에서 노이즈 제거 함수 $\mu_\psi (z_t, s_0, t)$를 학습하여 현재 확산된 잠재 변수 $z_t$와 확산 시간 단계 $t$를 기반으로 원래 잠재 변수 $z_0$을 예측한다. 여기서 손실함수로는 아래와 같은 식을 사용한다. 이 식이 작을수록 원래의 값 $z_0$과 가깝다는 뜻이다.

$$\mathcal{L}(\psi) = \mathbb{E}_{t \sim [1, T], \tau_H \sim D, z_0 \sim q_\phi(z|\tau_H),z_t \sim q(z_t|z_0)} \left [ min\{SNR(t), \gamma\}||z_0 - \mu_\psi(z_t, s_0, t)||^2 \right ]$$

 

아래의 그림은 Latent Diffusion Reinforcement Learning의 동작을 그림으로 간단하게 나타낸 것이다.

1. VAE를 이용하여 행동들의 경로 시퀀스를 잠재 모델 z로 변환한다.
2. 변환된 z들을 diffusion 모델을 이용하여 학습한다.
3. 학습된 모델을 통해 z를 샘플링하고 이것을 디코딩해서 행동을 진행한다.

---

Latent Diffusion-Constrained Q-Learning (LDCQ)

Batch-Constrained Q-Learning(BCQ)에서는 대상 Q-함수가 특정 상태에서 수행한 행동만을 사용하여 최적화되도록 제한된다.

$$\pi(s) = \arg\max_aQ(s,a) \quad (s, a) \in D$$

결정론적 MDP 환경(이산적인 환경)에서는 이론적으로 BCQ가 최적의 batch-constrained policy로 수렴한다. 하지만 현실에서는 특히 상태가 연속적인 경우 데이터셋에 존재하는 행동만을 정책(policy)로 제한하는 것이 어렵다. 따라서 데이터에서 학습된 정책 $\pi_psi(a|s)$이 필요하며 이 정책(policy)에서 샘플링된 행동들이 후보군으로 사용된다.

$$\pi(s) = \arg\max_{a_i \sim \pi_\psi(a|s)}Q(s,a_i)$$

 

본 연구에서는 잠재 공간에서의 Q-함수를 학습한다. 이는 주어진 상태 $s$에서 잠재 변수 $z$를 샘플링한 후 행동 시퀀스의 Q-값을 평가하는 방식이다. 확산 사전 확률(diffusion prior distribution) $p_\psi(z|s)$에서 여러 개의 후보 $z$를 생성하고 학습된 Q-함수를 이용해 최적의 $z$를 선택한다. 이후 선택된 z를 사용하여 저수준 정책(low-level policy) $\pi_\theta(a_i|s_i, z)$가 H-step 동안 행동(action)을 생성하게 만든다.

 

Training. 우선 데이터 셋 $D$에서 길이 $H$의 경로들을 수집하여 재생 버퍼(replay buffer) $\mathcal{B}$를 구축한다. 여기서 우리는 튜플 $(s_t, \textbf{z}, r_{t:t+H},s_{t+H})$ 로 이루어져 있으며 각각 뜻은 아래와 같다.

• $s_t$ : 초기 상태
• $\textbf{z}$ : 잠재 변수($q_phi(\textbf{z}|\tau_H)$에서 샘플링됨)
• $r_{t:t+H}$ : 감가된 보상의 합
• $s_{t+H}$ : 최종 상태

Q-함수는 TD(Temporal-Difference)를 이용하여 업데이트되며 확산 사전 확률에서 $\textbf{z}_i$를 샘플링하여 argmax를 계산한다.

$$Q(s_t, \textbf{z}) \leftarrow (r_{t:t+H} + \gamma^HQ(s_{t+H}, \underset{\textbf{z}_i \sim p_\psi(\textbf{z}|s_{t+H})}{\arg\max} (Q(s_{t+H}, \textbf{z}_i))))$$

 

Clipped Double Q-Learning을 사용하여 Q-함수의 과대추정을 방지하고 우선순위 경험 재생(Prioritized Experience Replay)을 사용하여 희소 보상 환경에서 학습 속도를 향상시켰다.

 

Policy Execution. 생성된 잠재 변수들은 학습된 Q-함수에 의해 점수가 매겨지며, 그중에 최적의 잠재 변수 $\textbf{z}_{max}$를 선택한다. 이 $\textbf{z}_{max}$를 VAE의 디코더 $\pi_\theta(a|s, \textbf{z}_{max})$를 통해 디코딩하여 H 단계의 행동 시퀀스(action sequences)를 생성하고 이를 순차적으로 실행한다.

---

Experimental Evaluation and Analysis

Temporal abstraction induces multi-modality in latent space

만약 데이터셋에 여러 행동(behavior modes)가 포함된 데이터셋을 사용할 경우, 비효율적인 행동 시퀀스가 생성될 수 있다. 따라서 이러한 행동들을 구별할 수 있어야 한다. 이를 실험적으로 분석하기 위해, kitchen-mixed-v0 작업을 사용하여 9 DoF 로봇 팔이 전자레인지, 주전자, 버너, 스위치를 순차적으로 조작하여 목표 구성을 달성하는 것을 목표 상태를 달성하는 환경을 설정했다.

 

시퀀스 길이 H가 증가하면 서로 다른 행동들이 분리될 것이라고 가정한다. 아래 사진을 보면 VAE 인코더의 잠재 공간을 2D PCA로 시각화한 결과를 확인할 수 있다. H가 증가할수록 정책이 특정 행동을 수행할 수 있도록 구별할 수 있게 되는 것을 볼 수 있다.

 

결과적으로 이 실험을 통해 시간적 추상화(temporal abstraction)을 활용하면 Q-함수 학습이 더 쉬워지고, 최적 성능도 향상됨을 확인할 수 있다.

---

LDMs address multi-modality in latent space

이 섹션에서는 잠재 확산 모델(latent diffusion model)이 다중 모드 분포(multi-modal distributions)를 VAE보다 더 효과적으로 모델링할 수 있음을 실험적으로 증명한다. 이를 위해 앞에서 사용했던 kitchen-mixed-v0에서 초기 상태에 대한 잠재 변수 분포를 PCA로 시각화해서 분석했다. 아래의 사진에서 그 결과를 볼 수 있다.

확산 모델(latent diffusion)에서는 각 모드를 정확하게 샘플링하는 반면, VAE에서는 분포가 넓게 퍼지면서 지원되지 않는 영역까지 샘플링하는 모습을 볼 수 있다. VAE를 사용할 경우 외삽 오류의 발생 가능성이 늘어나는 것이다. 그렇기 때문에 latent diffusion을 이용하여 Q-함수 학습 시 더 안정적인 정책을 학습이 가능하다.

---

Performance improvement with temporal abstraction

이 섹션에서는 시간적 추상화(temporal abstraction)이 강화학습 성능에 미치는 영향과 Diffusion 모델이 이러한 잠재 공간을 얼마나 효과적으로 모델링하는지를 분석한다. BCQ-H(시간적 추상화를 적용한 Batch-Constrained Q-Learning 변형 모델)과 LDCQ(본 연구에서 제안한 Latent Diffusion-Constrained Q-learning)을 비교한다. 두 모델은 동일한 VAE 구조를 사용하며 BCQ-H의 조건부 가우시안 prior를 diffusion 모델의 prior와 비교하는 방식으로 실험을 진행했다.

 

1. $H$가 증가할수록 둘다 성능이 향상되었다. 그러나 일정 수준 이상에서는 성능이 저하되기도 하였는데 이는 디코더의 한계일 가능성이 있다.
2. $H = 1$ 에서는 BCQ-H와 LDCQ의 성능이 유사했다. 잠재 공간 분포(latent distribution)가 정규 분포(normal distribution)에 가까워지기 때문이다.
3. $H$가 증가할수록 LDCQ가 더 나은 성능을 보였다. diffusion prior는 더 복잡한 잠재 분포를 효과적으로 샘플링이 가능하다.
4. BCQ-H도 $H$ 증가 시 성능이 향상되지만 LDCQ보다는 낮은 성능을 유지했다.

---

Offline RL benchmarks

이 섹션에서는 Latent Diffusion Reinforcement Learning(LDRL) 기법이 오프라인 RL 벤치마크에서 얼마나 효과적인지를 평가한다. 이를 행동 복제(Behavior Cloning) 및 최신 오프라인 RL 방법들과 비교한다. Batch Constrained Q-Learning(BCQ) (Fujimoto et al. 2019), Conservative Q-LearningCQL) (Kumar et al. 2020), Implicit Q-Learning(IQL) (Kostrikov et al.), Decision Transformer(DT) (Chen et al. 2021), Diffuser(Janner et al. 2022), Decision Diffuser(Ajay et al.) 의 모델들과 비교를 진행했다.

 

제안한 방법이 하이퍼파라미터 튜닝에 크게 의존하지 않는다는 사실을 발견했기 때문에 시퀀스 길이 $H$만을 조정하여 실험을 진행했다. 목표 조건이 있는 변형 모델(LDGC, Latent Diffusion Goal-Conditioning)도 같이 평가를 진행했다.

 

위의 표에서는 희소 보상을 가지는 D4RL 작업의 결과를 나타낸다. 특히 Maze2d, antmaze, frankaKitchen 환경은 D4RL에서 가장 어려운 환경으로 알려져 있으며 대부분의 알고리즘이 낮은 성능을 보인다. 

 

 

목표 기반 변형 모델(LDGC):

• Maze2d 및 AntMaze 환경에서 기존 방법 대비 우수한 성능을 기록했다.
• 특히 목표 기반 탐색(goal-conditioned navigation) 작업에서 효과적이었다.

일반 LDCQ 방식:

• 기존 Diffuser 방식보다 더 안정적인 성능을 보였다.
• 특히 Kitchen 환경에서 BCQ보다 높은 성능을 기록했다.

 

추가적으로 D4RL locomotion suite 및 Adroit robotics suite에서도 방법을 평가했다. 이 환경들은 경로 연결(trajectory stitching)이 주요 목표가 아니지만, 우리의 방법(LDCQ)은 기존 오프라인 RL 기법과 경쟁력 있는 성능을 보여주었다. Maze2D 및 AntMaze는 목표 도달(goal-reaching) 작업이므로 LDGC를 적용했고, Locomotion 및 Adroit robotics suite는 해당 작업이 아니므로 LDCQ만 실행했다.

---

Limitations

1. 추론 속도가 느림

• Diffusion 기반 RL 알고리즘은 반복적인 샘플링(iterative sampling) 과정을 포함하므로 추론 속도가 느리다.
• 특히 DDPM(Denoising Diffusion Probabilistic Model)을 단순 구현하여 사용했기 때문에 속도 저하가 더 두드러진다.

2. Locomotion 환경에서 평균적인 성능

• 희소 보상(sparse reward) 환경에서는 높은 성능을 보였지만, 보행(gait) 패턴이 중요한 Locomotion Task에서는 큰 성능 향상이 없었음.
• 이는 Locomotion 환경에서 보행 주기가 정해져 있어, 시간적 추상화(temporal abstraction)가 큰 효과를 발휘하지 못했기 때문으로 추정된다.

3. Perturbation Function을 사용하지 않음

• 기존 Q-learning 기법(Fujimoto et al., 2019)에서는 Perturbation Function(작은 변동 추가)을 활용하여 행동 정책을 더 세밀하게 조정하지만, 본 연구에서는 이를 사용하지 않아 중간 난이도의 Locomotion Task에서 성능 개선이 어려웠음.

4. 고정된 시퀀스 길이 $H$ 사용

• 실험 전체에서 고정된 시간적 추상화 수준($H$)을 사용했으며, 이는 환경마다 최적의 값이 다를 가능성이 있다.
• 기존 RL 기법들은 학습 도중 동적으로 시간적 추상화 수준을 조정할 수 있는 방법을 사용하기도 한다.

---

느낀점

여러가지가 섞인 방법을 제시하는 논문이라, 읽다가 이해가 안되는 부분을 찾다 보니 여러가지 개념을 같이 공부하게 된 것 같다. 특히 디퓨전 모델은 많이 찾아본 것 같다.

 

https://arxiv.org/abs/2309.06599

Reasoning with Latent Diffusion in Offline Reinforcement Learning