[Feature Request] Add batch support to CEM

[luisenp]

🚀 Feature Request

Add a batch dimension to all operations in CEM so that it can optimize several problems at once.

Motivation

The current implementation of our CEM optimizer only supports one optimization problem, which means that, for example, it cannot be used to compute action for several observations in parallel, which in turn slows down model evaluation.

Is your feature request related to a problem? Please describe.

Pitch

Add batch support to all operations involved in CEM, so that the return is an optimal value for each element in the batch.