What Must Generalist Agents Remember?

논문 제목: 범용 에이전트는 무엇을 기억해야 하는가? (What Must Generalist Agents Remember?)
저자: Khurram Yamin, Namrata Deka, Maitreyi Swaroop, Albert Ting, Jeff Schneider, Bryan Wilder (카네기 멜런 대학교, 조지아 공과대학교)

문제 정의

본 논문은 범용 강화 학습(Reinforcement Learning, RL) 에이전트의 발전에 있어 근본적인 질문, 즉 범용 에이전트가 여러 환경(도메인)과 다양한 목표에 걸쳐 최적에 가까운 행동을 수행하기 위해 내부적으로 어떤 정보를 저장해야 하는지를 다룹다.

기존 연구들이 에이전트에게 현재 작업이나 도메인에 대한 명시적인 인덱스 또는 라벨이 주어진다고 가정하는 것과 달리, 본 논문은 에이전트가 그러한 인덱스를 받지 못하는 설정을 연구한다. 에이전트는 오직 자신의 궤적(상태-행동 이력)만을 관찰하며, 적절한 행동 선택에 필요한 잠재적 도메인 정체성을 추론해야 한다. 저자들은 이를 잠재적 도메인이 포함된 부분 관측 가능한 마르코프 결정 과정(Partially Observable Markov Decision Process, POMDP)으로 프레임화하였으며, 여기서 숨겨진 변수는 전이 역학(transition dynamics)을 변화시키는 고정된 도메인 인덱스 $\sigma$이고, 보상 메커니즘은 모든 도메인에서 공유된다.

핵심 과제는 서로 다른 두 도메인이 동일한 관측값(상태)을 제시할 수 있지만, 각각 서로 양립할 수 없는 최적의 행동을 요구할 수 있다는 점이다. 본 논문은 메모리가 이러한 모호성을 해결하기 위해 엄격히 필요한 조건은 무엇인지, 그리고 충분히 유능한 메모리가 환경의 국소적 전이 역학을 재구성할 수 있게 하는지 묻는다.

방법론 및 이론적 프레임워크

저자들은 이산 시간(discrete time) 상에서 일련의 목표 $\Psi$와 도메인 집합 $\Sigma$를 사용하여 상호작용을 공식화하였다. 이들은 두 가지 주요 이론적 결과를 도입한다.

1. 표현 분리 (정리 1 - Representation Separation)
첫 번째 결과는 메모리의 필요성에 대한 필요조건을 확립한다. 저자들은 두 개의 서로 다른 도메인이 동일한 상태 $s^\dagger$에서 서로 분리된 최적 행동 집합을 요구하는 "공유 병목(shared bottleneck)" 시나리오(가정 1)를 정의한다.

• 주장: 에이전트가 모든 목표와 도메인에 대해 균등하게 최적에 가깝다면, 에이전트의 내부 메모리 상태 $M_t$는 병목 상태에서 이들 도메인을 구별해야 한다.
• 메커니즘: 이 정리는 에이전트의 정책이 $\epsilon$-근사 최적일 때, 도메인에 따른 병목 상태에서의 메모리 조건부 확률 분포 간의 총 변동(Total Variation, TV) 거리가 0보다 커야 함을 증명한다. 구체적으로, 두 도메인의 메모리 분포 간 TV 거리는 가치 격차(value gap), 병목 지점에 도달할 확률, 그리고 할인 계수(discount factor)의 함수에 의해 하한이 결정된다.
• 함의: 메모리는 단순히 최적 제어를 위한 충분 통계량(표준 POMDP의 믿음 상태처럼)이 아니라, 숨겨진 도메인 정체성이 행동에 영향을 미칠 때 균등한 능력을 달성하기 위한 필요한 조건이다.

2. 디코더 존재성 (정리 2 - Decoder Existence)
두 번째 결과는 분리에 필요한 메모리가 환경의 역학을 재구성하기에도 충분한지를 조사한다.

• 설정: 저자들은 일회성 예측 테스트 역할을 하는 보조 가족인 "프로브 목표(probe goals, $\Psi_F$)"를 도입한다. 이들은 "가치 충분 메모리(value-sufficient memory)"를 국소적 전이 커널을 구별할 수 있는 프로브 목표들에 대해 정확하게 최적 행동 가치를 예측할 수 있는 메모리로 정의한다.
• 주장: 만약 메모리가 국소 전이 커널을 구별할 수 있는(적분 확률 메트릭을 통해) 풍부한 프로브 목표 클래스에 대해 가치 충분하다면, 국소적 개입 전이 커널 $P^\sigma(\cdot | s, a)$를 근사적으로 재구성할 수 있는 디코더가 존재한다.
• 함의: 풍부한 목표 세트에 대해 최적의 계획(planning)을 지원하는 메모리 상태는 본질적으로 환경의 국소 역학을 디코딩할 수 있는 충분한 정보를 포함하고 있다. 이는 환경으로부터 명시적인 세계 모델을 제공받지 않고도 근사적인 계획을 가능하게 한다.

실험적 검증

저자들은 ForkWorld라는 커스텀 환경을 사용하여 이론적 주장을 검증하였다.

• 환경: T-접합 그리드 월드(T-junction gridworld)로, 복도가 갈림길(병목)로 이어진다. 도메인은 잠재적이다: "정상(normal)" 도메인에서는 '위(up)'가 북쪽으로 이동하고 '아래(down)'가 남쪽으로 이동하지만, "반전(swapped)" 도메인에서는 이 매핑이 뒤바뀐다. 목표(위 또는 아래 도달)는 관찰 가능하지만, 도메인은 관찰되지 않는다.
• 병목: 갈림길에서 최적의 행동은 숨겨진 도메인에 전적으로 의존한다. 메모리가 없는 에이전트는 도메인을 구분할 수 없으며, 따라서 추측에 의존하게 되어 성공률이 50%로 제한된다.
• 에이전트: 연구는 메모리 없는 에이전트(DQN), 명시적 히스토리 윈도우를 가진 에이전트(Stacked DQN), 그리고 순환 에이전트(DRQN)를 오라클(oracle) 베이스라인과 비교한다.

주요 결과:

1. 메모리의 필요성: 메모리 없는 에이전트는 50%의 베이스라인을 넘지 못하고 (19% 달성) 실패하였으며, 이는 메모리를 통한 도메인 분리가 없으면 정책이 안정적인 솔루션으로 수렴할 수 없음을 확인시켜 준다. 메모리를 갖춘 에이전트(Stacked DQN 및 DRQN)는 7588%의 성공률을 기록하며 오라클과의 격차를 좁혔다.
2. 표현 분리: 에이전트의 내부 상태로부터 도메인을 예측하도록 훈련된 선형 프로브(linear probe)를 사용한 결과, 메모리를 갖춘 에이전트는 갈림길을 재방문할 때 거의 완벽한 도메인 분류 성능(정확도 >97%)을 보였다. 반면 메모리 없는 에이전트는 여전히 무작위 수준(50%)에 머물렀다. 이는 정리 1을 경험적으로 확인해 준다.
3. 디코더 가능한 역학: 저자들은 에이전트의 메모리 표현으로부터 다음 상태를 예측하는 디코더를 훈련시켰다. 메모리를 갖춘 에이전트는 베이스라인에 비해 전이 예측 오차를 유의미하게 줄였으나, 메모리 없는 에이전트는 개선을 보이지 않았다. 이는 정리 2를 뒷받침하며, 계획에 사용되는 메모리가 국소 전이 역학을 인코딩한다는 것을 보여준다.

핵심 기여 및 의의

본 논문은 다음과 같은 기여를 한다:

1. 메모리의 형식적 필요성 입증: 잠재적 도메인 내에서 작동하는 범용 에이전트에게 메모리는 선택 사항이 아니라, 도메인이 관측을 공유하되 서로 다른 행동을 요구할 때 균등한 최적 근사성을 달기 위한 필수 조건임을 제어 이론적으로 증명하였다.
2. 세계 모델을 위한 기질로서의 메모리: 가치 예측을 위한 프로브 목표에 대해 충분히 풍부한 메모리 표현은 본질적으로 국소 전이 역학의 디코딩을 지원한다는 것을 보여주었다.
3. 분석의 전환: 에이전트에게 작업 인덱스가 주어진다고 가정하는 기존의 "좋은 조절자(good regulator)" 또는 인과적 RL 연구와 달리, 본 연구는 히스토리-투-메모리(history-to-memory) 매핑을 분석하여 에이전트가 궤적 데이터만으로 잠재적 도메인을 암묵적으로 분리해야 함을 보여주었다.

의의:
저자들은 이러한 결과가 메모리를 범용 에이전트를 위한 도메인 모호성 해소, 전이 모델 재구성, 그리고 계획을 지원하는 기질(substrate)로 규정한다고 주장한다. 이 연구는 잠재적 맥락을 추론하고 국소 역학을 재구성하는 능력이 단순히 추가할 수 있는 아키텍처적 특징이 아니라, 변화하는 환경 군(family of environments)에서 우수한 성능을 내기 위한 근본적인 요구사항임을 시사한다. 연구 결과는 성공적인 범용 RL 에이전트가 잠재적 도메인을 분리하는 표현을 내부에 인코딩해야 하며, 이러한 인코딩이 환경의 인과 구조 재구성을 자연스럽게 촉진한다는 점을 시사한다.