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