H-JEPA: 큰 목표와 작은 동작을 나누어 계획하는 계층형 세계 모델

배경과 문제

세계 모델은 에이전트가 경험으로부터 환경의 동역학을 학습하여 미래를 예측하고 계획을 세우게 하는 기술이다. 최근에는 픽셀을 재구성하는 대신 미래의 잠재 상태를 예측하는 JEPA(Joint-Embedding Predictive Architecture) 방식이 주목받고 있다. JEPA는 보상 신호나 픽셀 재구성 없이도 학습할 수 있어, 작업에 구애받지 않는 일반적인 세계 모델로 활용될 수 있다. 그러나 기존의 행동 조건부 JEPA 세계 모델은 단일 잠재공간에서 예측·계획하며, 여러 시간 지평을 지원하더라도 같은 표현 공간을 공유한다는 한계가 있다.

단일 시간 척도에서는 장기 예측을 위해 많은 세밀한 단계를 굴려야 하므로 예측 오차가 누적되고 행동 탐색 공간이 커진다. 또한 하나의 잠재공간이 저수준 동역학과 목표 매칭을 동시에 지원해야 하는데, 빠르게 변하는 세부 정보를 모두 담은 잠재 상태는 추상적인 목표(예: 위치 도달)와의 거리를 측정하는 데 적합하지 않을 수 있다. 예를 들어, 개미 로봇의 다리 자세까지 정확히 맞추는 것은 목표가 단순히 특정 위치에 도달하는 것일 때 불필요한 제약이 된다.

기존 접근의 한계

계층적 세계 모델은 이러한 문제를 해결하기 위해 각 계층이 서로 다른 시간 척도와 추상화 수준을 담당하도록 설계된다. 뇌의 피질 영역도 내재적 시간 척도의 계층을 이루며, 상위 영역일수록 느리게 변하는 표현을 가진다는 신경과학적 발견이 이러한 접근을 뒷받침한다. 그러나 기존의 계층적 잠재 비디오 모델은 픽셀을 재구성하지만 제어에는 사용되지 않았고, 계층적 모델 기반 강화학습은 보상으로부터 작업 특정 정책을 학습한다는 차이가 있다.

H-JEPA와 가장 유사한 HWM은 작업에 구애받지 않는 JEPA 세계 모델로 계층적으로 계획하지만, 서로 다른 예측 지평을 가진 예측기들이 하나의 잠재공간을 공유한다. 즉, 시간적 분해는 하지만 더 추상적인 공간에서 목표를 매칭할 수는 없다. H-JEPA는 각 계층이 자신만의 잠재공간을 학습하여, 데이터의 변화 속도가 충분히 분리된 환경에서는 상위 계층이 예측하기 어려운 빠른 세부 정보를 덜 담고 느린 상태를 유지하도록 학습한다. 이 선택적 추상화는 모든 환경에서 나타나는 보장이 아니다.

연구방법을 단계별로

H-JEPA는 관찰과 행동의 궤적으로부터 잠재 예측 모델의 계층을 학습한다. 계층은 시간적이다. 1계층은 가장 세밀한 관찰 스트림에서 작동하고, 각 상위 계층은 하위 계층이 생성한 잠재 상태를 더 거친 보폭(stride)으로 소비한다. 각 계층은 자신만의 관찰 인코더, 행동 인코더, 잠재 예측기를 가지며, 모든 계층이 과거 잠재 상태와 행동으로부터 미래 잠재 상태를 예측한다.

상위 계층의 상태 인코더는 하위 계층의 잠재 상태 창(window)을 하나의 추상 상태로 풀링하고, 행동 인코더는 하위 계층의 행동 임베딩을 집계하여 매크로 행동을 만든다. 각 계층은 JEPA로서 학습되며, 예측 손실과 SIGReg라는 붕괴 방지 정규화 항을 최소화한다. SIGReg는 임베딩 분포가 등방성 가우시안에 가까워지도록 장려하여 인코더가 상수 벡터로 붕괴하는 것을 막는다. 모든 계층은 처음부터 끝까지(end-to-end) 함께 학습되며, 상위 계층의 기울기가 하위 계층의 인코더로 전파된다.

계획은 학습된 계층을 통해 위에서 아래로 진행된다. 현재 관찰과 목표 관찰을 모든 계층에서 인코딩하여 초기 잠재 상태와 목표 잠재 상태를 얻는다. 최상위 계층은 목표까지의 거친 계획을 제안하고, 폐루프 설정에서는 각 하위 계층이 상위 계층 궤적의 첫 번째 예측 상태를 하위 목표로 삼아 더 세밀한 전이로 계획을 다듬는다. 1계층은 원시 행동을 출력하여 환경에서 실행한다. 각 계층의 행동 시퀀스는 경사 하강법으로 최적화되며, 폐루프 제어에서는 계획한 행동의 앞부분을 실행한 후 새로운 관찰을 인코딩하여 계획을 반복한다.

결과와 숫자의 의미

네 가지 시뮬레이션 환경(FourRoomDistractors, Visual AntMaze, Push-T, OGBench Cube)에서 H-JEPA는 단일 계층 LeWM보다 계획 성공률이 높았고, 특히 Visual AntMaze에서는 3계층 H-JEPA가 18%에서 73%로 성공률을 높였다. 이 성공률 비교는 환경별 50개 별도 평가 과제와 세 개 시드의 평균을 사용한 결과로, Visual AntMaze의 단일 계층 LeWM 18.0±3.5%와 3계층 H-JEPA 73.3±3.5%에 해당한다. ±는 표준오차이고, 계획 설정은 100 TFLOP 예산 아래에서 골랐다. 계산량별 비교에서도 H-JEPA가 더 나은 성공률·연산량 절충을 보였다. 그림 6의 파레토 프런트는 계층 추가가 대체로 계산 효율을 개선하지만, 깊이를 늘릴 때마다 항상 성능이 오르는 것은 아님을 보여준다. Push-T는 예외로, 3계층 모델이 최대 계산량에서도 성능이 나빴는데, 이는 에피소드가 짧아 3계층 모델을 훈련할 클립이 부족했기 때문으로 분석된다.

계층적 표현의 이점을 분리하기 위해, 상위 계층의 잠재공간에서 목표 비용만 계산하고 1계층 동역학은 그대로 사용하는 절제 실험을 수행했다. Visual AntMaze에서 상위 계층 투영 비용은 기본 1계층 비용보다 성공률을 크게 높였고, 이는 추상적인 목표 표현 자체가 계획에 도움이 됨을 보여준다. 또한 시간적 분해는 장기 작업을 더 짧은 하위 문제로 나누어 각 하위 계획자가 더 적은 순차적 롤아웃 단계로 탐색하게 하여 계산 효율을 높인다.

실제 로봇 비디오 데이터셋 DROID에서는 장면이 에피소드마다 변하기 때문에 순수 예측 목표만으로는 인코더가 배경을 유지하고 로봇을 버리는 '느린 특징 붕괴'가 발생한다. 이를 막기 위해 역동역학(IDM) 손실을 추가했다. IDM은 연속된 상태 표현으로부터 행동을 회귀하도록 하여, 인코더가 행동과 관련된 정보를 유지하도록 강제한다. DROID에서 2계층 H-JEPA는 단일 계층 모델보다 높은 Fréchet 충실도를 보였고, 더 낮은 플래너 예산에서도 우수했다.

구체적인 예시

Visual AntMaze 환경에서 3계층 H-JEPA의 계획 과정을 시각화하면, 3계층은 가장 추상적인 잠재공간에서 최종 목표를 향해 계획하고, 그 첫 번째 예측 상태가 2계층 플래너의 하위 목표가 된다. 이 분해는 1계층 플래너까지 재귀적으로 이어져 원시 행동을 출력한다. 각 계층의 시작과 목표 잠재를 디코딩한 이미지를 보면, 상위 계층의 잠재는 개미의 위치와 미로 레이아웃은 유지하면서 다리 자세 세부 정보는 잃는다.

FourRoomDistractors에서는 파란 방해물(distractor)이 무작위로 순간이동하는데, 상위 계층으로 갈수록 이 방해물의 위치 정보가 사라지고 빨간 에이전트의 위치는 유지된다. 이는 각 계층이 자신의 예측 지평에서 예측 가능한 특징만 유지하도록 학습되었기 때문이다. 방해물의 움직임은 짧은 지평에서는 예측 가능하지만, 긴 지평에서는 순간이동을 만날 가능성이 높아 예측하기 어렵다. 따라서 상위 계층은 이를 버리고 느리게 변하는 에이전트 위치만 남긴다.

한계와 오해하기 쉬운 점

선택적 추상화는 모든 환경에서 나타나지 않는다. Push-T와 OGBench Cube 같은 조작 데이터셋에서는 엔티티 주파수 간격이 작아 상위 계층이 빠른 특징을 버리지 않는다. 따라서 계층적 표현의 이점은 데이터에 느리게 변하는 요소와 빠르게 변하는 요소가 분리되어 있을 때 주로 나타난다. 또한 Push-T에서는 에피소드가 짧아 깊은 계층을 훈련할 유효 클립이 부족한 것이 3계층 성능 저하의 가능한 원인으로 제시된다.

DROID 실험은 오프라인 평가로, 시뮬레이터가 없어 계획을 실제로 실행하지 않고 전문가 궤적과의 Fréchet 거리로만 점수를 매긴다. 따라서 실제 로봇에서의 폐루프 제어 성능은 아직 검증되지 않았다. 또한 Fréchet 충실도는 경로의 모양을 측정하지만, 작업 완료 여부와 항상 일치하지는 않는다. 예를 들어, OGBench Cube에서 Fréchet 충실도가 특정 임계값을 넘어도 잡기에 실패할 수 있다.

실제로 어떻게 활용할까

H-JEPA는 보상 신호 없이 오프라인 데이터로부터 학습할 수 있어, 다양한 환경에서 제로샷 계획이 가능하다. 예를 들어, 로봇이 새로운 작업을 받았을 때 목표 이미지만 주어지면 학습된 계층적 세계 모델을 사용하여 행동 시퀀스를 계획할 수 있다. 다만 여기서 제로샷은 작업별 보상 학습 없이 목표를 계획한다는 의미이며, 임의의 새 환경으로의 일반화를 보장하지 않는다. DROID 결과는 기록된 영상으로 계획 경로를 평가한 것이어서 실제 로봇에 바로 적용할 수 있다는 증거로 확대하면 안 된다.

향후 연구 방향으로는 행동이 없는 비디오로부터 계층을 학습하는 것, 상위 계층의 잠재를 언어와 정렬하여 목표를 단어로 지정하는 것, 고정 보폭 대신 가변 길이 세그먼트로 상위 계층을 훈련하는 것 등이 제시된다. 이러한 확장은 H-JEPA를 더 일반적인 시각적 계획 시스템으로 발전시킬 수 있을 것이다.