# Path: envlogger/environment_wrapper.py
# class EnvironmentWrapper(dm_env.Environment):
#   def __init__(self, environment: dm_env.Environment):
#   def __getattr__(self, name):
#   def step(self, action) -> dm_env.TimeStep:
#   def reset(self) -> dm_env.TimeStep:
#   def action_spec(self):
#   def discount_spec(self):
#   def observation_spec(self):
#   def reward_spec(self):
#   def close(self):
# 
# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/backend_type.py
# class BackendType(enum.IntEnum):
#   RIEGELI = 0
#   IN_MEMORY = 1
# 
# 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/in_memory_backend.py
# class InMemoryBackendWriter(backend_writer.BackendWriter):
# class InMemoryBackendReader(backend_reader.BackendReader):
#   def __init__(self, **base_kwargs):
#   def _record_step(self, data: step_data.StepData,
#                    is_new_episode: bool) -> None:
#   def set_episode_metadata(self, data: Any) -> None:
#   def close(self) -> None:
#   def __init__(self, in_memory_backend_writer: InMemoryBackendWriter):
#   def _copy(self) -> 'InMemoryBackendReader':
#   def close(self) -> None:
#   def _get_nth_step(self, i: int) -> step_data.StepData:
#   def _get_nth_episode_info(self,
#                             i: int,
#                             include_metadata: bool = False
#                            ) -> episode_info.EpisodeInfo:
#   def _get_num_steps(self) -> int:
#   def _get_num_episodes(self) -> int:
#   def metadata(self):
# 
# Path: envlogger/backends/riegeli_backend_writer.py
# class RiegeliBackendWriter(backend_writer.BackendWriter):
#   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,
#   ):
#   def _record_step(self, data: step_data.StepData,
#                    is_new_episode: bool) -> None:
#   def set_episode_metadata(self, data: Any) -> None:
#   def close(self) -> None:
# 
# Path: envlogger/converters/spec_codec.py
# _ENVIRONMENT_SPEC_NAMES = [
#     'observation_spec',
#     'action_spec',
#     'reward_spec',
#     'discount_spec',
# ]
# def encode_environment_specs(
#     env: Optional[dm_env.Environment]) -> Dict[str, Any]:
# def decode_environment_specs(
#     encoded_specs: Dict[str, Any]) -> Dict[str, Optional[specs.Array]]:
# def _array_spec_to_dict(array_spec: specs.Array) -> Dict[str, Any]:
# def encode(
#     spec: Union[specs.Array, List[Any], Tuple[Any], Dict[str, Any]]
# ) -> Union[List[Any], Tuple[Any], Dict[str, Any]]:
# def decode(
#     spec: Union[List[Any], Tuple[Any], Dict[str, Any]]
# ) -> Union[specs.Array, List[Any], Tuple[Any], Dict[str, Any]]:


from typing import Any, Callable, Dict, Optional, Union
from envlogger import environment_wrapper
from envlogger import step_data
from envlogger.backends import backend_type
from envlogger.backends import backend_writer
from envlogger.backends import in_memory_backend
from envlogger.backends import riegeli_backend_writer
from envlogger.converters import spec_codec
import dm_env
        the environment itself and returns custom episodic data that's written
        when the current episode is over. If it is None or if it returns None,
        nothing is written.  This function is called at every step during the
        course of an episode, but only the last value it returns will actually
        be stored (all intermediate return values are ignored).
      metadata: Any dataset-level custom data to be written.
      backend: One of the following:
        * A `LoggingBackend` instance: `EnvLogger` will simply use this instance
          as is.
        * A `BackendType` enum indicating the backend to use: `EnvLogger` will
          construct a `LoggingBackend` from a list of predefined backends
          passing `backend_kwargs`.
        * A `Callable`: `EnvLogger` will call the given function passing
          `backend_kwargs`. The function _must_ return a `LoggingBackend`
          instance.
      **backend_kwargs: Extra arguments use to construct the backend. These will
        be handed to `backend` without any modification.
    """
    super().__init__(env)

    self._step_fn = step_fn
    self._episode_fn = episode_fn
    self._reset_next_step = True
    metadata = metadata or {}
    metadata['environment_specs'] = spec_codec.encode_environment_specs(env)
    backend_kwargs['metadata'] = metadata

    # Set backend.
    if isinstance(backend, backend_writer.BackendWriter):
      self._backend = backend

# Path: envlogger/environment_wrapper.py
# class EnvironmentWrapper(dm_env.Environment):
#   def __init__(self, environment: dm_env.Environment):
#   def __getattr__(self, name):
#   def step(self, action) -> dm_env.TimeStep:
#   def reset(self) -> dm_env.TimeStep:
#   def action_spec(self):
#   def discount_spec(self):
#   def observation_spec(self):
#   def reward_spec(self):
#   def close(self):
# 
# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/backend_type.py
# class BackendType(enum.IntEnum):
#   RIEGELI = 0
#   IN_MEMORY = 1
# 
# 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/in_memory_backend.py
# class InMemoryBackendWriter(backend_writer.BackendWriter):
# class InMemoryBackendReader(backend_reader.BackendReader):
#   def __init__(self, **base_kwargs):
#   def _record_step(self, data: step_data.StepData,
#                    is_new_episode: bool) -> None:
#   def set_episode_metadata(self, data: Any) -> None:
#   def close(self) -> None:
#   def __init__(self, in_memory_backend_writer: InMemoryBackendWriter):
#   def _copy(self) -> 'InMemoryBackendReader':
#   def close(self) -> None:
#   def _get_nth_step(self, i: int) -> step_data.StepData:
#   def _get_nth_episode_info(self,
#                             i: int,
#                             include_metadata: bool = False
#                            ) -> episode_info.EpisodeInfo:
#   def _get_num_steps(self) -> int:
#   def _get_num_episodes(self) -> int:
#   def metadata(self):
# 
# Path: envlogger/backends/riegeli_backend_writer.py
# class RiegeliBackendWriter(backend_writer.BackendWriter):
#   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,
#   ):
#   def _record_step(self, data: step_data.StepData,
#                    is_new_episode: bool) -> None:
#   def set_episode_metadata(self, data: Any) -> None:
#   def close(self) -> None:
# 
# Path: envlogger/converters/spec_codec.py
# _ENVIRONMENT_SPEC_NAMES = [
#     'observation_spec',
#     'action_spec',
#     'reward_spec',
#     'discount_spec',
# ]
# def encode_environment_specs(
#     env: Optional[dm_env.Environment]) -> Dict[str, Any]:
# def decode_environment_specs(
#     encoded_specs: Dict[str, Any]) -> Dict[str, Optional[specs.Array]]:
# def _array_spec_to_dict(array_spec: specs.Array) -> Dict[str, Any]:
# def encode(
#     spec: Union[specs.Array, List[Any], Tuple[Any], Dict[str, Any]]
# ) -> Union[List[Any], Tuple[Any], Dict[str, Any]]:
# def decode(
#     spec: Union[List[Any], Tuple[Any], Dict[str, Any]]
# ) -> Union[specs.Array, List[Any], Tuple[Any], Dict[str, Any]]:


# Path: envlogger/environment_logger.py
from typing import Any, Callable, Dict, Optional, Union
from envlogger import environment_wrapper
from envlogger import step_data
from envlogger.backends import backend_type
from envlogger.backends import backend_writer
from envlogger.backends import in_memory_backend
from envlogger.backends import riegeli_backend_writer
from envlogger.converters import spec_codec
import dm_env

        the environment itself and returns custom episodic data that's written
        when the current episode is over. If it is None or if it returns None,
        nothing is written.  This function is called at every step during the
        course of an episode, but only the last value it returns will actually
        be stored (all intermediate return values are ignored).
      metadata: Any dataset-level custom data to be written.
      backend: One of the following:
        * A `LoggingBackend` instance: `EnvLogger` will simply use this instance
          as is.
        * A `BackendType` enum indicating the backend to use: `EnvLogger` will
          construct a `LoggingBackend` from a list of predefined backends
          passing `backend_kwargs`.
        * A `Callable`: `EnvLogger` will call the given function passing
          `backend_kwargs`. The function _must_ return a `LoggingBackend`
          instance.
      **backend_kwargs: Extra arguments use to construct the backend. These will
        be handed to `backend` without any modification.
    """
    super().__init__(env)

    self._step_fn = step_fn
    self._episode_fn = episode_fn
    self._reset_next_step = True
    metadata = metadata or {}
    metadata['environment_specs'] = spec_codec.encode_environment_specs(env)
    backend_kwargs['metadata'] = metadata

    # Set backend.
    if isinstance(backend, backend_writer.BackendWriter):
      self._backend = backend

    elif isinstance(backend, backend_type.BackendType):