fixing envs specs.

PiperOrigin-RevId: 209242641

Author: blazejosinski

tensor2tensor/data_generators/gym_problems.py

@@ -54,7 +54,7 @@ def standard_atari_env_spec(env):
     env_lambda = lambda: gym.make(env)
   if callable(env):
     env_lambda = env
-  assert env is not None, "Unknown specification of environment"
+  assert env_lambda is not None, "Unknown specification of environment"
 
   return tf.contrib.training.HParams(
       env_lambda=env_lambda, wrappers=standard_wrappers, simulated_env=False)

tensor2tensor/models/research/rl.py

@@ -130,7 +130,7 @@ def simple_gym_spec(env):
     env_lambda = lambda: gym.make(env)
   if callable(env):
     env_lambda = env
-  assert env is not None, "Unknown specification of environment"
+  assert env_lambda is not None, "Unknown specification of environment"
 
   return tf.contrib.training.HParams(env_lambda=env_lambda,
                                      wrappers=standard_wrappers,

tensor2tensor/rl/rl_trainer_lib_test.py

@@ -45,7 +45,7 @@ def test_no_crash_cartpole(self):
         TrainTest.test_config)
 
     hparams.add_hparam(
-        "environment_spec", gym_problems.standard_atari_env_spec("CartPole-v0"))
+        "environment_spec", rl_models.simple_gym_spec("CartPole-v0"))
     rl_trainer_lib.train(hparams)
 
   # This test should successfully train pong.