DreamSAC: Learning Hamiltonian World Models via Symmetry Exploration

이 논문은 물리 법칙의 불변성을 학습하여 외삽 일반화 능력을 향상시키기 위해 해밀토니안 기반 호기심 보상을 활용한 대칭성 탐색 전략과 자체 지도 대조 학습을 결합한 'DreamSAC' 프레임워크를 제안합니다.

핵심

드림 SAC: 물리 법칙을 직접 탐험하며 배우는 AI 의 이야기

이 논문은 **"AI 가 세상을 어떻게 더 똑똑하게 이해할 수 있을까?"**라는 질문에 대한 새로운 답을 제시합니다. 기존 AI 들이 겪던 큰 실수를 지적하고, 이를 해결하는 창의적인 방법 '드림 SAC(DreamSAC)'을 소개합니다.

1. 기존 AI 의 문제: "무작위 TV 시청자" vs "물리학자"

기존의 AI(세계 모델) 들은 마치 무작위로 채널을 돌리며 TV 를 보는 사람과 같습니다.

• 무엇을 하나요? 화면에 나타나는 픽셀 (이미지) 패턴을 열심히 외웁니다. "공이 왼쪽으로 가면 오른쪽으로 튕긴다"는 식의 통계적 상관관계를 기억하죠.
• 어떤 문제가 있나요? 익숙한 상황에서는 잘하지만, 전혀 새로운 상황 (예: 공의 무게가 갑자기 2 배가 되거나, 중력이 변하는 상황) 이면 완전히 엉망이 됩니다. 왜냐하면 그들은 **'왜' 그렇게 움직이는지 (물리 법칙)**를 모르고, **'무엇이' 어떻게 움직였는지 (화면 패턴)**만 외웠기 때문입니다.

이들은 시각적으로는 다양해 보이지만, 물리적으로는 뻔한 데이터만 보고 학습해서, 실제 세상의 법칙을 깨닫지 못합니다.

2. 드림 SAC 의 해결책: "호기심 많은 물리학자"가 되다

저자들은 이 문제를 해결하기 위해 AI 를 능동적인 탐험가로 바꾸었습니다. 이를 위해 두 가지 핵심 장치를 도입했습니다.

① 대칭성 탐험 (Symmetry Exploration): "일으켜 세우는 호기심"

기존 AI 는 가만히 있는 데이터를 보지만, 드림 SAC 는 스스로 환경을 건드려보며 배웁니다.

• 비유: 아이가 장난감 자동차를 볼 때, 그냥 바라보는 게 아니라 바퀴를 굴리고, 벽에 부딪히고, 무거운 돌을 올려놓아 "어? 왜 이렇게 안 움직이지?"라고 궁금해하며 실험을 반복하는 것과 같습니다.
• 원리: AI 는 **'해밀토니안 (에너지 법칙) 호기심'**이라는 내재적 동기를 가집니다. "내가 무엇을 했을 때 에너지가 가장 크게 변할까?"를 찾아서 적극적으로 행동합니다. 이를 통해 AI 는 물리 법칙의 핵심 (에너지 보존 법칙 등) 을 깨닫는 데 필요한 가장 중요한 데이터를 스스로 수집합니다.

② 해밀토니안 세계 모델: "물리 법칙을 내장한 두뇌"

수집된 데이터를 학습할 때, 일반적인 AI 는 픽셀을 외우지만, 드림 SAC 는 **물리 법칙 (에너지, 운동량 보존)**을 수학적으로 내장한 구조를 사용합니다.

• 비유: 일반적인 AI 가 "이 그림은 공이 굴러가는 그림이야"라고 외우는 반면, 드림 SAC 는 "공의 질량과 속도를 계산해서 다음 위치를 물리 공식으로 예측하는" 두뇌를 가집니다.
• 핵심 기술: 카메라 각도가 달라져도 (시점 변화) 물리 법칙은 변하지 않는다는 점을 이용합니다. AI 는 카메라가 움직이는 '소음'을 제거하고, 사물 자체의 '물리 상태'만 남기는 **특수한 필터 (대조 학습)**를 통해 학습합니다.

3. 실제 효과: 새로운 세상에서도 당황하지 않는 AI

이 방법을 적용한 결과, 드림 SAC 는 놀라운 능력을 보여줍니다.

• 예측 능력: 같은 환경에서도 다른 AI 들보다 훨씬 정확하게 미래를 예측합니다. (오류가 10 배 이상 줄었습니다!)
• 적응력: 훈련할 때 보지 못했던 새로운 중력이나 미끄러운 바닥에 놓여도, 금방 적응하여 임무를 수행합니다. 마치 물리 법칙을 이해한 사람이라면 새로운 환경에서도 빠르게 움직일 수 있는 것과 같습니다.
• 제 0 회 (Zero-shot) 일반화: 아예 새로운 물리 법칙이 적용된 세상에서도, 기존에 배운 '물리 법칙'을 바탕으로 새로운 상황을 추론해냅니다.

4. 한 줄 요약

"기존 AI 는 TV 화면의 패턴을 외우는 '암기왕'이라면, 드림 SAC 는 직접 장난감을 부수고 조립하며 물리 법칙을 깨닫는 '호기심 많은 과학자'입니다."

이 연구는 AI 가 단순히 데이터를 외우는 것을 넘어, 세상의 근본적인 법칙 (물리 법칙) 을 스스로 발견하고 이해함으로써, 예측 불가능한 현실 세계에서도 robust(튼튼한) 하게 작동할 수 있는 길을 열었습니다.

기술 요약

1. 문제 정의 (Problem Statement)

기존의 학습된 월드 모델 (World Models) 은 고차원 픽셀 입력을 기반으로 환경의 예측을 수행하며, 익숙한 객체와 역학에 대한 보간 (Interpolation) 일반화에서는 탁월한 성과를 보입니다. 그러나 외삽 (Extrapolation) 일반화, 즉 훈련 데이터와 다른 물리적 속성 (예: 새로운 질량 비율, 마찰 계수, 중력 변화) 이나 복잡한 물리적 상호작용이 포함된 시나리오에서는 예측 능력이 급격히 저하됩니다.

• 근본 원인: 기존 모델들은 픽셀 수준의 통계적 상관관계 (Statistical Correlations) 를 학습하여 기술적 (Descriptive) 인 시스템이 되었을 뿐, 환경을 지배하는 근본적인 **생성 규칙 (Generative Rules)**이나 **물리 법칙 (보존 법칙, 대칭성)**을 학습하지 못했습니다.
• 한계: 수동적으로 수집된 데이터는 시각적으로는 다양할 수 있으나 물리적으로는 중복적 (Redundant) 인 경우가 많아, 모델이 우연한 통계적 패턴만 학습하게 만듭니다.

2. 제안 방법론 (Methodology: DreamSAC)

저자들은 **DreamSAC (Dream with Symmetry-Aware Curiosity)**라는 새로운 프레임워크를 제안합니다. 이는 물리 법칙에 기반한 월드 모델을 학습하여 외삽 능력을 극대화하는 것을 목표로 합니다. 핵심은 **대칭성 탐색 (Symmetry Exploration)**과 **해밀토니안 월드 모델 (Hamiltonian World Model)**의 결합입니다.

가. 해밀토니안 월드 모델 (Hamiltonian World Model)

• 구조: 기존의 비구조화된 recurrent state-space 모델 (RSSM) 을 수정하여, 물리적 대칭성을 강제하는 **해밀토니안 역학 (Hamiltonian Dynamics)**을 도입했습니다.
• 상태 표현: 관측된 픽셀 ($x_t$) 을 객체 중심의 잠재 상태 ($Z_t$) 로 매핑하며, 각 슬롯을 일반화 좌표 ($q_t$) 와 정준 운동량 ($p_t$) 으로 분해합니다.
• 대칭성 강제 (G-invariant): 내부 해밀토니안 $H_\phi(Z_t)$가 3D 물리 대칭군 (예: SE(3)) 에 대해 불변하도록 설계되었습니다. 이를 위해 Lie Transformer 아키텍처를 사용하여 구조적으로 대칭성을 보장합니다.
• 시점 강건성 (Viewpoint Robustness): 픽셀 입력은 시점에 의존적이지만 물리 법칙은 불변해야 한다는 모순을 해결하기 위해, 자기지도 학습 대비 학습 (Self-Supervised Contrastive Learning) 기반의 손실 함수 ($L_{vr}$) 를 도입했습니다. 이는 카메라 시점 변화에 무관한 물리 상태를 추출하도록 인코더를 훈련시킵니다.

나. 대칭성 탐색 (Symmetry Exploration)

• 동기 부여: 에이전트가 수동적으로 데이터를 수집하는 대신, **해밀토니안 기반의 호기심 (Hamiltonian-based Curiosity)**을 통해 능동적으로 환경을 탐색합니다.
• 내재적 보상 (Intrinsic Reward): 에이전트의 행동으로 인해 시스템의 해밀토니안 (에너지) 이 얼마나 변하는지 ($|\Delta H_\phi|$) 를 측정하여 보상을 부여합니다.
  • 수식: $r_{sym} \approx |\Delta H_\phi| - \lambda_s ||a_t - a_{t-1}||^2$
  • 의미: 시스템의 대칭성 (보존 법칙) 을 깨뜨리는, 즉 외부 힘에 의해 가장 많은 일 (Work) 을 수행하는 상호작용을 찾아내도록 유도합니다. 이는 모델이 물리 법칙의 구조적 속성 (강성, 퍼텐셜 장벽 등) 을 학습하는 데 가장 정보량이 많은 데이터를 수집하게 합니다.
• 학습 전략: 초기에는 안정적인 탐색을 위해 RND(Random Network Distillation) 보상을 사용하다가, 해밀토니안 모델이 학습됨에 따라 점진적으로 대칭성 탐색 보상 ($r_{sym}$) 으로 전환 (Annealing) 합니다.

다. 하류 작업 적응 (Downstream Task Adaptation)

• 차별화된 미세 조정 (Differentiated Fine-tuning): 새로운 작업 (OOD 환경) 에 적응할 때, 시각 인코더는 고정하고 해밀토니안 모델의 물리 파라미터 (질량, 마찰 등) 만 빠르게 미세 조정합니다. 이는 학습된 대칭성 구조를 유지하면서 새로운 물리 상수에 빠르게 적응할 수 있게 합니다.

3. 주요 기여 (Key Contributions)

1. 대칭성 탐색 (Symmetry Exploration): 해밀토니안 기반의 호기심 보상을 통해 물리적으로 정보량이 많은 데이터를 능동적으로 수집하는 비지도 탐색 전략을 제안했습니다.
2. 해밀토니안 월드 모델: 대비 학습 (Contrastive Learning) 을 통해 픽셀에서 시점 불변의 물리 상태를 추출하고, Lie Transformer 를 통해 물리 대칭성을 구조적으로 강제하는 월드 모델을 설계했습니다.
3. 외삽 일반화 성능 입증: 3D 물리 시뮬레이션 환경에서 기존 SOTA 모델 (DreamerV3 등) 을 압도하는 외삽 성능을 보여주었습니다.

4. 실험 결과 (Results)

DeepMind Control Suite (DMCS) 와 GymFetch 의 다양한 3D 물리 벤치마크에서 실험을 수행했습니다.

• 예측 정확도: DreamSAC 는 이미지 예측 오차 (MSE) 에서 기존 모델들보다 22%~163% 더 낮은 오차를 기록했습니다. 특히 긴 시간 단계 (Rollout Horizon) 에서도 안정적인 예측을 보여주었습니다.
• 외삽 일반화 (OOD Performance):
  • 구조적 OOD: 새로운 시점 (Unseen View), 새로운 객체 수 (Unseen Object), 새로운 목표 위치 (Unseen Goal) 에서 기존 모델 대비 월등히 높은 성공률과 보상을 달성했습니다.
  • 매개변수 OOD: 훈련과 다른 중력 (1.5 배), 마찰 (2.0 배), 물체 질량 등의 변화에 대해 빠른 적응을 보였습니다. 특히 "Unseen Distribution" (훈련 분포의 외곽 영역) 에서도 뛰어난 성능을 발휘하여 단순한 보간이 아닌 진정한 외삽 능력을 입증했습니다.
• 제거 실험 (Ablation Study):
  • $L_{vr}$ (시점 강건성 손실) 제거: 시점 변화에 대한 성능이 급격히 저하됨.
  • $H_\phi$ (해밀토니안 구조) 제거: 새로운 물리 파라미터 (중력 등) 에 대한 적응 능력이 크게 떨어짐.
  • SAVi (객체 중심 인코더) 제거: 물리 속성 일반화 성능 저하.
  • 이는 제안된 모든 구성 요소가 필수적임을 입증합니다.

5. 의의 및 결론 (Significance)

이 논문은 월드 모델 학습의 패러다임을 **"수동적인 통계적 학습"에서 "능동적인 물리 법칙 발견"**으로 전환합니다.

• 물리 기반의 일반화: 에이전트가 단순히 픽셀 패턴을 기억하는 것이 아니라, 에너지 보존 법칙과 같은 근본적인 물리 대칭성을 학습함으로써, 훈련되지 않은 새로운 물리적 환경에서도 robust 하게 작동할 수 있음을 증명했습니다.
• 데이터 효율성: 능동적인 탐색 전략을 통해 물리적으로 의미 있는 데이터를 효율적으로 수집함으로써, 적은 데이터로도 복잡한 물리 역학을 학습할 수 있는 가능성을 제시했습니다.
• 실제 적용 가능성: 로봇 제어 및 오픈 월드 환경과 같이 예측 불가능한 물리적 상호작용이 필요한 분야에서, 기존 모델의 한계를 극복할 수 있는 강력한 프레임워크를 제공합니다.

요약하자면, DreamSAC는 에이전트가 "호기심"을 통해 물리 법칙을 직접 탐구하게 함으로써, 단순한 패턴 인식을 넘어 물리적으로 grounded 된 진정한 일반화 능력을 갖춘 월드 모델을 구축하는 혁신적인 접근법입니다.