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 ParametricActionsModel(DistributionalQTFModel): def __init__(self, obs_space, action_space, num_outputs, model_config, name, true_obs_shape=(2, ), **kw): super(ParametricActionsModel, self).__init__(obs_space, action_space, num_outputs, model_config, name, **kw) self.action_value_model = FullyConnectedNetwork( Box(-1, 1, shape=true_obs_shape), action_space, num_outputs, model_config, name + "_action_values", ) self.register_variables(self.action_value_model.variables()) def forward(self, input_dict, state, seq_lens): action_mask = input_dict["obs"]["action_mask"] action_values, _ = self.action_value_model( {"obs": input_dict["obs"]["actual_obs"]}) inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min) return action_values + inf_mask, state
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 YetAnotherCentralizedCriticModel(TFModelV2): """Multi-agent model that implements a centralized value function. 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(YetAnotherCentralizedCriticModel, 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])
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, }
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 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) 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 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") 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"]["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 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) 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) @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 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()