code: 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)
      )
  
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 9 for calling a global function with example usage: SerializedPolicy(seq_model, n_controls, n_actions, observation_serializer, action_serializer) and returns: tl, tl, tl, jnp