[RLlib; Offline RL] Add CQLLearner and CQLTorchLearner.

[simonsays1980]

Why are these changes needed?

This PR proposes the Learner classes for CQL in the new API stack. It relies on the SACTorchModule (i.e. having a Q-network and Q-target-network) and implements the loss logic. Because three loss terms are optimized the gradients are collected inside the compute_loss_for_module method and are then just returned from compute_gradients. The reason for this are multiple passes through the networks where a follow-up path would erase the gradients computed in the pass before.

This PR is one of several PRs that should implement CQL in the new API stack.

Related issue number

Closes #37779

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[sven1977]

Can we explain this logic here and why we defend against different batch sizes coming in?

Don't we expect always the same batch size from the data pipeline?

If not, we should:
a) explain here why we are expecting various batch sizes
b) probably fix this logic here. What if we call compute_loss_for_module 10x with a batch of train_batch_size_per_learner - 1 (accumulating gradients for these, but not applying these) and then one batch of size train_batch_size_per_learner. In this case, the effective batch size would be roughly 11x the user configured one, correct?

c) Also, what if the incoming batch is larger than train_batch_size_per_learner?

[simonsays1980]

@sven1977 I fully agree on your comment.

a) This was also part of the old stack algorithm. The reason behind this is that when iterating over a dataset the last batch could have less than train_batch_size_per_learner samples in it.

b) The offline logic here is that after each call to compute_loss_for_module a compute_gradients and apply_gradients will occur. No SGD is run on the offline algorithms (yet). So the train batch size should be always as large as configured (without the last one being smaller).

c) This case should also not happen. iter_batches takes care of this and ensures that always the batch size is sampled, without the last one which can be avoided by setting a flag, but neglects data.

[sven1977]

Can we add more explanation here as to why caching any grads is necessary?

Maybe I don't understand the data pipelines properly yet, but should we try to fix these to always produce the exact same batch sizes (or number of episode timesteps) per iteration?

[simonsays1980]

This is not really caused by the batch sizes, but will only be executed if the batch size is as configured. The reason for the grad caching is merely to enable multiple passes through the network during loss calculation.

[sven1977]

Add: reduce="sum" here to make sure iterations are not EMA'd :) (the default).

[sven1977]

Actually, you don't need this here anymore. B/c you are logging it below, this entry will be created automatically.

[sven1977]

New that reduce="sum" is already defined above, you can just do:

log_value(
    (ALL_MODULES, TRAINING_ITERATION),
    1,
    reduce="sum",
)

[sven1977]

Change to:

self.metrics.peek(([your key]), default=0)

[sven1977]

To cover the case, in which this entry does not exist yet.
This way, you don't have to define it up front.

Just log into it as reduce="sum" (and value=1) and all will be good.

[sven1977]

Just some nits left, otherwise good to go! Awesome work @simonsays1980 .