{"record":{"id":"9d84a0e8418faa59","repo":"Unity-Technologies/ml-agents","slug":"registering-object-of-unsupported-type-to-model","errorCode":null,"errorMessage":"Registering Object of unsupported type {} to ModelSaver ","messagePattern":"Registering Object of unsupported type (.+?) to ModelSaver ","errorType":"exception","errorClass":"UnityPolicyException","httpStatus":null,"severity":"error","filePath":"ml-agents/mlagents/trainers/model_saver/torch_model_saver.py","lineNumber":40,"sourceCode":"\n    def __init__(\n        self, trainer_settings: TrainerSettings, model_path: str, load: bool = False\n    ):\n        super().__init__()\n        self.model_path = model_path\n        self.initialize_path = trainer_settings.init_path\n        self._keep_checkpoints = trainer_settings.keep_checkpoints\n        self.load = load\n\n        self.policy: Optional[TorchPolicy] = None\n        self.exporter: Optional[ModelSerializer] = None\n        self.modules: Dict[str, torch.nn.Modules] = {}\n\n    def register(self, module: Union[TorchPolicy, TorchOptimizer]) -> None:\n        if isinstance(module, TorchPolicy) or isinstance(module, TorchOptimizer):\n            self.modules.update(module.get_modules())  # type: ignore\n        else:\n            raise UnityPolicyException(\n                \"Registering Object of unsupported type {} to ModelSaver \".format(\n                    type(module)\n                )\n            )\n        if self.policy is None and isinstance(module, TorchPolicy):\n            self.policy = module\n            self.exporter = ModelSerializer(self.policy)\n\n    def save_checkpoint(self, behavior_name: str, step: int) -> Tuple[str, List[str]]:\n        if not os.path.exists(self.model_path):\n            os.makedirs(self.model_path)\n        checkpoint_path = os.path.join(self.model_path, f\"{behavior_name}-{step}\")\n        state_dict = {\n            name: module.state_dict() for name, module in self.modules.items()\n        }\n        pytorch_ckpt_path = f\"{checkpoint_path}.pt\"\n        export_ckpt_path = f\"{checkpoint_path}.onnx\"\n        torch.save(state_dict, f\"{checkpoint_path}.pt\")","sourceCodeStart":22,"sourceCodeEnd":58,"githubUrl":"https://github.com/Unity-Technologies/ml-agents/blob/3ecb446f75d1e7400eb404c562dc005d3164cffc/ml-agents/mlagents/trainers/model_saver/torch_model_saver.py#L22-L58","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Ensure the object passed to register() subclasses TorchPolicy (for policies) or TorchOptimizer (for optimizers).","Implement get_modules() returning the dict of nn.Modules to save if you have a custom subclass.","Register each component separately: register the TorchPolicy first, then the TorchOptimizer, instead of wrapping them in a container object."],"exampleFix":"# before\nsaver.register(my_custom_policy_class(env_behavior_spec))  # doesn't inherit TorchPolicy\n# after\nclass MyPolicy(TorchPolicy): ...\nsaver.register(MyPolicy(env_behavior_spec))","handlingStrategy":"type-guard","validationCode":"from mlagents.trainers.policy.torch_policy import TorchPolicy\nfrom mlagents.trainers.optimizer.torch_optimizer import TorchOptimizer\nif not isinstance(obj, (TorchPolicy, TorchOptimizer)):\n    raise TypeError(f\"ModelSaver.register expects TorchPolicy/TorchOptimizer, got {type(obj)}\")","typeGuard":"from mlagents.trainers.policy.torch_policy import TorchPolicy\nfrom mlagents.trainers.optimizer.torch_optimizer import TorchOptimizer\ndef is_registrable(obj) -> bool:\n    return isinstance(obj, (TorchPolicy, TorchOptimizer))","tryCatchPattern":"from mlagents.trainers.exception import UnityPolicyException\ntry:\n    saver.register(component)\nexcept UnityPolicyException as e:\n    logger.error(f\"Skipping unregistrable component: {e}\")","preventionTips":["Subclass TorchPolicy/TorchOptimizer for any custom components","Register policy and optimizer objects directly, never wrapper containers","Add an isinstance assert in custom trainer setup code"],"tags":["ml-agents","type-error","pytorch","api-misuse"],"backgroundTag":"unsupported-type-registration","analyzedSha":"3ecb446f75d1e7400eb404c562dc005d3164cffc","analyzedAt":"2026-09-02T16:33:12.832Z","contentChangedAt":null,"schemaVersion":2},"datasetVersion":"2026-09-09T21:17:11.164Z"}