le-wm

LeWorldModel (LeWM) es el código base oficial del artículo "LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels" de Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun y Randall Balestriero. Implementa una Arquitectura Predictiva de Incrustación Conjunta (JEPA) que, según el resumen, se entrena de manera estable de extremo a extremo desde píxeles crudos usando solo dos términos de pérdida: una pérdida de predicción de la siguiente incrustación y un regularizador que impone incrustaciones latentes con distribución gaussiana. Los autores afirman que esto reduce los hiperparámetros de pérdida ajustables de seis a uno en comparación con la única alternativa existente de extremo a extremo, y que el modelo tiene aproximadamente 15 millones de parámetros entrenables en una sola GPU en unas pocas horas. El repositorio se basa en dos proyectos externos: stable-worldmodel para la gestión del entorno, planificación y evaluación, y stable-pretraining para el entrenamiento. El objetivo declarado es reducir este repositorio a su contribución central, la arquitectura del modelo y el objetivo de entrenamiento, implementados en jepa.py como un módulo de PyTorch. La instalación utiliza uv con Python 3.10 e instala stable-worldmodel con las extensiones train y env. Los conjuntos de datos se distribuyen en formato HDF5 a través de Hugging Face y se descomprimen con tar --zstd; los archivos .h5 extraídos se colocan en $STABLEWM_HOME (por defecto ~/.stable-wm/), que se puede sobrescribir con una variable de entorno. Los nombres de los conjuntos de datos se referencian sin la extensión .h5 en los archivos de configuración de Hydra bajo config/train/. El entrenamiento se lanza con python train.py data=pusht después de configurar la entidad y el proyecto de WandB en config/train/lewm.yaml. Los checkpoints se guardan en $STABLEWM_HOME al finalizar. Las configuraciones de evaluación se encuentran en config/eval/, y el campo policy debe apuntar a una ruta de checkpoint relativa a $STABLEWM_HOME sin el sufijo _object.ckpt. Los checkpoints preentrenados de LeWM se reflejan en el Hugging Face Hub para los entornos pusht, cube, tworooms y reacher, junto con los conjuntos de datos en la misma colección. Una suite más amplia de checkpoints de referencia (PLDM, LeJEPA, IVL, IQL, GCBC, DINO-WM, DINO-WM-noprop) está disponible a través de Google Drive, con una tabla que indica qué métodos cubren qué entornos. Los checkpoints vienen en dos formas: un _object.ckpt que es un objeto de Python serializado usado por eval.py y la API stable_worldmodel, y un _weight.ckpt que es un diccionario de estado solo de pesos. La carga se realiza a través de swm.policy.AutoCostModel con un nombre de ejecución relativo a $STABLEWM_HOME; el módulo devuelto está en modo de evaluación y expone .state_dict(). Para los espejos de Hugging Face, el README proporciona un script de conversión que reconstruye el modelo JEPA desde config.json y weights.pt y guarda un checkpoint de objeto. El README también enlaza el artículo, un sitio web del proyecto y la colección de Hugging Face, e invita a reportar problemas y colaborar por correo electrónico. Informa que LeWM planifica hasta 48 veces más rápido que los modelos de mundo basados en modelos fundacionales, manteniéndose competitivo en diversas tareas de control 2D y 3D, que su espacio latente codifica estructura física como se muestra al sondear cantidades físicas, y que la evaluación de sorpresa detecta eventos físicamente implausibles. Estas son afirmaciones del resumen de los autores, no verificadas de forma independiente aquí.