from typing import Any
from typing import Dict


def _get_step_metadata(step: Any) -> Dict[str, Any]:
  """Extracts the custom metadata from the step.

  Note that it may return the dictionary contained in the step, so further
  modifications will affect the step itself.

  Args:
    step: step data (envlogger.StepData).

  Returns:
    Dictionary containing the step metadata.
  """
  if not isinstance(step.custom_data, dict):
    # We ignore step metadata if step.custom_data is not a dictionary
    return {}
  if 'episode_metadata' in step.custom_data:
    # Some datasets used an old version of EnvLogger that stores episode
    # metadata in the first step.
    return {
        k: v for k, v in step.custom_data.items() if k != 'episode_metadata'
    }
  else:
    return step.custom_data