Improving design of Learner.step: adding optional parameter sample

This could be a huge change in acme's API design but I would like to make a proposal or initiate the discussion about how we can change and improve the design of one of the acme's central APIs: `Learner.step()`.

In short, I think we could benefit **greatly** from adding an optional parameter `sample` to `Learner`. Specifically, we could design the method `Learner.step` such that it takes an optinal parameter `samples`, and it performs the usual learning step from the provided `sample` if given, or sample from `self._iterator` (the current behavior) otherwise.

```diff
class Learner(...):
@abc.abstractmethod
- def step(self):
+ def step(self, sample: Optional[reverb.ReplaySample]):
"""Perform an update step of the learner's parameters."""
```

Then, a typical implementation would look like:
```python
def step(self, sample: Optional[reverb.ReplaySample] = None):
sample = next(self._iterator)
transitions = types.Transition(*sample.data)

self._state, metrics = self._update_step(self._state, transitions)
```

### Why do we need it?

It will allow much easier extension and customization of learning behaviors.

One example is overriding reward functions; for example, intrinsic rewards. Some example sthat makes use of intrinsic reward functions are AIL (imitaiton learning) and RND, where the agent is built on top of underlying RL algorithm (called "direct RL" learner) but whose reward function is redefined by some other components: intrinsic reward computed by RND network or discriminator (AIL). In such cases, one needs to *process* and *override* the reward of a sample because the extrinsic/task reward may not be used.

However, doing this is very complicated with the current form of the API. Because there is **no way** to access or override the individual samples being sampled from the dataset iterator (usually reverb dataset) outside the learner. In my view [the current way how this is achieved](https://github.com/deepmind/acme/blob/master/acme/agents/jax/ail/learning.py#L161-L169) looks somewhat intimidating or overly complicated: the `iterator` of the underlying direct RL learner will be `tee`-ed and passed through `process_sample` which is able to arbitrarily process a sample data. To inject this iterator to the direct RL learner, AILLearner or RNDLearner would need to be provided with `direct_rl_learner_factory` that creates a direct RL learner when passed such a "processed" replay sample iterator. This approach works well with the current API design, but duplicating the iterator is error-prone to some errors about update-sample ratio when rate_limiter is engaged. Readability of code due to such encapsulation and difficult-to-track lambda functions would be another concern.

On the other hand, if we had this optional parameter, a design of nested Learners will be much simpler. For example:
```python
class AILLearner:
# ...
def step(self, sample=None):
sample: AILSample = self._iterator()
# ...
self._direct_rl_learner.step(self._process_sample(sample.rl_sample))
# ...

```

It will also enable implementing a custom learning process, for instance, (optional) on-policy learning that is not necessarily tied with the reverb replay buffer through `dataset` or `learner._iterator`.

### What would be concerns of making such changes?

- We need to ensure that the type of `sample` is compatible with that of `self._iterator`. Usually it is `replay.ReplaySample`, but in principle learner should not heavily coupled with `ReplaySample` itself. As in the above example, `AILLearner` expects `samples` to be `AILSample` rather than the raw `ReplaySample`; so the type of `sample` across different learners can be different and it actually needs to be generic.
- Since this is a breaking change of API requirement, some third-party or external use cases of Learner might not be compatible with this (see the alternative below).

In addition, I believe there might be other reasons why DM hasn't taken this approach in the first place. Any thoughts behind the design rationale?

### Alternatives?

An alternative, less-intrusive approach is to add an unified method like `def update_step(sample)` for `acme.Learner` or its specific (abstract) subclass, pretty similar to `SACLearner._update_step(learning_state, transitions)` or `self._sgd_step(learning_state, sample)`, etc. In this way we could keep the `acme.Learner.step(...)` intact. Currently all the subclass implementation of Learner do this all differently without a common interface --- protected method names are different, and signatures are different (some are taking `Transition` and others `ReplaySample`), and those functions are even hidden by JIT (`jax.jit` or `tf.function`) compilation, which makes overriding pretty difficult as usually are defined in the constructor.

If we take this approach, a refactoring will be required to unify types (with some use of Generics of course) and the signature of `sample`. We may want to add make this polymorphic extension engaged only for more specific, concrete subclass (e.g., `GenericLearner` , in a similar fashion as `GenericActor` extends `Actor`) to avoid a breaking API change in `acme.Learner`.

Example:

```python
Sample = acme.jax.types.Sample # TypeVar

class GenericLearner(Learner, Generic[Sample]):
@abstractmethod
def step(step, sample: Optional[Sample] = None):
...
```
It would be controversial whether subclass methods may not have different (additional) function signature, so this might not be a good idea. Instead, one could do like (apologizes for tentative naming and sketch designs)
```python
Sample = acme.jax.types.Sample # TypeVar
LearnerState = TypeVar('LearnerState')

class GenericLearner(Learner, Generic[LearnerState, Sample]):
def __init__(...):
self._logger = ...
self._state = self.make_initial_state(random_key)
self._counter = ...

def step(self):
"""An unified, common implementation of learner."""
sample = next(self._iterator)
transitions = types.Transition(*sample.data)
self._state, metrics = self.update_step(self._state, transitions)
# Optionally, measure elapsed time, etc.
counts = self._counter.increment(steps=1, walltime=elapsed_time)
self._logger.write({**metrics, **counts})

@abstractmethod
def make_initial_state(self, key: networks_lib.PRNGKey
) -> LearnerState:
...

@abstractmethod
def update_step(self, state: LearnerState, sample: Sample
) -> Tuple[LearnerState, Metrics]:
# TODO: what about jax.jit?
...
```