def Quality(
    body=None,
    normalizer=None,
    batch_axes=None,
    mode='train',
    n_actions=2,
    head_init_range=None,
):
  """The network takes as input an observation and outputs values of actions."""

  if body is None:
    body = lambda mode: []
  if normalizer is None:
    normalizer = lambda mode: []

  head_kwargs = {}
  if head_init_range is not None:
    head_kwargs['kernel_initializer'] = tl.RandomUniformInitializer(
        lim=head_init_range
    )

  return tl.Serial(
      _Batch(normalizer(mode=mode), batch_axes),
      _Batch(body(mode=mode), batch_axes),
      tl.Dense(n_actions, **head_kwargs),
  )
The network takes as input an observation and outputs values of actions.
trax/models/rl.py