A policy and value net function.

def policy_and_value_net(rng_key,
                         batch_observations_shape,
                         num_actions,
                         bottom_layers=None):
  """A policy and value net function."""

  # Layers.
  cur_layers = []
  if bottom_layers is not None:
    cur_layers.extend(bottom_layers)

  # Now, with the current logits, one head computes action probabilities and the
  # other computes the value function.
  # NOTE: The LogSoftmax instead of the Softmax because of numerical stability.
  cur_layers.extend([layers.Branch(), layers.Parallel(
      layers.Serial(layers.Dense(num_actions), layers.LogSoftmax()),
      layers.Dense(1)
  )])
  net = layers.Serial(*cur_layers)
  return net.initialize(batch_observations_shape, rng_key), net