def wrap_policy(seq_model, observation_space, action_space, vocab_size):  # pylint: disable=invalid-name
  """Wraps a sequence model in either RawPolicy or SerializedPolicy.

  Args:
    seq_model: Trax sequence model.
    observation_space: Gym observation space.
    action_space: Gym action space.
    vocab_size: Either the number of symbols for a serialized policy, or None.

  Returns:
    RawPolicy if vocab_size is None, else SerializedPolicy.
  """
  (n_controls, n_actions) = analyze_action_space(action_space)
  if vocab_size is None:
    policy_wrapper = RawPolicy
  else:
    obs_serializer = space_serializer.create(observation_space, vocab_size)
    act_serializer = space_serializer.create(action_space, vocab_size)
    policy_wrapper = functools.partial(SerializedPolicy,
                                       observation_serializer=obs_serializer,
                                       action_serializer=act_serializer)
  return policy_wrapper(seq_model, n_controls, n_actions)
Wraps a sequence model in either RawPolicy or SerializedPolicy.

  Args:
    seq_model: Trax sequence model.
    observation_space: Gym observation space.
    action_space: Gym action space.
    vocab_size: Either the number of symbols for a serialized policy, or None.

  Returns:
    RawPolicy if vocab_size is None, else SerializedPolicy.
  
trax/rl/serialization_utils.py