class_docstr: RND learner.
class_docstr_tok: ["RND","learner","."]
docstr: Uses the replay sample to train and update its reward.
  
  Args:
    sample: Replay sample to train on.
  
  Returns:
    The sample replay sample with an updated reward.
docstr_tok: ["Uses","the","replay","sample","to","train","and","update","its","reward",".","Args",":","sample",":","Replay","sample","to","train","on",".","Returns",":","The","sample","replay","sample","with","an","updated","reward","."]
code: def _process_sample(
      self, sample: reverb.ReplaySample
  ) -> reverb.ReplaySample:
      """Uses the replay sample to train and update its reward.
  
      Args:
        sample: Replay sample to train on.
  
      Returns:
        The sample replay sample with an updated reward.
      """
      transitions = reverb_utils.replay_sample_to_sars_transition(
          sample, is_sequence=self._is_sequence_based
      )
      self._state, metrics = self._update(self._state, transitions)
      rewards = self._get_reward(
          self._state.params, self._state.target_params, transitions
      )
  
      # Compute elapsed time.
      timestamp = time.time()
      elapsed_time = (
          timestamp - self._timestamp if self._timestamp else 0
      )
      self._timestamp = timestamp
  
      # Increment counts and record the current time
      counts = self._counter.increment(steps=1, walltime=elapsed_time)
  
      # Attempts to write the logs.
      self._logger.write({**metrics, **counts})
  
      return sample._replace(data=sample.data._replace(reward=rewards))
  
imports: import functools
  import time
  from typing import Any, Callable, Dict, Iterator, List, NamedTuple, Optional, Tuple
  import acme
  from acme import types
  from acme.agents.jax.rnd import networks
  from acme.jax import networks
  from acme.jax import utils
  from acme.utils import counting
  from acme.utils import loggers
  from acme.utils import reverb_utils
  import jax
  import jax.numpy
  import optax
  import reverb
label_desc: Use image node_id 2 for calling the RNDLearner obj's underlying member method code with example usage: obj._process_sample(sample) and returns: sample