World Model Implanting for Test-time Adaptation of Embodied Agents (ICML 2025)

A concise framework for training and evaluating multimodal agents.

Overview

WorMI seamlessly incorporates multiple "world models" into a main model, enabling both standard training and meta-learning approaches to tackle more complex scenarios.

Steps

STEP 1: Writing world model training curricula file

Create a Python file (e.g., "world_curricula.py") that describes how to train your world models.

Example:

from pathlib import Path

from wormi.curricula import WorldModelCurricula, WorldModelCurriculum
from wormi.trainer import SchedulerType, WorMITrainerConfig

dataset_root = Path("/path/to/dataset")
rooms = ["bathrooms", "bedrooms", "livingrooms", "kitchens"]

batch_size = 4

aux_trainer_args = WorMITrainerConfig(
    max_steps=200,
    save_steps=0,
    eval_steps=500,
    logging_steps=10,
    batch_size=batch_size,
    lr_scheduler_type=SchedulerType.COSINE,
)

curricula = WorldModelCurricula(
    curricula=[
        WorldModelCurriculum(
            name=room,
            base_model="meta-llama/Llama-3.2-1B",
            tokenizer="meta-llama/Llama-3.2-1B-Instruct",
            dataset=dataset_root / room,
            trainer_args=aux_trainer_args,
            behavior_cloning=True,
        )
        for room in rooms
    ],
)

STEP 2: Training world models

Use: wormi world train --curricula_path /path/to/world_curricula.py
to train each world model on the data defined in your curricula.

STEP 3 (optional): Evaluating world models

Use: wormi world eval --model-name /path/to/trained_world_model
to evaluate them on your test data.

STEP 4: Writing WorMI training curricula file

Create a Python file (e.g., "wormi_curricula.py") specifying how to integrate the world models with your main model.

Example:

from pathlib import Path

from wormi.curricula import WorldModel, WorMICurricula, WorMICurriculum
from wormi.model import WorMIIntegrateMethod
from wormi.trainer import SchedulerType, WorMITrainerConfig

dataset_root = Path("/path/to/dataset")
rooms = ["bathrooms", "bedrooms", "livingrooms", "kitchens"]

trainer_args = WorMITrainerConfig(
    max_steps=50,
    eval_steps=50,
    logging_steps=1,
    save_steps=100,
    learning_rate=1e-5,
    lr_scheduler_type=SchedulerType.COSINE,
    batch_size=4,
)

curricula = WorMICurricula(
    method=WorMIIntegrateMethod.WORLD_WISE_ATTENTION,
    meta_learning=True,
    base_model="meta-llama/Llama-3.2-3B-Instruct",
    connections=[13, 27],
    world_models=[
        WorldModel(
            model_name=f"/path/to/world/models/{room}/last",
            connections=[7, 15],
        )
        for room in rooms
    ],
    datasets=[str(dataset_root / room) for room in rooms],
    self_attention=False,
    num_iterations=10,
    train_curricula=[
        WorMICurriculum(
            name=f"combin-1",
            target_world_models=[0, 1],
            datasets=[0, 1],
            trainer_args=trainer_args,
            num_eval_samples=100,
        ),
        WorMICurriculum(
            name=f"combin-2",
            target_world_models=[0, 2],
            datasets=[0, 2],
            trainer_args=trainer_args,
            num_eval_samples=100,
        ),
        WorMICurriculum(
            name=f"combin-3",
            target_world_models=[1, 2],
            datasets=[1, 2],
            trainer_args=trainer_args,
            num_eval_samples=100,
        ),
    ],
    test_curricula=[
        WorMICurriculum(
            name=f"{i}",
            target_world_models=[0, 1, 2],
            datasets=[0, 1, 2],
            trainer_args=trainer_args,
            num_eval_samples=300,
        )
        for i in range(1, 6)
    ],
    test_continuously=False,
)

STEP 5: Training WorMI

Use: wormi train --curricula_path /path/to/wormi_curricula.py
to train the main model, optionally in a meta-learning fashion, depending on the curricula file.

STEP 6: Evaluating WorMI

Use: wormi eval --curricula_path /path/to/wormi_curricula.py
to evaluate the final integrated model.

Scripts

• wormi --help
  Shows available subcommands.
• wormi train
  Trains a WorMI model with the specified curricula.
• wormi eval
  Evaluates a WorMI model with the specified curricula.
• wormi world train
  Trains a world model with the specified curricula.
• wormi world eval
  Evaluates a world model with the specified curricula.