code: def build_ema_mean_std_normalizer(
      nested_spec: types.NestedSpec,
      tau: float = 0.995,
      epsilon: float = 1e-6,
  ) -> NormalizationFns:
      """Builds pure functions used for normalizing based on EMA mean and std.
  
      The built normalizer functions can be used to normalize nested arrays that
      have a structure corresponding to nested_spec. Currently only supports
      nested_spec where all leafs have float dtype.
  
      Arguments:
        nested_spec: A nested spec where all leaves have float dtype
        tau: tau parameter for exponential moving average
        epsilon: epsilon for avoiding division by zero std
  
      Returns:
        NormalizationFns to be used for normalization
      """
      nested_dims = jax.tree_util.tree_map(
          lambda x: len(x.shape), nested_spec
      )
  
      def init() -> EMAMeanStdNormalizerParams:
          first_moment = utils.zeros_like(nested_spec)
          second_moment = utils.zeros_like(nested_spec)
  
          return EMAMeanStdNormalizerParams(
              ema_counter=jnp.int32(0),
              biased_first_moment=first_moment,
              biased_second_moment=second_moment,
          )
  
      def _normalize_leaf(
          x: jnp.ndarray,
          ema_counter: jnp.int32,
          biased_first_moment: jnp.ndarray,
          biased_second_moment: jnp.ndarray,
      ) -> jnp.ndarray:
          zero_debias = 1.0 / (1.0 - jnp.power(tau, ema_counter))
          mean = biased_first_moment * zero_debias
          second_moment = biased_second_moment * zero_debias
          std = jnp.sqrt(jax.nn.relu(second_moment - mean**2))
  
          mean = jnp.broadcast_to(mean, x.shape)
          std = jnp.broadcast_to(std, x.shape)
          return (x - mean) / jnp.fmax(std, epsilon)
  
      def _normalize(
          nested_array: types.NestedArray,
          params: EMAMeanStdNormalizerParams,
      ) -> types.NestedArray:
          ema_counter = params.ema_counter
          normalized_nested_array = jax.tree_util.tree_map(
              lambda x, f, s: _normalize_leaf(x, ema_counter, f, s),
              nested_array,
              params.biased_first_moment,
              params.biased_second_moment,
          )
          return normalized_nested_array
  
      def normalize(
          nested_array: types.NestedArray,
          params: EMAMeanStdNormalizerParams,
      ) -> types.NestedArray:
          ema_counter = params.ema_counter
          norm_obs = jax.lax.cond(
              ema_counter > 0,
              _normalize,
              lambda o, p: o,
              nested_array,
              params,
          )
          return norm_obs
  
      def _compute_first_moment(x: jnp.ndarray, ndim: int):
          reduce_axes = tuple(range(len(x.shape) - ndim))
          first_moment = jnp.mean(x, axis=reduce_axes)
          return first_moment
  
      def _compute_second_moment(x: jnp.ndarray, ndim: int):
          reduce_axes = tuple(range(len(x.shape) - ndim))
          second_moment = jnp.mean(x**2, axis=reduce_axes)
          return second_moment
  
      def update(
          params: EMAMeanStdNormalizerParams,
          nested_array: types.NestedArray,
          pmap_axis_name: Optional[str] = None,
      ) -> EMAMeanStdNormalizerParams:
          # compute the stats
          first_moment = jax.tree_util.tree_map(
              _compute_first_moment, nested_array, nested_dims
          )
          second_moment = jax.tree_util.tree_map(
              _compute_second_moment, nested_array, nested_dims
          )
  
          # propagate across devices
          if pmap_axis_name is not None:
              first_moment, second_moment = jax.lax.pmean(
                  (first_moment, second_moment),
                  axis_name=pmap_axis_name,
              )
  
          # update running statistics
          new_first_moment = jax.tree_util.tree_map(
              lambda x, y: tau * x
              + (1.0 - tau) * y,  # pylint: disable=g-long-lambda
              params.biased_first_moment,
              first_moment,
          )
          new_second_moment = jax.tree_util.tree_map(
              lambda x, y: tau * x
              + (1.0 - tau) * y,  # pylint: disable=g-long-lambda
              params.biased_second_moment,
              second_moment,
          )
  
          # update ema_counter and return updated params
          new_params = EMAMeanStdNormalizerParams(
              ema_counter=params.ema_counter + 1,
              biased_first_moment=new_first_moment,
              biased_second_moment=new_second_moment,
          )
  
          return new_params
  
      return NormalizationFns(
          init=init,
          normalize=normalize,
          update=update,
      )
  
imports: import dataclasses
  from typing import Any, Callable, Generic, NamedTuple, Optional
  from acme import adders
  from acme import types
  from acme.agents.jax import actor_core
  from acme.agents.jax import actors
  from acme.jax import networks
  from acme.jax import running_statistics
  from acme.jax import utils
  from acme.jax import variable_utils
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: build_ema_mean_std_normalizer(nested_spec, tau, epsilon) and returns: NormalizationFns, EMAMeanStdNormalizerParams, unknown, normalized_nested_array, norm_obs, first_moment, second_moment, new_params