def __init__(self, env_spec, name='CategoricalMLPPolicy', hidden_sizes=(32, 32), hidden_nonlinearity=tf.nn.tanh, hidden_w_init=tf.glorot_uniform_initializer(), hidden_b_init=tf.zeros_initializer(), output_nonlinearity=tf.nn.softmax, output_w_init=tf.glorot_uniform_initializer(), output_b_init=tf.zeros_initializer(), layer_normalization=False):
    assert isinstance(env_spec.action_space, akro.Discrete), 'CategoricalMLPPolicy only works with akro.Discrete action space.'
    super().__init__(name, env_spec)
    self.obs_dim = env_spec.observation_space.flat_dim
    self.action_dim = env_spec.action_space.n
    self.model = MLPModel(output_dim=self.action_dim, hidden_sizes=hidden_sizes, hidden_nonlinearity=hidden_nonlinearity, hidden_w_init=hidden_w_init, hidden_b_init=hidden_b_init, output_nonlinearity=output_nonlinearity, output_w_init=output_w_init, output_b_init=output_b_init, layer_normalization=layer_normalization, name='MLPModel')
    self._initialize()
def __init__(self, env_spec, name='CategoricalMLPPolicy', hidden_sizes=(32, 32), hidden_nonlinearity=tf.nn.tanh, hidden_w_init=tf.glorot_uniform_initializer(), hidden_b_init=tf.zeros_initializer(), output_nonlinearity=tf.nn.softmax, output_w_init=tf.glorot_uniform_initializer(), output_b_init=tf.zeros_initializer(), layer_normalization=False):
    assert isinstance(env_spec.action_space, akro.Discrete), 'CategoricalMLPPolicy only works with akro.Discrete action space.'
    super().__init__(name, env_spec)
    self.obs_dim = env_spec.observation_space.flat_dim
    self.action_dim = env_spec.action_space.n
    self.model = MLPModel(output_dim=self.action_dim, hidden_sizes=hidden_sizes, hidden_nonlinearity=hidden_nonlinearity, hidden_w_init=hidden_w_init, hidden_b_init=hidden_b_init, output_nonlinearity=output_nonlinearity, output_w_init=output_w_init, output_b_init=output_b_init, layer_normalization=layer_normalization, name='MLPModel')
    self._initialize()