def Value(
    body=None,
    normalizer=None,
    inject_actions=False,
    inject_actions_n_layers=1,
    inject_actions_dim=64,
    batch_axes=None,
    mode='train',
    is_discrete=False,
    vocab_size=2,
    multiplicative_action_injection=False,
    head_init_range=None,
):
  """Attaches a value head to a model body."""
  if body is None:
    body = lambda mode: []
  if normalizer is None:
    normalizer = lambda mode: []

  def ActionInjector(mode):
    if inject_actions:
      if is_discrete:
        action_encoder = tl.Embedding(vocab_size, inject_actions_dim)
      else:
        action_encoder = tl.Dense(inject_actions_dim)
      encoders = tl.Parallel(
          tl.Dense(inject_actions_dim),
          action_encoder,
      )
      if multiplicative_action_injection:
        action_injector = tl.Serial(
            tl.Fn('TanhMulGate', lambda x, a: x * jnp.tanh(a)),
            tl.LayerNorm()  # compensate for reduced variance
        )
      else:
        action_injector = tl.Add()
      return tl.Serial(
          # Input: (body output, actions).
          encoders,
          action_injector,
          models.MLP(
              layer_widths=(inject_actions_dim,) * inject_actions_n_layers,
              out_activation=True,
              flatten=False,
              mode=mode,
          )
      )
    else:
      return []

  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),
      ActionInjector(mode=mode),
      tl.Dense(1, **head_kwargs),
  )
Attaches a value head to a model body.
trax/models/rl.py