def get_model(registry, inputs, num_outputs, options={}):
"""Returns a suitable model conforming to given input and output specs.
Args:
registry (obj): Registry of named objects (ray.tune.registry).
inputs (Tensor): The input tensor to the model.
num_outputs (int): The size of the output vector of the model.
options (dict): Optional args to pass to the model constructor.
Returns:
model (Model): Neural network model.
"""
if "custom_model" in options:
model = options["custom_model"]
print("Using custom model {}".format(model))
return registry.get(RLLIB_MODEL, model)(inputs, num_outputs,
options)
obs_rank = len(inputs.shape) - 1
# num_outputs > 1 used to avoid hitting this with the value function
if isinstance(
options.get("custom_options", {}).get(
"multiagent_fcnet_hiddens", 1), list) and num_outputs > 1:
return MultiAgentFullyConnectedNetwork(inputs, num_outputs,
options)
if obs_rank > 1:
return VisionNetwork(inputs, num_outputs, options)
return FullyConnectedNetwork(inputs, num_outputs, options)
def _get_model(inputs, num_outputs, options):
if "custom_model" in options:
model = options["custom_model"]
print("Using custom model {}".format(model))
return _global_registry.get(RLLIB_MODEL,
model)(inputs, num_outputs, options)
obs_rank = len(inputs.shape) - 1
# num_outputs > 1 used to avoid hitting this with the value function
if isinstance(
options.get("custom_options", {}).get(
"multiagent_fcnet_hiddens", 1), list) and num_outputs > 1:
return MultiAgentFullyConnectedNetwork(inputs, num_outputs,
options)
if obs_rank > 1:
return VisionNetwork(inputs, num_outputs, options)
return FullyConnectedNetwork(inputs, num_outputs, options)