# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/rlds_utils.py
# def to_rlds_step(prev_step: step_data.StepData,
#                  step: Optional[step_data.StepData]) -> Step:
# def _find_extra_shard(split_info: tfds.core.SplitInfo) -> Optional[Any]:
# def maybe_recover_last_shard(builder: tfds.core.DatasetBuilder):
# 
# Path: envlogger/backends/tfds_backend_testlib.py
# def generate_episode_data(
#     backend: backend_writer.BackendWriter,
#     num_episodes: int = 2,
# ) -> List[List[step_data.StepData]]:
# def catch_env_tfds_config(
#     name: str = 'catch_example') -> tfds.rlds.rlds_base.DatasetConfig:
# def tfds_backend_catch_env(
#     data_directory: str,
#     max_episodes_per_file: int = 1,
#     split_name: Optional[str] = None) -> tfds_backend_writer.TFDSBackendWriter:


from absl.testing import absltest
from envlogger import step_data
from envlogger.backends import rlds_utils
from envlogger.backends import tfds_backend_testlib
import dm_env
import tensorflow_datasets as tfds
        'is_terminal': False,
        'is_first': True,
        'is_last': False,
    }

    rlds_step = rlds_utils.to_rlds_step(prev_step, step)

    self.assertEqual(rlds_step, expected_step)

  def test_build_last_step(self):
    prev_step = step_data.StepData(
        timestep=dm_env.TimeStep(
            step_type=dm_env.StepType.LAST, reward=1, observation=2,
            discount=1),
        action=4)
    expected_step = {
        'observation': 2,
        'action': 0,
        'reward': 0,
        'discount': 0,
        'is_terminal': False,
        'is_first': False,
        'is_last': True,
    }

    rlds_step = rlds_utils.to_rlds_step(prev_step, None)

    self.assertEqual(rlds_step, expected_step)

  def test_build_terminal_step(self):

# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/rlds_utils.py
# def to_rlds_step(prev_step: step_data.StepData,
#                  step: Optional[step_data.StepData]) -> Step:
# def _find_extra_shard(split_info: tfds.core.SplitInfo) -> Optional[Any]:
# def maybe_recover_last_shard(builder: tfds.core.DatasetBuilder):
# 
# Path: envlogger/backends/tfds_backend_testlib.py
# def generate_episode_data(
#     backend: backend_writer.BackendWriter,
#     num_episodes: int = 2,
# ) -> List[List[step_data.StepData]]:
# def catch_env_tfds_config(
#     name: str = 'catch_example') -> tfds.rlds.rlds_base.DatasetConfig:
# def tfds_backend_catch_env(
#     data_directory: str,
#     max_episodes_per_file: int = 1,
#     split_name: Optional[str] = None) -> tfds_backend_writer.TFDSBackendWriter:


# Path: envlogger/backends/rlds_utils_test.py
from absl.testing import absltest
from envlogger import step_data
from envlogger.backends import rlds_utils
from envlogger.backends import tfds_backend_testlib
import dm_env
import tensorflow_datasets as tfds

        'is_terminal': False,
        'is_first': True,
        'is_last': False,
    }

    rlds_step = rlds_utils.to_rlds_step(prev_step, step)

    self.assertEqual(rlds_step, expected_step)

  def test_build_last_step(self):
    prev_step = step_data.StepData(
        timestep=dm_env.TimeStep(
            step_type=dm_env.StepType.LAST, reward=1, observation=2,
            discount=1),
        action=4)
    expected_step = {
        'observation': 2,
        'action': 0,
        'reward': 0,
        'discount': 0,
        'is_terminal': False,
        'is_first': False,
        'is_last': True,
    }

    rlds_step = rlds_utils.to_rlds_step(prev_step, None)

    self.assertEqual(rlds_step, expected_step)

  def test_build_terminal_step(self):

    prev_step = step_data.StepData(