Eval set of hparams for model-free PPO.
9309
def rlmf_eval():
  """"""
  hparams = rlmf_original()
  hparams.batch_size = 8
  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("attempt", 0)
  hparams.add_hparam("moe_loss_coef", 0)
  return hparams
434337
function hingeLoss(p, weights) {
  const [label, features] = p;
  const grad = [];
  let dotProd = 0;

  for (let i = 0; i < features.length; i++)
    dotProd += (features[i] || 0) * (weights[i] || 0);

  if (label * dotProd < 1)
    for (let i = 0; i < features.length; i++)
      grad[i] = -label * (features[i] || 0);
  else
    for (let i = 0; i < features.length; i++)
      grad[i] = 0;

  return grad;
}
132845
def next_frame_sv2p():
  """"""
  hparams = basic_stochastic.next_frame_basic_stochastic()
  hparams.optimizer = "true_adam"
  hparams.learning_rate_schedule = "constant"
  hparams.learning_rate_constant = 1e-3
  hparams.video_num_input_frames = 1
  hparams.video_num_target_frames = 3
  hparams.batch_size = 16
  hparams.bottom = {
      "inputs": modalities.video_raw_bottom,
      "targets": modalities.video_raw_targets_bottom,
  }
  hparams.loss = {
      "targets": modalities.video_l2_raw_loss,
  }
  hparams.top = {
      "targets": modalities.video_raw_top,
  }
  hparams.video_modality_loss_cutoff = 0.0
  hparams.scheduled_sampling_mode = "count"
  hparams.scheduled_sampling_k = 900.0
  hparams.add_hparam("reward_prediction", True)
  hparams.add_hparam("reward_prediction_stop_gradient", False)
  hparams.add_hparam("reward_prediction_buffer_size", 0)
  hparams.add_hparam("model_options", "CDNA")
  hparams.add_hparam("num_masks", 10)
  hparams.add_hparam("multi_latent", False)
  hparams.add_hparam("relu_shift", 1e-12)
  hparams.add_hparam("dna_kernel_size", 5)
  hparams.add_hparam("upsample_method", "conv2d_transpose")
  hparams.add_hparam("reward_model", "basic")
  hparams.add_hparam("visualize_logits_histogram", True)
  return hparams
1313481
def _define_train(
    train_env,
    ppo_hparams,
    eval_env_fn=None,
    sampling_temp=1.0,
    **collect_kwargs
):
  """"""
  memory, collect_summary, train_initialization = (
      _define_collect(
          train_env,
          ppo_hparams,
          "ppo_train",
          eval_phase=False,
          sampling_temp=sampling_temp,
          **collect_kwargs))
  ppo_summary = ppo.define_ppo_epoch(
      memory, ppo_hparams, train_env.action_space, train_env.batch_size)
  train_summary = tf.summary.merge([collect_summary, ppo_summary])

  if ppo_hparams.eval_every_epochs:
    # TODO(koz4k): Do we need this at all?
    assert eval_env_fn is not None
    eval_env = eval_env_fn(in_graph=True)
    (_, eval_collect_summary, eval_initialization) = (
        _define_collect(
            eval_env,
            ppo_hparams,
            "ppo_eval",
            eval_phase=True,
            sampling_temp=0.0,
            **collect_kwargs))
    return (train_summary, eval_collect_summary, (train_initialization,
                                                  eval_initialization))
  else:
    return (train_summary, None, (train_initialization,))
536499
def evaluate_world_model(
    agent_type, loop_hparams, planner_hparams, model_dir, policy_dir,
    random_starts_step_limit, debug_video_path, log_every_steps
):
  """"""
  if debug_video_path:
    debug_video_path = os.path.join(debug_video_path, "0.avi")

  storage_env = rl_utils.setup_env(loop_hparams, batch_size=1, max_num_noops=0)
  stacked_env = rl_utils.BatchStackWrapper(
      storage_env, loop_hparams.frame_stack_size
  )
  policy_hparams = trainer_lib.create_hparams(loop_hparams.base_algo_params)
  agent = make_agent_from_hparams(
      agent_type, storage_env, stacked_env, loop_hparams, policy_hparams,
      planner_hparams, model_dir, policy_dir,
      # TODO(koz4k): Loop over eval_sampling_temps?
      sampling_temp=loop_hparams.eval_sampling_temps[0],
  )
  collect_frames_for_random_starts(
      storage_env, stacked_env, agent, loop_hparams.frame_stack_size,
      random_starts_step_limit, log_every_steps
  )
  return rl_utils.evaluate_world_model(
      storage_env, loop_hparams, model_dir, debug_video_path, split=None
  )
1290059
def _evaluate(x, y, weights):
    '''
    
    '''
    i = _validI(x, y, weights)
    xx = x[i]
    y = y[i]

    try:
        fitParams = _fit(xx, y)
        # bound noise fn to min defined y value:
        minY = function(xx[0], *fitParams)
        fitParams = np.insert(fitParams, 0, minY)
        fn = lambda x, minY=minY: boundedFunction(x, *fitParams)
    except RuntimeError:
        print(
            "couldn't fit noise function with filtered indices, use polynomial fit instead")
        fitParams = None
        fn = smooth(xx, y, weights[i])
    return fitParams, fn, i
417062
def update_hparams_for_tpu(hparams):
  """"""

  # Adafactor uses less memory than Adam.
  # switch to Adafactor with its recommended learning rate scheme.
  hparams.optimizer = "Adafactor"
  hparams.learning_rate_schedule = "rsqrt_decay"
  hparams.learning_rate_warmup_steps = 10000

  # Avoid an expensive concat on TPU.
  # >1 shards helps with faster parameter distribution on multi-GPU machines
  hparams.symbol_modality_num_shards = 1

  # Adaptive batch sizes and sequence lengths are not supported on TPU.
  # Instead, every batch has the same sequence length and the same batch size.
  # Longer sequences are dropped and shorter ones are padded.
  #
  # It is therefore suggested to use a problem where examples have been combined
  # to a longer length, e.g. the "_packed" problems.
  #
  # For problems with variable sequence lengths, this parameter controls the
  # maximum sequence length.  Shorter sequences are dropped and longer ones
  # are padded.
  #
  # For problems with fixed sequence lengths - e.g. the "_packed" problems,
  # this hyperparameter is ignored.
  hparams.max_length = 64

  # TPUs have less memory than GPUs, so decrease the batch size
  hparams.batch_size = 2048

  # Using noise broadcast in the dropout layers saves memory during training.
  hparams.attention_dropout_broadcast_dims = "0,1"  # batch, heads
  hparams.relu_dropout_broadcast_dims = "1"  # length
  hparams.layer_prepostprocess_dropout_broadcast_dims = "1"  # length
  return hparams
1199339
def ppo_atari_base():
  """"""
  hparams = ppo_discrete_action_base()
  hparams.learning_rate_constant = 1e-4
  hparams.epoch_length = 200
  hparams.gae_gamma = 0.985
  hparams.gae_lambda = 0.985
  hparams.entropy_loss_coef = 0.003
  hparams.value_loss_coef = 1
  hparams.optimization_epochs = 3
  hparams.epochs_num = 1000
  hparams.policy_network = "feed_forward_cnn_small_categorical_policy"
  hparams.clipping_coef = 0.2
  hparams.optimization_batch_size = 20
  hparams.clip_grad_norm = 0.5
  return hparams
1228214
def _run_mplp(self, no_iterations):
        
        for niter in range(no_iterations):
            # We take the clusters in the order they were added in the model and update messages for all factors whose
            # scope is greater than 1
            for factor in self.model.get_factors():
                if len(factor.scope()) > 1:
                    self._update_message(self.cluster_set[frozenset(factor.scope())])
            # Find an integral solution by locally maximizing the single node beliefs
            self._local_decode()
            # If mplp converges to a global/local optima, we break.
            if self._is_converged(self.dual_threshold, self.integrality_gap_threshold) and niter >= 16:
                break
219907
def xmoe2_v1():
  
  hparams = xmoe2_dense(0)
  moe.set_default_moe_hparams(hparams)
  hparams.decoder_layers = (
      ["local_att", "local_att", "drd",
       "att", "drd", "local_att", "local_att", "hmoe"] * 4)[:-1]
  hparams.d_ff = 2048
  hparams.d_kv = 128
  hparams.moe_hidden_size = 32768
  hparams.mesh_shape = "b0:4;b1:8"
  hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
  hparams.outer_batch_size = 4
  hparams.moe_num_experts = [8, 4]
  hparams.num_heads = 4
  return hparams
1091571
def evaluate_world_model(
    real_env, hparams, world_model_dir, debug_video_path,
    split=tf.estimator.ModeKeys.EVAL,
):
  """"""
  frame_stack_size = hparams.frame_stack_size
  rollout_subsequences = []
  def initial_frame_chooser(batch_size):
    assert batch_size == len(rollout_subsequences)
    return np.stack([
        [frame.observation.decode() for frame in subsequence[:frame_stack_size]]    # pylint: disable=g-complex-comprehension
        for subsequence in rollout_subsequences
    ])

  env_fn = rl.make_simulated_env_fn_from_hparams(
      real_env, hparams, batch_size=hparams.wm_eval_batch_size,
      initial_frame_chooser=initial_frame_chooser, model_dir=world_model_dir
  )
  sim_env = env_fn(in_graph=False)
  subsequence_length = int(
      max(hparams.wm_eval_rollout_ratios) * hparams.simulated_rollout_length
  )
  rollouts = real_env.current_epoch_rollouts(
      split=split,
      minimal_rollout_frames=(subsequence_length + frame_stack_size)
  )

  video_writer = common_video.WholeVideoWriter(
      fps=10, output_path=debug_video_path, file_format="avi"
  )

  reward_accuracies_by_length = {
      int(ratio * hparams.simulated_rollout_length): []
      for ratio in hparams.wm_eval_rollout_ratios
  }
  for _ in range(hparams.wm_eval_num_batches):
    rollout_subsequences[:] = random_rollout_subsequences(
        rollouts, hparams.wm_eval_batch_size,
        subsequence_length + frame_stack_size
    )

    eval_subsequences = [
        subsequence[(frame_stack_size - 1):]
        for subsequence in rollout_subsequences
    ]

    # Check that the initial observation is the same in the real and simulated
    # rollout.
    sim_init_obs = sim_env.reset()
    def decode_real_obs(index):
      return np.stack([
          subsequence[index].observation.decode()
          for subsequence in eval_subsequences  # pylint: disable=cell-var-from-loop
      ])
    real_init_obs = decode_real_obs(0)
    assert np.all(sim_init_obs == real_init_obs)

    debug_frame_batches = []
    def append_debug_frame_batch(sim_obs, real_obs, sim_cum_rews,
                                 real_cum_rews, sim_rews, real_rews):
      """Add a debug frame."""
      rews = [[sim_cum_rews, sim_rews], [real_cum_rews, real_rews]]
      headers = []
      for j in range(len(sim_obs)):
        local_nps = []
        for i in range(2):
          img = PIL_Image().new("RGB", (sim_obs.shape[-2], 11),)
          draw = PIL_ImageDraw().Draw(img)
          draw.text((0, 0), "c:{:3}, r:{:3}".format(int(rews[i][0][j]),
                                                    int(rews[i][1][j])),
                    fill=(255, 0, 0))
          local_nps.append(np.asarray(img))
        local_nps.append(np.zeros_like(local_nps[0]))
        headers.append(np.concatenate(local_nps, axis=1))
      errs = absolute_hinge_difference(sim_obs, real_obs)
      headers = np.stack(headers)
      debug_frame_batches.append(  # pylint: disable=cell-var-from-loop
          np.concatenate([headers,
                          np.concatenate([sim_obs, real_obs, errs], axis=2)],
                         axis=1)
      )
    append_debug_frame_batch(sim_init_obs, real_init_obs,
                             np.zeros(hparams.wm_eval_batch_size),
                             np.zeros(hparams.wm_eval_batch_size),
                             np.zeros(hparams.wm_eval_batch_size),
                             np.zeros(hparams.wm_eval_batch_size))

    (sim_cum_rewards, real_cum_rewards) = (
        np.zeros(hparams.wm_eval_batch_size) for _ in range(2)
    )
    for i in range(subsequence_length):
      actions = [subsequence[i].action for subsequence in eval_subsequences]
      (sim_obs, sim_rewards, _) = sim_env.step(actions)
      sim_cum_rewards += sim_rewards

      real_rewards = np.array([
          subsequence[i + 1].reward for subsequence in eval_subsequences
      ])
      real_cum_rewards += real_rewards
      for (length, reward_accuracies) in six.iteritems(
          reward_accuracies_by_length
      ):
        if i + 1 == length:
          reward_accuracies.append(
              np.sum(sim_cum_rewards == real_cum_rewards) /
              len(real_cum_rewards)
          )

      real_obs = decode_real_obs(i + 1)
      append_debug_frame_batch(sim_obs, real_obs, sim_cum_rewards,
                               real_cum_rewards, sim_rewards, real_rewards)

    for debug_frames in np.stack(debug_frame_batches, axis=1):
      debug_frame = None
      for debug_frame in debug_frames:
        video_writer.write(debug_frame)

      if debug_frame is not None:
        # Append two black frames for aesthetics.
        for _ in range(2):
          video_writer.write(np.zeros_like(debug_frame))

  video_writer.finish_to_disk()

  return {
      "reward_accuracy/at_{}".format(length): np.mean(reward_accuracies)
      for (length, reward_accuracies) in six.iteritems(
          reward_accuracies_by_length
      )
  }
1367122
def parametrized_bottleneck(x, hparams):
  """"""
  if hparams.bottleneck_kind == "tanh_discrete":
    d, _ = tanh_discrete_bottleneck(
        x, hparams.bottleneck_bits, hparams.bottleneck_noise * 0.5,
        hparams.discretize_warmup_steps, hparams.mode)
    return d, 0.0
  if hparams.bottleneck_kind == "isemhash":
    return isemhash_bottleneck(
        x, hparams.bottleneck_bits, hparams.bottleneck_noise * 0.5,
        hparams.discretize_warmup_steps, hparams.mode,
        hparams.isemhash_noise_dev, hparams.isemhash_mix_prob)
  if hparams.bottleneck_kind == "vq":
    return vq_discrete_bottleneck(x, hparams.bottleneck_bits, hparams.vq_beta,
                                  hparams.vq_decay, hparams.vq_epsilon)
  if hparams.bottleneck_kind == "em":
    return vq_discrete_bottleneck(
        x,
        hparams.bottleneck_bits,
        hparams.vq_beta,
        hparams.vq_decay,
        hparams.vq_epsilon,
        soft_em=True,
        num_samples=hparams.vq_num_samples)
  if hparams.bottleneck_kind == "gumbel_softmax":
    return gumbel_softmax_discrete_bottleneck(
        x,
        hparams.bottleneck_bits,
        hparams.vq_beta,
        hparams.vq_decay,
        hparams.vq_epsilon,
        hparams.temperature_warmup_steps,
        hard=False,
        summary=True)

  raise ValueError(
      "Unsupported hparams.bottleneck_kind %s" % hparams.bottleneck_kind)
1331114
def adafactor_optimizer_from_hparams(hparams, lr):
  
  if hparams.optimizer_adafactor_decay_type == "adam":
    decay_rate = adafactor_decay_rate_adam(
        hparams.optimizer_adafactor_beta2)
  elif hparams.optimizer_adafactor_decay_type == "pow":
    decay_rate = adafactor_decay_rate_pow(
        hparams.optimizer_adafactor_memory_exponent)
  else:
    raise ValueError("unknown optimizer_adafactor_decay_type")
  if hparams.weight_dtype == "bfloat16":
    parameter_encoding = quantization.EighthPowerEncoding()
  else:
    parameter_encoding = None
  return AdafactorOptimizer(
      multiply_by_parameter_scale=(
          hparams.optimizer_adafactor_multiply_by_parameter_scale),
      learning_rate=lr,
      decay_rate=decay_rate,
      beta1=hparams.optimizer_adafactor_beta1,
      clipping_threshold=hparams.optimizer_adafactor_clipping_threshold,
      factored=hparams.optimizer_adafactor_factored,
      simulated_quantize_bits=getattr(
          hparams, "simulated_parameter_quantize_bits", 0),
      parameter_encoding=parameter_encoding,
      use_locking=False,
      name="Adafactor")
481984
def ppo_original_world_model_stochastic_discrete():
  """"""
  hparams = ppo_original_params()
  hparams.policy_network = "next_frame_basic_stochastic_discrete"
  hparams_keys = hparams.values().keys()
  video_hparams = basic_stochastic.next_frame_basic_stochastic_discrete()
  for (name, value) in six.iteritems(video_hparams.values()):
    if name in hparams_keys:
      hparams.set_hparam(name, value)
    else:
      hparams.add_hparam(name, value)
  # To avoid OOM. Probably way to small.
  hparams.optimization_batch_size = 1
  hparams.weight_decay = 0
  return hparams
885099
def lmx_base():
  """"""
  hparams = transformer.transformer_tpu()
  # sharing is counterproductive when underparameterized
  hparams.shared_embedding_and_softmax_weights = False
  # we judge by log-ppl, so label smoothing hurts.
  hparams.label_smoothing = 0.0
  # This makes the batch size on GPU the same as on TPU for a packed problem
  # with sequence length 256.
  # TODO(noam): fix the mess that is the data reading pipeline.
  hparams.max_length = 256
  # larger batch since we only have a decoder
  hparams.batch_size = 4096
  # save some memory so we can have a larger model
  hparams.activation_dtype = "bfloat16"
  return hparams
852694
def create(entropy_coefficient, value_coefficient, cliprange, max_grad_norm, discount_factor,
           normalize_advantage=True, gae_lambda=1.0):
    
    return PpoPolicyGradient(
        entropy_coefficient, value_coefficient, cliprange, max_grad_norm,
        discount_factor=discount_factor,
        normalize_advantage=normalize_advantage,
        gae_lambda=gae_lambda
    )
401205
def make_simulated_env_kwargs(real_env, hparams, **extra_kwargs):
  """"""
  objs_and_attrs = [
      (real_env, [
          "reward_range", "observation_space", "action_space", "frame_height",
          "frame_width"
      ]),
      (hparams, ["frame_stack_size", "intrinsic_reward_scale"])
  ]
  kwargs = {
      attr: getattr(obj, attr)  # pylint: disable=g-complex-comprehension
      for (obj, attrs) in objs_and_attrs for attr in attrs
  }
  kwargs["model_name"] = hparams.generative_model
  kwargs["model_hparams"] = trainer_lib.create_hparams(
      hparams.generative_model_params
  )
  if hparams.wm_policy_param_sharing:
    kwargs["model_hparams"].optimizer_zero_grads = True
  kwargs.update(extra_kwargs)
  return kwargs