internal merge of PR #1365

PiperOrigin-RevId: 228918597

Author: koz4k

tensor2tensor/rl/envs/simulated_batch_gym_env.py

@@ -20,9 +20,11 @@
 from __future__ import print_function
 
 from gym import Env
-from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
 
 import numpy as np
+
+from tensor2tensor.rl.envs.simulated_batch_env import SimulatedBatchEnv
+
 import tensorflow as tf
 
 

tensor2tensor/rl/evaluator.py

@@ -32,8 +32,8 @@
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.rl import trainer_model_based_params  # pylint: disable=unused-import
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
-from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import registry
+from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
 

tensor2tensor/rl/rl_utils.py

@@ -377,6 +377,7 @@ def __init__(
 
   def act(self, observations):
     def run_batch_from(observation, action):
+      """Run a batch of actions."""
       self._sim_env.initial_frames = np.array(
           [observation] * self._sim_env.batch_size
       )