Gin files for env_server and ppo binaries for learning rate tuning.

PiperOrigin-RevId: 261588529

Author: afrozenator

tensor2tensor/trax/rlax/configs/lr_learning_env_params_online_tune_wide_resnet_cifar10.gin

@@ -0,0 +1,41 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.rlax
+import tensor2tensor.trax.rlax.envs
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size = 32
+batch_fun.bucket_length = 32
+batch_fun.buckets = None
+batch_fun.eval_batch_size = 32
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 'cifar10'
+
+# Parameters for Momentum:
+# ==============================================================================
+Momentum.mass = 0.9
+
+# Parameters for shuffle_and_batch_data:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun = @trax.inputs.cifar10_no_augmentation_preprocess
+
+# Parameters for WideResnet:
+# ==============================================================================
+WideResnet.widen_factor = 2
+WideResnet.n_blocks = 3
+WideResnet.n_output_classes = 10
+
+# Parameters for OnlineTuneEnv:
+# ==============================================================================
+OnlineTuneEnv.inputs = @trax.inputs.inputs
+OnlineTuneEnv.model = @trax.models.WideResnet
+OnlineTuneEnv.optimizer = @trax.optimizers.Momentum
+OnlineTuneEnv.start_lr = 0.01
+OnlineTuneEnv.train_steps = 500
+OnlineTuneEnv.eval_steps = 50
+OnlineTuneEnv.env_steps = 128

tensor2tensor/trax/rlax/configs/lr_learning_ppo_params_online_tune_wide_resnet_cifar10.gin

@@ -0,0 +1,21 @@
+import tensor2tensor.trax.rlax
+
+# Parameters for ppo.training_loop:
+# ==============================================================================
+ppo.training_loop.n_optimizer_steps = 30
+ppo.training_loop.boundary = 128
+ppo.training_loop.max_timestep = 128
+ppo.training_loop.max_timestep_eval = 128
+ppo.training_loop.random_seed = 0
+ppo.training_loop.gamma = 0.99
+ppo.training_loop.lambda_ = 0.95
+ppo.training_loop.epsilon = 0.1
+ppo.training_loop.c1 = 1.0
+ppo.training_loop.c2 = 0.01
+ppo.training_loop.eval_every_n = 10
+ppo.training_loop.done_frac_for_policy_save = 0
+ppo.training_loop.enable_early_stopping = True
+ppo.training_loop.n_evals = 1
+ppo.training_loop.len_history_for_policy = 1  # this needs to be bumped up.
+ppo.training_loop.eval_temperatures = (1.0,)
+ppo.training_loop.epochs = 1000