[RL] Is Conditional Generative Modeling All You Need For Decision-Making?(Decision-Diffuser)

IS CONDITIONAL GENERATIVE MODELING ALL YOU NEED FOR DECISION-MAKING?

 

Link : https://arxiv.org/pdf/2211.15657.pdf

 

diffuser 논문을 재밌게 읽었는데 후속 논문이 나와서 읽어봄

Abstract

해당 논문에서는 Conditional Generative Model이 직접적으로 sequential decision-making 문제를 해결할 수 있는지에 대해 알아본다. 그리고 의사 결정을 강화 학습 관점이 아닌, 조건부 생성 모델링(conditional generative modeling)을 통해 해결해보려고 한다. 저자들은 해당 연구를 통해 policy를 return conditional diffusion model로 모델링함으로써, Dynamic programming의 필요성을 피하고, 결과적으로 기존 offline RL에서 발생하는 많은 복잡성을 제거하는 방법을 설명한다. 또한 2개의 condition variable (각각 constraint와 skills)로 conditional diffusion modeling을 하면 오는 이점도 함께 설명한다. 두가지 서로 다른 조건 변수인 constraint와 skill을 고려하여 conditional diffusion model로 policy를 모델링하는 것에 대한 이점을 보여주는데, training 중 하나의 제약 조건이나 스킬을 condition으로 설정하면 test time시, 여러 제약 조건을 함께 만족하거나 skill의 조합을 보여줄 수 있는 behavior로 이어지게 된다고 한다. 전체적으로 해당 논문에서는 conditional generative modeling이 decision-making을 위한 강력한 도구라는 점을 보여준다.

 

 

Introduction

최근 몇 년동안 conditional generative modeling은 놀라운 발전을 보여줬는데, 이미지 모델링에서 DALL-E나 ImageGen과 같은 모델은 high-resolution의 고품질 이미지를 text description만으로도 생성해낼 수 있음을 보인 것이 그 예이다. 또한 GPT같은 경우는 일부분의 text prompt만 condition으로 주여저도 전체 paragraph를 생성할 수 있으며 Minerva 모델은 고등학교 수준의 수학 문제를 step-by-step으로 생성할 수 있다. 하지만 이런 경우는 데이터를 인터넷에서 쉽게 구할 수 있는 경우이다. Decision-making의 문제 같은 경우는 환경으로부터 적극적으로 데이터를 수집해야 하기 때문에, Generative model을 deicision making문제에 직접적으로 적용하는 것은 그리 간단하지 않다. 하지만 offline dataset을 이용하면 의사 결정 문제는 이미 생성 모델이 성공을 거둔 경우와 유사해질 수 있다.

이전 offline decision-making 연구들은 RL tool에 의존하고 있고, TD(temporal difference) learning을 사용한다. 이는 function approximation, off-policy learning, bootstrapping을 함께 사용하는데서 오는 학습 불안정성과 divergence을 쉽게 야기한다. 이는 "deadly triad"로 불려진다.

우리는 제한된 양의 데이터셋을 사용해야 하기 때문에 데이터셋의 분포 안에서 policy가 유지되도록 하는 트릭이나 휴리스틱을 사용해야 하는데, 이 때문에 offline RL을 scaling하는 것이 어려워진다. 해당 논문에서는 offline decision-making문제를 conditional generative modeling관점으로 보고, 고정된 reward labeled trajectory 데이터셋이 주어지면, return-conditional diffusion model을 학습하는 방법을 제안한다. inference time에 low-temperature sampling을 하는 classifier-free guidance를 데이터셋 내의 best behaviors를 포착하기 위해 사용하고, return을 최대화하는 trajectory를 찾는다.

offline decision-making을 conditional generative modeling관점에서 보면, 단지 return을 conditioning하는 것 이상을 할 수 있다. Figure 1을 보면 스킬, 리워드, 제약조건이 라벨된 trajectory들을 통해서 여러 스킬이 섞였거나 여러가지 제약 조건을 만족하는 trajectory를 생성해낼 수 있음을 볼 수 있다.

동심원 모양의 linear dynamics navigates를 가진 환경이 있을 때, 로봇의 state-action trajectory가 주어지면, 각각은 두 가지 제약(더 큰 원 안에 로봇 포지션이 위치한 경우와 더 작은 원에 로봇 포지션이 위치한 경우) 사항 중 하나를 만족한다. constraint-conditioned model을 학습함으로써 모델은 어떤 constraint 셋을 만족하는 trajectory를 생성해낼 수 있고, inference시, 이 learned trajectory model은 데이터셋으로부터 constraints를 합성할 수 있고 combined constraint를 만족하는 trajectory를 생성해낼 수 있게 된다.

 

 

Background

Reinforcement Learning

먼저 standard RL은 Sequential Decision-making 문제는 <initial distribution, state, action, transition function, reward, discount factor>로 이루어진 discounted MDP로 formulate되고 return을 최대로 하는 policy pi를 찾는 것이다.

 

Temporal Difference Learning

TD method는 optimal Q function을 추정하는 방법으로, 현재 reward와 다음 step에 대한 estimated Q value의 합인 TD target과 실제 value의 차이인 TD error를 최소화하는 방향으로 parametrized Q function을 학습하며(eq (1)), 학습된 Q가 최대화되는 parametric policy(\phi)를 학습한다.(eq (2)) 이 때 dataset transition D는 에이전트가 환경과 상호 작용하고, Q_\theta와 \pi_\phi가 함께 학습함에 따라 진화한다. 이 방법은 deadly triad를 사용해야 하고 learning을 불안정하게 만든다.

 

Offline RL

Offline RL setting에서는 미지의 behavior policy \mu에 의해 수집된 transition의 고정된 데이터셋에서 return을 극대화하는 정책을 찾는 것이다. TD learning을 naively 사용하면, state visitation distributuion d\pi_\phi(s)가 데이터셋 d\mu(s)의 분포에서 멀어지게 된다. 결과적으로 policy pi_\phi는 데이터셋에서 보이는 것과는 상당히 다른 행동을 취하게 된다. offline RL 알고리즘은 TD learning 절차에서 직접 D(d\pi_\phi || d\mu) 형식의 제약 조건을 부과함으로써 distribution shift를 해결한다.(여기서 D는 KL과 같은 divergence metric이다.) 제약 최적화문제는 합리적인 성능을 위해 추가적인 하이퍼라마티 튜닝이나 구현단의 휴리스틱을 필요로 한다. 이에 비해 decision diffuser는 이러한 단점이 없다. 어떤 종류의 Q function도 추정할 필요가 없으므로 TD learning에서 사용하는 방법을 완전히 피할 수 있게 된다. 또한 generative model은 maximum likelihood estimation을 통해 학습되므로 distribution shift 리스크를 피할 수 있다.

 

 

Diffusion Probabilistic Models

Diffusion model은 generative model의 한 종류로 데이터 분포 q(x)를 데이터셋으로 부터 학습한다. data-generating procedure는 predefined forward noising process q(x_k+1|x_k) := N(x_k_1; \sqrt{\alpha_k}x_k, (1-\alpha_k)I)로 모델링되고, 이 forward process의 역인 reverse process p_\theta(x_k-1|x_k):=N(x_k-1|\mu_theta(x_k, k), \sigma_k)를 learning한다. 최정족으로는 가우시안 노이즈로부터 시작해서 trained reverse process를 통해 denoising step을 반복적으로 거친 후, x를 생성해낸다.

그리고 log p_theta의 tractable variational lower-bound를 통해 diffusion model을 학습할 수 있지만 DDPM 논문에서 제안되었듯 위 eq(3)처럼 간단하게 스텝의 noise를 예측하는 방식으로 diffusion model을 학습하게 된다. 즉 reverse process p를 직접적으로 parameterization하지 않고 noise를 neural net으로 parametrize하여 noise를 estimate한다. \mu_theta는 noise \epsilon_theta의 함수로 계산될 수 있기 때문에 p_theta의 mean을 예측하는 것과 noise를 예측하는 것은 같다고 볼 수 있는 것이다.

 

Guided Diffusion

그리고 conditional data distributuion q(x|y)를 모델링하는 것은 label y의 속성을 지닌 샘플을 생성해낼 수 있도록 한다. diffusion model과  score-matching은 사실상 같다는 것을 통해 \epsilon_theta가 log p(x_k)의 gradient와 비례하다는 것을 보이고, 이는 자연스럽게 conditioning을 위한 두 종류의 method가 제안되었다. 하나는 classifier-guided이고 다른 하나는 clssifier-free이다. 전자는 추가적인 clssifier p_\theta(y|x_k)를 nosiy data에 대해 학습하여 test-time에 y label에 가까운 데이터를 생성할 수 있도록 하는 방법이고, 후자는 classifier를 학습하는 것와 diffusion model을 학습하는 것을 따로 두지 않고 original training step 안에서, conditional \epsilon_\theta(x_k, y, k)와 unconditional \epsilon_\theta(x_k, k)를 함께 학습하도록 한 방법이다. unconditional noise는 dummy value Ø가 y를 대신하는 conditional noise \epslion_\theta(x_k, Ø, k)로 표현되며, 나중에 샘플을 생성할 때, perturb noise 가 사용된다.

 

 

Generative Modeling With The Decision Diffuser

제안된 방법은, offlien RL을 distribution shift를 유발하는 TD-learning 없이 해결하기 위해, 아래와 같은 conditional generative modeling 문제로 formulate한다.

우리의 목표는 conditional dta distribution p_theta를 추정하는 것이고, 이렇게 되면 information y(\tau)로부터 trajectory x_0(\tau)의 일부를 생성해낼 수 있게 된다. y의 한 예로는, trajectory에 있는 return 값이나, trajectory에 있는 skill 또는 trajectory에 의해 만족되는 constraint가 될 수 있다. 

우리는 이제 conditional diffusion process(eq. (5))에 따라 generative model을 구축하게 된다. 일반적으로 q는 forward noising process이고 p_\theta는 revese denoising process를 나타낸다. 이에 따라 먼저 어떻게 diffusion을 decision making에 사용할 수 있는지 논의해볼 것이다. 먼저 diffusion을 위한 modeling choice에 대해 논의하고, classifier-free guidance를 활용해 best trajectory를 포착할 수 있는 방법을 의논할 것이다. 그 다음, conditional diffusion model로 구현된 서로 다른 behavior들에 대해 얘기해볼 것이고 마지막으로 실제 training detail에 대해 설명한다.

 

 

Diffusing over states

이미지에서, diffusion process는 모든 픽셀 값에 대해 적용된다. trajectory의 state, action을 모델링하기 위해 비슷하게 적용해 볼 수 있지만, RL setting에서 action을 diffusion process로 직접적으로 모델링하는 것은 몇 가지 practical issue가 있다. 먼저, RL에서 일반적으로 state는 continuous한 반면, action은 더 다양하고, 종종 discrete하다. 더욱 action sequence는 종종 joint torques로 표현되는데, 이는 더 high-frequency하고 less smooth하기 때문에 이를 예측하고 모델링하는 것은 어렵다. 이러한 practical issue로 인해, 저자들은 아래와 같이 오직 state에 대해서만 diffusion process를 적용하였다.

k는 forward time step이며 t는 trajectory \tau안에서 방문된 state를 의미한다. x_k(\tau)는 length H만큼의 trajectory의 noisy state sequence이고, x_k(\tau)는 2차원의 array이다.

 

Acting with Inverse-Dynamics

diffusion model로부터state를 샘플링하는 것은, controller를 정의하는 것을 위해서는 충분하지 않다. 그러나, policy를 추론할 수 있는 방법은 x_0(\tau)안의 어떤 타임스텝 t라도 s_t가 s_t+1로 가게 하는 action a_t를 추정하는 것으로부터 추론할 수 있다. 두개의 연속적인 state, s_t, s_t+1이 주어진다면, 우리는 inverse dynamics model을 따르는 aciton을 생성해낼 수 있다.

reverse process p_\theta를 학습하는데 사용된 offline dataset을 이 inverse dynamics f_\phi를 학습하는데도 그대로 사용할 수 있다. 직접적으로 state 분포를 diffusing하고, action을 예칙하기 위해서는 inverse dynamics model을 사용한 이 디자인은 state와 action 둘다 diffusing한 것 보다 성능이 훨씬 향상되었다고 한다. 이에 대한 analyze는 논문의 Appendix에 나와 있다.

 

Planning with Classfier-Free Guidance

데이터셋의 서로 다른 trajectory를 나타내는 diffusion model이 주어지면, 우리는 이를 어떻게 planning에 활용할 것인지에 대해 얘기할 수 있다. 모델을 planning에 사용하기 위해서는 y(\tau)를 diffusion process에 추가적으로 condition하는 것은 필수적이다. 하나의 방법은, noisy trajectory x_k(\tau)로부터 y(\tau)를 예측하기 위한 classifier p_\phi(y(\tau)|x_k(\tau))를 학습하는 것이다. y(\tau)가 trajectory의 return이 되는 경우에는, Q-function estimating이 요구된다. 이는 복잡한 dynamic programming procedure가 필요하다. DP를 피하기 위한 또 다른 방법으로는 return y(\tau)를 condition한 conditional diffusion model을 직접적으로 학습하는 것이다. 그러나 데이터셋은 sub-optimal이므로 conditional diffusion model은 optimal behavior가 아닌 behavior로 인해 오염될 것이다.

이 이슈를 해결하기 위해 해당 논문에서는 classifier-free guidance with low-temperature sampling 기법을 활용해 높은 likelihood를 가진 trajectory를 추출하며 이렇게 되면 데이터셋 안에서 best set of behavior에 상응하는 trajectory를 찾을 수 있게 된다. classifier free gudance를 구현하기 위해 가우시안 노이즈 x_K(\tau)로 시작해서 각 중간 timestep 마다 perturb noise(\hat{\epsilon})와 함께 x_k(\tau)를 x_k-1(\tau)로 정제해 x_0(\tau)를 샘플링한다.

y(\tau)을 나타내는 데이터셋의 trajectory의 best portion을 augment하고 추출하기 위해 scalar omega가 적용되며, Decision Diffuser로부터 샘플링하는 것은 RL에서 planning하는 것과 비슷해진다. 먼저 environment에서 state가 관찰되면, y와 마지막으로 관찰된 C state의 history를 conditioned한 diffusion process를 사용해 특정 horizon의 state를 샘플링한다. 마지막으로 inverse dynamcis model을 사용해, 가장 즉각적인 예측된 state에 도달하기 위해 취해야할 action을 식별한다. 이 절차는 아래 알고리즘에 설명된, standard receding-horizon control loop에서 반복된다.

Conditioning Beyond Returns

지금까지는 explicitly conditioning variable y(\tau)를 정의했다. 이는 trajectory의 return이 될 수도 있으며 관련된 제약 조건을 충족하거나 특정 동작을 보여주는 state sequence로 diffusion process를 가이드하는 것도 고려할 수 있다.

 

Maximizing Returns

return을 최대화하는 trajectory를 생성해내기 위해서, noise model에 trajectory의 return을 condition할 수 있다. 이때 return R(\tau)는 0과1사이가 되게끔 normalize해주었다고 한다. 가장 최대의 return을 conditioning하기 위해서는 R(\tau)를 1로 설정하고, 이때 return이 되는 값은 Dynamic programming이 필요한 어떠한 Q value도 사용하지 않았다. 

Satisfying Constraints

Trajectory들은 특정 goal을 만족한다던지, 특정 순서로 state를 방문한다던지, 또는 특정 state space를 피한다던지 하는 다양한 constriants를 만족할 것이다. 이 각각의 constraints들을 set C_i라고 할 때, C_i가 주어지면, noise model에 one-hot encoding을 통해 이러한 costraint를 줄 수 있다. 

trajectory들이 사용 가능한 제약 조건 중 하나만 만족하는 offline dataset으로 모델을 학습함에도, 추론시에는 여러 제약조건을 함께 만족할 수 있다.

 

Composing Skills

Skill i는 demonstrains B_i의 셋으로부터 특정되고, 특정 skill을 하는 trajectory를 생성하기 위해서 우리는 noise model에 마찬가지로 one-hot encoding으로 skill을 줄 수 있다.

우리가 학습시 individual skill을 사용하더라도, 마찬가지로 skill들이 compose하도록 추론시 조정할 수 있다. 아래와 같이 n개의 서로 다른 conditioning variable을 위한 학습된 data distribution이 있을 때, 

우리는 perturbed noise(eq (9))를 사용해서

composed data distribution:

으로부터 샘플할 수 있다. 이러한 특성은 y(\tau)가 state trajectory x_0(\tau)가 주어졌을 때 conditionally independent하다는 것을 가정한다. 하지만 저자들은 이와 같은 가정을 conditioning variable의 composition이 가능하다면, 엄격하게 충족할 필요가 없다는 것을 empirically 관찰했다고 한다.

하지만 이를 위해서는 데이터셋이 아래와 같은 것들을 만족해야 한다.

먼저, 데이터셋은 각 조건 변수 y_i(\tau)를 만족시키는 다양한 demo가 있어야 하며 이를 통해 decision diffuser는 각 조건 변수 y를 만족시키는 다양한 방법을 학습 할 수 있다. action을 추출하기 위해 inverse dynamics를 사용하기 때문에, 서로 다른 조건 변수의 구성으로인한 state trajectory x_0(\tau)에는 데모를 생성한 동일한 분포에서 나온 연속적인 state pair(s_t, s_t+1)이 포함되어 있다고 가정한다. 그렇지 않으면 inverse dynamics는 잘못된 예측을 할 수 있다.

 

저자들은 이러한 속성을 사용해 테스트시 둘 이상의 제약 조건 또는 스킬을 함께 구성할 수 있다고 얘기한다. 그리고 더 나아가 Appendix에서는 특정 constraint나 skill을 피하도록 할 수 있음도 보였다.

 

Training The Decision Diffuser

decision making을 위한 conditional generative model은 supervised manner로 학습되었으며, 각 달성한 return, 만족하는 제약 조건 또는 demo가능한 skill로 레이블이 지정된 trajectory 데이터셋이 주어지면, noise model \theta를 통해 prarameterized reverse diffusion process p_\theta와 inverse model f_\phi를 아래 loss로 동시에 학습한다.

각 trajectory에 대해, 먼저 가우시안 노이즈와 timestep을 샘플하고, noisy state array x_k(\tau)를 구성하고, 마지막으로 noise \epsilon 을 예측한다. 확률 p를 사용하면 조건 정보를 무시하고 inverse dynamics model은 trajectory가 아닌 개별 transition으로 학습된다.

 

Architecture

\theta를 학습하기 위해 Convolutional residual block이 반복적으로 구성된 temporal U-net 를 사용했다. 이렇게 하면 높이가 단일 상태의 차원을 나타내고, 너비가 trajectory의 길이를 나타내는 이미지처럼 생각될 수 있고, state sequence x_k(\tau)를 효과적으로 처리할 수 있다. conditioning information y(\tau)는 scalar 또는 one hot vector로 인코딩하고 이를 MLP를 사용해 latent variable z \in R^h에 project한다.

 

Low-temperature Sampling

알고리즘의 denoising step에서, 우리는 state의 noisy sequence와 예측된 noise를 가지고 \mu_k-1을 계산하고 \sigma_k-1를 계산할 수 있다. 분산이 [0,1]로 스케일링 된 x_k-1 ~ N(\mu_k-1, \alpha \sigma_k-1)을 샘플링하면 더 나은 품질의 sequence를 얻을 수 있다.(lower temperature로 sampling하는 것과 일치한다.)

 

 

Discussion

해당 논문에서는 decision making을 위한 conditional generative model인 Decision diffuser를 제안했다. offline sequence decision making을 조건부 생성 모델링으로구성하고, TD learning의 필요성을 회피하여 의사 결정 파이프라인을 더 단순하게 만들 수 있음을 보였다. 제안된 방법은 high return에 대한 샘플링을 통해 데이터셋 내의 best behavior를 포착할 수 있으며 좋은 성능을 보여주었다. 또 return 외에도 constraint나 skill을 컨디셔닝으로 설정할 수 있으며 테스트 타임에는 제약 조건과 스킬을 유연하게 조합하여 새로운 behavior도 구성할 수 있음을 보였다. 그리고, exploration을 위해 state sequence model의 엔트로피를 활용해 decision diffuser를 온라인 finetuning하는 것을 future work으로 제안한다. 그리고 해당 논문에서는 state를 기반한 environment에 초점을 두었지만, 이전 연구들이 수행한 것 처럼 observation space가 아닌 latent space에서의 diffusion을 수행해 이미지 기반 환경으로 확장할 수도 있을 것이라고 한다. 

 

해당 논문의 Appendix 맨 끝에서는 한계점도 명시해두었는데,

먼저 Deicison diffuser을 학습하기 위해서는 fully observable MDP가 필요하므로 partial observability한 데이터셋(POMDP)으로는 해당 모델을 학습할 수 없으며 self-delusions를 야기할 가능성이 있다고 한다.

그리고 future work으로 언급했던 것 처럼 online setting에서 exploration을 위한 update가 불가능하다는 점과,

실험은 모두 state-based environment에서 수행했다는 점,

그리고 오직 AND와 NOT compositon만 support한다는 점(명시적으로 density를 추정하지 않기 때문에),

stochastic dynamics환경에서는 성능이 저하된다는 점 등이 있다.

 

후속 연구들에 대해서 상세하게 언급해 놓아서 관련 연구들을 진행해도 재밌을 것 같다. 비록 나는 생성 모델을 잘 몰라서 뭐가 뭔지 잘모르겠지마는..