code: 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)
  
imports: import functools
  import gym
  import numpy
  from trax import layers
  from trax.fastmath import numpy
  from trax.rl import space_serializer
label_desc: Use image node_id 12 for calling a global function with example usage: wrap_policy(seq_model, observation_space, action_space, vocab_size) and returns: policy_wrapper