code: def RawPolicy(seq_model, n_controls, n_actions):
      """Wraps a sequence model in a policy interface.
  
      The resulting model takes as input observation anc action sequences, but only
      uses the observations. Adds output heads for action logits and value
      predictions.
  
      Args:
        seq_model: Trax sequence model taking as input and outputting a sequence of
          continuous vectors.
        n_controls: Number of controls.
        n_actions: Number of action categories in each control.
  
      Returns:
        A model of signature (obs, act) -> (act_logits, values), with shapes:
          obs: (batch_size, length + 1, obs_depth)
          act: (batch_size, length, n_controls)
          act_logits: (batch_size, length, n_controls, n_actions)
          values: (batch_size, length)
      """
  
      def SplitControls():  # pylint: disable=invalid-name
          """Splits logits for actions in different controls."""
  
          def f(x):
              return jnp.reshape(
                  x, x.shape[:2] + (n_controls, n_actions)
              )
  
          return tl.Fn("SplitControls", f)
  
      action_head = [
          # Predict all action logits at the same time.
          tl.Dense(n_controls * n_actions),
          # Then group them into separate controls, adding a new dimension.
          SplitControls(),
          tl.LogSoftmax(),
      ]
      return tl.Serial(  # (obs, act)
          tl.Select([0], n_in=2),  # (obs,)
          seq_model,  # (obs_hidden,)
          tl.Dup(),  # (obs_hidden, obs_hidden)
          tl.Parallel(
              action_head, [tl.Dense(1), tl.Flatten()]
          ),  # (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 7 for calling a global function with example usage: RawPolicy(seq_model, n_controls, n_actions) and returns: tl, tl, jnp