def PolicyAndValue(
    policy_distribution,
    body=None,
    policy_top=Policy,
    value_top=Value,
    normalizer=None,
    joint=True,
    mode='train',
):
  """Attaches policy and value heads to a model body."""
  if normalizer is None:
    normalizer = lambda mode: []
  if body is None:
    body = lambda mode: []

  common_kwargs = {'body': None, 'normalizer': None, 'mode': mode}
  policy_top = policy_top(
      policy_distribution=policy_distribution, **common_kwargs
  )
  value_top = value_top(**common_kwargs)

  layers = [normalizer(mode=mode)]
  if joint:
    layers.extend([
        body(mode=mode),
        tl.Branch(policy_top, value_top),
    ])
  else:
    layers.append(tl.Branch(
        tl.Serial(body(mode=mode), policy_top),
        tl.Serial(body(mode=mode), value_top),
    ))
  return tl.Serial(layers)
Attaches policy and value heads to a model body.
trax/models/rl.py