Unity-Technologies/ml-agents · error · UnityPolicyException

Registering Object of unsupported type

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.

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))
Defensive patterns

Strategy: type-guard

Validate before calling

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)}")

Type guard

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))

Try / catch

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

Prevention

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

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

Appendix: 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)