<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 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