DOCSEARCHbyAlgolia
Submit Docs

LeWorldModel (LeWM)

repository·main·Indexed 23 days ago

https://github.com/lucas-maes/le-wm

A stable, end-to-end Joint-Embedding Predictive Architecture (JEPA) that learns world models from raw pixels. Designed for efficiency and fast planning in 2D and 3D control tasks, LeWM utilizes stable-worldmodel for environment management and planning, and stable-pretraining for training. It supports HDF5 dataset formats and provides the AutoCostModel API for loading checkpoints for MPC and planning.

Tokens
788
Snippets
6
Records
6
Agent score
38%

What's inside LeWorldModel

  1. Train LeWorldModel

    main

    Training is implemented in jepa.py and configured via Hydra.

    1. Configure WandB: Before training, set your entity and project in config/train/lewm.yaml:
    wandb:
  config:
    entity: your_entity
    project: your_project
    1. Launch Training: Run the training script specifying the data:
    python train.py data=pusht

    Checkpoints are automatically saved to $STABLEWM_HOME upon completion.

  2. Configure and manage LeWorldModel data

    main

    Datasets are stored in HDF5 format.

    1. Download data from HuggingFace.
    2. Decompress archives using tar --zstd -xvf archive.tar.zst.
    3. Place .h5 files in the directory specified by $STABLEWM_HOME.

    By default, $STABLEWM_HOME is ~/.stable-wm/. You can override this by setting the environment variable:

    export STABLEWM_HOME=/path/to/your/storage

    Note: Dataset names in configuration files should be specified without the .h5 extension (e.g., pusht_expert_train resolves to $STABLEWM_HOME/pusht_expert_train.h5).

  3. Install LeWorldModel

    main

    Install the LeWorldModel environment using uv. This project relies on stable-worldmodel for environment management, planning, and evaluation, and stable-pretraining for training.

    uv venv --python=3.10
source .venv/bin/activate
uv pip install stable-worldmodel[train,env]

  4. Evaluate LeWorldModel via Planning

    main

    Evaluation configurations are located in config/eval/. When using the policy flag, provide the checkpoint path relative to $STABLEWM_HOME, and ensure you omit the _object.ckpt suffix.

    Correct usage:

    python eval.py --config-name=pusht.yaml policy=pusht/lewm

    Incorrect usage (will fail):

    python eval.py --config-name=pusht.yaml policy=pusht/lewm_object.ckpt

  5. Convert Hugging Face checkpoints to object checkpoints

    main

    Hugging Face model repositories provide weights.pt (state dict) and config.json. To use these with eval.py or AutoCostModel, you must convert them into an _object.ckpt file.

    1. Download the files:
    hf download quentinll/lewm-pusht --local-dir $STABLEWM_HOME/hf_pusht
    1. Run the conversion script (provided below) to generate the _object.ckpt in the expected directory structure under $STABLEWM_HOME.

  6. Load a checkpoint using AutoCostModel

    main

    To load a checkpoint (specifically for MPC/planning), use the stable_worldmodel.policy.AutoCostModel API. This works with checkpoints extracted from the Drive archive.

    AutoCostModel accepts:

    • run_name: The checkpoint path relative to $STABLEWM_HOME, without the _object.ckpt suffix.
    • cache_dir: (Optional) An override for the checkpoint root. Defaults to $STABLEWM_HOME.
    import stable_worldmodel as swm

# Load the cost model (for MPC)
cost = swm.policy.AutoCostModel('pusht/lewm')