from typing import Any
from typing import AbstractSet
from typing import Dict


def get_step_metadata(
    step: Any,
    step_metadata_skip_list: AbstractSet[str] = frozenset()
) -> Dict[str, Any]:
  """Extracts the custom metadata from the step.

  Args:
    step: step data (envlogger.StepData).
    step_metadata_skip_list: Set of metadata keys to filter.

  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 {}
  return {
      k: v
      for k, v in step.custom_data.items()
      if k != 'episode_metadata' and k not in step_metadata_skip_list
  }