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)
  )
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)
  
trax/rl/serialization_utils.py