(
    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,
)
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.
Attaches a value head to a model body.
Attaches a value head to a model body .
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),
  )
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 , ) : 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 ) , )