def policy_net(rng_key,
               batch_observations_shape,
               num_actions,
               bottom_layers=None):
  """A policy net function."""
  # Use the bottom_layers as the bottom part of the network and just add the
  # required layers on top of it.
  if bottom_layers is None:
    bottom_layers = []

  # NOTE: The LogSoftmax instead of the Softmax.
  bottom_layers.extend([layers.Dense(num_actions), layers.LogSoftmax()])
  net = layers.Serial(*bottom_layers)

  return net.initialize(batch_observations_shape, rng_key), net
A policy net function.