# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/backend_writer.py
# class BackendWriter(metaclass=abc.ABCMeta):
#   def __init__(self,
#                metadata: Optional[Dict[str, Any]] = None,
#                scheduler: Optional[schedulers.Scheduler] = None):
#   def record_step(self, data: step_data.StepData, is_new_episode: bool) -> None:
#   def set_episode_metadata(self, data: Any) -> None:
#   def _record_step(self, data: step_data.StepData,
#                    is_new_episode: bool) -> None:
#   def close(self) -> None:
#   def __del__(self):
#   def metadata(self):
# 
# 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/converters/codec.py
# _SCALAR_DIM_SIZE = -438
# def _python_int_to_bytes(py_int: int) -> bytes:
# def _set_datum_values_from_scalar(scalar: Union[ScalarNumber, bool, str, bytes],
#                                   datum: storage_pb2.Datum) -> bool:
# def _set_datum_values_from_array(array: np.ndarray,
#                                  values: storage_pb2.Datum.Values) -> None:
# def encode(
#     user_data: Union[np.ndarray, List[Any], Tuple[Any, ...], Dict[str, Any]]
# ) -> storage_pb2.Data:
# def decode_datum(
#     datum: storage_pb2.Datum
# ) -> Union[np.ndarray, ScalarNumber, bool, str, bytes]:
#     def from_bigint(int_bytes):
# def decode(
#     user_data: storage_pb2.Data
# ) -> Optional[Union[ScalarNumber, bool, str, bytes, np.ndarray, List[Any],


from typing import Any, Optional
from absl import logging
from envlogger import step_data
from envlogger.backends import backend_writer
from envlogger.backends import schedulers
from envlogger.backends.python import riegeli_dataset_writer
from envlogger.converters import codec
from pybind11_abseil import status
  def __init__(
      self,
      data_directory: str,
      max_episodes_per_file: int = 10000,
      writer_options: str = 'transpose,brotli:6,chunk_size:1M',
      flush_scheduler: Optional[schedulers.Scheduler] = None,
      **base_kwargs,
  ):
    """Constructor.

    Calling `close()` will flush the trajectories and the index to disk and will
    ensure that they can be read later on. If it isn't called, there is a large
    risk of losing data. This is particularly common in some RL frameworks that
    do not clean up their environments. If the environment runs for a very long
    time, this can happen only to the last shard, but if the instance is
    short-lived, then a large portion of the trajectories can disappear.

    Args:
      data_directory: Destination for the episode data.
      max_episodes_per_file: maximum number of episodes stored in one file.
      writer_options: Comma-separated list of options that are passed to the
        Riegeli RecordWriter as is.
      flush_scheduler: This controls when data is flushed to permanent storage.
        If `None`, it defaults to a step-wise Bernoulli scheduler with 1/5000
        chances of flushing.
      **base_kwargs: arguments for the base class.
    """
    super().__init__(**base_kwargs)
    self._data_directory = data_directory
    if flush_scheduler is None:

# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/backend_writer.py
# class BackendWriter(metaclass=abc.ABCMeta):
#   def __init__(self,
#                metadata: Optional[Dict[str, Any]] = None,
#                scheduler: Optional[schedulers.Scheduler] = None):
#   def record_step(self, data: step_data.StepData, is_new_episode: bool) -> None:
#   def set_episode_metadata(self, data: Any) -> None:
#   def _record_step(self, data: step_data.StepData,
#                    is_new_episode: bool) -> None:
#   def close(self) -> None:
#   def __del__(self):
#   def metadata(self):
# 
# 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/converters/codec.py
# _SCALAR_DIM_SIZE = -438
# def _python_int_to_bytes(py_int: int) -> bytes:
# def _set_datum_values_from_scalar(scalar: Union[ScalarNumber, bool, str, bytes],
#                                   datum: storage_pb2.Datum) -> bool:
# def _set_datum_values_from_array(array: np.ndarray,
#                                  values: storage_pb2.Datum.Values) -> None:
# def encode(
#     user_data: Union[np.ndarray, List[Any], Tuple[Any, ...], Dict[str, Any]]
# ) -> storage_pb2.Data:
# def decode_datum(
#     datum: storage_pb2.Datum
# ) -> Union[np.ndarray, ScalarNumber, bool, str, bytes]:
#     def from_bigint(int_bytes):
# def decode(
#     user_data: storage_pb2.Data
# ) -> Optional[Union[ScalarNumber, bool, str, bytes, np.ndarray, List[Any],


# Path: envlogger/backends/riegeli_backend_writer.py
from typing import Any, Optional
from absl import logging
from envlogger import step_data
from envlogger.backends import backend_writer
from envlogger.backends import schedulers
from envlogger.backends.python import riegeli_dataset_writer
from envlogger.converters import codec
from pybind11_abseil import status

  def __init__(
      self,
      data_directory: str,
      max_episodes_per_file: int = 10000,
      writer_options: str = 'transpose,brotli:6,chunk_size:1M',
      flush_scheduler: Optional[schedulers.Scheduler] = None,
      **base_kwargs,
  ):
    """Constructor.

    Calling `close()` will flush the trajectories and the index to disk and will
    ensure that they can be read later on. If it isn't called, there is a large
    risk of losing data. This is particularly common in some RL frameworks that
    do not clean up their environments. If the environment runs for a very long
    time, this can happen only to the last shard, but if the instance is
    short-lived, then a large portion of the trajectories can disappear.

    Args:
      data_directory: Destination for the episode data.
      max_episodes_per_file: maximum number of episodes stored in one file.
      writer_options: Comma-separated list of options that are passed to the
        Riegeli RecordWriter as is.
      flush_scheduler: This controls when data is flushed to permanent storage.
        If `None`, it defaults to a step-wise Bernoulli scheduler with 1/5000
        chances of flushing.
      **base_kwargs: arguments for the base class.
    """
    super().__init__(**base_kwargs)
    self._data_directory = data_directory
    if flush_scheduler is None:

      self._flush_scheduler = schedulers.bernoulli_step_scheduler(