Distributional Reinforcement Learning in the Brain

Lowet et al. 2020 리뷰논문 요약.
이 논문이 나오게 된 계기 또는 이전의 흐름은 Deep mind에서 발표한 distributional RL 논문 두개. Bellamare et al. 2017, Dabney et al. 2017, 그리고 이 두 논문의 computational 효과, 그리고 biological implementation에 대한 논문. Dabney et al. 2020. 이 세 논문이 전제되어 있다.

Biological and Artificial Intelligence

위에서 말한 세 논문을 인용하면서 이것이 neuroscience과 AI사이의 훌륭한 상호작용이라 칭하고 있다.

Development of Distributional Reinforcement Learning in AI

RL, 그 중에서도 DQN을 소개하고, 여기서 distributional RL(이제 distRL로 줄인다.)이 개발되었음을 설명하고 있다.
distRL의 핵심 아이디어는 미래 보상, 상태 이전이 필연적으로 stochastic하다는 것을 전제로 기대 보상(가치)를 평균값으로 산출하지 않고, 아예 확률 분포로 추정하도록 학습한다는 것이다.

How Distributional RL works

distRL에서 중요한 점은 1) 어떻게 보상의 분포가 표상되는가? 2) 어떻게 그것을 학습하는가? ¹
먼저 개발된 categorical dist RL은 histogram을 학습한다. 하지만 그 다음에 개발된 quantile dist RL은 population coding 개념을 이용했고, 이것의 학습 알고리듬이 좀 더 biological plausible하다고 주장한다.

Learning from Prediction Errors

Rescorla-Wagner rule에 따라 가치가 갱신된다는 것을 생각해보자.

$V \leftarrow V + \alpha \cdot \delta$

여기서 $\alpha$ 는 learning rate parameter이고, $\delta$ 는 prediction error다. 즉 $\delta := R - V$ 이다. 이 RW rule은 잘 알려져 있다시피 dopamine firing( $\delta$ )이 prediction error에 반응한다는 것으로도 유명하다. 또한 RL algorithm중 temporal difference learning으로 발전되었다.
하지만 둘은 미묘하게 다른데, RW rule의 경우 $\delta$ 가 실제 보상 R과 차이를 통해 계산된다. 하지만 TD learning에서는 다음 스텝에서 얻는 보상에 더해 $\gamma V(s_{t+1})$ (discounted value of next state)가 포함된다. 즉 TD learning은 $\delta := r_{t} + \gamma V(s_{t+1}) - V(s_{t})$ 이다. 이건 boostrapping으로 알려진 절차이고 그렇기 때문에 정확하지 않다. DQN으로 접목되기 힘들었던 이유가 parameter를 업데이트하면 $V(s_{t+1})$ 까지 변해버리기 때문이다.² 물론 그럼에도 불구하고 실질적으로 TD learning은 local optimum으로 잘 수렴하긴 한다.

Toward Distributional RL

보상이나 가치를 평균이라는 하나의 값으로 뭉쳐버리면 실제 현실을 충분히 반영하지 못 한다. 실제 현실이 아주 낮은 확률을 가진 높은 보상과, 높은 확률을 가진 낮은 보상 등이 섞여 있다면 이것의 평균은 낮을 것이고 그러면 비록 확률은 낮지만 높은 보상의 기회를 놓치게 된다. 그래서 dist RL이 등장하는데, 확률 분포 자체를 학습하는 것은 까다로워 보이나, 연구자들은 이를 단순한 수식 변환으로 해결했다.

$V \leftarrow V + \alpha \cdot \left\lbrace \begin{array}{l l} -1 &\text{if} \quad \delta \leq 0 \\ \ \ 1 &\text{if} \quad \delta > 0 \end{array} \right.$

그 결과 value는보상의 평균이 아니라, median으로 수렴하게 된다. ³
여기서 quantile regression으로 갈 때는 $\alpha$ 라는 learning rate를 추가한다. 그러면 무조건 50% 순위로 수렴하던 아까와는 달리, $\alpha$ 값의 비율에 따라 quantile rank로 수렴하게 된다.

$V \leftarrow V + \left\lbrace \begin{array}{l l} \alpha^{-}_{i} \cdot (-1) &\text{if} \quad \delta \leq 0 \\ \alpha^{+}_{i} \cdot (+1) &\text{if} \quad \delta > 0 \end{array} \right.$

여기서 이면 의 수치가 곧 수렴하게 되는 순위이고, 그렇지 않으면 두 값에 가 차지하는 비율이 수렴하게 되는 순위가 된다. 이걸 논문에서는 quantile을 이라고 표현하고 있다.
- 따라서 $\alpha$ 값이 서로 다른 value predictor를 동시에 학습 시키면 서로 다른 $\tau_{i}$ 순위의 predictor가 학습되고 이걸 마치 population code와 같다고 한다.
이제 위의 quantile rule에서 $+1,\, -1$ 값 대신 prediction error, 즉 $\delta$ 값을 대입하면 median을 비롯한 순위에 따른 quantile이 아니라, 보장에 대한 기대값(가치)에 수렴하므로, 이걸 expectile이라고 한다. 만약 $\alpha^{+}$ 값의 비중이 0.5라면 mean 값에 수렴할 것이다. ⁴

Distributional RL as the Process of Minimizing Estimation Errors

이제 위 알고리듬을 보편적인 선형회귀 문제로 살펴보자. N번 시도했을 때 보상을 얻고, 그 때 기대하고 있던 가치(V)와의 mean squared error(MSE)는 다음 과 같다.

$MSE(V) = \frac{1}{N} \sum^{N}_{n=1} (r_{n} - V)^{2}$

MSE 가 V에 의존적인 함수로 변환된다. V를 변화시켜 MSE를 최소화 시키는 것이 머신러닝의 기본이다. 제일 좋은 방법은 보상 값 전체를 저장했다가 한 번에 계산하는 것이지만, 실용적인 방법은 아니다. 실용적인 방법은 잘 알려져있는 Stochastic gradient descent(SGD)다. 이걸 수학적으론 MSE를 미분하고 그 값에 따라 V를 업데이트하는 것인데, 이건 결국 RW rule과 동일해진다.
그리고 SGD에 입력하는 error function을 MSE로 할 것인지, absolute error로 할 것인지, 아니면 quantile $\tau$ 값으로 할 것인지, 여기 error 값을 곱해 expectile로 할 것이지에 따라 위에서 다뤘던 learning과 일치하게 된다.

parametrized problem

이 section제목은 원래 논문에 없지만, 임의로 넣음.
V 값이 function approximator를 통해 구해질 경우 우리는 function approximator에 있는 coefficients 또는 parameter를 갱신해야 한다. 이는 다음과 같이 표현 할 수 있다.

$\theta \leftarrow \theta + \alpha \sum^{M}_{i=1} \bigtriangledown_{\theta}q_{i}(\theta) \cdot \left\lbrace \begin{array}{l l} -(1-\tau) &\text{if} \quad \delta_{i,n} \leq 0 \\ \quad \; \tau &\text{if} \quad \delta_{i,n} > 0 \end{array} \right.$

위 수식에서 expectile로 변경하려면 $q_{i}$ 를 $e_{i}$ 값으로 변경하고 prediction error term $\delta_{i,n}$ 도 추가하면 된다. 다시 정리하면 위에서 직접 update했던 value들 대신 parameter $\theta$ 에 gradient form를 추가해 $\theta$ 를 update하는 방법이다. DNN이 들어간다면 복잡해지긴 하겠지만, 기본 아이디어는 같다.

Traditional and Distributional RL in the Brain

VTA의 dopaminergic neuron의 다양성에 대해 강조하고 있다. 다른 영역의 뉴런에 비해 꽤 homogeneous하긴 하지만, 최근 연구 결과에서 Kim et al, 2019은 VTA의 dopaminergic neuron이 꽤 diverse하다는 것을 보고 했다. 그리고 이것이 RPE 신호의 체계적인 variation의 근거라고 주장한다. Distributional RL은 이 다양성에 대한 설명을 제공할 수 있다.

Empirically Testing Distributional RL

비대칭적 스케일 factor인 $\tau = \frac{\alpha^{+}}{\alpha^{+} + \alpha^{-}}$ 가 distRL의 핵심 아이디어다. 그럼 $\alpha$ 는 무엇인가? 아마 dopaminergic neuron이 RPE에 대해 상대적으로 firing rate를 증가시키는지 감소시키는지에 의해 결정될 것이다.
dopaminergic neuron에서도 optimistic, pessimistic predictor가 따로 있다고 가정하는 것이다. optimistic dopamine neuron이라면 high value predictor이고, 따라서 큰 reward가 주어져야 RPE가 matching될 것이다. $\tau$ 가 높은 값으로 매겨져 있다면, RPE에 반응하는 reversal point가 높은 쪽으로 치우쳐져 있다는 의미다.

이 가설을 시험하기 위해 Dabney et al. 2020에서는 optogenetical 방법을 사용했다. 그 결과 어떤 뉴런은 주어지는 보상의 평균 보다 아래에서도 firing했고, 어떤 뉴런은 보상의 평균보다 높아야 firing했다.
이런 가설을 바탕으로 dopaminergic neuron firing rate를 decoding했더니 각기 다른 reversal point of expectile을 가진 것으로 추정되었다. 이전에는 이런 결과를 단순히 noise로 해석할 수밖에 없었다.
DistRL 가설, 이론은 몇몇 실험적 예측을 제시한다. 예를 들어 도파민 뉴런은 다른 보상 분포에서 비교적 일정한 비대칭적 $\tau$ 를 보여야 한다. 낙관적 세포는 음의 RPE에서 느리게 학습되어야 한다. 또한 downstream target에서도 quantile-like 분포가 관찰되어야 한다. 조작적 조건화 실험에서도 DistRL에 준하는 행동이 관찰되어야 한다. 낙관적 도파민 신경세포에 의해 risk-taking behavior가 유도되어야 한다.

Is Distributional RL Biologically Plausible?

Diversity in Asymmetric Scaling and Independent Loops

Positive, negative RPE 신호가 분리되어 처리된다는 증거가 있다. Lateral habenula에 lesion을 줄 경우 negative RPE에 둔감해지고, 결과적으로 optimistic하게 된다.
Optimistic, pessimistic dopamine neuron이 topographically organized 되어 있다.

Learning Rate Parameters in Striatum and Cortex

Positive, negative RPE signal이 striatum에서 D1, D2 type dopamine receptor를 발현하고 있는 target neuron에 따라 다르게 처리된다.

How Does the Brian Benefit from Distributional Representations?

DistRL의 concept은 multilayer neural net에 적합하다. 단순히 reward나 value의 평균 scalar을 전달하는 것보다 distribution을 backpropagation하기 때문에 hidden layer에 rich information이 표상된다.
- 이런 distribution coding이 보상이나 확률에 대해서 뿐만 아니라, delay interval이나 “distributions in the common currency of value”를 표상할 수 있는지는 더 연구가 필요하다.
Quantile code는 nonparametric code이다. population coding에 대해서는 probabilistic population code(PPC)나 distributed distributional code가 있고, 이것들은 parametric code가 된다. 예를 들어 PPC는 Gaussian distribution을 따른다. 하지만 dopamine neuron의 RPE도 이럴지는 모른다.
- PCC가 Bayesian inference를 지지하고, Quantile code는 cumulative prospect theory의 단순 버전을 지지할 수 있다.

Distributional TD Updates in the Brain

전통적 RL에서는 RPE( $\delta$ )가 single, local 추정치로 계산될 수 있다.
하지만 distRL에서는 당연히 reward distribution을 알 수 있을만큼의 sample이 필요하다.
- 그래서 single unit이 아니라 neuronal population 개념으로 접근해야 한다.

—

최근 내가 관심을 하고 있는 free-energy principle이나 Bayesian brain theory와 관련해서 확률분포를 어떻게 신경세포 또는 neural unit이 학습하는가? 는 중요한 이슈이다. ↩
그래서 deep mind에서 DQN을 처음 이용할 때 replay buffer를 이용해 마치 stochastic images를 DQN이 학습하는 것처럼 만들었다. ↩
Rank에서 50%로 향하게 되므로. ↩
Expectile에 대한 부연 설명. quantile은 상위 몇%에 있는 해당 single value를 반환한다. 하지만 expectile은 해당 %를 기준으로 두 집단으로 나눈 다음 해당 집단의 평균 값을 구하고, 두 값의 평균을 구하는 작업이다. 그래서 cumulative density function을 적용할 수 없는 것이다. 만약 상위 1% quantile이라면 해당 수치가 이미 발생 확률을 내포하고 있다.(1/100이므로) 하지만 1% expectile은 하위 99%의 평균값에 의해 수치가 조정되어 버리므로 발생 확률과 상관이 없다. ↩