code: def get_normalization_stats(
      iterator: Iterator[types.Transition],
      num_normalization_batches: int = 50,
  ) -> running_statistics.RunningStatisticsState:
      """Precomputes normalization statistics over a fixed number of batches.
  
      The iterator should contain batches of 3-transitions, i.e. with two leading
      dimensions, the first one denoting the batch dimension and the second one the
      previous, current and next timesteps. The statistics are calculated using the
      data of the previous timestep.
  
      Args:
        iterator: Iterator of batchs of 3-transitions.
        num_normalization_batches: Number of batches to calculate the statistics.
  
      Returns:
        RunningStatisticsState containing the normalization statistics.
      """
      # Set up normalization:
      example = next(iterator)
      unbatched_single_example = jax.tree_map(
          lambda x: x[0, PREVIOUS, :], example
      )
      mean_std = running_statistics.init_state(unbatched_single_example)
  
      for batch in itertools.islice(
          iterator, num_normalization_batches - 1
      ):
          example = jax.tree_map(lambda x: x[:, PREVIOUS, :], batch)
          mean_std = running_statistics.update(mean_std, example)
  
      return mean_std
  
file_path: acme/acme/agents/jax/mbop/dataset.py
imports: import functools
  import itertools
  from typing import Iterator, Optional
  from acme import types
  from acme.jax import running_statistics
  import jax
  import jax.numpy
  import rlds
  import tensorflow
  import tree
label_desc: Use image node_id 7 for calling a global function with example usage: get_normalization_stats(iterator, num_normalization_batches) and returns: mean_std