# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/backend_reader.py
# T = TypeVar('T')
# class _SequenceAdapter(Generic[T], Sequence[T]):
# class BackendReader(metaclass=abc.ABCMeta):
#   def __init__(self, count: int, get_nth_item: Callable[[int], T]):
#   def __getitem__(self, index: Union[int, slice]) -> Union[T, List[T]]:
#   def __len__(self) -> int:
#   def __iter__(self) -> Iterator[T]:
#   def __next__(self) -> T:
#   def __init__(self):
#   def copy(self) -> 'BackendReader':
#   def _copy(self) -> 'BackendReader':
#   def _init_visitors(self):
#   def _get_nth_step(self, i: int) -> step_data.StepData:
#   def _get_num_steps(self) -> int:
#   def _get_num_episodes(self) -> int:
#   def _get_nth_episode_info(self,
#                             i: int,
#                             include_metadata: bool = False
#                            ) -> episode_info.EpisodeInfo:
#   def _get_nth_episode(self, i: int) -> Sequence[step_data.StepData]:
#     def get_nth_step_from_episode(j: int):
#   def _get_nth_episode_metadata(self, i: int) -> Optional[Any]:
#   def __enter__(self):
#   def __exit__(self, exc_type, exc_value, tb):
#   def __del__(self):
#   def close(self) -> None:
#   def metadata(self) -> Dict[str, Any]:
#   def episodes(self) -> Sequence[Sequence[step_data.StepData]]:
#   def episode_metadata(self) -> Sequence[Optional[Any]]:
#   def steps(self) -> Sequence[step_data.StepData]:
# 
# Path: envlogger/backends/backend_type.py
# class BackendType(enum.IntEnum):
#   RIEGELI = 0
#   IN_MEMORY = 1
# 
# 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_reader.py
# class RiegeliBackendReader(backend_reader.BackendReader):
#   def __init__(self, data_directory: str):
#   def _copy(self) -> 'RiegeliBackendReader':
#   def close(self):
#   def _decode_step_data(self, data: Tuple[Any, Any, Any]) -> step_data.StepData:
#   def _get_num_steps(self):
#   def _get_num_episodes(self):
#   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 metadata(self):
# 
# 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]]:


import copy
from typing import Any, Optional, Sequence, Union
from absl import logging
from envlogger import step_data
from envlogger.backends import backend_reader
from envlogger.backends import backend_type
from envlogger.backends import in_memory_backend
from envlogger.backends import riegeli_backend_reader
from envlogger.converters import spec_codec
# coding=utf-8
# Copyright 2022 DeepMind Technologies Limited..
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Reader of EnvironmentLogger data."""




class Reader:
  """Reader of trajectories generated by EnvLogger."""

  def __init__(self,
               *backend_args,
               backend: Union[
                   backend_reader.BackendReader,

# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/backend_reader.py
# T = TypeVar('T')
# class _SequenceAdapter(Generic[T], Sequence[T]):
# class BackendReader(metaclass=abc.ABCMeta):
#   def __init__(self, count: int, get_nth_item: Callable[[int], T]):
#   def __getitem__(self, index: Union[int, slice]) -> Union[T, List[T]]:
#   def __len__(self) -> int:
#   def __iter__(self) -> Iterator[T]:
#   def __next__(self) -> T:
#   def __init__(self):
#   def copy(self) -> 'BackendReader':
#   def _copy(self) -> 'BackendReader':
#   def _init_visitors(self):
#   def _get_nth_step(self, i: int) -> step_data.StepData:
#   def _get_num_steps(self) -> int:
#   def _get_num_episodes(self) -> int:
#   def _get_nth_episode_info(self,
#                             i: int,
#                             include_metadata: bool = False
#                            ) -> episode_info.EpisodeInfo:
#   def _get_nth_episode(self, i: int) -> Sequence[step_data.StepData]:
#     def get_nth_step_from_episode(j: int):
#   def _get_nth_episode_metadata(self, i: int) -> Optional[Any]:
#   def __enter__(self):
#   def __exit__(self, exc_type, exc_value, tb):
#   def __del__(self):
#   def close(self) -> None:
#   def metadata(self) -> Dict[str, Any]:
#   def episodes(self) -> Sequence[Sequence[step_data.StepData]]:
#   def episode_metadata(self) -> Sequence[Optional[Any]]:
#   def steps(self) -> Sequence[step_data.StepData]:
# 
# Path: envlogger/backends/backend_type.py
# class BackendType(enum.IntEnum):
#   RIEGELI = 0
#   IN_MEMORY = 1
# 
# 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_reader.py
# class RiegeliBackendReader(backend_reader.BackendReader):
#   def __init__(self, data_directory: str):
#   def _copy(self) -> 'RiegeliBackendReader':
#   def close(self):
#   def _decode_step_data(self, data: Tuple[Any, Any, Any]) -> step_data.StepData:
#   def _get_num_steps(self):
#   def _get_num_episodes(self):
#   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 metadata(self):
# 
# 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/reader.py
import copy
from typing import Any, Optional, Sequence, Union
from absl import logging
from envlogger import step_data
from envlogger.backends import backend_reader
from envlogger.backends import backend_type
from envlogger.backends import in_memory_backend
from envlogger.backends import riegeli_backend_reader
from envlogger.converters import spec_codec

# coding=utf-8
# Copyright 2022 DeepMind Technologies Limited..
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Reader of EnvironmentLogger data."""




class Reader:
  """Reader of trajectories generated by EnvLogger."""

  def __init__(self,
               *backend_args,
               backend: Union[
                   backend_reader.BackendReader,

                   backend_type.BackendType] = backend_type.BackendType.RIEGELI,