class_docstr: Computes PWIL rewards along a trajectory.
  
  The rewards measure similarity to the demonstration transitions and are based
  on a greedy approximation to the Wasserstein distance between trajectories.
docstr: Computes reward and updates state, advancing it along a trajectory.
  
  Subsequent calls to append_and_compute_reward assume inputs are subsequent
  trajectory points.
  
  Args:
    observation: observation on a trajectory, to compare with the expert
      demonstration(s).
    action: the action following the observation on the trajectory.
  
  Returns:
    the reward value: the return contribution from the trajectory point.
code: def append_and_compute_reward(
      self, observation: jnp.ndarray, action: jnp.ndarray
  ) -> np.float32:
      """Computes reward and updates state, advancing it along a trajectory.
  
      Subsequent calls to append_and_compute_reward assume inputs are subsequent
      trajectory points.
  
      Args:
        observation: observation on a trajectory, to compare with the expert
          demonstration(s).
        action: the action following the observation on the trajectory.
  
      Returns:
        the reward value: the return contribution from the trajectory point.
  
      """
      # If we run out of demonstrations, penalize further action.
      if self._all_expert_weights_zero:
          return np.float32(0.0)
  
      # Scale observation and action.
      if self._use_actions_for_distance:
          agent_atom = np.concatenate([observation, action])
      else:
          agent_atom = observation
      agent_atom /= self._std
  
      cost = 0.0
      # A special marker for records with zero expert weight. Has to be large so
      # that argmin will not return it.
      DELETED = 1e10  # pylint: disable=invalid-name
      # As we match the expert's weights with the agent's weights, we might
      # raise an error due to float precision, we substract a small epsilon from
      # the agent's weights to prevent that.
      weight = 1.0 / self._episode_length - 1e-6
      norms = np.array(
          self._compute_norm(self.expert_atoms, agent_atom)
      )
      # We need to mask out states with zero weight, so that 'argmin' would not
      # return them.
      adjusted_norms = (
          1 - np.sign(self.expert_weights)
      ) * DELETED + norms
      while weight > 0:
          # Get closest expert state action to agent's state action.
          argmin = adjusted_norms.argmin()
          effective_weight = min(weight, self.expert_weights[argmin])
  
          if adjusted_norms[argmin] >= DELETED:
              self._all_expert_weights_zero = True
              break
  
          # Update cost and weights.
          weight -= effective_weight
          self.expert_weights[argmin] -= effective_weight
          cost += effective_weight * norms[argmin]
          adjusted_norms[argmin] = DELETED
  
      if weight > 0:
          # We have a 'partial' cost if we ran out of demonstrations in the reward
          # computation loop. We assign a high cost (infinite) in this case which
          # makes the reward equal to 0.
          reward = np.array(0.0)
      else:
          reward = self._reward_scale * np.exp(
              -self._reward_sigma * cost
          )
  
      return reward.astype("float32")
  
file_path: acme/acme/agents/jax/pwil/rewarder.py
imports: from typing import Iterator
  from acme import types
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 4 for calling the WassersteinDistanceRewarder obj's underlying member method code with example usage: obj.append_and_compute_reward(observation, action) and returns: reward, np