Unity-Technologies/ml-agents · error · UnityPolicyException

Registering Object of unsupported type {} to ModelSaver

Error message

Registering Object of unsupported type {} to ModelSaver 

What it means

UnityPolicyException thrown by TorchModelSaver.register when the object passed is neither a TorchPolicy nor a TorchOptimizer. The model saver only knows how to extract save/restore modules via get_modules() on those two types, so any other object is rejected before training can checkpoint.

Source

Thrown at ml-agents/mlagents/trainers/model_saver/torch_model_saver.py:40

    def __init__(
        self, trainer_settings: TrainerSettings, model_path: str, load: bool = False
    ):
        super().__init__()
        self.model_path = model_path
        self.initialize_path = trainer_settings.init_path
        self._keep_checkpoints = trainer_settings.keep_checkpoints
        self.load = load

        self.policy: Optional[TorchPolicy] = None
        self.exporter: Optional[ModelSerializer] = None
        self.modules: Dict[str, torch.nn.Modules] = {}

    def register(self, module: Union[TorchPolicy, TorchOptimizer]) -> None:
        if isinstance(module, TorchPolicy) or isinstance(module, TorchOptimizer):
            self.modules.update(module.get_modules())  # type: ignore
        else:
            raise UnityPolicyException(
                "Registering Object of unsupported type {} to ModelSaver ".format(
                    type(module)
                )
            )
        if self.policy is None and isinstance(module, TorchPolicy):
            self.policy = module
            self.exporter = ModelSerializer(self.policy)

    def save_checkpoint(self, behavior_name: str, step: int) -> Tuple[str, List[str]]:
        if not os.path.exists(self.model_path):
            os.makedirs(self.model_path)
        checkpoint_path = os.path.join(self.model_path, f"{behavior_name}-{step}")
        state_dict = {
            name: module.state_dict() for name, module in self.modules.items()
        }
        pytorch_ckpt_path = f"{checkpoint_path}.pt"
        export_ckpt_path = f"{checkpoint_path}.onnx"
        torch.save(state_dict, f"{checkpoint_path}.pt")

View on GitHub (pinned to 3ecb446f75)

Solutions

1. Ensure the object passed to register() subclasses TorchPolicy (for policies) or TorchOptimizer (for optimizers).
2. Implement get_modules() returning the dict of nn.Modules to save if you have a custom subclass.
3. Register each component separately: register the TorchPolicy first, then the TorchOptimizer, instead of wrapping them in a container object.

Example fix

# before
saver.register(my_custom_policy_class(env_behavior_spec))  # doesn't inherit TorchPolicy
# after
class MyPolicy(TorchPolicy): ...
saver.register(MyPolicy(env_behavior_spec))

from mlagents.trainers.policy.torch_policy import TorchPolicy
from mlagents.trainers.optimizer.torch_optimizer import TorchOptimizer
if not isinstance(obj, (TorchPolicy, TorchOptimizer)):
    raise TypeError(f"ModelSaver.register expects TorchPolicy/TorchOptimizer, got {type(obj)}")

from mlagents.trainers.policy.torch_policy import TorchPolicy
from mlagents.trainers.optimizer.torch_optimizer import TorchOptimizer
def is_registrable(obj) -> bool:
    return isinstance(obj, (TorchPolicy, TorchOptimizer))

from mlagents.trainers.exception import UnityPolicyException
try:
    saver.register(component)
except UnityPolicyException as e:
    logger.error(f"Skipping unregistrable component: {e}")

When it happens

Trigger: Calling torch_model_saver.register(obj) with any object that is not a TorchPolicy or TorchOptimizer instance — e.g. passing a bare nn.Module, a custom Policy subclass that doesn't inherit TorchPolicy, or an optimizer from a different framework.

Common situations: Writing a custom trainer that wires its own policy/optimizer into the saver; migrating from TF to PyTorch trainers and passing the legacy policy class; refactoring where the custom policy forgot to subclass TorchPolicy.

Related errors

• shape and dimensionProperties must have the same length.
• Unsupported config {d} for {t.__name__}.
• Index out of bounds, expected a number between 0 and {Length
• Enumerator not started.
• Enumerator has reached the end already.

AI-assisted analysis of Unity-Technologies/ml-agents@3ecb446f75 (2026-09-02). Data as JSON: /api/errors/9d84a0e8418faa59. Report an issue: GitHub.