def SerializedPolicy(
    seq_model, n_controls, n_actions, observation_serializer, action_serializer
):
  """Wraps a policy in serialization machinery for training.

  The resulting model takes as input observation and action sequences, and
  serializes them into one sequence similar to SerializedModel, before passing
  to the given sequence model. Adds output heads for action logits and value
  predictions.

  Args:
    seq_model: Trax sequence model taking as input a sequence of symbols and
      outputting a sequence of continuous vectors.
    n_controls: Number of controls.
    n_actions: Number of action categories in each control.
    observation_serializer: Serializer to use for observations.
    action_serializer: Serializer to use for actions.

  Returns:
    A model of signature (obs, act) -> (act_logits, values), same as in
    RawPolicy.
  """
  if action_serializer.representation_length != n_controls:
    raise ValueError(
        'Action symbols should correspond 1-1 to controls, but got {} '
        'controls and {} symbols.'.format(
            n_controls, action_serializer.representation_length
        )
    )

  def FirstSymbol():
    return tl.Fn('FirstSymbol', lambda x: x[:, :, 0])

  def PadRight(n_to_pad):
    def pad_right(x):
      pad_widths = [(0, 0), (0, n_to_pad)] + [(0, 0)] * (x.ndim - 2)
      return jnp.pad(
          x, pad_widths, mode='constant', constant_values=x.dtype.type(0))
    return tl.Fn(f'PadRight({n_to_pad})', pad_right)

  action_head = [
      tl.Dense(n_actions),
      tl.LogSoftmax(),
  ]
  value_head = [
      # Take just the vectors corresponding to the first action symbol.
      FirstSymbol(),
      # Predict values.
      tl.Dense(1),
      # Get rid of the singleton dimension.
      tl.Flatten(),
  ]
  return tl.Serial(
      # (obs, act)
      tl.Parallel(Serialize(observation_serializer),
                  Serialize(action_serializer)),
      # (obs_repr, act_repr)
      Interleave(),
      # (obs_act_repr,)

      # Add one dummy action to the right - we'll use the output at its first
      # symbol to predict the value for the last observation.
      PadRight(action_serializer.representation_length),

      # Shift one symbol to the right, so we predict the n-th action symbol
      # based on action symbols 1..n-1 instead of 1..n.
      tl.ShiftRight(),
      seq_model,
      # (obs_act_hidden,)
      Deinterleave(observation_serializer.representation_length,
                   action_serializer.representation_length),
      # (obs_hidden, act_hidden)
      tl.Select([1, 1]),
      # (act_hidden, act_hidden)
      tl.Parallel(action_head, value_head),
      # (act_logits, values)
  )
Wraps a policy in serialization machinery for training.

  The resulting model takes as input observation and action sequences, and
  serializes them into one sequence similar to SerializedModel, before passing
  to the given sequence model. Adds output heads for action logits and value
  predictions.

  Args:
    seq_model: Trax sequence model taking as input a sequence of symbols and
      outputting a sequence of continuous vectors.
    n_controls: Number of controls.
    n_actions: Number of action categories in each control.
    observation_serializer: Serializer to use for observations.
    action_serializer: Serializer to use for actions.

  Returns:
    A model of signature (obs, act) -> (act_logits, values), same as in
    RawPolicy.
  
trax/rl/serialization_utils.py