def Policy(
    policy_distribution,
    body=None,
    normalizer=None,
    head_init_range=None,
    batch_axes=None,
    mode='train',
):
  """Attaches a policy head to a model body."""
  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(policy_distribution.n_inputs, **head_kwargs),
  )
Attaches a policy head to a model body.
trax/models/rl.py