Visualizing the World Model in an Agent's Value Function

Could model-free reinforcement learning algorithms actually be learning a model of the environment? In this work, we apply counterfactual interventions which hold everything about a state fixed, besides a targeted change, to explore if the value function of a Proximal Policy Optimization agent forms a world model of the Crafter environment. Despite only being trained to optimize reward, we discover that the value function learns concepts such as the rules governing when a diamond is collectible, water and food having an impact on health, heuristics regarding inventory, and more. Finally, we discuss the safety implications of our visualizations and some shortcomings of current reinforcement learning agents that will prevent us from building truly advanced machine intelligence.

한 줄 요약

PPO 에이전트의 가치 함수가 환경 모델을 암묵적으로 학습한다는 것을 반사적 분석을 통해 밝혀냄.

핵심 기여도

PPO 에이전트의 가치 함수가 Crafter 환경의 규칙(다이아몬드 수집 조건, 물과 음식의 건강 영향 등)을 학습함.
반사적 분석을 통해 가치 함수 내 암묵적 세계 모델을 시각화함.
100억 환경 스텝 동안 훈련된 PPO 에이전트의 가치 함수에서 3가지 핵심 환경 개념을 확인함.
가치 함수가 인벤토리 구성 요소를 에피소드 진행 지표로 활용함.

핵심 아이디어

기존 강화 학습 에이전트는 보상 최적화에 집중하지만, 이 연구는 가치 함수가 환경의 내재적 규칙을 암묵적으로 학습할 수 있음을 제시한다.
PPO 에이전트의 가치 함수(크리틱)가 환경의 미래 상태와 보상을 예측하기 위해 내부적으로 세계 모델을 형성한다는 가정을 세웠다.
이를 검증하기 위해 반사적 분석(counterfactual analysis)을 사용하여 특정 관측 요소만 변경하고 나머지 상태를 고정한 후, 가치 예측의 변화를 관찰했다.
예를 들어, 다이아몬드가 보이더라도 철 곡괭이가 없으면 수집 불가능하다는 환경 규칙이 가치 함수에 인코딩되어 있음을 확인했다.

기술적 접근법

**알고리즘**: Proximal Policy Optimization (PPO)
**환경**: Crafter (2D Minecraft 유사 환경)
**분석 방법**: 반사적 분석(counterfactual interventions)
**관측 요소 조작**: 에이전트의 관측 상태 중 특정 요소(예: 인벤토리, 건강, 자원)만 변경하고 나머지를 고정
**훈련 스텝**: 100억 환경 스텝
**가치 함수**: 신경망 기반 크리틱(critic)을 사용하여 미래 보상의 기대치를 예측

주요 결과

**Crafter 환경에서 가치 함수가 3가지 핵심 개념을 학습함**:

1. 환경 내 유한 보상의 존재
2. 특정 업적(예: 다이아몬드 수집) 달성 조건
3. 건강 증감 요인(예: 물, 음식)

**인벤토리 구성 요소가 에피소드 진행 지표로 활용됨** → 가치 예측에 큰 영향
**다이아몬드 수집 가능성 판단**: 다이아몬드가 보이더라도 철 곡괭이가 없으면 수집 불가능하다는 규칙을 인식
**식물은 익었을 때만 먹을 수 있음** → 가치 함수가 이 조건을 반영함
**물과 소는 에이전트의 수분/음식 수준에 따라 유용성 달라짐** → 환경 내부와 외부 상태의 상호작용을 이해

의의 및 한계

**의의**: 가치 함수가 암묵적으로 세계 모델을 형성한다는 사실은 강화 학습 에이전트의 내부 작동 메커니즘을 이해하는 데 기여하며, 안전성 확보와 인간 가치 정렬에 중요한 단서를 제공.
**한계**: 현재 강화 학습 에이전트는 복잡한 환경에서의 일반화 능력이 부족하며, 암묵적 모델이 항상 정확하게 학습되는 것은 아님.
**추가 한계**: 반사적 분석은 특정 요소만 조작하기 때문에, 전체 환경 상호작용을 완전히 이해하기는 어려움.

실용적 활용

**로봇 제어**: 에이전트가 환경 내 규칙을 이해하면, 실제 물리 환경에서 안정적인 제어가 가능해짐.
**보안 및 안전성 분석**: 가치 함수 내 암묵적 모델을 분석하여 악의적 에이전트의 행동을 예측하고 방지할 수 있음.
**게임 AI 개발**: 게임 내 복잡한 규칙을 이해하는 AI를 개발하여 플레이어 경험을 향상시킬 수 있음.