(seq_model, n_controls, n_actions)
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)
Wraps a sequence model in a policy interface.
Wraps a sequence model in a policy interface .
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)
  )
def RawPolicy ( seq_model , n_controls , n_actions ) : 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) )