# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# 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:


from absl.testing import absltest
from absl.testing import parameterized
from envlogger import step_data
from envlogger.backends import schedulers
import dm_env
import numpy as np
# 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.

"""Tests for schedulers."""



def _create_episode(num_transitions: int):
  """Creates an episode with `num_transition` transitions."""
  episode = [step_data.StepData(dm_env.restart(observation=None), action=None)]
  for _ in range(num_transitions):
    episode.append(
        step_data.StepData(
            dm_env.transition(observation=None, reward=None), action=None))
  episode.append(

# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# 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/backends/schedulers_test.py
from absl.testing import absltest
from absl.testing import parameterized
from envlogger import step_data
from envlogger.backends import schedulers
import dm_env
import numpy as np

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

"""Tests for schedulers."""



def _create_episode(num_transitions: int):
  """Creates an episode with `num_transition` transitions."""
  episode = [step_data.StepData(dm_env.restart(observation=None), action=None)]
  for _ in range(num_transitions):
    episode.append(
        step_data.StepData(
            dm_env.transition(observation=None, reward=None), action=None))
  episode.append(

      step_data.StepData(