전통적인 강화학습에서는 연속 행동 공간(continuous action domain)에서 확률적 정책(stochastic policy)을 학습할 때, 정책의 parameterization을 보통 Gaussian distribution으로 둔다. 예를 들어 상태 s가 주어졌을 때 정책은
$$a \sim \mathcal N(\mu_\theta(s), \Sigma_\theta(s))$$
와 같이 표현되며, policy gradient나 PPO/TRPO 계열 알고리즘은 이 Gaussian likelihood를 기반으로 정책을 업데이트한다.
이러한 parameterization은 계산이 안정적이고 likelihood 및 entropy를 쉽게 계산할 수 있다는 장점이 있다. 특히 reparameterization trick과 결합했을 때 gradient estimation이 효율적이며, 고차원 연속 제어 문제에서도 비교적 안정적으로 동작한다는 이유로 오랫동안 표준적인 선택으로 사용되어 왔다.
다만 Gaussian policy는 표현 가능한 action distribution의 형태가 제한적이다. 특히 하나의 상태에서 여러 행동 mode를 동시에 유지해야 하거나, 탐색 과정에서 서로 떨어진 action region을 함께 고려해야 하는 경우에는 단일 Gaussian의 unimodal 구조가 제약으로 작용할 수 있다. 이러한 한계는 offline RL이나 imitation learning에서 mode averaging 문제로 더 명확히 드러나지만, online RL에서도 복잡한 exploration 구조를 표현하는 데에는 부족할 수 있다.
이런 관점에서 flow-based policy는 Gaussian policy를 대체한다기보다, Gaussian policy가 제공하는 계산적 편의성을 어느 정도 유지하면서 더 유연한 action distribution을 표현하려는 시도로 볼 수 있다. 단순 Gaussian을 복잡한 invertible transformation으로 확장함으로써, multimodal하거나 highly non-Gaussian한 action distribution을 표현할 수 있다는 점에서 주목받고 있다.
그러한 점에서 본 논문은 continuous normalizing flow(CNF)를 policy parameterization에 도입하여, 기존 Gaussian policy보다 더 expressive한 action distribution을 구성하고자 한다. 구체적으로는 latent variable z를 base distribution에서 샘플링한 뒤, 이를 flow dynamics를 통해 변환하여 action distribution의 평균을 생성한다. 이후 해당 평균 주변에서 Gaussian sampling을 수행함으로써, 기존 정책 최적화 기법과의 호환성을 유지하면서도 multimodal하거나 highly non-Gaussian한 정책 분포를 표현할 수 있도록 한다.
특히 본 논문은 flow를 직접 density estimator로 사용하는 대신, flow가 생성한 representation 위에 conditional Gaussian policy를 정의하는 구조를 사용한다. 이를 통해 복잡한 flow 기반 표현력을 활용하면서도, PPO 계열 알고리즘에서 필요한 likelihood ratio 계산을 보다 안정적으로 수행할 수 있도록 설계하였다.
Main Method
본 논문의 핵심은 Gaussian policy의 평균 $\mu_\theta(s)$ 를 단순한 neural network output으로 두는 대신, continuous normalizing flow의 terminal output $\phi_1(z;s)$ 로 정의하는 것이다.
기존 Gaussian policy는 보통
$$a \sim \mathcal N(\mu_\theta(s), \sigma^2)$$
처럼 상태 s에서 하나의 평균을 만들고, 그 주변에서 action을 샘플링한다. 반면 이 논문에서는 먼저 latent variable
$$z \sim p_z(z)$$
를 샘플링한 뒤, z를 상태 s에 조건화된 flow를 통해 이동시킨다.
$$\frac{d}{dt}\phi_t(z;s) = v_t(\phi_t(z;s);s), \qquad \phi_0(z;s)=z$$
여기서 $v_t$는 neural network로 parameterize된 velocity field다. 이 ODE를 $t=0$부터 $t=1$까지 적분하면 terminal point
$$\phi_1(z;s)$$
를 얻는다. 논문은 이 값을 action distribution의 평균으로 사용한다.
$$a=\phi_1(z;s)+n, \qquad n\sim \mathcal N(0,\sigma^2)$$
즉 z가 고정되면 action은
$$\pi(a\mid z,s) = \mathcal N(a;\phi_1(z;s),\sigma^2)$$
를 따른다.
하지만 실제 policy는 z를 모르는 상태에서 action만 보는 분포이므로, z에 대해 적분한 marginal distribution이다.
$$\pi(a\mid s) = \int \pi(a\mid z,s)p_z(z)\,dz$$
따라서 전체 policy는 하나의 Gaussian이 아니라, 여러 z가 만드는 Gaussian들의 연속 mixture가 된다. 이 때문에 기존 Gaussian policy보다 더 복잡한 multimodal action distribution을 표현할 수 있다.
문제는 PPO-style objective를 쓰려면 importance ratio가 필요하다는 점이다.
$$\frac{\pi(a\mid z,s)} {\hat\pi(a\mid z,s)}$$
여기서 $\hat\pi$는 rollout을 생성한 old policy이고, $\pi$는 현재 업데이트하려는 new policy다.
정확히 계산하면
$$\frac{ \mathcal N(a;\phi_1(z;s),\sigma^2) }{ \mathcal N(a;\hat\phi_1(z;s),\hat\sigma^2) }$$
가 된다.
그런데 $\phi_1(z;s)$와 $\hat\phi_1(z;s)$를 얻으려면 각각 ODE를 끝까지 풀어야 한다. 특히 학습 중에 new flow $\phi_1$에 대해 gradient까지 흘리려면 비용이 크고 불안정해진다. 그래서 논문은 ratio를 직접 계산하지 않고 근사한다.
핵심 관찰은 Gaussian likelihood ratio가 shift-invariant하다는 것이다.
$$\frac{ \mathcal N(a;\phi_1,\sigma^2) }{ \mathcal N(a;\hat\phi_1,\hat\sigma^2) } = \frac{ \mathcal N(a-\hat\phi_1;\phi_1-\hat\phi_1,\sigma^2) }{ \mathcal N(a-\hat\phi_1;0,\hat\sigma^2) }$$
즉 중요한 것은 new terminal point 자체가 아니라,
$$\delta\phi_1(z;s) = \phi_1(z;s)-\hat\phi_1(z;s)$$
이다. 이는 old flow의 terminal point에서 new flow가 얼마나 이동했는지를 의미한다.
하지만 $\delta\phi_1$도 정확히 계산하려면 new ODE를 풀어야 한다. 그래서 논문은 terminal shift $\delta\phi_1$를 직접 계산하지 않고, velocity field의 차이로 근사한다.
old flow의 시작점은 z, 끝점은 $\hat\phi_1$다. 이 둘을 잇는 직선 경로를
$$x_t=(1-t)z+t\hat\phi_1(z;s)$$
로 정의한다. 그리고 이 경로 위에서 new velocity와 old velocity의 차이를 본다.
$$\delta v_t(x_t;s) = v_t(x_t;s)-\hat v_t(x_t;s)$$
직관적으로는, terminal output의 차이를 직접 계산하는 대신, 중간 경로에서의 velocity 차이로 대체하는 것이다.
그래서 importance ratio를 다음과 같이 근사한다:
$$\rho = \frac{ \mathcal N(a-\hat\phi_1(z;s); v_t(x_t;s,\theta)-\hat v_t(x_t;s), \sigma^2) }{ \mathcal N(a-\hat\phi_1(z;s);0,\hat\sigma^2) }$$
이제 학습 시에는 new flow의 ODE를 끝까지 적분할 필요가 없다. 대신 $t\sim Unif[0,1]$를 샘플링하고, interpolation point $x_t$에서 velocity field만 평가하면 된다.
논문에서는 이렇게 proxy로 importance ratio를 계산하는 것에 대해서, approximate bound를 제시함으로써 정당성을 확인하니 확인해보면 좋을 것 같다.
Policy Entropy Regularization
PPO나 SAC 같은 알고리즘에서는 정책 entropy를 키워 exploration을 유도하는 것이 일반적이다. Gaussian policy에서는 entropy가 닫힌형태로 계산된다. 예를 들어 diagonal Gaussian의 경우 entropy는 분산만 알면 바로 계산할 수 있다.
하지만 flow-based policy에서는 상황이 더 복잡하다. 정책 분포가
$$\pi(a\mid s) = \int \mathcal N(a;\phi_1(z;s),\sigma^2)p_z(z)\,dz$$
와 같은 implicit mixture 형태를 가지기 때문에, 전체 policy entropy
$$\mathcal H(\pi(\cdot\mid s))$$
를 직접 계산하기 어렵다. 정통 CNF 방식으로 density를 계산하려면 divergence trace를 ODE와 함께 적분해야 하므로 계산 비용도 커진다. 따라서 기존 방식처럼 policy entropy를 직접 계산해서 regularization하는 것은 이 구조에서는 부담스럽다.
이를 해결하기 위해 논문은 Brownian motion에서 영감을 받은 entropy regularizer를 제안한다. 핵심 아이디어는 entropy를 직접 계산하지 않고, flow trajectory가 자연스럽게 퍼지는 방향으로 velocity field를 유도하는 것이다.
Brownian motion은 시간이 지날수록 particle들이 공간상에 퍼지는 확산 과정이다. 이때 확률밀도 $p_t(x)$는 heat equation을 따른다.
$$\frac{\partial p_t(x)}{\partial t} = \nabla_x^2 p_t(x)$$
한편 flow로 induced되는 density의 변화는 continuity equation으로 표현된다.
$$\frac{\partial p_t(x)}{\partial t} = -\nabla_x\cdot(p_t(x)v_t(x))$$
여기서 velocity field를
$$v_t(x) = -\nabla_x\log p_t(x)$$
로 두면 continuity equation이 heat equation과 연결된다. 즉, negative score 방향의 velocity field는 density를 퍼뜨리는 확산적 움직임과 관련된다.
이 관찰을 이용해 논문은 learned velocity field가 reference flow의 negative score 방향을 따르도록 regularization한다. 다만 score $\nabla_x\log \hat p_t(x_t;s)$를 직접 계산하기 어렵기 때문에, reference velocity field $\hat v_t$와 score 사이의 관계를 사용한다.
$$\nabla_x\log \hat p_t(x_t;s) = \frac{1}{1-t} \left( t\hat v_t(x_t;s)-x_t \right)$$
따라서 negative score는 대략
$$-\nabla_x\log \hat p_t(x_t;s) = \frac{1}{1-t} \left( x_t-t\hat v_t(x_t;s) \right)$$
가 된다.
즉 현재 velocity v$v_t(x_t;s,\theta)$가 이 방향과 비슷해지도록 만들고 싶다. 하지만 위 식은 $t\to 1$일 때 $1/(1-t)$ 때문에 불안정해진다. 그래서 논문은 양변에 $1-t$를 곱한 형태로 regularizer를 정의한다.
$$\eta_t(x_t;s,\theta) = (1-t)v_t(x_t;s,\theta) - \left( x_t-t\hat v_t(x_t;s) \right)$$
그리고 이 값을 작게 만드는 항을 objective에 추가한다.
$$-w_b \|\eta_t(x_t;s,\theta)\|_2^2$$
최대화 objective이므로 음수 부호가 붙는다. 즉 $\eta_t$가 작아질수록 objective가 커진다. 결과적으로 learned velocity field는 reference flow의 negative score 방향과 정렬된다.
직관적으로 말하면, 이 regularizer는 flow가 좁은 영역으로 collapse하지 않고 Brownian motion처럼 퍼지는 방향을 갖도록 유도한다. 그래서 여러 latent z에서 출발한 trajectory들이 비슷한 action mode로 몰리는 것을 완화하고, 더 넓은 action region을 탐색하게 만든다.
또한 논문은 terminal Gaussian noise의 entropy도 함께 regularization한다.
$$\frac{w_g}{2} \sum_{i=1}^d \log(2\pi e\sigma_i^2)$$
이는 action을
$$a=\phi_1(z;s)+n, \qquad n\sim\mathcal N(0,\sigma^2)$$
로 샘플링할 때 들어가는 Gaussian noise n의 entropy다. 이 항은 분산이 너무 작아져 deterministic policy처럼 되는 것을 막고, terminal action sampling 단계에서도 stochasticity를 유지하도록 한다.
따라서 최종 regularizer는 두 부분으로 구성된다.
$$J_{\text{Reg}}(\theta,\sigma) = \mathbb E \left[ -w_b\|\eta_t(x_t;s,\theta)\|_2^2 + \frac{w_g}{2} \sum_{i=1}^d \log(2\pi e\sigma_i^2) \right]$$
첫 번째 항은 flow trajectory 자체가 퍼지는 방향을 갖게 만들고, 두 번째 항은 terminal Gaussian noise의 entropy를 키운다.
전체적으로 PolicyFlow의 학습 objective는 PPO-style clipped surrogate와 Brownian regularizer의 합이다.
$$J_{\text{Flow}} + J_{\text{Reg}}$$
즉 이 방법은 다음 두 가지를 동시에 달성하려고 한다.
- PPO처럼 advantage가 높은 action의 probability를 증가시킨다.
- Flow trajectory가 collapse하지 않고 다양한 action mode를 유지하도록 한다.
정리하면, Brownian regularizer는 정책 entropy를 직접 계산하지 않고도, flow의 velocity field를 확산적인 방향으로 유도하여 exploration과 trajectory diversity를 증가시키는 장치다. 다만 논문에서도 말하듯이 이것은 엄밀한 Brownian dynamics의 exact derivation은 아니고, rectified flow에서 알려진 velocity-score 관계를 policy flow에 가져온 practical regularization으로 보는 것이 정확하다.
Experiment는 굳이 확인하지 않도록 하겠다. 꽤 복잡해서 이렇게까지 해야 하나 싶기는 한데... 아무튼 여기서 마무리하겠다!