def test_correct_piece(self, batch_size):
env = pbe.PiecewiseBernoulliPyEnvironment(
[[0.1, 0.2, 0.3], [0.3, 0.2, 0.1], [0.1, 0.12, 0.14]],
self.deterministic_duration_generator(), batch_size)
for t in range(100):
env.reset()
self.assertEqual(int(t / 10) % 3, env._current_piece)
_ = env.step([0])
def test_out_of_bound_parameter(self):
with self.assertRaisesRegexp(
ValueError, r'All parameters should be floats in \[0, 1\]\.'):
pbe.PiecewiseBernoulliPyEnvironment(
[[0.1, 1.2, 0.3]], self.deterministic_duration_generator())