# 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
# 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.

"""Wrapper that logs `TimeStep`s and actions, metadata and episodes.

Data can be read back using:

- The python reader API (environment_logger.reader).

"""



_DEFAULT_BACKEND = backend_type.BackendType.RIEGELI


class EnvLogger(environment_wrapper.EnvironmentWrapper):
  """Wrapper that logs timestep and actions."""

  def __init__(
      self,
      env: dm_env.Environment,
      step_fn: Optional[Callable[[dm_env.TimeStep, Any, dm_env.Environment],
                                 Any]] = None,
      episode_fn: Optional[Callable[[dm_env.TimeStep, Any, dm_env.Environment],
                                    Any]] = None,
      metadata: Optional[Dict[str, Any]] = None,
      backend: Union[backend_writer.BackendWriter, backend_type.BackendType,
                     Callable[...,

# 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

# 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.

"""Wrapper that logs `TimeStep`s and actions, metadata and episodes.

Data can be read back using:

- The python reader API (environment_logger.reader).

"""



_DEFAULT_BACKEND = backend_type.BackendType.RIEGELI


class EnvLogger(environment_wrapper.EnvironmentWrapper):
  """Wrapper that logs timestep and actions."""

  def __init__(
      self,
      env: dm_env.Environment,
      step_fn: Optional[Callable[[dm_env.TimeStep, Any, dm_env.Environment],
                                 Any]] = None,
      episode_fn: Optional[Callable[[dm_env.TimeStep, Any, dm_env.Environment],
                                    Any]] = None,
      metadata: Optional[Dict[str, Any]] = None,
      backend: Union[backend_writer.BackendWriter, backend_type.BackendType,
                     Callable[...,

                              backend_writer.BackendWriter]] = _DEFAULT_BACKEND,