# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/schedulers.py
# def n_step_scheduler(step_interval: int) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def bernoulli_step_scheduler(keep_probability: float) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def n_episode_scheduler(episode_interval: int) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:
# def bernoulli_episode_scheduler(keep_probability: float) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:
# def list_step_scheduler(
#     desired_steps: Union[List[int], np.ndarray]) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def list_episode_scheduler(
#     desired_episodes: Union[List[int], np.ndarray]) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:


import abc
from typing import Any, Dict, Optional
from envlogger import step_data
from envlogger.backends import schedulers
  """Abstract trajectory logging interface."""

  def __init__(self,
               metadata: Optional[Dict[str, Any]] = None,
               scheduler: Optional[schedulers.Scheduler] = None):
    """BackendWriter base class.

    Args:
      metadata: Any dataset-level custom data to be written.
      scheduler: A callable that takes the current timestep, current
        action, the environment itself and returns True if the current step
        should be logged, False otherwise. This function is called _before_
        `step_fn`, meaning that if it returns False, `step_fn` will not be
        called at all. NOTE: This scheduler should NEVER skip the first timestep
        in the episode, otherwise `EnvLogger` will not know that such episode
        really exists.
    """
    self._scheduler = scheduler
    self._metadata = metadata

  def record_step(self, data: step_data.StepData, is_new_episode: bool) -> None:
    if (self._scheduler is not None and not self._scheduler(data)):
      return
    self._record_step(data, is_new_episode)

  @abc.abstractmethod
  def set_episode_metadata(self, data: Any) -> None:
    pass

  @abc.abstractmethod

# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/schedulers.py
# def n_step_scheduler(step_interval: int) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def bernoulli_step_scheduler(keep_probability: float) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def n_episode_scheduler(episode_interval: int) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:
# def bernoulli_episode_scheduler(keep_probability: float) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:
# def list_step_scheduler(
#     desired_steps: Union[List[int], np.ndarray]) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def list_episode_scheduler(
#     desired_episodes: Union[List[int], np.ndarray]) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:


# Path: envlogger/backends/backend_writer.py
import abc
from typing import Any, Dict, Optional
from envlogger import step_data
from envlogger.backends import schedulers

  """Abstract trajectory logging interface."""

  def __init__(self,
               metadata: Optional[Dict[str, Any]] = None,
               scheduler: Optional[schedulers.Scheduler] = None):
    """BackendWriter base class.

    Args:
      metadata: Any dataset-level custom data to be written.
      scheduler: A callable that takes the current timestep, current
        action, the environment itself and returns True if the current step
        should be logged, False otherwise. This function is called _before_
        `step_fn`, meaning that if it returns False, `step_fn` will not be
        called at all. NOTE: This scheduler should NEVER skip the first timestep
        in the episode, otherwise `EnvLogger` will not know that such episode
        really exists.
    """
    self._scheduler = scheduler
    self._metadata = metadata

  def record_step(self, data: step_data.StepData, is_new_episode: bool) -> None:
    if (self._scheduler is not None and not self._scheduler(data)):
      return
    self._record_step(data, is_new_episode)

  @abc.abstractmethod
  def set_episode_metadata(self, data: Any) -> None:
    pass

  @abc.abstractmethod

  def _record_step(self, data: step_data.StepData,