Belief-State RWKV for Reinforcement Learning under Partial Observability

Each language version is independently generated for its own context, not a direct translation.

이 논문은 인공지능 (AI) 이 "모르는 상황"에서 어떻게 더 똑똑하게 행동할 수 있는지에 대한 새로운 아이디어를 제시합니다.

핵심 아이디어를 일상적인 비유로 설명해 드릴게요.

🧐 문제: "안개 속의 운전"

상상해 보세요. 안개가 자욱한 도로를 운전하고 있다고 칩시다.

기존 AI (RWKV 의 일반적 방식): 운전사는 앞이 잘 안 보이지만, 과거의 경험과 현재 보이는 것들을 머릿속에 **'한 줄의 요약 메모'**로 저장합니다. "아까 차가 있었지, 지금도 차가 있을 거야"라고 생각하며 운전합니다. 하지만 이 메모에는 **"내가 얼마나 확신하는가?"**라는 정보가 빠져 있습니다. 안개가 짙을수록 AI 는 자신이 틀릴 수도 있다는 사실을 모른 채 무작정 결정을 내릴 수 있습니다.

💡 해결책: "불확실성을 아는 운전사" (Belief-State RWKV)

이 논문은 AI 의 머릿속 메모를 단순히 "사실"만 적는 것이 아니라, "사실 + 확신도" 두 가지로 나누어 저장하자고 제안합니다.

위치 (µt): "지금 차가 어디쯤 있을 것 같은가?" (예: 오른쪽 차선)
불확실성 (Σt): "내가 이 추정에 얼마나 확신하는가?" (예: 안개가 너무 짙어서 50% 만 확신함)

이제 AI 는 **"차량이 오른쪽에 있을 것 같지만, 확신이 50% 라니, 너무 빨리 추월하지 말고 잠시 기다려보자"**라고 더 현명한 결정을 내릴 수 있게 됩니다.

🛠️ 어떻게 작동할까요? (간단한 비유)

RWKV 란? AI 가 긴 이야기를 기억하면서도 메모리 용량을 적게 쓰는 아주 효율적인 '기억 장치'입니다.
기존 방식: 이 기억 장치에서 나온 정보를 그냥 "정답"처럼 받아들였습니다.
이 논문의 방식: 이 기억 장치에서 나온 정보를 **"추측"**으로 해석합니다.
- "내 기억이 흔들리고 있어 (불확실성 높음) → 조심스럽게 행동해."
- "내 기억이 선명해 (불확실성 낮음) → 과감하게 행동해."

📊 실험 결과: 언제 도움이 될까요?

연구팀은 AI 를 훈련시켜서 안개 (노이즈) 가 심한 상황에서 결정을 내리게 했습니다.

평범한 상황 (안개가 얇을 때): 기존 방식과 별 차이가 없거나, 오히려 기존 방식이 조금 더 빨랐습니다. (불필요한 고민을 안 해도 되니까요.)
어려운 상황 (안개가 짙을 때): 이 새로운 방식이 훨씬 잘했습니다.
- 안개가 너무 짙어서 정보가 흐릿해지면, 기존 AI 는 당황해서 실수를 했지만, 이 새로운 AI 는 "아, 내가 지금 잘 모르겠네"라고 인정하고 더 신중하게 행동하거나, 기다리는 전략을 썼습니다.
- 특히 훈련할 때 보지 못했던 새로운 안개 패턴이 등장했을 때도 가장 잘 견뎌냈습니다.

🚀 왜 중요한가요?

이 연구의 가장 큰 장점은 성능뿐만 아니라 이해의 용이성입니다.
기존 AI 는 "왜 그런 결정을 내렸지?"라고 물어보면 "그냥 그렇게 계산했어"라고 대답하지만, 이 새로운 AI 는 **"내가 그 선택에 대해 60% 만 확신해서 신중하게 행동했어"**라고 설명할 수 있습니다.

📝 요약

이 논문은 **"AI 가 무조건 모든 것을 다 안다고 착각하지 말고, 자신이 모르는 부분 (불확실성) 을 스스로 인식하게 만들자"**고 말합니다.

쉬운 말: AI 에게 "내가 얼마나 확신하는지"를 알려주는 나침반을 추가했습니다.
효과: 안개가 짙은 (정보를 알 수 없는) 상황일수록, 이 나침반을 가진 AI 가 더 안전하고 똑똑하게 행동합니다.
미래: 앞으로 AI 가 더 복잡한 미로나 예측 불가능한 세상에서 살아남을 수 있도록, 이 '불확실성 인식' 기술을 더 발전시킬 예정입니다.

결론적으로, 이 기술은 AI 가 **"무조건 빠른 것"보다 "상황을 잘 파악한 후 올바른 결정을 내리는 것"**을 더 중요하게 생각하게 만드는 첫걸음입니다.

Each language version is independently generated for its own context, not a direct translation.

1. 문제 정의 (Problem)

강화 학습 (RL) 에서 에이전트가 환경의 완전한 상태를 관측하지 못하는 부분 관측성 (Partial Observability) 상황은 중요한 도전 과제입니다.

기존 접근법의 한계: RWKV 와 같은 재귀적 (Recurrent) 시퀀스 모델을 RL 에 적용할 때, 기존의 고정된 크기의 숨겨진 상태 (hidden state, $h_t$ ) 는 단순히 과거의 정보를 압축하여 저장하는 '불투명한 벡터'로만 작용합니다.
핵심 결여: 이러한 방식은 에이전트가 환경의 잠재 상태에 대해 얼마나 불확실한지 (uncertainty) 를 명시적으로 표현하지 못합니다. 즉, 에이전트는 증거는 기억할 수 있지만, 그 증거에 대한 '신뢰도'나 '불확실성'을 정책 (Policy) 이나 가치 함수 (Value) 에 반영하지 못합니다.
목표: 불투명한 숨겨진 상태를, 신념 상태 (Belief State) 로 재해석하여, 에이전트가 무엇을 믿는지 ( $\mu_t$ ) 와 그 믿음에 대한 불확실성 ( $\Sigma_t$ ) 을 모두 명시적으로 관리할 수 있는 구조를 제안합니다.

2. 방법론 (Methodology)

2.1. 신념 상태 RWKV 구조 (Belief-State RWKV Recurrence)

기존의 단일 숨겨진 상태 $h_t$ 대신, 고정된 크기의 구조화된 신념 상태 $b_t = (\mu_t, \Sigma_t)$ 를 도입합니다.

$\mu_t$ (위치 통계량): 에이전트가 현재 환경 상태에 대해 가지고 있는 믿음의 중심 (예: 평균).
$\Sigma_t$ (불확실성 통계량): 해당 믿음에 대한 불확실성의 정도 (예: 분산).
구현: RWKV 의 시간 혼합 (Time-Mix) 레이어에서 생성된 재귀적 상태 ( $s_t$ $s_{t}$ ) 를 기반으로, 경량화된 읽기 헤드 (Readout Heads) 를 통해 $\mu_t$ $μ_{t}$ 와 $\Sigma_t$ $Σ_{t}$ 를 계산합니다.
- $z_t = \psi(u_t, s_t)$
- $\mu_t = W_\mu z_t$
- $\log \Sigma_t = W_\Sigma z_t$

2.2. 정책 및 가치 조건부 학습 (Policy and Value Conditioning)

기존의 Actor-Critic 프레임워크를 유지하되, 정책 ( $\pi$ ) 과 가치 함수 ( $V$ ) 가 단순히 $h_t$ 가 아닌 $(\mu_t, \Sigma_t)$ 쌍에 조건부로 학습됩니다.

이는 불확실성을 고려한 의사결정을 가능하게 하며, RWKV 의 상수 공간 복잡도 (constant-space inference) 이점을 유지합니다.

2.3. 확장 가능성 (Extensions)

논문에서는 두 가지 추가적인 확장 아이디어를 제안하지만, 초기 실험에서는 이를 사용하지 않았습니다.

신념 기반 메모리 제어 (Belief-Conditioned Memory Control): 불확실성 ( $\Sigma_t$ ) 이 높을 때 메모리 쓰기 (Write) 를 증가시키고, 불확실성이 낮을 때 유지 (Retention) 를 증가시키는 게이트 메커니즘.
저랭크 어댑터 (Low-Rank Adapters): 보상과 관련된 하위 공간에 특화되도록 신념 상태를 저랭크 행렬로 변환.
특권적 신념 감독 (Privileged Belief Supervision): 학습 시에만 접근 가능한 시뮬레이터의 잠재 변수를 사용하여 신념 상태를 지도 학습 (Supervised Learning) 하도록 하는 보조 손실 함수.

3. 주요 기여 (Key Contributions)

신념 상태 기반 RWKV RL 프레임워크 도입: 정책과 가치 함수를 불투명한 숨겨진 상태가 아닌, 명시적인 불확실성 통계량 $(\mu_t, \Sigma_t)$ 에 조건부인 모델로 설계.
이론적 기반 마련: 근사적 충분성 (Approximate Sufficiency), 안정성 (Stability), 그리고 저랭크 보상 관련성 (Low-rank Reward Relevance) 에 대한 명제 (Proposition) 와 증명 스케치를 제시하여 이론적 타당성을 입증.
초기 실험 및 검증: 숨겨진 노이즈가 있는 부분 관측 환경에서 실험을 수행하여, 분포 외 (OOD) 상황과 가장 어려운 환경에서 신념 상태 기반 정책의 우월성을 입증.
간단한 읽기 헤드의 효과 입증: 복잡한 게이트 제어나 특권적 감독 없이, 단순한 '신념 읽기 (Belief Readout)' 만으로도 분포 이동 (Distribution Shift) 에 대한 강건성이 향상됨을 확인.

4. 실험 결과 (Results)

4.1. 실험 환경

작업: 'Stop-or-Guess' 환경. 에이전트는 숨겨진 라벨 ( $z \in \{-1, +1\}$ ) 과 숨겨진 노이즈 수준 ( $\sigma$ ) 하에서 관측값을 받습니다.
목표: 증거가 충분할 때까지 기다리거나 (Wait), 정답을 추측 (Guess) 하는 것. 노이즈가 높을수록 추측을 미루는 것이 유리합니다.

4.2. 주요 성과

일반적 성능: MLP(메모리 없음) 와 기존 RWKV 요약 상태 (Summary State) 모델에 비해 전체 평균 보상은 비슷하거나 약간 낮을 수 있으나, 가장 어려운 환경 (High Noise) 과 분포 외 (OOD) 환경에서 더 나은 성능을 보였습니다.
- 표 2: 훈련 데이터보다 더 높은 노이즈 범위 (OOD) 에서 테스트 시, 신념 상태 RWKV 가 가장 높은 평균 보상 (0.650) 을 기록했습니다.
강건성 (Robustness): 숨겨진 노이즈가 증가함에 따라, 기존 요약 상태 모델은 성능이 급격히 떨어지는 반면, 신념 상태 모델은 더 완만하게 성능이 저하되며 높은 노이즈 구간에서 우위를 점했습니다.
보정 (Calibration): 신념 상태 모델은 결정 시점의 기대 오차 (ECE) 가 더 낮아, 불확실성을 더 정확하게 반영하여 더 신중한 결정을 내리는 경향이 있었습니다.

4.3. 애블레이션 (Ablation) 연구

간단한 읽기 헤드가 최강: '게이트된 메모리 (Gated Memory)'나 '특권적 목표 (Privileged Targets)'를 추가한 복잡한 모델들은 오히려 OOD 성능이 저하되거나 과적합되는 경향을 보였습니다.
결론: 복잡한 구조보다는 불확실성을 명시적으로 표현하는 것 자체가 분포 이동에 대한 강건성을 높이는 데 가장 효과적이었습니다.

5. 의의 및 결론 (Significance & Conclusion)

인터페이스의 명확성: RWKV 의 고정된 재귀 메모리를 '불투명한 벡터'가 아닌 '구조화된 신념 상태'로 재해석함으로써, 에이전트의 내부 상태가 무엇을 의미하는지 연구자가 더 명확하게 이해하고 제어할 수 있게 되었습니다.
불확실성 활용의 중요성: 부분 관측성 하의 강화 학습에서, 에이전트는 단순히 과거 정보를 압축하는 것을 넘어 '얼마나 확신하는지' 를 명시적으로 표현해야 합니다. 이는 특히 훈련 데이터와 다른 환경 (Distribution Shift) 에서 에이전트의 강건성을 결정하는 핵심 요소입니다.
향후 방향: 단순한 신념 읽기만으로도 분포 외 일반화에 큰 도움이 되지만, 메모리 제어와 신념의 통합, 그리고 더 복잡한 벤치마크를 통한 검증이 필요함을 시사합니다.

요약하자면, 이 논문은 RWKV 기반 RL 에 불확실성 (Uncertainty) 을 명시적으로 통합함으로써, 부분 관측 환경에서 에이전트의 의사결정 강건성을 획기적으로 개선할 수 있는 새로운 패러다임을 제시합니다.