class ParametricActionsModel(DistributionalQTFModel): def __init__(self, obs_space, action_space, num_outputs, model_config, name, **kw): print("{} : [INFO] ParametricActionsModel {}, {}, {}, {}, {}".format( datetime.now(), action_space, obs_space, num_outputs, name, model_config)) super(ParametricActionsModel, self).__init__(obs_space, action_space, num_outputs, model_config, name, **kw) # print("####### obs_space {}".format(obs_space)) # raise Exception("END") self.action_param_model = FullyConnectedNetwork( FLAT_OBSERVATION_SPACE, action_space, num_outputs, model_config, name + "_action_param") self.register_variables(self.action_param_model.variables()) def forward(self, input_dict, state, seq_lens): # Extract the available actions tensor from the observation. action_mask = input_dict["obs"]["action_mask"] # Compute the predicted action embedding action_param, _ = self.action_param_model( {"obs": input_dict["obs"]["state"]}) # Mask out invalid actions (use tf.float32.min for stability) inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min) return action_param + inf_mask, state def value_function(self): return self.action_param_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 FCMaskedActionsModelTF(DistributionalQTFModel, TFModelV2): def __init__(self, obs_space, action_space, num_outputs, model_config, name, **kw): super(FCMaskedActionsModelTF, self).__init__(obs_space, action_space, num_outputs, model_config, name, **kw) true_obs_space = gym.spaces.MultiBinary(n=obs_space.shape[0] - action_space.n) self.action_embed_model = FullyConnectedNetwork( obs_space=true_obs_space, action_space=action_space, num_outputs=action_space.n, model_config=model_config, name=name + "action_model") self.register_variables(self.action_embed_model.variables()) def forward(self, input_dict, state, seq_lens): action_mask = input_dict["obs"]["action_mask"] # Compute the predicted action embedding raw_actions, _ = self.action_embed_model( {"obs": input_dict["obs"]["real_obs"]}) #inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min) logits = tf.where(tf.math.equal(action_mask, 1), raw_actions, tf.float32.min) return logits, state def value_function(self): return self.action_embed_model.value_function()
class ActionMaskModel(TFModelV2): 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) low = np.asarray(model_config["custom_model_config"]["obs_space_low"]) high = np.asarray(model_config["custom_model_config"]["obs_space_high"]) self.policy = FullyConnectedNetwork( spaces.Box(low, high, shape=(len(low),)), action_space, num_outputs, model_config, "policy_network", ) def forward(self, input_dict, state, seq_lens): obs = input_dict["obs"]["real_obs"] action_mask = input_dict["obs"]["action_mask"] action_logits, _ = self.policy({"obs": obs}, state, seq_lens) if self.num_outputs == 1: return action_logits, state # Mask out invalid actions (use tf.float32.min for stability) inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min) return action_logits + inf_mask, state def value_function(self): return self.policy.value_function()
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") @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 value_function(self): # Delegate to our FCNet. return self.fcnet.value_function() @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 metrics(self): return { "policy_loss": self.policy_loss, "imitation_loss": self.imitation_loss, }
class ActionMaskModel(TFModelV2): """Model that handles simple discrete action 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, **kwargs): orig_space = getattr(obs_space, "original_space", obs_space) assert (isinstance(orig_space, Dict) and "action_mask" in orig_space.spaces and "observations" in orig_space.spaces) super().__init__(obs_space, action_space, num_outputs, model_config, name) self.internal_model = FullyConnectedNetwork( orig_space["observations"], action_space, num_outputs, model_config, name + "_internal", ) # disable action masking --> will likely lead to invalid actions self.no_masking = model_config["custom_model_config"].get( "no_masking", False) def forward(self, input_dict, state, seq_lens): # Extract the available actions tensor from the observation. action_mask = input_dict["obs"]["action_mask"] # Compute the unmasked logits. logits, _ = self.internal_model( {"obs": input_dict["obs"]["observations"]}) # If action masking is disabled, directly return unmasked logits if self.no_masking: return logits, state # Convert action_mask into a [0.0 || -inf]-type mask. inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min) masked_logits = logits + inf_mask # Return masked logits. return masked_logits, state def value_function(self): return self.internal_model.value_function()
class CustomModel(TFModelV2): """Example of a keras custom model that just delegates to an 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) 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 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=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", ) 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"]["cart"]}) # 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.math.log(action_mask), tf.float32.min) return action_logits + inf_mask, state def value_function(self): return self.action_embed_model.value_function()
class CustomModel(TFModelV2): 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 = 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 OwnershipActionMaskingModel(FullyConnectedNetwork): """ Parametric action model that handles the dot product and masking. This assumes the outputs are logits for a single Categorical action dist. """ def __init__(self, obs_space, action_space, num_outputs, model_config, name, **kw): super(OwnershipActionMaskingModel, self).__init__(obs_space, action_space, num_outputs, model_config, name, **kw) self.true_obs_shape = model_config['custom_model_config'][ 'true_obs_shape'] self.action_embed_size = model_config['custom_model_config'][ 'action_embed_size'] self.action_embed_model = FullyConnectedNetwork( self.true_obs_shape, action_space, self.action_embed_size, model_config, name + "_action_embed") # Box(-1, 0, shape=true_obs_shape) 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"]["obs"]}) # 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=1) # Mask out invalid actions (use tf.float32.min for stability) inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min) return action_logits + inf_mask, state def value_function(self): return self.action_embed_model.value_function()
class CentralizedCriticModel(TFModelV2): """Multi-agent model that implements a centralized value function.""" 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) # 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) @override(ModelV2) 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(tf.cast(opponent_actions, tf.int32), 2) ]), [-1], ) @override(ModelV2) def value_function(self): return self.model.value_function() # not used
class CustomModel(TFModelV2, selection_DelayedImpactEnv): """Example of a keras custom model that just delegates to an fc-net.""" obs_space = selection_DelayedImpactEnv.observation_space action_space = selection_DelayedImpactEnv.action_space # num_outputs=169 model_config = {} name = 'My_model' 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 KP0ActionMaskModel(TFModelV2): def __init__(self, obs_space, action_space, num_outputs, model_config, name, true_obs_shape=(11, ), action_embed_size=5, *args, **kwargs): super(KP0ActionMaskModel, self).__init__(obs_space, action_space, num_outputs, model_config, name, *args, **kwargs) self.action_embed_model = FullyConnectedNetwork( CatanatronEnv.observation_space, 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 ParametricActionsModelThatLearnsEmbeddings(DistributionalQTFModel): """Same as the above ParametricActionsModel. However, this version also learns the action embeddings. """ def __init__(self, obs_space, action_space, num_outputs, model_config, name, true_obs_shape=(4, ), action_embed_size=2, **kw): super(ParametricActionsModelThatLearnsEmbeddings, self).__init__(obs_space, action_space, num_outputs, model_config, name, **kw) action_ids_shifted = tf.constant(list(range(1, num_outputs + 1)), dtype=tf.float32) obs_cart = tf.keras.layers.Input(shape=true_obs_shape, name="obs_cart") valid_avail_actions_mask = tf.keras.layers.Input( shape=(num_outputs), name="valid_avail_actions_mask") self.pred_action_embed_model = FullyConnectedNetwork( Box(-1, 1, shape=true_obs_shape), action_space, action_embed_size, model_config, name + "_pred_action_embed", ) # Compute the predicted action embedding pred_action_embed, _ = self.pred_action_embed_model({"obs": obs_cart}) _value_out = self.pred_action_embed_model.value_function() # 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(pred_action_embed, 1) valid_avail_actions = action_ids_shifted * valid_avail_actions_mask # Embedding for valid available actions which will be learned. # Embedding vector for 0 is an invalid embedding (a "dummy embedding"). valid_avail_actions_embed = tf.keras.layers.Embedding( input_dim=num_outputs + 1, output_dim=action_embed_size, name="action_embed_matrix", )(valid_avail_actions) # Batch dot product => shape of logits is [BATCH, MAX_ACTIONS]. action_logits = tf.reduce_sum(valid_avail_actions_embed * intent_vector, axis=2) # Mask out invalid actions (use tf.float32.min for stability) inf_mask = tf.maximum(tf.math.log(valid_avail_actions_mask), tf.float32.min) action_logits = action_logits + inf_mask self.param_actions_model = tf.keras.Model( inputs=[obs_cart, valid_avail_actions_mask], outputs=[action_logits, _value_out], ) self.param_actions_model.summary() def forward(self, input_dict, state, seq_lens): # Extract the available actions mask tensor from the observation. valid_avail_actions_mask = input_dict["obs"][ "valid_avail_actions_mask"] action_logits, self._value_out = self.param_actions_model( [input_dict["obs"]["cart"], valid_avail_actions_mask]) return action_logits, state def value_function(self): return self._value_out