프로젝트 소개
LeWorldModel(LeWM)은 Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun, Randall Balestriero의 논문 "LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels"의 공식 코드베이스입니다. 이는 JEPA(Joint-Embedding Predictive Architecture)를 구현하며, 초록에 따르면 원시 픽셀에서 종단 간 안정적으로 학습하며 다음 임베딩 예측 손실과 가우시안 분포 잠재 임베딩을 강제하는 정규화기의 두 가지 손실 항만 사용합니다. 저자들은 기존 종단 간 대안과 비교해 조정 가능한 손실 하이퍼파라미터를 6개에서 1개로 줄였으며, 모델은 약 1500만 개의 파라미터로 단일 GPU에서 몇 시간 내에 학습 가능하다고 밝힙니다.
저장소는 환경 관리, 계획 및 평가를 위한 stable-worldmodel과 학습을 위한 stable-pretraining이라는 두 외부 프로젝트를 기반으로 합니다. 목표는 이 저장소를 핵심 기여인 모델 아키텍처와 학습 목표로 축소하고, 이를 jepa.py에 PyTorch 모듈로 구현하는 것입니다.
설치는 Python 3.10과 함께 uv를 사용하며, train 및 env 엑스트라로 stable-worldmodel을 설치합니다. 데이터셋은 Hugging Face를 통해 HDF5 형식으로 배포되며 tar --zstd로 압축 해제됩니다. 추출된 .h5 파일은 $STABLEWM_HOME(기본값 ~/.stable-wm/) 아래에 배치되며 환경 변수로 재정의할 수 있습니다. 데이터셋 이름은 config/train/ 아래의 Hydra 구성 파일에서 .h5 확장자 없이 참조됩니다.
학습은 config/train/lewm.yaml에서 WandB 엔티티와 프로젝트를 설정한 후 python train.py data=pusht로 시작합니다. 체크포인트는 완료 시 $STABLEWM_HOME에 저장됩니다. 평가 구성은 config/eval/ 아래에 있으며, policy 필드는 $STABLEWM_HOME 기준 체크포인트 경로를 _object.ckpt 접미사 없이 가리켜야 합니다.
사전 학습된 LeWM 체크포인트는 pusht, cube, tworooms, reacher 환경에 대해 Hugging Face Hub에 미러링되어 있으며, 동일 컬렉션에 데이터셋도 포함됩니다. 더 넓은 기준 체크포인트 제품군(PLDM, LeJEPA, IVL, IQL, GCBC, DINO-WM, DINO-WM-noprop)은 Google Drive에서 제공되며, 어떤 방법이 어떤 환경을 지원하는지 표로 표시됩니다. 체크포인트는 eval.py와 stable_worldmodel API에서 사용하는 직렬화된 Python 객체인 _object.ckpt와 가중치 전용 상태 사전인 _weight.ckpt의 두 가지 형태로 제공됩니다. 로딩은 swm.policy.AutoCostModel을 통해 $STABLEWM_HOME 기준 실행 이름으로 수행되며, 반환된 모듈은 eval 모드이며 .state_dict()를 노출합니다. Hugging Face 미러의 경우 README는 config.json과 weights.pt에서 JEPA 모델을 재구성하고 객체 체크포인트를 저장하는 변환 스크립트를 제공합니다.
README는 또한 논문, 프로젝트 웹사이트, Hugging Face 컬렉션을 링크하고 이메일을 통한 이슈 및 협업을 초대합니다. LeWM이 기반 모델 기반 월드 모델보다 최대 48배 빠르게 계획하면서 다양한 2D 및 3D 제어 작업에서 경쟁력을 유지하고, 잠재 공간이 물리량 프로빙으로 물리적 구조를 인코딩하며, 놀라움 평가가 물리적으로 불가능한 이벤트를 감지한다고 보고합니다. 이는 저자 초록의 주장이며 여기서 독립적으로 검증되지 않았습니다.
Comments
0 Rating appears after 10 ratings
Sign in to join the discussion.