malib.registration module

class malib.registration.Registry

Bases: object

Global registry of algorithms, models, preprocessors and environments

Examples:

>>> # register custom model
>>> Registry.register_custom_model("MyCustomModel", model_class)
>>> # register custom policy
>>> Registry.register_custom_policy("MyCustomPolicy", policy_class)
>>> # register custom environment
>>> Registry.register_custom_env("MyCustomEnvironment", environment_class)
>>> # register custom algorithm
>>> Registry.register_custom_algorithm(
...     name="MyCustomAlgo",
...     policy="registered_policy_name_or_cls",
...     trainer="registered_trainer_name_or_cls",
...     loss="registered_loss_name_or_cls")
>>>

static register_custom_algorithm(name: str, policy: Union[type, str], trainer: Union[type, str], loss: Optional[Union[type, str]] = None) → None

Register a custom algorithm by name.

Parameters

• name – str, Name to register the algorithm under.

• policy – Union[type, str], Python class or registered name of policy.

• trainer – Union[type, str], Python class or registered name of trainer.

• loss – Union[type, str], Python class or registered name of loss function.

Returns

static register_custom_env(name: str, env_class: type) → None

Register a custom environment by name.

Parameters

• name – str, Name to register the environment under.

• env_class – type, Python class of the environment.

static register_custom_model(name: str, model_class: type) → None

Register a custom model by name.

Parameters

• name – str, Name to register the model under.

• model_class – type, Python class of the model.

Returns

static register_custom_policy(name: str, policy_class: type) → None

Register a custom policy by name.

Parameters

• name – str, Name to register the policy under.

• policy_class – type, Python class of the policy.