code: def policy_net(rng_key,
                 batch_observations_shape,
                 num_actions,
                 bottom_layers=None):
    
    # 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
docs: A policy net function.