internal merge of PR #1391

PiperOrigin-RevId: 230279058

tensor2tensor/models/research/rl.py

@@ -21,7 +21,6 @@
 import gym
 import six
 
-from tensor2tensor.data_generators import gym_env
 from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import video_utils
 from tensor2tensor.layers import common_hparams
@@ -379,27 +378,22 @@ def rlmf_base():
 
 
 @registry.register_hparams
-def rlmf_final_eval():
-  """Base set of hparams for model-free PPO."""
+def rlmf_eval():
+  """Eval set of hparams for model-free PPO."""
   hparams = rlmf_original()
   hparams.batch_size = 8
-  hparams.eval_sampling_temps=[0.0, 1.0]
+  hparams.eval_sampling_temps = [0.0, 0.5, 1.0]
   hparams.eval_rl_env_max_episode_steps = -1
   hparams.add_hparam("ppo_epoch_length", 128)
   hparams.add_hparam("ppo_optimization_batch_size", 32)
   hparams.add_hparam("ppo_epochs_num", 10000)
   hparams.add_hparam("ppo_eval_every_epochs", 500)
+  hparams.add_hparam("ppo_eval_every_epochs", 500)
   hparams.add_hparam("attempt", 0)
+  hparams.add_hparam("moe_loss_coef", 0)
   return hparams
 
 
-@registry.register_ranged_hparams
-def rlmf_human_score_games(rhp):
-  rhp.set_categorical("game",
-                      gym_env.ATARI_GAMES_WITH_HUMAN_SCORE_NICE)
-  rhp.set_discrete("attempt", list(range(5)))
-
-
 @registry.register_hparams
 def rlmf_tiny():
   hparams = rlmf_base()

tensor2tensor/rl/trainer_model_free.py

@@ -27,8 +27,8 @@
 from __future__ import division
 from __future__ import print_function
 
-import pprint
 import os
+import pprint
 
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_utils