Reward Prediction with Factorized World States

이 논문은 언어 모델을 활용하여 비구조화된 관측을 계층적 객체 - 속성 구조로 변환하는 'StateFactory'를 제안함으로써, 훈련 데이터의 편향 없이 다양한 도메인에서 목표 상태와의 의미적 유사성을 기반으로 한 정확한 보상 예측과 향상된 에이전트 계획 성능을 달성함을 보여줍니다.

핵심

이 논문은 "로봇이나 AI 가 목표를 달성할 때, 얼마나 잘하고 있는지 스스로 판단하는 방법" 에 대한 연구입니다.

기존의 AI 는 "정답"을 알려주는 선생님 (지도 학습) 을 통해 reward(보상) 를 배우는데, 이는 새로운 상황에서는 잘 못 할 수 있습니다. 이 논문은 "세상을 잘 이해하는 것만으로도, 스스로 보상을 예측할 수 있다" 는 아이디어를 제시합니다.

핵심 내용을 일상적인 비유로 설명해 드릴게요.

---

1. 문제: "선생님이 없으면 어떻게 해?"

상상해 보세요. 아이가 요리 배울 때, 매번 "잘했어! 100 점!"이라고 말해주는 선생님이 옆에 있다면요? 아이는 그 선생님이 가르친 레시피만 잘 따라할 뿐, 새로운 재료가 나오면 당황합니다.

기존 AI 도 비슷합니다. 특정 게임이나 환경에서 "어떤 행동을 하면 점수가 오른다"고 가르쳐주면 그 게임만 잘합니다. 하지만 환경이 조금만 바뀌어도 (예: 주방이 달라지거나, 목표가 조금 변하면) AI 는 "내가 지금 잘하고 있는 건가?"를 판단하지 못해 길을 잃습니다.

2. 해결책: "StateFactory(상태 공장)"

저자들은 AI 가 스스로 상황을 이해하고 보상을 계산할 수 있게 해주는 'StateFactory' 라는 새로운 방법을 만들었습니다.

비유: "산더미 같은 뉴스 기사 정리하기"

• 기존 방법 (Unstructured): AI 가 세상을 볼 때, "주방에 냄비가 있고, 가스레인지가 켜져 있고, 바닥에 물이 있고, 창문은 열려 있고..." 하는 식의 산더미 같은 텍스트를 그대로 봅니다. 이걸 보면 AI 는 "아, 냄비가 뜨거운 건가? 아니면 바닥이 젖은 게 문제인가?" 헷갈려 합니다.
• StateFactory 방법: 이 산더미 같은 정보를 가공 (Factorize) 해서 정리합니다.
  • 객체 (Object): 냄비, 가스레인지, 바닥
  • 속성 (Attribute): 냄비 (뜨거움), 가스레인지 (켜짐), 바닥 (젖음)

이렇게 "누가 (객체) 어떤 상태 (속성) 에 있다" 는 식으로 레고 블록처럼 쪼개서 정리하면, AI 는 혼란 없이 상황을 파악할 수 있습니다.

3. 보상 예측: "목표와 현재 상태 비교하기"

이제 AI 는 이렇게 정리된 정보를 바탕으로 "내가 얼마나 잘하고 있지?"를 계산합니다.

• 목표 (Goal): "뜨거운 냄비를 식탁에 올려라."
• 현재 상태 (State): "냄비가 가스레인지 위에 있고 뜨겁다."

StateFactory 는 목표와 현재 상태를 비교합니다.

• 냄비가 식탁에 있나? -> 아님 (0 점)
• 냄비가 뜨겁긴 하나? -> 맞음 (점수 상승)
• 냄비가 가스레인지에서 내려왔나? -> 아님 (점수 상승)

이처럼 목표와 현재 상태의 '의미적 거리'를 계산해서, 단계별로 점수를 매겨줍니다. 선생님이 점수를 줘서 가르치는 게 아니라, AI 가 스스로 "아, 냄비가 식탁에 가까워졌네? 그럼 점수가 올라가겠구나!" 하고 판단하는 것입니다.

4. 실험 결과: "새로운 환경에서도 잘한다"

저자들은 5 가지 다른 환경 (요리, 과학 실험, 웹 쇼핑, 퍼즐 등) 에서 이 방법을 테스트했습니다.

• 기존 AI (지도 학습): 특정 환경 (예: 요리) 에서만 점수를 잘 맞췄지만, 다른 환경 (예: 과학 실험) 으로 가면 완전히 엉뚱한 점수를 매겼습니다. (선생님의 말만 믿다가 새로운 상황에 당황한 셈)
• StateFactory: 어떤 환경이든 처음 보는 상황 (Zero-shot) 에서도 목표와 현재 상태를 비교해 정확한 점수를 매겼습니다.
  • 결과: 기존 방법보다 60%~8% 더 정확한 보상 예측을 했습니다.
  • 실제 효과: 이 정확한 점수 덕분에 AI 가 미로를 헤매지 않고 목표에 더 빨리 도달했습니다. (성공률 20% 이상 향상)

5. 요약: 왜 이것이 중요한가?

이 연구는 "AI 에게 정답을 외우게 하는 게 아니라, 세상을 구조적으로 이해하는 법을 가르쳐주면, AI 는 스스로 길을 찾을 수 있다" 는 것을 증명했습니다.

• StateFactory: 세상을 '객체'와 '속성'으로 쪼개서 정리하는 정리 정돈 기술.
• 효과: 새로운 환경에서도 스스로 "내가 잘하고 있구나"를 판단하여, 더 똑똑하고 유연하게 행동하게 됩니다.

마치 아이가 요리할 때, "선생님이 점수를 매겨주는 게 아니라, '냄비가 식탁에 가까워졌으니 성공에 한 걸음 더 다가갔다'는 논리를 스스로 깨우치는 것"과 같습니다. 이렇게 되면 AI 는 어떤 새로운 집, 새로운 주방에서도 요리할 수 있는 것입니다.

기술 요약

이 논문은 **Factorized World States(분해된 세계 상태)**를 활용한 **보상 예측 (Reward Prediction)**에 대한 연구로, 에이전트가 다양한 목표와 환경에서 일반화될 수 있도록 돕는 새로운 프레임워크인 StateFactory와 이를 평가하기 위한 벤치마크 RewardPrediction를 제안합니다.

주요 내용은 다음과 같습니다.

1. 문제 정의 (Problem)

• 배경: 에이전트는 목표에 도달하기 위해 행동의 결과를 추론하고 보상을 최대화하는 행동을 선택해야 합니다.
• 한계: 기존의 지도 학습 기반 보상 모델은 훈련 데이터에 내재된 편향을 도입하여 새로운 목표나 환경으로의 일반화 (Zero-shot) 능력이 제한적입니다.
• 과제: 명시적인 보상 모델링 없이, 잘 정의된 세계 상태 (World State) 표현만으로도 정확한 보상 예측이 가능한지, 그리고 이를 통해 에이전트의 계획 (Planning) 성능을 향상시킬 수 있는지 확인하는 것이 핵심 문제입니다.
• 현재 기술의 병목 현상:
  1. 작업 진행도를 정확히 반영하는 추상적인 상태 표현 공간 구축의 어려움 (특히 고수준 언어 기반 에이전트에게).
  2. 텍스트 공간에서 작동하는 에이전트를 위한 세밀한 단계별 (step-wise) 보상 품질 평가를 위한 벤치마크 부재.

2. 방법론 (Methodology)

A. RewardPrediction 벤치마크

연구진은 5 가지 다양한 도메인 (AlfWorld, ScienceWorld, TextWorld, WebShop, BlocksWorld) 을 아우르는 RewardPrediction 벤치마크를 구축했습니다.

• 데이터: 총 2,454 개의 고유한 행동 - 관찰 궤적 (Trajectory) 을 포함하며, 각 단계마다 Ground-truth 보상 (0~1 사이의 진행도 점수) 이 존재합니다.
• 평가 지표: 예측된 보상과 실제 보상의 차이를 정량화하기 위해 **EPIC 거리 (Equivalent Policy-Invariant Comparison distance)**를 사용합니다. EPIC 거리가 낮을수록 보상 함수의 정렬도가 높음을 의미합니다.

B. StateFactory 프레임워크

기존의 비구조화된 텍스트 관찰을 계층적 객체 - 속성 (Hierarchical Object-Attribute) 구조로 변환하여 보상을 예측하는 새로운 방법론입니다.

1. 상태 추출 (State Extraction):

   • 비구조화된 관찰 (Observation) 을 LLM 을 통해 **원자적 의미 요소 (Atomic Semantic Factors)**로 분해합니다.
   • 결과물은 객체 (Identity)와 동적 속성 (Attributes)의 집합으로 구조화됩니다 (예: "컵"이라는 객체에 "온도: 뜨거움", "위치: 테이블 위" 등의 속성).
   • 이는 작업과 무관한 노이즈를 제거하고 핵심 상태 변화만 포착하도록 설계되었습니다.

2. 목표 해석 (Goal Interpretation):

   • 고정된 목표 표현이 아닌, 상호작용 역사와 현재 상황에 따라 동적으로 업데이트되는 목표 상태를 생성합니다.
   • 이는 작업 수행 중 환경 변화에 적응하여 "진행의 착시 (Illusion of progress)"를 방지합니다.

3. 계층적 라우팅 (Hierarchical Routing) 및 보상 계산:

   • 현재 상태와 목표 상태 간의 **의미적 유사성 (Semantic Similarity)**을 계층적으로 계산합니다.
   • 객체 매칭: 목표 객체와 현재 상태의 객체를 식별합니다.
   • 속성 매칭: 해당 객체의 속성 값들 간의 유사도를 계산합니다.
   • 보상 도출: 모든 목표 객체의 만족도 점수를 집계하여 최종 보상 (0~1) 을 산출합니다.
   • 특징: 이 과정은 별도의 학습 (Supervised Training) 없이 Zero-shot으로 작동하며, 상태 표현의 구조적 이점을 통해 일반화 능력을 극대화합니다.

3. 주요 기여 (Key Contributions)

1. RewardPrediction 벤치마크 제안: 5 가지 텍스트 기반 환경에서 2,454 개의 궤적과 단계별 Ground-truth 보상을 포함하는 새로운 평가 기준을 마련했습니다.
2. StateFactory 방법론 개발: 평탄한 텍스트 설명을 구조화된 계층적 표현으로 변환하여 일반화 가능한 보상 예측을 가능하게 하는 새로운 표현 기법을 제시했습니다.
3. 에이전트 계획 성능 향상 검증: StateFactory 가 세계 모델 (World Model) 및 행동 제안과 결합될 때, 에이전트의 성공률을 유의미하게 높인다는 것을 실험적으로 증명했습니다.

4. 실험 결과 (Results)

A. 보상 예측 정확도 (Reward Prediction Accuracy)

• Supervised 모델의 한계: 특정 도메인에서 훈련된 보상 모델은 새로운 도메인으로 전이될 때 EPIC 거리가 138% 증가하여 일반화 실패를 보였습니다.
• StateFactory 의 우위:
  • Zero-shot 성능: VLWM-critic(비구조화) 과 LLM-as-a-Judge(프롬팅 기반) 대비 각각 60% 와 8% 낮은 EPIC 거리를 기록했습니다.
  • 전체 평균 EPIC: StateFactory 는 0.297 을 기록하여, 단일 도메인에서 훈련된 감독 모델의 일반화 성능을 능가했습니다.
  • 구조의 중요성: 비구조화된 텍스트 (0.57) 나 단순 객체 중심 (0.35) 표현보다 **객체 - 속성 분해 (0.30)**가 가장 우수한 성능을 보였습니다.

B. 에이전트 계획 성능 (Agent Planning Performance)

StateFactory 가 생성한 보상 신호를 ReAct 에이전트에 적용한 결과, 성공률이 크게 향상되었습니다.

• AlfWorld: 성공률 +21.64% 증가 (34.33% → 55.97%).
• ScienceWorld: 성공률 +12.40% 증가 (22.63% → 35.03%).
• BlocksWorld: 성공률 +8.00% 증가 (85.00% → 93.00%).
• 의미: 희소 보상 (Sparse Reward) 환경에서 에이전트가 "사고의 고리 (Deadlock)"를 깨고 장기 계획 (Long-horizon planning) 을 성공적으로 수행할 수 있게 되었습니다.

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

• 일반화 가능성: 명시적인 보상 모델 학습 없이, 구조화된 세계 상태 표현만으로도 다양한 도메인에서 정확한 보상 예측이 가능함을 입증했습니다.
• 시스템 1 vs 시스템 2: StateFactory 는 반응형 (System-1) 에이전트의 성능을 획기적으로 높일 뿐만 아니라, MCTS(몬테카를로 트리 탐색) 기반의 계획형 (System-2) 에이전트에게도 강력한 휴리스틱 신호를 제공하여 탐색 효율성을 높입니다.
• 미래 방향: 이 연구는 에이전트가 복잡한 환경에서 물리적 증거에 기반한 구조화된 추론을 통해 목표를 달성하는 데 있어, **상태 표현의 구조화 (State Factorization)**가 핵심 요소임을 강조합니다.

요약하자면, 이 논문은 **"비구조화된 텍스트를 구조화된 객체 - 속성 형태로 분해하면, 학습 데이터에 의존하지 않고도 다양한 환경에서 정확한 보상 예측이 가능하며, 이는 에이전트의 계획 성공률을 크게 향상시킨다"**는 것을 증명했습니다.