content: # -*- coding: utf-8 -*-
  """
  An implementation of the policyValueNet in Theano and Lasagne
  
  @author: Junxiao Song
  """
  
  from __future__ import print_function
  import theano
  import theano.tensor as T
  import lasagne
  import pickle
  
  
  class PolicyValueNet():
      """policy-value network """
      def __init__(self, board_width, board_height, model_file=None):
          self.board_width = board_width
          self.board_height = board_height
          self.learning_rate = T.scalar('learning_rate')
          self.l2_const = 1e-4  # coef of l2 penalty
          self.create_policy_value_net()
          self._loss_train_op()
          if model_file:
              try:
                  net_params = pickle.load(open(model_file, 'rb'))
              except:
                  # To support loading pretrained model in python3
                  net_params = pickle.load(open(model_file, 'rb'),
                                           encoding='bytes')
              lasagne.layers.set_all_param_values(
                      [self.policy_net, self.value_net], net_params
                      )
  
      def create_policy_value_net(self):
          """create the policy value network """
          self.state_input = T.tensor4('state')
          self.winner = T.vector('winner')
          self.mcts_probs = T.matrix('mcts_probs')
          network = lasagne.layers.InputLayer(
                  shape=(None, 4, self.board_width, self.board_height),
                  input_var=self.state_input
                  )
          # conv layers
          network = lasagne.layers.Conv2DLayer(
                  network, num_filters=32, filter_size=(3, 3), pad='same')
          network = lasagne.layers.Conv2DLayer(
                  network, num_filters=64, filter_size=(3, 3), pad='same')
          network = lasagne.layers.Conv2DLayer(
                  network, num_filters=128, filter_size=(3, 3), pad='same')
          # action policy layers
          policy_net = lasagne.layers.Conv2DLayer(
                  network, num_filters=4, filter_size=(1, 1))
          self.policy_net = lasagne.layers.DenseLayer(
                  policy_net, num_units=self.board_width*self.board_height,
                  nonlinearity=lasagne.nonlinearities.softmax)
          # state value layers
          value_net = lasagne.layers.Conv2DLayer(
                  network, num_filters=2, filter_size=(1, 1))
          value_net = lasagne.layers.DenseLayer(value_net, num_units=64)
          self.value_net = lasagne.layers.DenseLayer(
                  value_net, num_units=1,
                  nonlinearity=lasagne.nonlinearities.tanh)
          # get action probs and state score value
          self.action_probs, self.value = lasagne.layers.get_output(
                  [self.policy_net, self.value_net])
          self.policy_value = theano.function([self.state_input],
                                              [self.action_probs, self.value],
                                              allow_input_downcast=True)
  
      def policy_value_fn(self, board):
          """
          input: board
          output: a list of (action, probability) tuples for each available
              action and the score of the board state
          """
          legal_positions = board.availables
          current_state = board.current_state()
          act_probs, value = self.policy_value(
              current_state.reshape(-1, 4, self.board_width, self.board_height)
              )
          act_probs = zip(legal_positions, act_probs.flatten()[legal_positions])
          return act_probs, value[0][0]
  
      def _loss_train_op(self):
          """
          Three loss terms：
          loss = (z - v)^2 - pi^T * log(p) + c||theta||^2
          """
          params = lasagne.layers.get_all_params(
                  [self.policy_net, self.value_net], trainable=True)
          value_loss = lasagne.objectives.squared_error(
                  self.winner, self.value.flatten())
          policy_loss = lasagne.objectives.categorical_crossentropy(
                  self.action_probs, self.mcts_probs)
          l2_penalty = lasagne.regularization.apply_penalty(
                  params, lasagne.regularization.l2)
          self.loss = self.l2_const*l2_penalty + lasagne.objectives.aggregate(
                  value_loss + policy_loss, mode='mean')
          # policy entropy，for monitoring only
          self.entropy = -T.mean(T.sum(
                  self.action_probs * T.log(self.action_probs + 1e-10), axis=1))
          # get the train op
          updates = lasagne.updates.adam(self.loss, params,
                                         learning_rate=self.learning_rate)
          self.train_step = theano.function(
              [self.state_input, self.mcts_probs, self.winner, self.learning_rate],
              [self.loss, self.entropy],
              updates=updates,
              allow_input_downcast=True
              )
  
      def get_policy_param(self):
          net_params = lasagne.layers.get_all_param_values(
                  [self.policy_net, self.value_net])
          return net_params
  
      def save_model(self, model_file):
          """ save model params to file """
          net_params = self.get_policy_param()  # get model params
          pickle.dump(net_params, open(model_file, 'wb'), protocol=2)
  