def __init__(self, obs_space, action_space, num_outputs, model_config, name): super(CustomModel, self).__init__(obs_space, action_space, num_outputs, model_config, name) self.model = FullyConnectedNetwork(obs_space, action_space, num_outputs, model_config, name) self.register_variables(self.model.variables())
def __init__(self, obs_space, action_space, num_outputs, model_config, name, true_obs_shape=(4, ), action_embed_size=6, **kw): super(ParametricActionsModel, self).__init__( obs_space, action_space, num_outputs, model_config, name, **kw) if model_config['custom_options']['spy']: true_obs_space = make_spy_space(model_config['custom_options']['parties'], model_config['custom_options']['blocks']) else: true_obs_space = make_blind_space(model_config['custom_options']['parties'], model_config['custom_options']['blocks']) if model_config['custom_options']['extended']: action_embed_size = 6 else: action_embed_size = 4 total_dim = 0 for space in true_obs_space: total_dim += get_preprocessor(space)(space).size self.action_embed_model = FullyConnectedNetwork( Box(-1, 1, shape = (total_dim,)), action_space, action_embed_size, model_config, name + "_action_embed") self.register_variables(self.action_embed_model.variables())
class MaskedActionsMLP(DistributionalQModel, TFModelV2): """Tensorflow model for Envs that provide action masks with observations.""" def __init__(self, obs_space, action_space, num_outputs, model_config, name, **kwargs): super().__init__(obs_space, action_space, num_outputs, model_config, name, **kwargs) # DictFlatteningPreprocessor, combines all obs components together # obs.shape for MLP should be a flattened game board obs original_space = obs_space.original_space['board'] flat_obs_space = spaces.Box(low=np.min(original_space.low), high=np.max(original_space.high), shape=(np.prod(original_space.shape), )) self.mlp = FullyConnectedNetwork(flat_obs_space, action_space, num_outputs, model_config, name) self.register_variables(self.mlp.variables()) def forward(self, input_dict, state, seq_lens): obs = flatten(input_dict['obs']['board']) action_mask = tf.maximum(tf.log(input_dict['obs']['action_mask']), tf.float32.min) model_out, _ = self.mlp({'obs': obs}) return action_mask + model_out, state def value_function(self): return self.mlp.value_function()
class VMActionMaskModel(TFModelV2): def __init__(self, obs_space, action_space, num_outputs, model_config, name, true_obs_shape=(51, 3), action_embed_size=50, *args, **kwargs): super(VMActionMaskModel, self).__init__(obs_space, action_space, num_outputs, model_config, name, *args, **kwargs) self.action_embed_model = FullyConnectedNetwork( spaces.Box(0, 1, shape=true_obs_shape), action_space, action_embed_size, model_config, name + "_action_embedding") self.register_variables(self.action_embed_model.variables()) def forward(self, input_dict, state, seq_lens): avail_actions = input_dict["obs"]["avail_actions"] action_mask = input_dict["obs"]["action_mask"] action_embedding, _ = self.action_embed_model( {"obs": input_dict["obs"]["state"]}) intent_vector = tf.expand_dims(action_embedding, 1) action_logits = tf.reduce_sum(avail_actions * intent_vector, axis=1) inf_mask = tf.maximum(tf.log(action_mask), tf.float32.min) return action_logits + inf_mask, state def value_function(self): return self.action_embed_model.value_function()
class CustomTFRPGModel(TFModelV2): """Example of interpreting repeated observations.""" def __init__(self, obs_space, action_space, num_outputs, model_config, name): super().__init__(obs_space, action_space, num_outputs, model_config, name) self.model = TFFCNet(obs_space, action_space, num_outputs, model_config, name) self.register_variables(self.model.variables()) def forward(self, input_dict, state, seq_lens): # The unpacked input tensors, where M=MAX_PLAYERS, N=MAX_ITEMS: # { # 'items', <tf.Tensor shape=(?, M, N, 5)>, # 'location', <tf.Tensor shape=(?, M, 2)>, # 'status', <tf.Tensor shape=(?, M, 10)>, # } print("The unpacked input tensors:", input_dict["obs"]) print() print("Unbatched repeat dim", input_dict["obs"].unbatch_repeat_dim()) print() if tf.executing_eagerly(): print("Fully unbatched", input_dict["obs"].unbatch_all()) print() return self.model.forward(input_dict, state, seq_lens) def value_function(self): return self.model.value_function()
class CentralizedCriticModel(TFModelV2): """Multi-agent model that implements a centralized VF.""" def __init__(self, obs_space, action_space, num_outputs, model_config, name): super(CentralizedCriticModel, self).__init__( obs_space, action_space, num_outputs, model_config, name) # Base of the model self.model = FullyConnectedNetwork(obs_space, action_space, num_outputs, model_config, name) self.register_variables(self.model.variables()) # Central VF maps (obs, opp_ops, opp_act) -> vf_pred obs = tf.keras.layers.Input(shape=(6, ), name="obs") opp_obs = tf.keras.layers.Input(shape=(6, ), name="opp_obs") opp_act = tf.keras.layers.Input(shape=(2, ), name="opp_act") concat_obs = tf.keras.layers.Concatenate(axis=1)( [obs, opp_obs, opp_act]) central_vf_dense = tf.keras.layers.Dense( 16, activation=tf.nn.tanh, name="c_vf_dense")(concat_obs) central_vf_out = tf.keras.layers.Dense( 1, activation=None, name="c_vf_out")(central_vf_dense) self.central_vf = tf.keras.Model( inputs=[obs, opp_obs, opp_act], outputs=central_vf_out) self.register_variables(self.central_vf.variables) def forward(self, input_dict, state, seq_lens): return self.model.forward(input_dict, state, seq_lens) def central_value_function(self, obs, opponent_obs, opponent_actions): return tf.reshape( self.central_vf( [obs, opponent_obs, tf.one_hot(opponent_actions, 2)]), [-1])
def __init__(self, obs_space, action_space, num_outputs, model_config, name, true_obs_shape=(24, ), action_embed_size=None): super(ParametricActionsModel, self).__init__(obs_space, action_space, num_outputs, model_config, name) if action_embed_size is None: action_embed_size = action_space.n # this works for Discrete() action # we get the size of the output of the preprocessor automatically chosen by rllib for the real_obs space real_obs = obs_space.original_space['real_obs'] true_obs_shape = get_preprocessor(real_obs)( real_obs).size # this will we an integer # true_obs_shape = obs_space.original_space['real_obs'] self.action_embed_model = FullyConnectedNetwork( obs_space=Box(-1, 1, shape=(true_obs_shape, )), action_space=action_space, num_outputs=action_embed_size, model_config=model_config, name=name + "_action_embed") self.base_model = self.action_embed_model.base_model self.register_variables(self.action_embed_model.variables())
def __init__(self, obs_space, action_space, num_outputs, model_config, name): super(CentralizedCriticModel, self).__init__(obs_space, action_space, num_outputs, model_config, name) # Base of the model self.model = FullyConnectedNetwork(obs_space, action_space, num_outputs, model_config, name) self.register_variables(self.model.variables()) # Central VF maps (obs, opp_ops, opp_act) -> vf_pred self.max_num_agents = model_config['custom_options']['max_num_agents'] self.obs_space_shape = obs_space.shape[0] other_obs = tf.keras.layers.Input(shape=(obs_space.shape[0] * self.max_num_agents, ), name="opp_obs") central_vf_dense = tf.keras.layers.Dense( model_config['custom_options']['central_vf_size'], activation=tf.nn.tanh, name="c_vf_dense")(other_obs) central_vf_out = tf.keras.layers.Dense( 1, activation=None, name="c_vf_out")(central_vf_dense) self.central_vf = tf.keras.Model(inputs=[other_obs], outputs=central_vf_out) self.register_variables(self.central_vf.variables)
def __init__(self, obs_space, action_space, num_outputs, model_config, name, **kwargs): super(HanabiFullyConnected, self).__init__(obs_space, action_space, num_outputs, model_config, name, **kwargs) self.fc = FullyConnectedNetwork(obs_space.original_space["board"], action_space, num_outputs, model_config, name + "fc") self.register_variables(self.fc.variables())
class ParametricActionsModel(DistributionalQTFModel): """Parametric action model that handles the dot product and masking. This assumes the outputs are logits for a single Categorical action dist. Getting this to work with a more complex output (e.g., if the action space is a tuple of several distributions) is also possible but left as an exercise to the reader. """ def __init__(self, obs_space, action_space, num_outputs, model_config, name, true_obs_shape=(4, ), action_embed_size=6, **kw): super(ParametricActionsModel, self).__init__( obs_space, action_space, num_outputs, model_config, name, **kw) if model_config['custom_options']['spy']: true_obs_space = make_spy_space(model_config['custom_options']['parties'], model_config['custom_options']['blocks']) else: true_obs_space = make_blind_space(model_config['custom_options']['parties'], model_config['custom_options']['blocks']) if model_config['custom_options']['extended']: action_embed_size = 6 else: action_embed_size = 4 total_dim = 0 for space in true_obs_space: total_dim += get_preprocessor(space)(space).size self.action_embed_model = FullyConnectedNetwork( Box(-1, 1, shape = (total_dim,)), action_space, action_embed_size, model_config, name + "_action_embed") self.register_variables(self.action_embed_model.variables()) def forward(self, input_dict, state, seq_lens): # Extract the available actions tensor from the observation. avail_actions = input_dict["obs"]["avail_actions"] action_mask = input_dict["obs"]["action_mask"] # Compute the predicted action embedding action_embed, _ = self.action_embed_model({ "obs": input_dict["obs"]["bitcoin"] }) # Expand the model output to [BATCH, 1, EMBED_SIZE]. Note that the # avail actions tensor is of shape [BATCH, MAX_ACTIONS, EMBED_SIZE]. intent_vector = tf.expand_dims(action_embed, 1) # Batch dot product => shape of logits is [BATCH, MAX_ACTIONS]. action_logits = tf.reduce_sum(avail_actions * intent_vector, axis=2) # Mask out invalid actions (use tf.float32.min for stability) inf_mask = tf.maximum(tf.log(action_mask), tf.float32.min) return action_logits + inf_mask, state def value_function(self): return self.action_embed_model.value_function()
def __init__(self, obs_space, action_space, num_outputs, model_config, name): super().__init__(obs_space, action_space, num_outputs, model_config, name) self.fcnet = FullyConnectedNetwork(self.obs_space, self.action_space, num_outputs, model_config, name="fcnet") self.register_variables(self.fcnet.variables())
def __init__(self, obs_space, action_space, num_outputs, model_config, name, **kwargs): super().__init__(obs_space, action_space, num_outputs, model_config, name, **kwargs) # DictFlatteningPreprocessor, combines all obs components together # obs.shape for MLP should be a flattened game board obs original_space = obs_space.original_space['board'] flat_obs_space = spaces.Box(low=np.min(original_space.low), high=np.max(original_space.high), shape=(np.prod(original_space.shape), )) self.mlp = FullyConnectedNetwork(flat_obs_space, action_space, num_outputs, model_config, name) self.register_variables(self.mlp.variables())
class CustomModel(TFModelV2): """Example of a custom model that just delegates to a fc-net.""" def __init__(self, obs_space, action_space, num_outputs, model_config, name): super(CustomModel, self).__init__(obs_space, action_space, num_outputs, model_config, name) self.model = FullyConnectedNetwork(obs_space, action_space, num_outputs, model_config, name) self.register_variables(self.model.variables()) def forward(self, input_dict, state, seq_lens): return self.model.forward(input_dict, state, seq_lens) def value_function(self): return self.model.value_function()
class FCModel(TFModelV2): '''Fully Connected Model''' def __init__(self, obs_space, action_space, num_outputs, model_config, name): super(FCModel, self).__init__(obs_space, action_space, num_outputs, model_config, name) self.model = FullyConnectedNetwork(obs_space, action_space, num_outputs, model_config, name) self.register_variables(self.model.variables()) def forward(self, input_dict, state, seq_lens): return self.model.forward(input_dict, state, seq_lens) def value_function(self): return self.model.value_function()
def __init__(self, obs_space, action_space, num_outputs, model_config, name, true_obs_shape=(4, ), action_embed_size=2, **kw): super(ParametricActionsModel, self).__init__( obs_space, action_space, num_outputs, model_config, name, **kw) self.action_embed_model = FullyConnectedNetwork( Box(-1, 1, shape=true_obs_shape), action_space, action_embed_size, model_config, name + "_action_embed") self.register_variables(self.action_embed_model.variables())
def __init__(self, obs_space, action_space, num_outputs, model_config, name): super(CentralizedCriticModel, self).__init__( obs_space, action_space, num_outputs, model_config, name) self.action_model = FullyConnectedNetwork( Box(low=0, high=1, shape=(6, )), # one-hot encoded Discrete(6) action_space, num_outputs, model_config, name + "_action") self.register_variables(self.action_model.variables()) self.value_model = FullyConnectedNetwork(obs_space, action_space, 1, model_config, name + "_vf") self.register_variables(self.value_model.variables())
def _build_layers_v2(self, input_dict, num_outputs, options): self.obs_in = input_dict["obs"] with tf.variable_scope("shared", reuse=tf.AUTO_REUSE): self.fcnet = FullyConnectedNetwork(input_dict, self.obs_space, self.action_space, num_outputs, options) return self.fcnet.outputs, self.fcnet.last_layer
def __init__(self, obs_space, action_space, num_outputs, model_config, name): super().__init__(obs_space, action_space, num_outputs, model_config, name) self.model = TFFCNet(obs_space, action_space, num_outputs, model_config, name) self.register_variables(self.model.variables())
class ParametricActionsModel(TFModelV2): """ Parametric model that handles varying action spaces""" def __init__(self, obs_space, action_space, num_outputs, model_config, name, true_obs_shape=(24, ), action_embed_size=None): super(ParametricActionsModel, self).__init__(obs_space, action_space, num_outputs, model_config, name) if action_embed_size is None: action_embed_size = action_space.n # this works for Discrete() action # we get the size of the output of the preprocessor automatically chosen by rllib for the real_obs space real_obs = obs_space.original_space['real_obs'] true_obs_shape = get_preprocessor(real_obs)( real_obs).size # this will we an integer # true_obs_shape = obs_space.original_space['real_obs'] self.action_embed_model = FullyConnectedNetwork( obs_space=Box(-1, 1, shape=(true_obs_shape, )), action_space=action_space, num_outputs=action_embed_size, model_config=model_config, name=name + "_action_embed") self.base_model = self.action_embed_model.base_model self.register_variables(self.action_embed_model.variables()) def forward(self, input_dict, state, seq_lens): # Compute the predicted action probabilties # input_dict["obs"]["real_obs"] is a list of 1d tensors if the observation space is a Tuple while # it should be a tensor. When it is a list we concatenate the various 1d tensors obs_concat = input_dict["obs"]["real_obs"] if isinstance(obs_concat, list): obs_concat = tf.concat(values=flatten_list(obs_concat), axis=1) action_embed, _ = self.action_embed_model({"obs": obs_concat}) # Mask out invalid actions (use tf.float32.min for stability) action_mask = input_dict["obs"]["action_mask"] inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min) return action_embed + inf_mask, state def value_function(self): return self.action_embed_model.value_function()
def __init__(self, obs_space, action_space, num_outputs, model_config, name, true_obs_shape=(51, 3), action_embed_size=50, *args, **kwargs): super(VMActionMaskModel, self).__init__(obs_space, action_space, num_outputs, model_config, name, *args, **kwargs) self.action_embed_model = FullyConnectedNetwork( spaces.Box(0, 1, shape=true_obs_shape), action_space, action_embed_size, model_config, name + "_action_embedding") self.register_variables(self.action_embed_model.variables())
class CentralizedCriticModel(TFModelV2): """Multi-agent model that implements a centralized VF. It assumes the observation is a dict with 'own_obs' and 'opponent_obs', the former of which can be used for computing actions (i.e., decentralized execution), and the latter for optimization (i.e., centralized learning). This model has two parts: - An action model that looks at just 'own_obs' to compute actions - A value model that also looks at the 'opponent_obs' / 'opponent_action' to compute the value (it does this by using the 'obs_flat' tensor). """ def __init__(self, obs_space, action_space, num_outputs, model_config, name): super(CentralizedCriticModel, self).__init__(obs_space, action_space, num_outputs, model_config, name) self.action_model = FullyConnectedNetwork( Box(low=0, high=1, shape=(6, )), # one-hot encoded Discrete(6) action_space, num_outputs, model_config, name + "_action", ) self.register_variables(self.action_model.variables()) self.value_model = FullyConnectedNetwork( obs_space, action_space, 1, model_config, name + "_vf", ) self.register_variables(self.value_model.variables()) def forward(self, input_dict, state, seq_lens): self._value_out, _ = self.value_model({"obs": input_dict["obs_flat"]}, state, seq_lens) return self.action_model({"obs": input_dict["obs"]["own_obs"]}, state, seq_lens) def value_function(self): return tf.reshape(self._value_out, [-1])
def __init__(self, obs_space, action_space, num_outputs, model_config, name): super(CentralizedCriticModel, self).__init__(obs_space, action_space, num_outputs, model_config, name) observation_length = 12 obs_lower_bound_list = [0, 0, 20] * observation_length obs_upper_bound_list = [650, 2, 40] * observation_length self.action_model = FullyConnectedNetwork( Box(np.array(obs_lower_bound_list), np.array(obs_upper_bound_list)), # one-hot encoded Discrete(6) action_space, num_outputs, model_config, name + "_action") self.register_variables(self.action_model.variables()) self.value_model = FullyConnectedNetwork(obs_space, action_space, 1, model_config, name + "_vf") self.register_variables(self.value_model.variables())
def __init__( self, obs_space, action_space, num_outputs, model_config, name, ): name = "Pa_model" super(ParametricActionsModel, self).__init__(obs_space, action_space, num_outputs, model_config, name) # get real obs space, discarding action mask real_obs_space = obs_space.original_space.spaces['array_obs'] # define action embed model self.action_embed_model = FullyConnectedNetwork( real_obs_space, action_space, num_outputs, model_config, name + "_action_embed") self.register_variables(self.action_embed_model.variables())
class CustomLossModel(TFModelV2): """Custom model that adds an imitation loss on top of the policy loss.""" def __init__(self, obs_space, action_space, num_outputs, model_config, name): super().__init__(obs_space, action_space, num_outputs, model_config, name) self.fcnet = FullyConnectedNetwork( self.obs_space, self.action_space, num_outputs, model_config, name="fcnet") self.register_variables(self.fcnet.variables()) @override(ModelV2) def forward(self, input_dict, state, seq_lens): # Delegate to our FCNet. return self.fcnet(input_dict, state, seq_lens) @override(ModelV2) def custom_loss(self, policy_loss, loss_inputs): # Create a new input reader per worker. reader = JsonReader( self.model_config["custom_model_config"]["input_files"]) input_ops = reader.tf_input_ops() # Define a secondary loss by building a graph copy with weight sharing. obs = restore_original_dimensions( tf.cast(input_ops["obs"], tf.float32), self.obs_space) logits, _ = self.forward({"obs": obs}, [], None) # You can also add self-supervised losses easily by referencing tensors # created during _build_layers_v2(). For example, an autoencoder-style # loss can be added as follows: # ae_loss = squared_diff( # loss_inputs["obs"], Decoder(self.fcnet.last_layer)) print("FYI: You can also use these tensors: {}, ".format(loss_inputs)) # Compute the IL loss. action_dist = Categorical(logits, self.model_config) self.policy_loss = policy_loss self.imitation_loss = tf.reduce_mean( -action_dist.logp(input_ops["actions"])) return policy_loss + 10 * self.imitation_loss def custom_stats(self): return { "policy_loss": self.policy_loss, "imitation_loss": self.imitation_loss, }
def __init__(self, obs_space, action_space, num_outputs, model_config, name): super(CentralizedCriticModel, self).__init__( obs_space, action_space, num_outputs, model_config, name) # Base of the model self.model = FullyConnectedNetwork(obs_space, action_space, num_outputs, model_config, name) self.register_variables(self.model.variables()) # Central VF maps (obs, opp_obs, opp_act) -> vf_pred obs = tf.keras.layers.Input(shape=(6, ), name="obs") opp_obs = tf.keras.layers.Input(shape=(6, ), name="opp_obs") opp_act = tf.keras.layers.Input(shape=(2, ), name="opp_act") concat_obs = tf.keras.layers.Concatenate(axis=1)( [obs, opp_obs, opp_act]) central_vf_dense = tf.keras.layers.Dense( 16, activation=tf.nn.tanh, name="c_vf_dense")(concat_obs) central_vf_out = tf.keras.layers.Dense( 1, activation=None, name="c_vf_out")(central_vf_dense) self.central_vf = tf.keras.Model( inputs=[obs, opp_obs, opp_act], outputs=central_vf_out) self.register_variables(self.central_vf.variables)
def __init__(self, observation_space, action_space, num_outputs, model_config, name): super().__init__(observation_space, action_space, num_outputs, model_config, name) inputs = tf.keras.layers.Input(shape=observation_space.shape) self.fcnet = FullyConnectedNetwork(obs_space=self.obs_space, action_space=self.action_space, num_outputs=self.num_outputs, model_config=self.model_config, name="fc1") out, value_out = self.fcnet.base_model(inputs) def lambda_(x): eager_out = tf.py_function(self.forward_eager, [x], tf.float32) with tf.control_dependencies([eager_out]): eager_out.set_shape(x.shape) return eager_out out = tf.keras.layers.Lambda(lambda_)(out) self.base_model = tf.keras.models.Model(inputs, [out, value_out]) self.register_variables(self.base_model.variables)
class CentralizedCriticModel(TFModelV2): """Multi-agent model that implements a centralized VF.""" # TODO(@evinitsky) make this work with more than boxes def __init__(self, obs_space, action_space, num_outputs, model_config, name): super(CentralizedCriticModel, self).__init__(obs_space, action_space, num_outputs, model_config, name) # Base of the model self.model = FullyConnectedNetwork(obs_space, action_space, num_outputs, model_config, name) self.register_variables(self.model.variables()) # Central VF maps (obs, opp_ops, opp_act) -> vf_pred self.max_num_agents = model_config['custom_options']['max_num_agents'] self.obs_space_shape = obs_space.shape[0] other_obs = tf.keras.layers.Input(shape=(obs_space.shape[0] * self.max_num_agents, ), name="opp_obs") central_vf_dense = tf.keras.layers.Dense( model_config['custom_options']['central_vf_size'], activation=tf.nn.tanh, name="c_vf_dense")(other_obs) central_vf_out = tf.keras.layers.Dense( 1, activation=None, name="c_vf_out")(central_vf_dense) self.central_vf = tf.keras.Model(inputs=[other_obs], outputs=central_vf_out) self.register_variables(self.central_vf.variables) def forward(self, input_dict, state, seq_lens): return self.model.forward(input_dict, state, seq_lens) def central_value_function(self, obs, opponent_obs): return tf.reshape(self.central_vf([opponent_obs]), [-1]) def value_function(self): return self.model.value_function() # not used
def __init__(self, obs_space, action_space, num_outputs, model_config, name, **kwargs): super(HanabiHandInference, self).__init__(obs_space, action_space, model_config["custom_options"]["q_module_hiddens"][-1], model_config, name, **kwargs) self.obs_module = FullyConnectedNetwork(obs_space.original_space["board"], None, model_config["custom_options"]["obs_module_hiddens"][-1], { "fcnet_activation": model_config["fcnet_activation"], "fcnet_hiddens": model_config["custom_options"][ "obs_module_hiddens"], "no_final_linear": True, "vf_share_layers": True}, name + "obs_module") obs_module_output_dummy = numpy.zeros(model_config["custom_options"]["obs_module_hiddens"][-1]) self.q_module = FullyConnectedNetwork(obs_module_output_dummy, None, model_config["custom_options"]["q_module_hiddens"][-1], {"fcnet_activation": model_config["fcnet_activation"], "fcnet_hiddens": model_config["custom_options"]["q_module_hiddens"], "no_final_linear": True, "vf_share_layers": True}, name + "q_module") self.aux_module = FullyConnectedNetwork(obs_module_output_dummy, None, model_config["custom_options"]["aux_module_hiddens"][-1], {"fcnet_activation": model_config["fcnet_activation"], "fcnet_hiddens": model_config["custom_options"][ "aux_module_hiddens"], "no_final_linear": True, "vf_share_layers": True}, name + "aux_module") aux_head_input_dummy = numpy.zeros(model_config["custom_options"]["aux_module_hiddens"][-1]) self.aux_head = FullyConnectedNetwork(aux_head_input_dummy, None, numpy.prod(obs_space.original_space["hidden_hand"].shape), {"fcnet_activation": model_config["fcnet_activation"], "fcnet_hiddens": model_config["custom_options"][ "aux_head_hiddens"], "no_final_linear": False, "vf_share_layers": True}, name + "aux_head") self.register_variables(self.obs_module.variables()) self.register_variables(self.q_module.variables()) self.register_variables(self.aux_module.variables()) self.register_variables(self.aux_head.variables()) self.aux_loss_formula = get_aux_loss_formula(model_config["custom_options"].get("aux_loss_formula", "sqrt"))
class HanabiFullyConnected(LegalActionsDistributionalQModel): def __init__(self, obs_space, action_space, num_outputs, model_config, name, **kwargs): super(HanabiFullyConnected, self).__init__(obs_space, action_space, num_outputs, model_config, name, **kwargs) self.fc = FullyConnectedNetwork(obs_space.original_space["board"], action_space, num_outputs, model_config, name + "fc") self.register_variables(self.fc.variables()) def forward(self, input_dict, state, seq_lens): model_out, state = self.fc({"obs": input_dict["obs"]["board"]}, state, seq_lens) self.calculate_and_store_q(input_dict, model_out) return model_out, state
class EagerModel(TFModelV2): """Example of using embedded eager execution in a custom model. This shows how to use tf.py_function() to execute a snippet of TF code in eager mode. Here the `self.forward_eager` method just prints out the intermediate tensor for debug purposes, but you can in general perform any TF eager operation in tf.py_function(). """ def __init__(self, observation_space, action_space, num_outputs, model_config, name): super().__init__(observation_space, action_space, num_outputs, model_config, name) inputs = tf.keras.layers.Input(shape=observation_space.shape) self.fcnet = FullyConnectedNetwork(obs_space=self.obs_space, action_space=self.action_space, num_outputs=self.num_outputs, model_config=self.model_config, name="fc1") out, value_out = self.fcnet.base_model(inputs) def lambda_(x): eager_out = tf.py_function(self.forward_eager, [x], tf.float32) with tf.control_dependencies([eager_out]): eager_out.set_shape(x.shape) return eager_out out = tf.keras.layers.Lambda(lambda_)(out) self.base_model = tf.keras.models.Model(inputs, [out, value_out]) self.register_variables(self.base_model.variables) @override(ModelV2) def forward(self, input_dict, state, seq_lens): out, self._value_out = self.base_model(input_dict["obs"], state, seq_lens) return out, [] @override(ModelV2) def value_function(self): return tf.reshape(self._value_out, [-1]) def forward_eager(self, feature_layer): assert tf.executing_eagerly() if random.random() > 0.99: print("Eagerly printing the feature layer mean value", tf.reduce_mean(feature_layer)) return feature_layer