Method _call

acme/tf/networks/recurrence.py:202–243

Computes a forward step for a single element. The observation and state are packed together in order to use `tf.vectorized_map` to handle batches of observations. See this module's __call__() function. Args: observation_and_state: the observation and state packed in a tuple.

(
      self, observation_and_state: Tuple[types.NestedTensor,
                                         PolicyCriticRNNState]
  )

    return tf.vectorized_map(self._call, (observation, prev_state))

  def _call(
      self, observation_and_state: Tuple[types.NestedTensor,
                                         PolicyCriticRNNState]
  ) -> Tuple[types.NestedTensor, PolicyCriticRNNState]:
    """Computes a forward step for a single element.

    The observation and state are packed together in order to use
    `tf.vectorized_map` to handle batches of observations.
    See this module's __call__() function.

    Args:
      observation_and_state: the observation and state packed in a tuple.

    Returns:
      The selected action and the corresponding state.
    """
    observation, prev_state = observation_and_state

    # Tile input observations and states to allow multiple policy predictions.
    tiled_observation, tiled_prev_state = utils.tile_nested(
        (observation, prev_state), self._num_action_samples)
    actions, policy_states = self._policy_network(
        tiled_observation, tiled_prev_state.policy)

    # Evaluate multiple critic predictions with the sampled actions.
    value_distribution, critic_states = self._critic_network(
        tiled_observation, actions, tiled_prev_state.critic)
    value_estimate = value_distribution.mean()

    # Resample a single action of the sampled actions according to logits given
    # by the tempered Q-values.
    selected_action_idx = tfp.distributions.Categorical(
        probs=tf.nn.softmax(value_estimate / self._temperature_beta)).sample()
    selected_action = actions[selected_action_idx]

    # Select and return the RNN state that corresponds to the selected action.
    states = PolicyCriticRNNState(
        policy=policy_states, critic=critic_states)
    selected_state = tree.map_structure(
        lambda x: x[selected_action_idx], states)

    return selected_action, selected_state

  def initial_state(self, batch_size: int) -> PolicyCriticRNNState:
    return PolicyCriticRNNState(