':param input_shape: Shape of the input data.
:param output_dim: Dimension of output.
:param hidden_sizes: Number of hidden units of each layer of the mean network.
:param hidden_nonlinearity: Non-linearity used for each layer of the mean network.
:param optimizer: Optimizer for minimizing the negative log-likelihood.
:param use_trust_region: Whether to use trust region constraint.
:param step_size: KL divergence constraint for each iteration'

def __init__(self, input_shape, output_dim, prob_network=None, hidden_sizes=(32, 32), hidden_nonlinearity=NL.rectify, optimizer=None, use_trust_region=True, step_size=0.01, normalize_inputs=True, name=None):

    Serializable.quick_init(self, locals())
    if (optimizer is None):
        if use_trust_region:
            optimizer = PenaltyLbfgsOptimizer()
        else:
            optimizer = LbfgsOptimizer()
    self.output_dim = output_dim
    self._optimizer = optimizer
    if (prob_network is None):
        prob_network = MLP(input_shape=input_shape, output_dim=output_dim, hidden_sizes=hidden_sizes, hidden_nonlinearity=hidden_nonlinearity, output_nonlinearity=NL.softmax)
    l_prob = prob_network.output_layer
    LasagnePowered.__init__(self, [l_prob])
    xs_var = prob_network.input_layer.input_var
    ys_var = TT.imatrix('ys')
    old_prob_var = TT.matrix('old_prob')
    x_mean_var = theano.shared(np.zeros(((1,) + input_shape)), name='x_mean', broadcastable=((True,) + ((False,) * len(input_shape))))
    x_std_var = theano.shared(np.ones(((1,) + input_shape)), name='x_std', broadcastable=((True,) + ((False,) * len(input_shape))))
    normalized_xs_var = ((xs_var - x_mean_var) / x_std_var)
    prob_var = L.get_output(l_prob, {prob_network.input_layer: normalized_xs_var})
    old_info_vars = dict(prob=old_prob_var)
    info_vars = dict(prob=prob_var)
    dist = self._dist = Categorical(output_dim)
    mean_kl = TT.mean(dist.kl_sym(old_info_vars, info_vars))
    loss = (- TT.mean(dist.log_likelihood_sym(ys_var, info_vars)))
    predicted = special.to_onehot_sym(TT.argmax(prob_var, axis=1), output_dim)
    self._f_predict = ext.compile_function([xs_var], predicted)
    self._f_prob = ext.compile_function([xs_var], prob_var)
    self._prob_network = prob_network
    self._l_prob = l_prob
    optimizer_args = dict(loss=loss, target=self, network_outputs=[prob_var])
    if use_trust_region:
        optimizer_args['leq_constraint'] = (mean_kl, step_size)
        optimizer_args['inputs'] = [xs_var, ys_var, old_prob_var]
    else:
        optimizer_args['inputs'] = [xs_var, ys_var]
    self._optimizer.update_opt(**optimizer_args)
    self._use_trust_region = use_trust_region
    self._name = name
    self._normalize_inputs = normalize_inputs
    self._x_mean_var = x_mean_var
    self._x_std_var = x_std_var
