Plugin Registry #

TorchWM’s plugin registry lets you register custom world models and environment backends so they are discoverable through the standard factory API alongside built-in models.

Why use the registry?#

Discovery: create_model("your-model") and list_models() work for registered models — no code change needed in experiment scripts.
Interop: Registered models appear in get_model_spec() and support the same config override flow as Dreamer, JEPA, IRIS, etc.
Aliases: Attach short or alternative names to your model.

Registering a world model #

Use register_world_model as a decorator on your model class:

from torchwm import register_world_model

@register_world_model(
    "my-agent",
    import_path="my_package.models:MyAgent",
    config_path="my_package.configs:MyConfig",
    description="My custom world model agent",
    aliases=("my_agent", "custom"),
)
class MyAgent:
    def __init__(self, config):
        self.config = config

After registration:

import torchwm

# Standard factory API works:
cfg = torchwm.create_config("my-agent", learning_rate=1e-4)
agent = torchwm.create_model("my-agent", cfg)

# Discovery:
print(torchwm.list_models())       # includes "my-agent"
spec = torchwm.get_model_spec("my-agent")
print(spec.description)            # "My custom world model agent"

Without a config class #

If your model does not need a config object, omit config_path:

@register_world_model("simple-model", import_path="my_package.models:create_simple")
class SimpleModel:
    ...

create_config("simple-model") returns None; keyword overrides are forwarded directly to the factory:

model = torchwm.create_model("simple-model", hidden_size=128)

Without a class (direct registration)#

Passing a Python import-path string skips the decorator pattern entirely:

from torchwm import register_world_model

register_world_model(
    "my-agent",
    import_path="my_package.models:create_my_agent",
    config_path="my_package.configs:MyAgentConfig",
)

This is useful when the model is built by a factory function instead of a class constructor.

Aliases #

Aliases let users refer to your model by multiple names:

@register_world_model(
    "my-agent",
    import_path="...",
    aliases=("my_agent", "custom", "ma"),
)
class MyAgent:
    ...

All of the following resolve to the same model:

torchwm.create_model("my-agent")
torchwm.create_model("my_agent")
torchwm.create_model("custom")
torchwm.create_model("ma")

Overriding an existing model #

Pass override=True to replace a previously registered model (including built-in models — use with caution):

@register_world_model(
    "dreamer",
    import_path="my_package.models:MyDreamer",
    override=True,
)
class MyDreamer:
    ...

This shadows the built-in dreamer spec. All existing scripts calling create_model("dreamer") will now use your replacement.

Environment backends #

The same registry pattern works for environment backends:

from torchwm import register_env_backend

register_env_backend(
    "custom-env",
    factory_path="my_package.envs:make_custom_env",
    description="My custom environment backend",
    aliases=("ce",),
)

Registered backends appear in list_env_backends() and work with make_env(backend="custom-env").

Discovering registered models #

import torchwm

# All models (built-in + registered):
all_models = torchwm.list_models()

# Only externally registered models:
ext_models = torchwm.list_registered_models()

# Spec lookup:
spec = torchwm.get_model_spec("my-agent")
print(spec.name, spec.import_path, spec.config_path, spec.description)

# Remove a registration:
torchwm.deregister_world_model("my-agent")

Deprecating a model #

Use the deprecated_class decorator to mark an older model as deprecated and point users to its replacement:

from torchwm import deprecated_class

@deprecated_class(version="0.6.0", alternative="MyNewAgent")
@register_world_model("old-agent", import_path="my_package.models:OldAgent")
class OldAgent:
    ...

Instantiating OldAgent (or using create_model("old-agent")) emits a DeprecationWarning:

DeprecationWarning: 'OldAgent' is deprecated since v0.6.0 — Use MyNewAgent instead

Lower-level deprecated_function and generic deprecated decorators are also available from torchwm:

from torchwm import deprecated, deprecated_function

Complete example #

"""my_package/models.py"""

import torchwm
from torchwm import register_world_model, deprecated_class


class MyConfig:
    def __init__(self):
        self.learning_rate = 1e-4
        self.hidden_size = 256


@register_world_model(
    "research-model",
    import_path="my_package.models:ResearchModel",
    config_path="my_package.models:MyConfig",
    description="Research world model with config overrides.",
    aliases=("research", "rm"),
)
class ResearchModel:
    def __init__(self, config):
        self.config = config

    def forward(self, obs):
        ...

Usage in any experiment script:

import torchwm

model = torchwm.create_model(
    "research-model",
    learning_rate=3e-4,    # overrides the default in MyConfig
)