From ca133e269975292ec46fb7e7094341b6ae861a67 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Mon, 17 Aug 2020 11:12:20 -0700 Subject: [PATCH] [rllib] Remove extra model config kwargs passed incorrectly for Torch models (#10055) --- rllib/examples/custom_loss.py | 2 +- rllib/models/catalog.py | 46 ++++++++++++++++++++------- rllib/policy/dynamic_tf_policy.py | 3 +- rllib/policy/torch_policy_template.py | 3 +- 4 files changed, 38 insertions(+), 16 deletions(-) diff --git a/rllib/examples/custom_loss.py b/rllib/examples/custom_loss.py index b39f407fb..61e77b7dd 100644 --- a/rllib/examples/custom_loss.py +++ b/rllib/examples/custom_loss.py @@ -31,7 +31,7 @@ parser.add_argument( type=str, default=os.path.join( os.path.dirname(os.path.abspath(__file__)), - "../tests/data/cartpole/small")) + "../tests/data/cartpole/small.json")) if __name__ == "__main__": ray.init() diff --git a/rllib/models/catalog.py b/rllib/models/catalog.py index f7c91801c..75a377a68 100644 --- a/rllib/models/catalog.py +++ b/rllib/models/catalog.py @@ -82,8 +82,10 @@ MODEL_DEFAULTS: ModelConfigDict = { # === Options for custom models === # Name of a custom model to use "custom_model": None, - # Extra options to pass to the custom classes. - # These will be available in the Model's + # Extra options to pass to the custom classes. These will be available to + # the Model's constructor in the model_config field. Also, they will be + # attempted to be passed as **kwargs to ModelV2 models. For an example, + # see rllib/models/[tf|torch]/attention_net.py. "custom_model_config": {}, # Name of a custom action distribution to use. "custom_action_dist": None, @@ -302,6 +304,11 @@ class ModelCatalog: model_config["custom_model_config"] = \ model_config.pop("custom_options") + # Allow model kwargs to be overriden / augmented by + # custom_model_config. + customized_model_kwargs = dict( + model_kwargs, **model_config.get("custom_model_config", {})) + if isinstance(model_config["custom_model"], type): model_cls = model_config["custom_model"] else: @@ -329,19 +336,19 @@ class ModelCatalog: # accept these as kwargs, not get them from # config["custom_model_config"] anymore). try: - instance = model_cls(obs_space, action_space, - num_outputs, model_config, - name, **model_kwargs) + instance = model_cls( + obs_space, action_space, num_outputs, + model_config, name, **customized_model_kwargs) except TypeError as e: # Keyword error: Try old way w/o kwargs. if "__init__() got an unexpected " in e.args[0]: + instance = model_cls(obs_space, action_space, + num_outputs, model_config, + name, **model_kwargs) logger.warning( "Custom ModelV2 should accept all custom " "options as **kwargs, instead of expecting" " them in config['custom_model_config']!") - instance = model_cls(obs_space, action_space, - num_outputs, model_config, - name) # Other error -> re-raise. else: raise e @@ -361,9 +368,26 @@ class ModelCatalog: else: # PyTorch automatically tracks nn.Modules inside the parent # nn.Module's constructor. - # TODO(sven): Do this for TF as well. - instance = model_cls(obs_space, action_space, num_outputs, - model_config, name, **model_kwargs) + # Try calling with kwargs first (custom ModelV2 should + # accept these as kwargs, not get them from + # config["custom_model_config"] anymore). + try: + instance = model_cls(obs_space, action_space, + num_outputs, model_config, name, + **customized_model_kwargs) + except TypeError as e: + # Keyword error: Try old way w/o kwargs. + if "__init__() got an unexpected " in e.args[0]: + instance = model_cls(obs_space, action_space, + num_outputs, model_config, + name, **model_kwargs) + logger.warning( + "Custom ModelV2 should accept all custom " + "options as **kwargs, instead of expecting" + " them in config['custom_model_config']!") + # Other error -> re-raise. + else: + raise e return instance # TODO(sven): Hard-deprecate Model(V1). This check will be # superflous then. diff --git a/rllib/policy/dynamic_tf_policy.py b/rllib/policy/dynamic_tf_policy.py index ec1872a92..f4dc78192 100644 --- a/rllib/policy/dynamic_tf_policy.py +++ b/rllib/policy/dynamic_tf_policy.py @@ -200,8 +200,7 @@ class DynamicTFPolicy(TFPolicy): action_space=action_space, num_outputs=logit_dim, model_config=self.config["model"], - framework="tf", - **self.config["model"].get("custom_model_config", {})) + framework="tf") # Create the Exploration object to use for this Policy. self.exploration = self._create_exploration() diff --git a/rllib/policy/torch_policy_template.py b/rllib/policy/torch_policy_template.py index f7ce8ed2f..a7a0c2001 100644 --- a/rllib/policy/torch_policy_template.py +++ b/rllib/policy/torch_policy_template.py @@ -201,8 +201,7 @@ def build_torch_policy( action_space=action_space, num_outputs=logit_dim, model_config=self.config["model"], - framework="torch", - **self.config["model"].get("custom_model_config", {})) + framework="torch") # Make sure, we passed in a correct Model factory. assert isinstance(self.model, TorchModelV2), \