(self, obs, use_eval_model=True)
return values
Runs value model.
Runs value model.
Runs value model .
def _run_value_model(self, obs, use_eval_model=True):
    """Runs value model."""
    n_devices = fastmath.local_device_count()
    if use_eval_model:
      weights = tl.for_n_devices(self._value_eval_model.weights, n_devices)
      state = tl.for_n_devices(self._value_eval_model.state, n_devices)
      rng = self._value_eval_model.rng
    else:
      # TODO(henrykm): this strangely looking solution address the problem that
      # value_batches_stream calls _run_value_model _once_ before
      # the trainer is initialized.
      try:
        weights = tl.for_n_devices(self._value_trainer.model_weights, n_devices)
        state = tl.for_n_devices(self._value_trainer.model_state, n_devices)
        rng = self._value_trainer._rng  # pylint: disable=protected-access
      except AttributeError:
        weights = tl.for_n_devices(self._value_eval_model.weights, n_devices)
        state = tl.for_n_devices(self._value_eval_model.state, n_devices)
        rng = self._value_eval_model.rng
    # TODO(henrykm): the line below fails on TPU with the error
    # ValueError: Number of devices (8) does not evenly divide batch size (1).
    obs_batch = obs.shape[0]
    if n_devices > obs_batch:
      obs = jnp.repeat(obs, int(n_devices / obs_batch), axis=0)
    values, _ = self._value_eval_jit(obs, weights, state, rng)
    values = values[:obs_batch]
    values *= self._value_network_scale
    return values
def _run_value_model ( self , obs , use_eval_model = True ) : n_devices = fastmath . local_device_count ( ) if use_eval_model : weights = tl . for_n_devices ( self . _value_eval_model . weights , n_devices ) state = tl . for_n_devices ( self . _value_eval_model . state , n_devices ) rng = self . _value_eval_model . rng else : # TODO(henrykm): this strangely looking solution address the problem that # value_batches_stream calls _run_value_model _once_ before # the trainer is initialized. try : weights = tl . for_n_devices ( self . _value_trainer . model_weights , n_devices ) state = tl . for_n_devices ( self . _value_trainer . model_state , n_devices ) rng = self . _value_trainer . _rng # pylint: disable=protected-access except AttributeError : weights = tl . for_n_devices ( self . _value_eval_model . weights , n_devices ) state = tl . for_n_devices ( self . _value_eval_model . state , n_devices ) rng = self . _value_eval_model . rng # TODO(henrykm): the line below fails on TPU with the error # ValueError: Number of devices (8) does not evenly divide batch size (1). obs_batch = obs . shape [ 0 ] if n_devices > obs_batch : obs = jnp . repeat ( obs , int ( n_devices / obs_batch ) , axis = 0 ) values , _ = self . _value_eval_jit ( obs , weights , state , rng ) values = values [ : obs_batch ] values *= self . _value_network_scale return values