def rlmb_base_stochastic_discrete():
  """Base setting with stochastic discrete model."""
  hparams = rlmb_base()
  hparams.learning_rate_bump = 1.0
  hparams.grayscale = False
  hparams.generative_model = "next_frame_basic_stochastic_discrete"
  hparams.generative_model_params = "next_frame_basic_stochastic_discrete"
  # The parameters below are the same as base, but repeated for easier reading.
  hparams.ppo_epoch_length = 50
  hparams.simulated_rollout_length = 50
  hparams.simulated_batch_size = 16
  return hparams
Base setting with stochastic discrete model.