H-JEPA가 계층별 잠재공간 예측으로 장기 시각 계획 성공률을 73%로 높인다

무엇인가

장기 호라이즌 시각 계획에서 잠재 세계 모델은 여러 시간 스케일과 추상화 수준에 걸쳐 추론해야 한다. 기존의 태스크 비의존적 JEPA 세계 모델은 하나의 잠재공간, 대개 하나의 시간 스케일에서 예측하고 계획한다. 논문은 여기서 두 가지 한계를 짚는다. 첫째, 단일 시간 스케일에서는 장기 예측이 세밀한 스텝을 많이 롤아웃해야 하므로 예측 오차가 누적되고 행동 탐색 공간이 커진다. 둘째, 하나의 공유 잠재공간이 저수준 동역학과 목표 매칭을 동시에 떠받쳐야 하는데, 빠르게 변하는 디테일을 모두 담은 잠재는 저수준 동역학은 잘 모델링해도 목표가 상태보다 추상적일 때(포즈를 맞추는 게 아니라 도달할 위치를 찾는 경우) 목표 매칭 비용으로는 형편없어진다. 뇌의 피질 영역이 내재적 시간 스케일의 계층을 이루고 느린 표현이 위층에 온다는 관찰이 동기로 제시된다.

어떻게 동작하나

H-JEPA는 행동 조건부 JEPA들을 계층으로 쌓아 end-to-end로 학습하는 레시피다. 각 레벨은 자기 잠재공간 안에서, 아래 레벨보다 더 긴 호라이즌과 더 거친 시간 stride로 미래를 예측한다. 레벨 1은 가장 세밀한 관측 스트림을 다루고, 상위 레벨은 아래 레벨이 만든 잠재 상태를 더 큰 stride로 소비한다. 레벨마다 자체 관측 인코더, 행동 인코더, 잠재 예측기를 갖는다. 상위 레벨의 시간 하이퍼파라미터는 두 개다. stride s는 상위 시간 t를 하위 시간 t·s로 사상하는 서브샘플링 인자이고, 윈도 크기 w는 상위 상태 하나가 요약하는 하위 스텝 수다. 상태 인코더는 w스텝 하위 윈도를 하나의 추상 상태로 풀링하고, 행동 인코더는 w와 무관하게 s개의 하위 행동 임베딩을 집계해 상위 상태에서 다음 상위 상태로 가는 전이 전체를 덮는다. 모든 실험에서 w=1이라 상위 상태 인코더는 pointwise다. 각 레벨은 자기 잠재공간에서 JEPA로 훈련된다. 컨텍스트 잠재 상태들과 행동 임베딩이 주어지면 예측기가 미래 잠재 상태를 예측하고, teacher-forcing 방식의 잠재 예측 손실에 SIGReg라는 스케치된 정규성 정규화 항을 더한다. SIGReg는 임베딩 분포를 등방성 가우시안에 가깝게 밀어 표현 붕괴를 막는다. 레벨 1 JEPA는 LeWM으로, ViT-Tiny 인코더가 [CLS] 토큰으로 이미지를 전역 표현으로 요약하고 인과 트랜스포머 예측기가 뒤따른다. 상위 레벨은 2층 MLP 인코더와 인과 트랜스포머 예측기를 쓰고 stride 2, 윈도 1을 쓴다. 계획은 top-down이다. 현재와 목표 관측을 모든 레벨에서 인코딩한 뒤, 최상위 레벨이 목표까지의 거친 계획을 세운다. 최상위 플래너는 최종 예측 상태와 목표 잠재 사이의 L2 거리를 최소화하는 매크로 행동을 찾는다. 그 롤아웃이 아래 레벨의 서브골이 되고, 각 하위 레벨은 상위 레벨 잠재공간으로 인코딩된 서브골에 자기 예측 상태를 맞추도록 행동을 최적화한다. 이 과정이 레벨 1까지 내려가 원시 행동을 낸다. 서브골 개수 K를 정하면 하위 호라이즌이 K·s+w-1로 결정된다. 폐루프 계획은 K=1로 첫 서브골만 맞추고, 개루프는 K=H로 중간 서브골까지 맞춘다. 행동 시퀀스 최적화는 모든 레벨에서 경사하강법으로 한다.

무엇과 다른가

실험은 내비게이션용 FourRoomDistractors와 Visual AntMaze, 조작용 Push-T와 OGBench Cube 네 개의 시뮬레이션 환경에서 최대 4레벨까지 학습해 이뤄졌다. 핵심 수치는 Visual AntMaze다. 3레벨 계층이 성공률을 18%에서 73%로 끌어올리면서 플래너 연산은 오히려 줄였다. FourRoom, AntMaze, Cube에서는 레벨을 하나씩 추가할 때마다 성공률-연산 Pareto frontier가 위와 왼쪽으로 이동했다. Push-T는 예외로, 2레벨은 모든 예산에서 LeWM과 같거나 앞섰지만 3레벨은 최대 연산에서도 성능이 나빴다. 저자들은 원인을 훈련 데이터로 본다. 클립이 한 에피소드 안에 있어야 하는데 Push-T 에피소드가 3레벨 모델 학습에 필요한 최소 클립 수에 비해 짧아, 에폭당 LeWM 전이의 58%만 보게 된다. 다른 환경들은 3레벨에서 75~88%를 유지한다.

어떻게 쓰나

시간 계층만으로 얻는 이득과 표현 계층이 주는 이득을 분리하기 위해, 하나의 잠재공간에 갇힌 계층 플래너 HWM과 깊이 대 깊이로 비교한다. HWM은 같은 구현, 데이터, 학습 설정을 쓰되 레벨 1 위 인코더를 항등맵으로 바꿔 모든 레벨이 같은 잠재공간에서 예측·계획하게 한 것이다. H-JEPA의 우위는 AntMaze에서 가장 두드러져, 3레벨 H-JEPA가 HWM 성공률의 거의 두 배에 도달한다. Cube에서도 높은 깊이에서 H-JEPA가 앞서지만 AntMaze보다 완만하고, FourRoom과 Push-T에서는 두 계층이 비슷하다. HWM 자체도 네 환경 모두에서 LeWM보다 평균 성공률이 높아, 표현 계층이 없어도 시간 분해만으로 이득이 생긴다는 점이 확인된다. 논문은 계획 개선을 두 메커니즘으로 분해한다. 하나는 시간 분해, 다른 하나는 더 추상적인 잠재공간에서 미래를 채점하는 것이다. 두 번째 요인만 떼어내기 위해, 레벨 1 세계 모델과 롤아웃 동역학은 그대로 두고 목표 비용만 상위 레벨 잠재공간에서 계산하는 실험을 한다. 상위 인코더들이 pointwise이므로 레벨 1 잠재 하나가 모든 레벨로 투영된다. 결과적으로 테스트한 모든 다중 레벨 H-JEPA에서 최소 하나의 상위 레벨 비용이 레벨 1 비용보다 계획 성공률을 높였다. 이득은 레벨 1 공간 계획이 잘 안 되는 곳에서 가장 컸고, 최고의 투영 레벨 1 플래너는 AntMaze의 모든 테스트 깊이에서 별도로 학습한 LeWM 베이스라인(18.0±3.5%)을 평균 성공률로 앞섰다. 즉 상위 플래너가 서브골을 만들어주지 않아도, 표현 계층만으로 더 나은 목표 거리 척도를 얻는다. FourRoom에서도 추상 목표 비용이 flat planning을 개선했지만, 선택적 추상화가 관찰되지 않은 Push-T와 Cube에서는 투영 비용이 뚜렷한 이득을 주지 못했다. 시간 분해 쪽 근거로는, 레벨 1 플래너가 레벨 2 서브골을 추적할 때의 비용이 전체 호라이즌 목표 비용보다 일관되게 감소한다는 점이 제시된다. 서브골 추적의 단조성(비용과 시간의 음의 Spearman 상관)은 0.89~1.00인 반면 전체 호라이즌 목표 비용은 0.48~0.75다.

전제와 한계

표현 분석에서는 상위 레벨이 빠르게 변하는 특징을 버리고 느린 특징을 남긴다는 증거가 나온다. AntMaze의 몸체 상태와 FourRoom의 방해물 위치는 계층이 깊어질수록 복구 가능성이 떨어지지만, 에이전트 위치는 두 계층 모두에서 끝까지 정확히 복구된다. 버려지는 디테일은 각 레벨의 예측 호라이즌과 일치한다. FourRoom 방해물은 텔레포트 사이에서는 국소적으로 예측 가능하지만 호라이즌이 길어지면 랜덤 텔레포트를 지날 확률이 커져 미래 위치를 맞추기 어려워지고, AntMaze의 개미 관절 구성은 전역 위치보다 훨씬 빠르게 진동한다. 저자들은 이를 JEPA의 인코더-예측기 공동 최적화로 설명한다. 예측 오차 그래디언트가 인코더를 통과하면서 각 레벨이 자기 시간 스케일에서 예측 가능한 특징은 남기고 그렇지 않은 것은 추상화해 버리도록 유도된다는 것이다. 다만 이런 선택적 추상화가 모든 환경에서 나타나지는 않는다. 데이터셋의 엔티티 주파수 격차, 즉 가장 빠르게 변하는 상태 성분과 가장 느리게 변하는 성분의 스펙트럼 중심 비율이 클수록 레벨 2에서 빠른 엔티티 정보 손실이 커지고 느린 엔티티 복구 가능성은 레벨 1 수준을 유지한다. 이 분리가 AntMaze, Humanoid, FourRoom에서는 나타나고 조작 데이터셋들은 격차가 작아 선택적 추상화가 거의 없다.

실제 로봇으로 확장한 부분도 있다. DROID는 장면, 조명, 조작 물체가 에피소드마다 바뀌는 실제 원격조작 조작 영상 데이터셋이다. 이런 비정상적 배경에서는 순수 예측 목적만으로도 느린 특징 해법이 성립한다. 인코더가 호라이즌 내내 예측 가능한 장면 정체성을 유지하고, 행동이 있어야만 예측 가능한 움직이는 에이전트를 버리는 것이다. 논문은 정규화 항이 임베딩의 주변 분포만 볼 때 이 해법이 항상 실현 가능함을 보이고, 이를 막으려면 시간에 걸친 표현의 결합 분포에 작용하는 손실이 필요하다고 말한다. 그래서 연속한 두 상태 표현에서 그 사이 행동을 회귀하는 역동역학(IDM) 손실을 레벨 목적함수에 추가한다. IDM 없이는 모델이 붕괴해 성능이 0이다. OGBench Cube는 장면 변화가 거의 없어 이 문제를 피해 가는데, 에피소드 간 잠재 분산이 DROID 75% 대 Cube 28%다. IDM을 넣으면 LeWM이 강한 베이스라인이 되고 그 위에 두 번째 레벨이 성능을 더한다. DROID에서도 H-JEPA가 HWM보다 뚜렷하게 앞서 표현 레버와 시간 레버가 모두 작동함을 보이며, flat 모델보다 낮은 계획 예산에서 더 높은 성능에 도달한다. 평가는 Fréchet fidelity로 한다. 계획된 엔드이펙터의 3차원 누적 경로가 전문가 경로를 얼마나 따르는지 재는 값으로, 0%는 팔이 전혀 움직이지 않은 경우, 100%는 전문가 경로, 음수는 아무것도 하지 않는 것보다 나쁘다는 뜻이다.

저자들이 밝힌 한계와 전제는 이렇다. Push-T처럼 에피소드가 짧은 환경에서는 계층 깊이를 늘릴 훈련 데이터가 부족해 3레벨 이상이 무너진다. 선택적 추상화는 데이터의 시간 스케일 분리가 뚜렷할 때만 나타나므로, 조작 데이터셋처럼 격차가 작으면 표현 계층의 이득이 사라진다. DROID 실험은 오프라인 개루프 계획 충실도까지만 검증했고, 실제 로봇에서의 폐루프 제어 전이는 다음 단계로 남겨 두었다. 향후 방향으로는 행동 없는 비디오에서 계층을 학습하는 것, 상위 잠재를 언어와 정렬해 목표를 단어로 지정하는 것, 고정 stride 대신 가변 길이 세그먼트로 상위 레벨을 훈련하는 것이 제안된다.

실무 관점에서 이 논문은 보상 없이 오프라인 영상만으로 장기 호라이즌 시각 계획을 세우려는 로봇 조작·내비게이션 파이프라인에 직접 쓸 수 있는 레시피다. 도입 전에 확인할 것은 세 가지다. 대상 데이터에 시간 스케일 분리가 있는지(없으면 표현 계층이 무의미하다), 에피소드 길이가 원하는 계층 깊이를 감당하는지, 그리고 장면이 에피소드마다 바뀌는 환경이라면 IDM 손실을 반드시 넣고 SIGReg 가중치와 함께 교차검증해야 한다는 점이다.