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 = 100
OnlineTuneEnv.env_steps = 100

# Parameters for ppo.training_loop:
# ==============================================================================
ppo.training_loop.n_optimizer_steps = 30
ppo.training_loop.boundary = 20
ppo.training_loop.max_timestep = 128
ppo.training_loop.max_timestep_eval = 20000
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
ppo.training_loop.eval_temperatures = (1.0,)
