def test_train_mode_reset(self):
"""Test reset() in train mode."""
train_env = Env(mode="train")
x, x_, theta, theta_ = train_env.reset()
assert abs(x) <= 2.3
assert x_ == 0
assert abs(theta) <= 0.3
assert theta_ == 0
def test_eval_mode_reset(self):
"""Test reset() in eval mode."""
eval_env = Env(mode="eval")
x, x_, theta, theta_ = eval_env.reset()
assert abs(x) <= 1.0
assert x_ == 0
assert abs(theta) <= 0.3
assert theta_ == 0