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: Initializes the rewarder.
  
  Args:
    demonstrations_it: An iterator over acme.types.Transition.
    episode_length: a target episode length (policies will be encouraged by
      the imitation reward to have that length).
    use_actions_for_distance: whether to use action to compute reward.
    alpha: float scaling the reward function.
    beta: float controling the kernel size of the reward function.
code: def __init__(
      self,
      demonstrations_it: Iterator[types.Transition],
      episode_length: int,
      use_actions_for_distance: bool = False,
      alpha: float = 5.0,
      beta: float = 5.0,
  ):
      """Initializes the rewarder.
  
      Args:
        demonstrations_it: An iterator over acme.types.Transition.
        episode_length: a target episode length (policies will be encouraged by
          the imitation reward to have that length).
        use_actions_for_distance: whether to use action to compute reward.
        alpha: float scaling the reward function.
        beta: float controling the kernel size of the reward function.
      """
      self._episode_length = episode_length
  
      self._use_actions_for_distance = use_actions_for_distance
      self._vectorized_demonstrations = self._vectorize(
          demonstrations_it
      )
  
      # Observations and actions are flat.
      atom_dims = self._vectorized_demonstrations.shape[1]
      self._reward_sigma = (
          beta * self._episode_length / np.sqrt(atom_dims)
      )
      self._reward_scale = alpha
  
      self._std = np.std(
          self._vectorized_demonstrations, axis=0, dtype="float64"
      )
      # The std is set to 1 if the observation values are below a threshold.
      # This prevents normalizing observation values that are constant (which can
      # be problematic with e.g. demonstrations coming from a different version
      # of the environment and where the constant values are slightly different).
      self._std = (self._std < 1e-6) + self._std
  
      self.expert_atoms = self._vectorized_demonstrations / self._std
      self._compute_norm = jax.jit(
          lambda a, b: jnp.linalg.norm(a - b, axis=1),
          device=jax.local_devices(backend="cpu")[0],
      )
  
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 1 to create a new WassersteinDistanceRewarder object  with example: obj = WassersteinDistanceRewarder(demonstrations_it, episode_length, use_actions_for_distance, alpha, beta)