Add a unit test for default batch* properties.

tensorflow · Apr 9, 2019 · dd0da02 · dd0da02
1 parent 545d2a7
commit dd0da02
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/tf_agents/environments/wrappers_test.py b/tf_agents/environments/wrappers_test.py
@@ -51,6 +51,12 @@ def test_batch_properties(self, batch_size):
     env = wrappers.PyEnvironmentBaseWrapper(nested_env)
     self.assertEqual(env.batched, nested_env.batched)
     self.assertEqual(nested_env.batch_size, env.batch_size)
+
+  def test_default_batch_properties(self):
+    cartpole_env = gym.spec('CartPole-v1').make()
+    env = gym_wrapper.GymWrapper(cartpole_env)
+    self.assertFalse(env.batched)
+    self.assertEqual(env.batch_size, None)
 
 
 class TimeLimitWrapperTest(absltest.TestCase):