def substitute_inner_policy(wrapped_policy, inner_policy, vocab_size):  # pylint: disable=invalid-name
  """Substitutes the inner weights/state in a {Raw,Serialized}Policy.

  Args:
    wrapped_policy (pytree): Weights or state of a wrapped policy.
    inner_policy (pytree): Weights or state of an inner policy.
    vocab_size (int or None): Vocabulary size of a serialized policy, or None
      in case of a raw policy.

  Returns:
    New weights or state of wrapped_policy, with the inner weights/state
      copied from inner_policy.
  """
  if vocab_size is None:
    substitute_fn = substitute_inner_policy_raw
  else:
    substitute_fn = substitute_inner_policy_serialized
  return substitute_fn(wrapped_policy, inner_policy)
Substitutes the inner weights/state in a {Raw,Serialized}Policy.

  Args:
    wrapped_policy (pytree): Weights or state of a wrapped policy.
    inner_policy (pytree): Weights or state of an inner policy.
    vocab_size (int or None): Vocabulary size of a serialized policy, or None
      in case of a raw policy.

  Returns:
    New weights or state of wrapped_policy, with the inner weights/state
      copied from inner_policy.
  
trax/rl/serialization_utils.py